XCEngine/tests/RHI/Vulkan/unit/test_compute.cpp

#if defined(XCENGINE_SUPPORT_VULKAN)

#include "fixtures/VulkanTestFixture.h"
#include "XCEngine/RHI/RHIDescriptorPool.h"
#include "XCEngine/RHI/RHIDescriptorSet.h"
#include "XCEngine/RHI/RHIPipelineLayout.h"
#include "XCEngine/RHI/RHIPipelineState.h"
#include "XCEngine/RHI/RHIResourceView.h"
#include "XCEngine/RHI/Vulkan/VulkanTexture.h"

#include <cstring>
#include <vector>

using namespace XCEngine::RHI;

namespace {

constexpr uint32_t kWriteRedComputeSpirv[] = {
    0x07230203, 0x00010000, 0x0008000B, 0x00000017, 0x00000000, 0x00020011, 0x00000001, 0x0006000B,
    0x00000001, 0x4C534C47, 0x6474732E, 0x3035342E, 0x00000000, 0x0003000E, 0x00000000, 0x00000001,
    0x0005000F, 0x00000005, 0x00000004, 0x6E69616D, 0x00000000, 0x00060010, 0x00000004, 0x00000011,
    0x00000001, 0x00000001, 0x00000001, 0x00030003, 0x00000002, 0x000001C2, 0x00040005, 0x00000004,
    0x6E69616D, 0x00000000, 0x00040005, 0x00000009, 0x616D4975, 0x00006567, 0x00030047, 0x00000009,
    0x00000019, 0x00040047, 0x00000009, 0x00000021, 0x00000000, 0x00040047, 0x00000009, 0x00000022,
    0x00000000, 0x00040047, 0x00000016, 0x0000000B, 0x00000019, 0x00020013, 0x00000002, 0x00030021,
    0x00000003, 0x00000002, 0x00030016, 0x00000006, 0x00000020, 0x00090019, 0x00000007, 0x00000006,
    0x00000001, 0x00000000, 0x00000000, 0x00000000, 0x00000002, 0x00000004, 0x00040020, 0x00000008,
    0x00000000, 0x00000007, 0x0004003B, 0x00000008, 0x00000009, 0x00000000, 0x00040015, 0x0000000B,
    0x00000020, 0x00000001, 0x00040017, 0x0000000C, 0x0000000B, 0x00000002, 0x0004002B, 0x0000000B,
    0x0000000D, 0x00000000, 0x0005002C, 0x0000000C, 0x0000000E, 0x0000000D, 0x0000000D, 0x00040017,
    0x0000000F, 0x00000006, 0x0000000004, 0x0004002B, 0x00000006, 0x00000010, 0x3F800000, 0x0004002B,
    0x00000006, 0x00000011, 0x00000000, 0x0007002C, 0x0000000F, 0x00000012, 0x00000010, 0x00000011,
    0x00000011, 0x00000010, 0x00040015, 0x00000013, 0x00000020, 0x00000000, 0x00040017, 0x00000014,
    0x00000013, 0x00000003, 0x0004002B, 0x00000013, 0x00000015, 0x00000001, 0x0006002C, 0x00000014,
    0x00000016, 0x00000015, 0x00000015, 0x00000015, 0x00050036, 0x00000002, 0x00000004, 0x00000000,
    0x00000003, 0x000200F8, 0x00000005, 0x0004003D, 0x00000007, 0x0000000A, 0x00000009, 0x00040063,
    0x0000000A, 0x0000000E, 0x00000012, 0x000100FD, 0x00010038
};

ShaderCompileDesc MakeWriteRedComputeShaderDesc() {
    ShaderCompileDesc shaderDesc = {};
    shaderDesc.sourceLanguage = ShaderLanguage::SPIRV;
    shaderDesc.profile = L"cs_6_0";
    shaderDesc.source.resize(sizeof(kWriteRedComputeSpirv));
    std::memcpy(shaderDesc.source.data(), kWriteRedComputeSpirv, sizeof(kWriteRedComputeSpirv));
    return shaderDesc;
}

ShaderCompileDesc MakeWriteRedComputeShaderFromGlslDesc() {
    static const char* computeSource = R"(#version 450
layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0, rgba8) uniform writeonly image2D uImage;
void main() {
    imageStore(uImage, ivec2(0, 0), vec4(1.0, 0.0, 0.0, 1.0));
}
)";

    ShaderCompileDesc shaderDesc = {};
    shaderDesc.sourceLanguage = ShaderLanguage::GLSL;
    shaderDesc.source.assign(computeSource, computeSource + std::strlen(computeSource));
    return shaderDesc;
}

