src/libANGLE/renderer/vulkan/CLKernelVk.cpp - third_party/angle - Git at Google

 //
 // Copyright 2021 The ANGLE Project Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 //
 // CLKernelVk.cpp: Implements the class methods for CLKernelVk.
 //

 #include "libANGLE/renderer/vulkan/CLKernelVk.h"
 #include "libANGLE/renderer/vulkan/CLContextVk.h"
 #include "libANGLE/renderer/vulkan/CLDeviceVk.h"
 #include "libANGLE/renderer/vulkan/CLMemoryVk.h"
 #include "libANGLE/renderer/vulkan/CLProgramVk.h"
 #include "libANGLE/renderer/vulkan/cl_types.h"
 #include "libANGLE/renderer/vulkan/vk_wrapper.h"

 #include "libANGLE/CLBuffer.h"
 #include "libANGLE/CLContext.h"
 #include "libANGLE/CLDevice.h"
 #include "libANGLE/CLKernel.h"
 #include "libANGLE/CLProgram.h"

 #include "spirv/unified1/NonSemanticClspvReflection.h"

 #include <algorithm>

 namespace rx
 {

 // Setup the CLKernelArgument with info from ClspvKernelArgument
 CLKernelArgument::CLKernelArgument(const CLContextVk *context, const ClspvKernelArgument &info)
     : mCompiledInfo(info), mSize(0), mHandle(nullptr)
 {
     mArgInfo.typeQualifier    = mCompiledInfo.info.typeQualifier;
     mArgInfo.addressQualifier = mCompiledInfo.info.addressQualifier;
     mArgInfo.accessQualifier  = mCompiledInfo.info.accessQualifier;
     mArgInfo.typeName         = mCompiledInfo.info.typeName;
     mArgInfo.name             = mCompiledInfo.info.name;
 }

 cl::Memory *CLKernelArgument::getMemoryHandle() const
 {
     if (!mCompiledInfo.used)
     {
         return nullptr;
     }
     return cl::Memory::Cast(static_cast<cl_mem>(mHandle));
 }

 cl::Sampler *CLKernelArgument::getSamplerHandle() const
 {
     ASSERT(mCompiledInfo.type == NonSemanticClspvReflectionArgumentSampler);
     return cl::Sampler::Cast(*static_cast<const cl_sampler *>(mHandle));
 }

 bool CLKernelArgument::isReadOnly() const
 {
     // if not used, can safely assume readonly
     if (!mCompiledInfo.used)
     {
         return true;
     }
     switch (getReflectionType())
     {
         case NonSemanticClspvReflectionArgumentPodUniform:
         case NonSemanticClspvReflectionArgumentUniform:
         case NonSemanticClspvReflectionArgumentPointerUniform:
         case NonSemanticClspvReflectionArgumentUniformTexelBuffer:
         case NonSemanticClspvReflectionConstantDataStorageBuffer:
         {
             return true;
         }
         case NonSemanticClspvReflectionArgumentStorageBuffer:
         case NonSemanticClspvReflectionArgumentStorageTexelBuffer:
         case NonSemanticClspvReflectionArgumentStorageImage:
         case NonSemanticClspvReflectionArgumentSampledImage:
         {
             return mArgInfo.accessQualifier == CL_KERNEL_ARG_ACCESS_READ_ONLY;
         }
         default:
         {
             return false;
         }
     }
 }

 CLKernelVk::CLKernelVk(const cl::Kernel &kernel, std::string &name, std::string &attributes)
     : CLKernelImpl(kernel),
       mProgram(&kernel.getProgram().getImpl<CLProgramVk>()),
       mContext(&kernel.getProgram().getContext().getImpl<CLContextVk>()),
       mName(name),
       mAttributes(attributes),
       mArgs(),
       mPodBuffer(nullptr),
       mLocalMemoryArgSizes(
           mProgram->getDeviceProgramData(mName.c_str())->getKernelArguments(mName.c_str()).size(),
           0)
 {
     mShaderProgramHelper.setShader(gl::ShaderType::Compute,
                                    mKernel.getProgram().getImpl<CLProgramVk>().getShaderModule());
 }

 CLKernelVk::~CLKernelVk()
 {
     mComputePipelineCache.destroy(mContext);
     mShaderProgramHelper.destroy(mContext->getRenderer());
 }

 angle::Result CLKernelVk::init()
 {
     const CLProgramVk::DeviceProgramData *deviceProgramData =
         mProgram->getDeviceProgramData(mName.c_str());

     // Literal sampler handling
     for (const ClspvLiteralSampler &literalSampler :
          deviceProgramData->reflectionData.literalSamplers)
     {
         mDescriptorSetLayoutDescs[DescriptorSetIndex::LiteralSampler].addBinding(
             literalSampler.binding, VK_DESCRIPTOR_TYPE_SAMPLER, 1, VK_SHADER_STAGE_COMPUTE_BIT,
             nullptr);
     }

     // Setup the arguments
     for (auto &clspvArg : deviceProgramData->getKernelArguments(mName.c_str()))
     {
         mArgs.push_back(std::make_shared<CLKernelArgument>(mContext, clspvArg));
     }

     vk::DescriptorSetLayoutDesc &descriptorSetLayoutDesc =
         mDescriptorSetLayoutDescs[DescriptorSetIndex::KernelArguments];
     VkPushConstantRange pcRange = deviceProgramData->pushConstRange;
     size_t podBufferSize        = 0;

     bool podFound = false;
     for (const auto arg : getArgs())
     {
         VkDescriptorType descType = VK_DESCRIPTOR_TYPE_MAX_ENUM;
         switch (arg->getReflectionType())
         {
             case NonSemanticClspvReflectionArgumentStorageBuffer:
                 descType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
                 break;
             case NonSemanticClspvReflectionArgumentUniform:
                 descType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
                 break;
             case NonSemanticClspvReflectionArgumentPodUniform:
             case NonSemanticClspvReflectionArgumentPodStorageBuffer:
             case NonSemanticClspvReflectionArgumentPointerUniform:
             {
                 uint32_t newPodBufferSize =
                     arg->getPodStorageBufferOffset() + arg->getPodStorageBufferSize();
                 podBufferSize = newPodBufferSize > podBufferSize ? newPodBufferSize : podBufferSize;
                 if (podFound)
                 {
                     continue;
                 }
                 descType =
                     arg->getReflectionType() == NonSemanticClspvReflectionArgumentPodStorageBuffer
                         ? VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
                         : VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
                 podFound = true;
                 break;
             }
             case NonSemanticClspvReflectionArgumentPodPushConstant:
             case NonSemanticClspvReflectionArgumentPointerPushConstant:
                 // Get existing push constant range and see if we need to update
                 if (arg->getPushConstantOffset() + arg->getPushConstantSize() >
                     pcRange.offset + pcRange.size)
                 {
                     pcRange.size =
                         arg->getPushConstantOffset() + arg->getPushConstantSize() - pcRange.offset;
                 }
                 continue;
             case NonSemanticClspvReflectionArgumentSampledImage:
                 descType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE;
                 break;
             case NonSemanticClspvReflectionArgumentStorageImage:
                 descType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
                 break;
             case NonSemanticClspvReflectionArgumentSampler:
                 descType = VK_DESCRIPTOR_TYPE_SAMPLER;
                 break;
             case NonSemanticClspvReflectionArgumentStorageTexelBuffer:
                 descType = VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
                 break;
             case NonSemanticClspvReflectionArgumentUniformTexelBuffer:
                 descType = VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER;
                 break;
             default:
                 continue;
         }
         if (descType != VK_DESCRIPTOR_TYPE_MAX_ENUM)
         {
             descriptorSetLayoutDesc.addBinding(arg->getDescriptorBinding(), descType, 1,
                                                VK_SHADER_STAGE_COMPUTE_BIT, nullptr);
         }
     }

     if (podBufferSize > 0)
     {
         mPodBuffer = cl::BufferPtr::Create(
             const_cast<cl::Context &>(mKernel.getProgram().getContext()), cl::Memory::PropArray{},
             cl::MemFlags(CL_MEM_READ_ONLY), podBufferSize, nullptr);
     }

     if (usesPrintf() && !usesPrintfBufferPointerPushConstant())
     {
         mDescriptorSetLayoutDescs[DescriptorSetIndex::Printf].addBinding(
             deviceProgramData->reflectionData.printfBufferStorage.binding,
             VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr);
     }

     // Get pipeline layout from cache (creates if missed)
     // A given kernel need not have resulted in use of all the descriptor sets. Unless the
     // graphicsPipelineLibrary extension is supported, the pipeline layout need all the descriptor
     // set layouts to be valide. So set them up in the order of their occurrence.
     mPipelineLayoutDesc = {};
     for (DescriptorSetIndex index : angle::AllEnums<DescriptorSetIndex>())
     {
         if (!mDescriptorSetLayoutDescs[index].empty())
         {
             mPipelineLayoutDesc.updateDescriptorSetLayout(index, mDescriptorSetLayoutDescs[index]);
         }
     }

     // push constant range size and offset both need to be a multiple of 4
     pcRange.size   = roundUpPow2(pcRange.size, 4u);
     pcRange.offset = roundDownPow2(pcRange.offset, 4u);
     mPipelineLayoutDesc.updatePushConstantRange(pcRange.stageFlags, pcRange.offset, pcRange.size);

     mPodArgumentPushConstants.resize(pcRange.size);

     angle::EnumIterator<DescriptorSetIndex> layoutIndex(DescriptorSetIndex::LiteralSampler);
     for (DescriptorSetIndex index : angle::AllEnums<DescriptorSetIndex>())
     {
         if (!getDescriptorSetLayoutDesc(index).empty())
         {
             ANGLE_CL_IMPL_TRY_ERROR(mContext->getDescriptorSetLayoutCache()->getDescriptorSetLayout(
                                         mContext, getDescriptorSetLayoutDesc(index),
                                         &getDescriptorSetLayouts()[*layoutIndex]),
                                     CL_INVALID_OPERATION);
             ASSERT(getDescriptorSetLayouts()[*layoutIndex]->valid());
             ++layoutIndex;
         }
     }
     ANGLE_CL_IMPL_TRY_ERROR(initPipelineLayout(), CL_INVALID_OPERATION);

     // initialize the descriptor pools
     // descriptor pools are setup as per their indices
     return mContext->initializeDescriptorPools(this);
 }

 angle::Result CLKernelVk::setArg(cl_uint argIndex, size_t argSize, const void *argValue)
 {
     auto &arg = mArgs.at(argIndex);
     if (arg->getUsed())
     {
         switch (arg->getReflectionType())
         {
             case NonSemanticClspvReflectionArgumentPodPushConstant:
             {
                 ASSERT(mPodArgumentPushConstants.size() >=
                        arg->getPushConstantSize() + arg->getPushConstantOffset());
                 void *handle = &mPodArgumentPushConstants[arg->getPushConstantOffset()];
                 size_t handleSize =
                     (argSize > arg->getPushConstantSize() ? arg->getPushConstantSize() : argSize);
                 if (argSize > 0 && argValue != nullptr)
                 {
                     // Copy the contents since app is free to delete/reassign the contents after
                     ANGLE_UNSAFE_TODO(memcpy(handle, argValue, handleSize));
                 }
                 arg->set(handleSize, handle);
                 break;
             }
             case NonSemanticClspvReflectionArgumentPodUniform:
             case NonSemanticClspvReflectionArgumentPodStorageBuffer:
             {
                 ASSERT(mPodBuffer->getSize() >= argSize + arg->getPodUniformOffset());
                 if (argSize > 0 && argValue != nullptr)
                 {
                     ANGLE_TRY(mPodBuffer->getImpl<CLBufferVk>().copyFrom(
                         argValue, arg->getPodStorageBufferOffset(), argSize));
                 }
                 break;
             }
             case NonSemanticClspvReflectionArgumentUniform:
             case NonSemanticClspvReflectionArgumentStorageBuffer:
             {
                 ASSERT(argSize == sizeof(cl_mem *));
                 size_t handleSize = argSize;
                 void *handle      = nullptr;
                 // Argument can be null ptr
                 if (argValue && *static_cast<const cl_mem *>(argValue))
                 {
                     handle = *static_cast<const cl_mem *>(argValue);
                 }
                 arg->set(handleSize, handle);
                 break;
             }
             case NonSemanticClspvReflectionArgumentStorageImage:
             case NonSemanticClspvReflectionArgumentSampledImage:
             case NonSemanticClspvReflectionArgumentUniformTexelBuffer:
             case NonSemanticClspvReflectionArgumentStorageTexelBuffer:
             case NonSemanticClspvReflectionArgumentPointerPushConstant:
             case NonSemanticClspvReflectionArgumentPointerUniform:
             {
                 ASSERT(argSize == sizeof(cl_mem *));
                 cl_mem memHandle = *static_cast<const cl_mem *>(argValue);

                 if (cl::Memory::Cast(memHandle)->getFlags().intersects(CL_MEM_READ_ONLY) &&
                     !arg->isReadOnly())
                 {
                     WARN() << "The cl_mem argument at index:" << argIndex
                            << " is setup as read-only, and the kernel is writing to it -- This "
                               "could lead to undefined behavior";
                 }
                 arg->set(argSize, memHandle);
                 break;
             }
             case NonSemanticClspvReflectionArgumentWorkgroup:
                 ASSERT(arg->getWorkGroupElementSize() != 0);
                 mLocalMemoryArgSizes[argIndex] = argSize;
                 arg->set(argSize, const_cast<void *>(argValue));
                 break;
             default:
                 // Just store ptr and size (if we end up here)
                 arg->set(argSize, const_cast<void *>(argValue));
                 break;
         }
     }

     return angle::Result::Continue;
 }

 angle::Result CLKernelVk::createInfo(CLKernelImpl::Info *info) const
 {
     info->functionName = mName;
     info->attributes   = mAttributes;
     info->numArgs      = static_cast<cl_uint>(mArgs.size());
     for (const auto arg : mArgs)
     {
         info->args.push_back(arg->getArgInfo());
     }

     auto &ctx = mKernel.getProgram().getContext();
     info->workGroups.resize(ctx.getDevices().size());
     const CLProgramVk::DeviceProgramData *deviceProgramData = nullptr;
     for (auto i = 0u; i < ctx.getDevices().size(); ++i)
     {
         auto &workGroup     = info->workGroups[i];
         const auto device   = &ctx.getDevices()[i];
         const auto deviceVk = &device->get()->getImpl<CLDeviceVk>();
         deviceProgramData   = mProgram->getDeviceProgramData(ctx.getDevices()[i]->getNative());
         if (deviceProgramData == nullptr)
         {
             continue;
         }

         // TODO: http://anglebug.com/42267005
         ANGLE_TRY(
             deviceVk->getInfoSizeT(cl::DeviceInfo::MaxWorkGroupSize, &workGroup.workGroupSize));

         // TODO: http://anglebug.com/42267004
         workGroup.privateMemSize = 0;

         workGroup.localMemSize = deviceProgramData->reflectionData.workgroupVariableSize.size;
         workGroup.prefWorkGroupSizeMultiple = deviceVk->getWorkGroupSizeMultiple();
         workGroup.globalWorkSize            = {0, 0, 0};
         if (deviceProgramData->reflectionData.kernelCompileWorkgroupSize.contains(mName))
         {
             workGroup.compileWorkGroupSize = {
                 deviceProgramData->reflectionData.kernelCompileWorkgroupSize.at(mName)[0],
                 deviceProgramData->reflectionData.kernelCompileWorkgroupSize.at(mName)[1],
                 deviceProgramData->reflectionData.kernelCompileWorkgroupSize.at(mName)[2]};
         }
         else
         {
             workGroup.compileWorkGroupSize = {0, 0, 0};
         }

         if (device->get()->getInfo().khrSubgroups)
         {
             workGroup.subGroupSizeForNDRange =
                 mContext->getRenderer()->getPhysicalDeviceSubgroupProperties().subgroupSize;
         }
         else
         {
             workGroup.subGroupSizeForNDRange = 0;
         }
     }

     return angle::Result::Continue;
 }

 angle::Result CLKernelVk::initPipelineLayout()
 {
     PipelineLayoutCache *pipelineLayoutCache = mContext->getPipelineLayoutCache();
     return pipelineLayoutCache->getPipelineLayout(mContext, mPipelineLayoutDesc,
                                                   mDescriptorSetLayouts, &mPipelineLayout);
 }

 angle::Result CLKernelVk::getOrCreateComputePipeline(vk::PipelineCacheAccess *pipelineCache,
                                                      const cl::NDRange &ndrange,
                                                      const cl::Device &device,
                                                      vk::PipelineHelper **pipelineOut)
 {
     const CLProgramVk::DeviceProgramData *devProgramData =
         getProgram()->getDeviceProgramData(device.getNative());
     ASSERT(devProgramData != nullptr);

     // Populate program specialization constants (if any)
     uint32_t constantDataOffset = 0;
     std::vector<uint32_t> specConstantData;
     std::vector<VkSpecializationMapEntry> mapEntries;
     for (const auto specConstantUsed : devProgramData->reflectionData.specConstantsUsed)
     {
         switch (specConstantUsed)
         {
             case SpecConstantType::WorkDimension:
                 specConstantData.push_back(ndrange.workDimensions);
                 break;
             case SpecConstantType::WorkgroupSizeX:
                 specConstantData.push_back(ndrange.localWorkSize[0]);
                 break;
             case SpecConstantType::WorkgroupSizeY:
                 specConstantData.push_back(ndrange.localWorkSize[1]);
                 break;
             case SpecConstantType::WorkgroupSizeZ:
                 specConstantData.push_back(ndrange.localWorkSize[2]);
                 break;
             case SpecConstantType::GlobalOffsetX:
                 specConstantData.push_back(ndrange.globalWorkOffset[0]);
                 break;
             case SpecConstantType::GlobalOffsetY:
                 specConstantData.push_back(ndrange.globalWorkOffset[1]);
                 break;
             case SpecConstantType::GlobalOffsetZ:
                 specConstantData.push_back(ndrange.globalWorkOffset[2]);
                 break;
             case SpecConstantType::SubgroupMaxSize:
             {
                 // We should only be here if cl_khr_subgroups is supported
                 ASSERT(device.getInfo().khrSubgroups);
                 specConstantData.push_back(
                     mContext->getRenderer()->getPhysicalDeviceSubgroupProperties().subgroupSize);
                 break;
             }
             default:
                 UNIMPLEMENTED();
                 continue;
         }
         mapEntries.push_back(VkSpecializationMapEntry{
             .constantID = devProgramData->reflectionData.specConstantIDs[specConstantUsed],
             .offset     = constantDataOffset,
             .size       = sizeof(uint32_t)});
         constantDataOffset += sizeof(uint32_t);
     }
     // Populate kernel specialization constants (if any)
     for (const auto arg : mArgs)
     {
         if (arg->getUsed() &&
             arg->getReflectionType() == NonSemanticClspvReflectionArgumentWorkgroup)
         {
             ASSERT(arg->getWorkGroupElementSize() != 0);

             specConstantData.push_back(
                 static_cast<uint32_t>(arg->getSize() / arg->getWorkGroupElementSize()));
             mapEntries.push_back(
                 VkSpecializationMapEntry{.constantID = arg->getWorkGroupBufferSpecId(),
                                          .offset     = constantDataOffset,
                                          .size       = sizeof(uint32_t)});
             constantDataOffset += sizeof(uint32_t);
         }
     }
     VkSpecializationInfo computeSpecializationInfo{
         .mapEntryCount = static_cast<uint32_t>(mapEntries.size()),
         .pMapEntries   = mapEntries.data(),
         .dataSize      = specConstantData.size() * sizeof(uint32_t),
         .pData         = specConstantData.data(),
     };

     // Now get or create (on compute pipeline cache miss) compute pipeline and return it
     vk::ComputePipelineOptions options = vk::GetComputePipelineOptions(
         vk::PipelineRobustness::NonRobust, vk::PipelineProtectedAccess::Unprotected);
     return mShaderProgramHelper.getOrCreateComputePipeline(
         mContext, &mComputePipelineCache, pipelineCache, getPipelineLayout(), options,
         PipelineSource::Draw, pipelineOut, mName.c_str(), &computeSpecializationInfo);
 }

 bool CLKernelVk::usesPrintf() const
 {
     return mProgram->getDeviceProgramData(mName.c_str())->getKernelFlags(mName) &
            NonSemanticClspvReflectionMayUsePrintf;
 }

 bool CLKernelVk::usesPrintfBufferPointerPushConstant() const
 {
     return mProgram->getDeviceProgramData(mName.c_str())
         ->reflectionData.pushConstants.contains(
             NonSemanticClspvReflectionPrintfBufferPointerPushConstant);
 }

 angle::Result CLKernelVk::initializeDescriptorPools()
 {
     for (DescriptorSetIndex index : angle::AllEnums<DescriptorSetIndex>())
     {
         if (!mDescriptorSetLayoutDescs[index].empty())
         {
             ANGLE_TRY(mContext->getMetaDescriptorPool().bindCachedDescriptorPool(
                 mContext, mDescriptorSetLayoutDescs[index], 1,
                 mContext->getDescriptorSetLayoutCache(), &mDynamicDescriptorPools[index]));
         }
     }
     return angle::Result::Continue;
 }

 angle::Result CLKernelVk::allocateDescriptorSet(
     DescriptorSetIndex index,
     angle::EnumIterator<DescriptorSetIndex> layoutIndex,
     vk::OutsideRenderPassCommandBufferHelper *computePassCommands)
 {
     if (mDescriptorSets[index] && mDescriptorSets[index]->valid())
     {
         // Safe to reuse: descriptor set is no longer in use by the GPU.
         if (mContext->getRenderer()->hasResourceUseFinished(
                 mDescriptorSets[index]->getResourceUse()))
         {
             // Set DS serial to current CB serial upon reuse.
             mDescriptorSets[index]->setQueueSerial(computePassCommands->getQueueSerial());
             return angle::Result::Continue;
         }
         else
         {
             mDescriptorSets[index].reset();
         }
     }

     if (mDynamicDescriptorPools[index]->valid())
     {
         ANGLE_TRY(mDynamicDescriptorPools[index]->allocateDescriptorSet(
             mContext, *mDescriptorSetLayouts[*layoutIndex], &mDescriptorSets[index]));
         computePassCommands->retainResource(mDescriptorSets[index].get());
     }

     return angle::Result::Continue;
 }

 cl_ulong CLKernelVk::getLocalMemSizeUsed(const cl::Device &device) const
 {
     return getAllArgLocalMemSize() + getCompiledLocalMemSize(device);
 }

 cl_ulong CLKernelVk::getAllArgLocalMemSize() const
 {
     return std::reduce(mLocalMemoryArgSizes.begin(), mLocalMemoryArgSizes.end(), 0,
                        std::plus<size_t>());
 }

 cl_ulong CLKernelVk::getCompiledLocalMemSize(const cl::Device &device) const
 {
     return mProgram->getDeviceProgramData(const_cast<cl_device_id>(device.getNative()))
         ->reflectionData.workgroupVariableSize.size;
 }

 }  // namespace rx
	//
	// Copyright 2021 The ANGLE Project Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.
	//
	// CLKernelVk.cpp: Implements the class methods for CLKernelVk.
	//

	#include "libANGLE/renderer/vulkan/CLKernelVk.h"
	#include "libANGLE/renderer/vulkan/CLContextVk.h"
	#include "libANGLE/renderer/vulkan/CLDeviceVk.h"
	#include "libANGLE/renderer/vulkan/CLMemoryVk.h"
	#include "libANGLE/renderer/vulkan/CLProgramVk.h"
	#include "libANGLE/renderer/vulkan/cl_types.h"
	#include "libANGLE/renderer/vulkan/vk_wrapper.h"

	#include "libANGLE/CLBuffer.h"
	#include "libANGLE/CLContext.h"
	#include "libANGLE/CLDevice.h"
	#include "libANGLE/CLKernel.h"
	#include "libANGLE/CLProgram.h"

	#include "spirv/unified1/NonSemanticClspvReflection.h"

	#include <algorithm>

	namespace rx
	{

	// Setup the CLKernelArgument with info from ClspvKernelArgument
	CLKernelArgument::CLKernelArgument(const CLContextVk *context, const ClspvKernelArgument &info)
	: mCompiledInfo(info), mSize(0), mHandle(nullptr)
	{
	mArgInfo.typeQualifier = mCompiledInfo.info.typeQualifier;
	mArgInfo.addressQualifier = mCompiledInfo.info.addressQualifier;
	mArgInfo.accessQualifier = mCompiledInfo.info.accessQualifier;
	mArgInfo.typeName = mCompiledInfo.info.typeName;
	mArgInfo.name = mCompiledInfo.info.name;
	}

	cl::Memory *CLKernelArgument::getMemoryHandle() const
	{
	if (!mCompiledInfo.used)
	{
	return nullptr;
	}
	return cl::Memory::Cast(static_cast<cl_mem>(mHandle));
	}

	cl::Sampler *CLKernelArgument::getSamplerHandle() const
	{
	ASSERT(mCompiledInfo.type == NonSemanticClspvReflectionArgumentSampler);
	return cl::Sampler::Cast(static_cast<const cl_sampler >(mHandle));
	}

	bool CLKernelArgument::isReadOnly() const
	{
	// if not used, can safely assume readonly
	if (!mCompiledInfo.used)
	{
	return true;
	}
	switch (getReflectionType())
	{
	case NonSemanticClspvReflectionArgumentPodUniform:
	case NonSemanticClspvReflectionArgumentUniform:
	case NonSemanticClspvReflectionArgumentPointerUniform:
	case NonSemanticClspvReflectionArgumentUniformTexelBuffer:
	case NonSemanticClspvReflectionConstantDataStorageBuffer:
	{
	return true;
	}
	case NonSemanticClspvReflectionArgumentStorageBuffer:
	case NonSemanticClspvReflectionArgumentStorageTexelBuffer:
	case NonSemanticClspvReflectionArgumentStorageImage:
	case NonSemanticClspvReflectionArgumentSampledImage:
	{
	return mArgInfo.accessQualifier == CL_KERNEL_ARG_ACCESS_READ_ONLY;
	}
	default:
	{
	return false;
	}
	}
	}

	CLKernelVk::CLKernelVk(const cl::Kernel &kernel, std::string &name, std::string &attributes)
	: CLKernelImpl(kernel),
	mProgram(&kernel.getProgram().getImpl<CLProgramVk>()),
	mContext(&kernel.getProgram().getContext().getImpl<CLContextVk>()),
	mName(name),
	mAttributes(attributes),
	mArgs(),
	mPodBuffer(nullptr),
	mLocalMemoryArgSizes(
	mProgram->getDeviceProgramData(mName.c_str())->getKernelArguments(mName.c_str()).size(),
	0)
	{
	mShaderProgramHelper.setShader(gl::ShaderType::Compute,
	mKernel.getProgram().getImpl<CLProgramVk>().getShaderModule());
	}

	CLKernelVk::~CLKernelVk()
	{
	mComputePipelineCache.destroy(mContext);
	mShaderProgramHelper.destroy(mContext->getRenderer());
	}

	angle::Result CLKernelVk::init()
	{
	const CLProgramVk::DeviceProgramData *deviceProgramData =
	mProgram->getDeviceProgramData(mName.c_str());

	// Literal sampler handling
	for (const ClspvLiteralSampler &literalSampler :
	deviceProgramData->reflectionData.literalSamplers)
	{
	mDescriptorSetLayoutDescs[DescriptorSetIndex::LiteralSampler].addBinding(
	literalSampler.binding, VK_DESCRIPTOR_TYPE_SAMPLER, 1, VK_SHADER_STAGE_COMPUTE_BIT,
	nullptr);
	}

	// Setup the arguments
	for (auto &clspvArg : deviceProgramData->getKernelArguments(mName.c_str()))
	{
	mArgs.push_back(std::make_shared<CLKernelArgument>(mContext, clspvArg));
	}

	vk::DescriptorSetLayoutDesc &descriptorSetLayoutDesc =
	mDescriptorSetLayoutDescs[DescriptorSetIndex::KernelArguments];
	VkPushConstantRange pcRange = deviceProgramData->pushConstRange;
	size_t podBufferSize = 0;

	bool podFound = false;
	for (const auto arg : getArgs())
	{
	VkDescriptorType descType = VK_DESCRIPTOR_TYPE_MAX_ENUM;
	switch (arg->getReflectionType())
	{
	case NonSemanticClspvReflectionArgumentStorageBuffer:
	descType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
	break;
	case NonSemanticClspvReflectionArgumentUniform:
	descType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
	break;
	case NonSemanticClspvReflectionArgumentPodUniform:
	case NonSemanticClspvReflectionArgumentPodStorageBuffer:
	case NonSemanticClspvReflectionArgumentPointerUniform:
	{
	uint32_t newPodBufferSize =
	arg->getPodStorageBufferOffset() + arg->getPodStorageBufferSize();
	podBufferSize = newPodBufferSize > podBufferSize ? newPodBufferSize : podBufferSize;
	if (podFound)
	{
	continue;
	}
	descType =
	arg->getReflectionType() == NonSemanticClspvReflectionArgumentPodStorageBuffer
	? VK_DESCRIPTOR_TYPE_STORAGE_BUFFER
	: VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
	podFound = true;
	break;
	}
	case NonSemanticClspvReflectionArgumentPodPushConstant:
	case NonSemanticClspvReflectionArgumentPointerPushConstant:
	// Get existing push constant range and see if we need to update
	if (arg->getPushConstantOffset() + arg->getPushConstantSize() >
	pcRange.offset + pcRange.size)
	{
	pcRange.size =
	arg->getPushConstantOffset() + arg->getPushConstantSize() - pcRange.offset;
	}
	continue;
	case NonSemanticClspvReflectionArgumentSampledImage:
	descType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE;
	break;
	case NonSemanticClspvReflectionArgumentStorageImage:
	descType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
	break;
	case NonSemanticClspvReflectionArgumentSampler:
	descType = VK_DESCRIPTOR_TYPE_SAMPLER;
	break;
	case NonSemanticClspvReflectionArgumentStorageTexelBuffer:
	descType = VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
	break;
	case NonSemanticClspvReflectionArgumentUniformTexelBuffer:
	descType = VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER;
	break;
	default:
	continue;
	}
	if (descType != VK_DESCRIPTOR_TYPE_MAX_ENUM)
	{
	descriptorSetLayoutDesc.addBinding(arg->getDescriptorBinding(), descType, 1,
	VK_SHADER_STAGE_COMPUTE_BIT, nullptr);
	}
	}

	if (podBufferSize > 0)
	{
	mPodBuffer = cl::BufferPtr::Create(
	const_cast<cl::Context &>(mKernel.getProgram().getContext()), cl::Memory::PropArray{},
	cl::MemFlags(CL_MEM_READ_ONLY), podBufferSize, nullptr);
	}

	if (usesPrintf() && !usesPrintfBufferPointerPushConstant())
	{
	mDescriptorSetLayoutDescs[DescriptorSetIndex::Printf].addBinding(
	deviceProgramData->reflectionData.printfBufferStorage.binding,
	VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr);
	}

	// Get pipeline layout from cache (creates if missed)
	// A given kernel need not have resulted in use of all the descriptor sets. Unless the
	// graphicsPipelineLibrary extension is supported, the pipeline layout need all the descriptor
	// set layouts to be valide. So set them up in the order of their occurrence.
	mPipelineLayoutDesc = {};
	for (DescriptorSetIndex index : angle::AllEnums<DescriptorSetIndex>())
	{
	if (!mDescriptorSetLayoutDescs[index].empty())
	{
	mPipelineLayoutDesc.updateDescriptorSetLayout(index, mDescriptorSetLayoutDescs[index]);
	}
	}

	// push constant range size and offset both need to be a multiple of 4
	pcRange.size = roundUpPow2(pcRange.size, 4u);
	pcRange.offset = roundDownPow2(pcRange.offset, 4u);
	mPipelineLayoutDesc.updatePushConstantRange(pcRange.stageFlags, pcRange.offset, pcRange.size);

	mPodArgumentPushConstants.resize(pcRange.size);

	angle::EnumIterator<DescriptorSetIndex> layoutIndex(DescriptorSetIndex::LiteralSampler);
	for (DescriptorSetIndex index : angle::AllEnums<DescriptorSetIndex>())
	{
	if (!getDescriptorSetLayoutDesc(index).empty())
	{
	ANGLE_CL_IMPL_TRY_ERROR(mContext->getDescriptorSetLayoutCache()->getDescriptorSetLayout(
	mContext, getDescriptorSetLayoutDesc(index),
	&getDescriptorSetLayouts()[*layoutIndex]),
	CL_INVALID_OPERATION);
	ASSERT(getDescriptorSetLayouts()[*layoutIndex]->valid());
	++layoutIndex;
	}
	}
	ANGLE_CL_IMPL_TRY_ERROR(initPipelineLayout(), CL_INVALID_OPERATION);

	// initialize the descriptor pools
	// descriptor pools are setup as per their indices
	return mContext->initializeDescriptorPools(this);
	}

	angle::Result CLKernelVk::setArg(cl_uint argIndex, size_t argSize, const void *argValue)
	{
	auto &arg = mArgs.at(argIndex);
	if (arg->getUsed())
	{
	switch (arg->getReflectionType())
	{
	case NonSemanticClspvReflectionArgumentPodPushConstant:
	{
	ASSERT(mPodArgumentPushConstants.size() >=
	arg->getPushConstantSize() + arg->getPushConstantOffset());
	void *handle = &mPodArgumentPushConstants[arg->getPushConstantOffset()];
	size_t handleSize =
	(argSize > arg->getPushConstantSize() ? arg->getPushConstantSize() : argSize);
	if (argSize > 0 && argValue != nullptr)
	{
	// Copy the contents since app is free to delete/reassign the contents after
	ANGLE_UNSAFE_TODO(memcpy(handle, argValue, handleSize));
	}
	arg->set(handleSize, handle);
	break;
	}
	case NonSemanticClspvReflectionArgumentPodUniform:
	case NonSemanticClspvReflectionArgumentPodStorageBuffer:
	{
	ASSERT(mPodBuffer->getSize() >= argSize + arg->getPodUniformOffset());
	if (argSize > 0 && argValue != nullptr)
	{
	ANGLE_TRY(mPodBuffer->getImpl<CLBufferVk>().copyFrom(
	argValue, arg->getPodStorageBufferOffset(), argSize));
	}
	break;
	}
	case NonSemanticClspvReflectionArgumentUniform:
	case NonSemanticClspvReflectionArgumentStorageBuffer:
	{
	ASSERT(argSize == sizeof(cl_mem *));
	size_t handleSize = argSize;
	void *handle = nullptr;
	// Argument can be null ptr
	if (argValue && static_cast<const cl_mem >(argValue))
	{
	handle = static_cast<const cl_mem >(argValue);
	}
	arg->set(handleSize, handle);
	break;
	}
	case NonSemanticClspvReflectionArgumentStorageImage:
	case NonSemanticClspvReflectionArgumentSampledImage:
	case NonSemanticClspvReflectionArgumentUniformTexelBuffer:
	case NonSemanticClspvReflectionArgumentStorageTexelBuffer:
	case NonSemanticClspvReflectionArgumentPointerPushConstant:
	case NonSemanticClspvReflectionArgumentPointerUniform:
	{
	ASSERT(argSize == sizeof(cl_mem *));
	cl_mem memHandle = static_cast<const cl_mem >(argValue);

	if (cl::Memory::Cast(memHandle)->getFlags().intersects(CL_MEM_READ_ONLY) &&
	!arg->isReadOnly())
	{
	WARN() << "The cl_mem argument at index:" << argIndex
	<< " is setup as read-only, and the kernel is writing to it -- This "
	"could lead to undefined behavior";
	}
	arg->set(argSize, memHandle);
	break;
	}
	case NonSemanticClspvReflectionArgumentWorkgroup:
	ASSERT(arg->getWorkGroupElementSize() != 0);
	mLocalMemoryArgSizes[argIndex] = argSize;
	arg->set(argSize, const_cast<void *>(argValue));
	break;
	default:
	// Just store ptr and size (if we end up here)
	arg->set(argSize, const_cast<void *>(argValue));
	break;
	}
	}

	return angle::Result::Continue;
	}

	angle::Result CLKernelVk::createInfo(CLKernelImpl::Info *info) const
	{
	info->functionName = mName;
	info->attributes = mAttributes;
	info->numArgs = static_cast<cl_uint>(mArgs.size());
	for (const auto arg : mArgs)
	{
	info->args.push_back(arg->getArgInfo());
	}

	auto &ctx = mKernel.getProgram().getContext();
	info->workGroups.resize(ctx.getDevices().size());
	const CLProgramVk::DeviceProgramData *deviceProgramData = nullptr;
	for (auto i = 0u; i < ctx.getDevices().size(); ++i)
	{
	auto &workGroup = info->workGroups[i];
	const auto device = &ctx.getDevices()[i];
	const auto deviceVk = &device->get()->getImpl<CLDeviceVk>();
	deviceProgramData = mProgram->getDeviceProgramData(ctx.getDevices()[i]->getNative());
	if (deviceProgramData == nullptr)
	{
	continue;
	}

	// TODO: http://anglebug.com/42267005
	ANGLE_TRY(
	deviceVk->getInfoSizeT(cl::DeviceInfo::MaxWorkGroupSize, &workGroup.workGroupSize));

	// TODO: http://anglebug.com/42267004
	workGroup.privateMemSize = 0;

	workGroup.localMemSize = deviceProgramData->reflectionData.workgroupVariableSize.size;
	workGroup.prefWorkGroupSizeMultiple = deviceVk->getWorkGroupSizeMultiple();
	workGroup.globalWorkSize = {0, 0, 0};
	if (deviceProgramData->reflectionData.kernelCompileWorkgroupSize.contains(mName))
	{
	workGroup.compileWorkGroupSize = {
	deviceProgramData->reflectionData.kernelCompileWorkgroupSize.at(mName)[0],
	deviceProgramData->reflectionData.kernelCompileWorkgroupSize.at(mName)[1],
	deviceProgramData->reflectionData.kernelCompileWorkgroupSize.at(mName)[2]};
	}
	else
	{
	workGroup.compileWorkGroupSize = {0, 0, 0};
	}

	if (device->get()->getInfo().khrSubgroups)
	{
	workGroup.subGroupSizeForNDRange =
	mContext->getRenderer()->getPhysicalDeviceSubgroupProperties().subgroupSize;
	}
	else
	{
	workGroup.subGroupSizeForNDRange = 0;
	}
	}

	return angle::Result::Continue;
	}

	angle::Result CLKernelVk::initPipelineLayout()
	{
	PipelineLayoutCache *pipelineLayoutCache = mContext->getPipelineLayoutCache();
	return pipelineLayoutCache->getPipelineLayout(mContext, mPipelineLayoutDesc,
	mDescriptorSetLayouts, &mPipelineLayout);
	}

	angle::Result CLKernelVk::getOrCreateComputePipeline(vk::PipelineCacheAccess *pipelineCache,
	const cl::NDRange &ndrange,
	const cl::Device &device,
	vk::PipelineHelper **pipelineOut)
	{
	const CLProgramVk::DeviceProgramData *devProgramData =
	getProgram()->getDeviceProgramData(device.getNative());
	ASSERT(devProgramData != nullptr);

	// Populate program specialization constants (if any)
	uint32_t constantDataOffset = 0;
	std::vector<uint32_t> specConstantData;
	std::vector<VkSpecializationMapEntry> mapEntries;
	for (const auto specConstantUsed : devProgramData->reflectionData.specConstantsUsed)
	{
	switch (specConstantUsed)
	{
	case SpecConstantType::WorkDimension:
	specConstantData.push_back(ndrange.workDimensions);
	break;
	case SpecConstantType::WorkgroupSizeX:
	specConstantData.push_back(ndrange.localWorkSize[0]);
	break;
	case SpecConstantType::WorkgroupSizeY:
	specConstantData.push_back(ndrange.localWorkSize[1]);
	break;
	case SpecConstantType::WorkgroupSizeZ:
	specConstantData.push_back(ndrange.localWorkSize[2]);
	break;
	case SpecConstantType::GlobalOffsetX:
	specConstantData.push_back(ndrange.globalWorkOffset[0]);
	break;
	case SpecConstantType::GlobalOffsetY:
	specConstantData.push_back(ndrange.globalWorkOffset[1]);
	break;
	case SpecConstantType::GlobalOffsetZ:
	specConstantData.push_back(ndrange.globalWorkOffset[2]);
	break;
	case SpecConstantType::SubgroupMaxSize:
	{
	// We should only be here if cl_khr_subgroups is supported
	ASSERT(device.getInfo().khrSubgroups);
	specConstantData.push_back(
	mContext->getRenderer()->getPhysicalDeviceSubgroupProperties().subgroupSize);
	break;
	}
	default:
	UNIMPLEMENTED();
	continue;
	}
	mapEntries.push_back(VkSpecializationMapEntry{
	.constantID = devProgramData->reflectionData.specConstantIDs[specConstantUsed],
	.offset = constantDataOffset,
	.size = sizeof(uint32_t)});
	constantDataOffset += sizeof(uint32_t);
	}
	// Populate kernel specialization constants (if any)
	for (const auto arg : mArgs)
	{
	if (arg->getUsed() &&
	arg->getReflectionType() == NonSemanticClspvReflectionArgumentWorkgroup)
	{
	ASSERT(arg->getWorkGroupElementSize() != 0);

	specConstantData.push_back(
	static_cast<uint32_t>(arg->getSize() / arg->getWorkGroupElementSize()));
	mapEntries.push_back(
	VkSpecializationMapEntry{.constantID = arg->getWorkGroupBufferSpecId(),
	.offset = constantDataOffset,
	.size = sizeof(uint32_t)});
	constantDataOffset += sizeof(uint32_t);
	}
	}
	VkSpecializationInfo computeSpecializationInfo{
	.mapEntryCount = static_cast<uint32_t>(mapEntries.size()),
	.pMapEntries = mapEntries.data(),
	.dataSize = specConstantData.size() * sizeof(uint32_t),
	.pData = specConstantData.data(),
	};

	// Now get or create (on compute pipeline cache miss) compute pipeline and return it
	vk::ComputePipelineOptions options = vk::GetComputePipelineOptions(
	vk::PipelineRobustness::NonRobust, vk::PipelineProtectedAccess::Unprotected);
	return mShaderProgramHelper.getOrCreateComputePipeline(
	mContext, &mComputePipelineCache, pipelineCache, getPipelineLayout(), options,
	PipelineSource::Draw, pipelineOut, mName.c_str(), &computeSpecializationInfo);
	}

	bool CLKernelVk::usesPrintf() const
	{
	return mProgram->getDeviceProgramData(mName.c_str())->getKernelFlags(mName) &
	NonSemanticClspvReflectionMayUsePrintf;
	}

	bool CLKernelVk::usesPrintfBufferPointerPushConstant() const
	{
	return mProgram->getDeviceProgramData(mName.c_str())
	->reflectionData.pushConstants.contains(
	NonSemanticClspvReflectionPrintfBufferPointerPushConstant);
	}

	angle::Result CLKernelVk::initializeDescriptorPools()
	{
	for (DescriptorSetIndex index : angle::AllEnums<DescriptorSetIndex>())
	{
	if (!mDescriptorSetLayoutDescs[index].empty())
	{
	ANGLE_TRY(mContext->getMetaDescriptorPool().bindCachedDescriptorPool(
	mContext, mDescriptorSetLayoutDescs[index], 1,
	mContext->getDescriptorSetLayoutCache(), &mDynamicDescriptorPools[index]));
	}
	}
	return angle::Result::Continue;
	}

	angle::Result CLKernelVk::allocateDescriptorSet(
	DescriptorSetIndex index,
	angle::EnumIterator<DescriptorSetIndex> layoutIndex,
	vk::OutsideRenderPassCommandBufferHelper *computePassCommands)
	{
	if (mDescriptorSets[index] && mDescriptorSets[index]->valid())
	{
	// Safe to reuse: descriptor set is no longer in use by the GPU.
	if (mContext->getRenderer()->hasResourceUseFinished(
	mDescriptorSets[index]->getResourceUse()))
	{
	// Set DS serial to current CB serial upon reuse.
	mDescriptorSets[index]->setQueueSerial(computePassCommands->getQueueSerial());
	return angle::Result::Continue;
	}
	else
	{
	mDescriptorSets[index].reset();
	}
	}

	if (mDynamicDescriptorPools[index]->valid())
	{
	ANGLE_TRY(mDynamicDescriptorPools[index]->allocateDescriptorSet(
	mContext, mDescriptorSetLayouts[layoutIndex], &mDescriptorSets[index]));
	computePassCommands->retainResource(mDescriptorSets[index].get());
	}

	return angle::Result::Continue;
	}

	cl_ulong CLKernelVk::getLocalMemSizeUsed(const cl::Device &device) const
	{
	return getAllArgLocalMemSize() + getCompiledLocalMemSize(device);
	}

	cl_ulong CLKernelVk::getAllArgLocalMemSize() const
	{
	return std::reduce(mLocalMemoryArgSizes.begin(), mLocalMemoryArgSizes.end(), 0,
	std::plus<size_t>());
	}

	cl_ulong CLKernelVk::getCompiledLocalMemSize(const cl::Device &device) const
	{
	return mProgram->getDeviceProgramData(const_cast<cl_device_id>(device.getNative()))
	->reflectionData.workgroupVariableSize.size;
	}

	} // namespace rx