//
// 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