HlslComputeShaderInterface.h

#pragma once
 
// Standard.
#include <unordered_map>
#include <format>
 
// Custom.
#include "shader/ComputeShaderInterface.h"
#include "render/directx/resources/DirectXResource.h"
#include "render/directx/descriptors/DirectXDescriptorHeap.h"
 
// External.
#include "directx/d3dx12.h"
 
namespace ne {
    class HlslComputeShaderInterface : public ComputeShaderInterface {
        // Only parent class can create instances of this class because there are some specific things
        // we need to do when creating objects of this class and parent class handles these things.
        friend class ComputeShaderInterface;
 
    public:
        HlslComputeShaderInterface() = delete;
        virtual ~HlslComputeShaderInterface() override = default;
 
        HlslComputeShaderInterface(const HlslComputeShaderInterface&) = delete;
        HlslComputeShaderInterface& operator=(const HlslComputeShaderInterface&) = delete;
 
        [[nodiscard]] virtual std::optional<Error> bindResource(
            GpuResource* pResource,
            const std::string& sShaderResourceName,
            ComputeResourceUsage usage,
            bool bUpdateOnlyCurrentFrameResourceDescriptors = false) override;
 
        inline void dispatchOnGraphicsQueue(ID3D12GraphicsCommandList* pCommandList) {
            // Bind CBVs.
            for (const auto& [iRootParameterIndex, pResource] : cbvResources) {
                pCommandList->SetComputeRootConstantBufferView(
                    iRootParameterIndex, pResource->getInternalResource()->GetGPUVirtualAddress());
            }
 
            // Bind UAVs.
            for (const auto& [iRootParameterIndex, pResource] : uavResources) {
                pCommandList->SetComputeRootUnorderedAccessView(
                    iRootParameterIndex, pResource->getInternalResource()->GetGPUVirtualAddress());
            }
 
            // Bind SRVs.
            for (const auto& [iRootParameterIndex, pResource] : srvResources) {
                pCommandList->SetComputeRootShaderResourceView(
                    iRootParameterIndex, pResource->getInternalResource()->GetGPUVirtualAddress());
            }
 
            // Bind table views.
            for (const auto& [iRootParameterIndex, pDescriptor] : tableResources) {
                // Set table.
                pCommandList->SetComputeRootDescriptorTable(
                    iRootParameterIndex,
                    D3D12_GPU_DESCRIPTOR_HANDLE{
                        pCbvSrvUavHeap->getInternalHeap()->GetGPUDescriptorHandleForHeapStart().ptr +
                        pDescriptor->getDescriptorOffsetInDescriptors() * iCbvSrvUavDescriptorSize});
            }
 
            // Add a dispatch command.
            pCommandList->Dispatch(getThreadGroupCountX(), getThreadGroupCountY(), getThreadGroupCountZ());
        }
 
    protected:
        HlslComputeShaderInterface(
            Renderer* pRenderer,
            const std::string& sComputeShaderName,
            ComputeExecutionStage executionStage,
            ComputeExecutionGroup executionGroup);
 
        std::unordered_map<UINT, DirectXResource*> cbvResources;
 
        std::unordered_map<UINT, DirectXResource*> uavResources;
 
        std::unordered_map<UINT, DirectXResource*> srvResources;
 
        std::unordered_map<UINT, DirectXDescriptor*> tableResources;
 
        DirectXDescriptorHeap* pCbvSrvUavHeap = nullptr;
 
        UINT iCbvSrvUavDescriptorSize = 0;
    };
}