TEST_F(VulkanGraphicsFixture, CreateUnorderedAccessViewProducesValidView) {
    RHITexture* texture = m_device->CreateTexture(CreateColorTextureDesc(4, 4));
    ASSERT_NE(texture, nullptr);

    ResourceViewDesc viewDesc = {};
    viewDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    viewDesc.dimension = ResourceViewDimension::Texture2D;

    RHIResourceView* uav = m_device->CreateUnorderedAccessView(texture, viewDesc);
    ASSERT_NE(uav, nullptr);
    EXPECT_TRUE(uav->IsValid());
    EXPECT_EQ(uav->GetViewType(), ResourceViewType::UnorderedAccess);

    uav->Shutdown();
    delete uav;
    texture->Shutdown();
    delete texture;
}

TEST_F(VulkanGraphicsFixture, DispatchWritesUavTexture) {
    RHITexture* texture = m_device->CreateTexture(CreateColorTextureDesc(4, 4));
    ASSERT_NE(texture, nullptr);

    ResourceViewDesc uavDesc = {};
    uavDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    uavDesc.dimension = ResourceViewDimension::Texture2D;
    RHIResourceView* uav = m_device->CreateUnorderedAccessView(texture, uavDesc);
    ASSERT_NE(uav, nullptr);

    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 1;
    poolDesc.shaderVisible = true;
    RHIDescriptorPool* pool = m_device->CreateDescriptorPool(poolDesc);
    ASSERT_NE(pool, nullptr);

    DescriptorSetLayoutBinding uavBinding = {};
    uavBinding.binding = 0;
    uavBinding.type = static_cast<uint32_t>(DescriptorType::UAV);
    uavBinding.count = 1;
    uavBinding.visibility = static_cast<uint32_t>(ShaderVisibility::All);

    DescriptorSetLayoutDesc setLayout = {};
    setLayout.bindings = &uavBinding;
    setLayout.bindingCount = 1;

    RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.setLayouts = &setLayout;
    pipelineLayoutDesc.setLayoutCount = 1;
    RHIPipelineLayout* pipelineLayout = m_device->CreatePipelineLayout(pipelineLayoutDesc);
    ASSERT_NE(pipelineLayout, nullptr);

    RHIDescriptorSet* descriptorSet = pool->AllocateSet(setLayout);
    ASSERT_NE(descriptorSet, nullptr);
    descriptorSet->Update(0, uav);

    ComputePipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.computeShader = MakeWriteRedComputeShaderDesc();
    RHIPipelineState* pipelineState = m_device->CreateComputePipelineState(pipelineDesc);
    ASSERT_NE(pipelineState, nullptr);
    EXPECT_TRUE(pipelineState->HasComputeShader());
    EXPECT_EQ(pipelineState->GetType(), PipelineType::Compute);

    RHICommandList* commandList = CreateCommandList();
    ASSERT_NE(commandList, nullptr);

    commandList->Reset();
    commandList->TransitionBarrier(uav, ResourceStates::Common, ResourceStates::UnorderedAccess);
    commandList->SetPipelineState(pipelineState);
    RHIDescriptorSet* descriptorSets[] = { descriptorSet };
    commandList->SetComputeDescriptorSets(0, 1, descriptorSets, pipelineLayout);
    commandList->Dispatch(1, 1, 1);
    SubmitAndWait(commandList);

    const std::vector<uint8_t> pixels = ReadTextureRgba8(static_cast<VulkanTexture*>(texture));
    ASSERT_GE(pixels.size(), 4u);
    EXPECT_EQ(pixels[0], 255u);
    EXPECT_EQ(pixels[1], 0u);
    EXPECT_EQ(pixels[2], 0u);
    EXPECT_EQ(pixels[3], 255u);

    commandList->Shutdown();
    delete commandList;
    pipelineState->Shutdown();
    delete pipelineState;
    descriptorSet->Shutdown();
    delete descriptorSet;
    pipelineLayout->Shutdown();
    delete pipelineLayout;
    pool->Shutdown();
    delete pool;
    uav->Shutdown();
    delete uav;
    texture->Shutdown();
    delete texture;
}

TEST_F(VulkanGraphicsFixture, DispatchWritesUavTextureWithGlslComputeShader) {
    RHITexture* texture = m_device->CreateTexture(CreateColorTextureDesc(4, 4));
    ASSERT_NE(texture, nullptr);

    ResourceViewDesc uavDesc = {};
    uavDesc.format = static_cast<uint32_t>(Format::R8G8B8A8_UNorm);
    uavDesc.dimension = ResourceViewDimension::Texture2D;
    RHIResourceView* uav = m_device->CreateUnorderedAccessView(texture, uavDesc);
    ASSERT_NE(uav, nullptr);

    DescriptorPoolDesc poolDesc = {};
    poolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    poolDesc.descriptorCount = 1;
    poolDesc.shaderVisible = true;
    RHIDescriptorPool* pool = m_device->CreateDescriptorPool(poolDesc);
    ASSERT_NE(pool, nullptr);

    DescriptorSetLayoutBinding uavBinding = {};
    uavBinding.binding = 0;
    uavBinding.type = static_cast<uint32_t>(DescriptorType::UAV);
    uavBinding.count = 1;
    uavBinding.visibility = static_cast<uint32_t>(ShaderVisibility::All);

    DescriptorSetLayoutDesc setLayout = {};
    setLayout.bindings = &uavBinding;
    setLayout.bindingCount = 1;

    RHIPipelineLayoutDesc pipelineLayoutDesc = {};
    pipelineLayoutDesc.setLayouts = &setLayout;
    pipelineLayoutDesc.setLayoutCount = 1;
    RHIPipelineLayout* pipelineLayout = m_device->CreatePipelineLayout(pipelineLayoutDesc);
    ASSERT_NE(pipelineLayout, nullptr);

    RHIDescriptorSet* descriptorSet = pool->AllocateSet(setLayout);
    ASSERT_NE(descriptorSet, nullptr);
    descriptorSet->Update(0, uav);

    ComputePipelineDesc pipelineDesc = {};
    pipelineDesc.pipelineLayout = pipelineLayout;
    pipelineDesc.computeShader = MakeWriteRedComputeShaderFromGlslDesc();
    RHIPipelineState* pipelineState = m_device->CreateComputePipelineState(pipelineDesc);
    ASSERT_NE(pipelineState, nullptr);
    EXPECT_TRUE(pipelineState->HasComputeShader());
    EXPECT_EQ(pipelineState->GetType(), PipelineType::Compute);

    RHICommandList* commandList = CreateCommandList();
    ASSERT_NE(commandList, nullptr);

    commandList->Reset();
    commandList->TransitionBarrier(uav, ResourceStates::Common, ResourceStates::UnorderedAccess);
    commandList->SetPipelineState(pipelineState);
    RHIDescriptorSet* descriptorSets[] = { descriptorSet };
    commandList->SetComputeDescriptorSets(0, 1, descriptorSets, pipelineLayout);
    commandList->Dispatch(1, 1, 1);
    SubmitAndWait(commandList);

    const std::vector<uint8_t> pixels = ReadTextureRgba8(static_cast<VulkanTexture*>(texture));
    ASSERT_GE(pixels.size(), 4u);
    EXPECT_EQ(pixels[0], 255u);
    EXPECT_EQ(pixels[1], 0u);
    EXPECT_EQ(pixels[2], 0u);
    EXPECT_EQ(pixels[3], 255u);

    commandList->Shutdown();
    delete commandList;
    pipelineState->Shutdown();
    delete pipelineState;
    descriptorSet->Shutdown();
    delete descriptorSet;
    pipelineLayout->Shutdown();
    delete pipelineLayout;
    pool->Shutdown();
    delete pool;
    uav->Shutdown();
    delete uav;
    texture->Shutdown();
    delete texture;
}

} // namespace

#endif