KSG-Framework/DX12Shader.h

#pragma once

#include "Array.h"
#include "DX12Buffer.h"

struct ID3D12Device5;
struct ID3D12GraphicsCommandList;
struct D3D12_INPUT_ELEMENT_DESC;
struct D3D12_ROOT_PARAMETER1;
struct D3D12_CONSTANT_BUFFER_VIEW_DESC;

namespace Framework
{
    class Texture;
    class DX12TLAS;

    enum DX12ShaderRegister
    {
        DX12_SHADER_REGISTER_B_CONST_BUFFER = 0,
        DX12_SHADER_REGISTER_T_SHADER_RESOURCE = 1,
        DX12_SHADER_REGISTER_U_UNORDERED_ACCESS = 2,
        // TODO: Do we need Sampler?
    };

    struct DX12ShaderRegisterUsage
    {
        DX12ShaderRegister registerType;
        int registerIndex;
        int spaceIndex;
        int descriptorHeapIndex;
    };

    class DX12ShaderHeap;

    class DX12ShaderSignature : public ReferenceCounter
    {
    private:
        ID3D12RootSignature* signature;
        Array<DX12ShaderRegisterUsage> registerUsages;
        bool changed;

    public:
        DX12ShaderSignature();
        ~DX12ShaderSignature();
        /**
         * for each datastructure with : register(...) in the shader code,
         * either this function or addRegisterUsageLinkedToDescriptorHeap must
         * be called to link the register usage to the shader binding table or
         * descriptor heap.
         *
         * \param registerType the register type e.g.
         * DX12_SHADER_REGISTER_B_CONST_BUFFER for : register(b0)
         * \param registerIndex the register index e.g. 1 for : register(b1)
         * \param spaceIndex the optional space index e.g. 3 for : register(b1,
         * space3)
         */
        void addRegisterUsageLinkedToShaderBindingTable(
            DX12ShaderRegister registerType,
            int registerIndex,
            int spaceIndex = 0);
        /**
         * for each datastructure with : register(...) in the shader code,
         * either this function or addRegisterUsageLinkedToDescriptorHeap must
         * be called to link the register usage to the shader binding table or
         * descriptor heap.
         *
         * \param descriptorHeapIndex the index in the descriptor heap witch
         * contains the resource for this register usage
         * \param registerType the register type e.g.
         * DX12_SHADER_REGISTER_B_CONST_BUFFER for : register(b0)
         * \param registerIndex the register index e.g. 1 for : register(b1)
         * \param spaceIndex the optional space index e.g. 3 for : register(b1,
         * space3)
         */
        void addRegisterUsageLinkedToDescriptorHeap(int descriptorHeapIndex,
            DX12ShaderRegister registerType,
            int registerIndex,
            int spaceIndex = 0);
        /**
         * Creates the root signature.
         */
        void createSignature(ID3D12Device5* zDevice,
            PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature);
        ID3D12RootSignature* zSignature() const;
        const Array<DX12ShaderRegisterUsage>& getRegisterUsages() const;
    };

    class DX12ShaderFunction : public ReferenceCounter
    {
    private:
        Text functionName;
        DX12ShaderSignature* signature;
        D3D12_EXPORT_DESC* exportDesc;

    public:
        DX12ShaderFunction(
            const Text& functionName, DX12ShaderSignature* signature);
        ~DX12ShaderFunction();
        const Text& getFunctionName() const;
        DX12ShaderSignature* zSignature() const;
        D3D12_EXPORT_DESC* zExportDesc() const;
    };

    class DX12Shader : public ReferenceCounter
    {
    private:
        RCArray<DX12ShaderFunction> functions;
        const unsigned char* shaderBytes;
        int shaderBytesSize;
        D3D12_DXIL_LIBRARY_DESC* libraryDesc;

    public:
        DX12Shader(const unsigned char* shaderBytes, int shaderBytesSize);
        ~DX12Shader();
        void addFunction(DX12ShaderFunction* function);
        int getShaderBytesSize() const;
        const unsigned char* getShaderBytes() const;
        const RCArray<DX12ShaderFunction>& getFunctions() const;
        D3D12_DXIL_LIBRARY_DESC* zLibraryDesc() const;
    };

    class DX12ShaderHitGroup : public ReferenceCounter
    {
    private:
        Text name;
        DX12ShaderFunction* closestHitShaderFunction;
        DX12ShaderFunction* anyHitShaderFunction;
        DX12ShaderFunction* intersectionShaderFunction;
        int payloadSize;
        int attributeSize;
        D3D12_HIT_GROUP_DESC* hitGroupDesc;

    public:
        DX12ShaderHitGroup(const Text name);
        ~DX12ShaderHitGroup();
        void setClosestHitShaderFunction(
            DX12ShaderFunction* zClosestHitShaderFunction);
        void setAnyHitShaderFunction(DX12ShaderFunction* zAnyHitShaderFunction);
        void setIntersectionShaderFunction(
            DX12ShaderFunction* zIntersectionShaderFunction);
        void setPayloadSize(int payloadSize);
        void setAttributeSize(int attributeSize);
        const Text& getName() const;
        DX12ShaderFunction* zClosestHitShaderFunction() const;
        DX12ShaderFunction* zAnyHitShaderFunction() const;
        DX12ShaderFunction* zIntersectionShaderFunction() const;
        int getPayloadSize() const;
        int getAttributeSize() const;
        D3D12_HIT_GROUP_DESC* zHitGroupDesc() const;
    };

    class DX12ShaderBindingTable;
    class DX12GlobalDescriptorHeap;

    class DX12Pipeline : public ReferenceCounter
    {
    private:
        RCArray<DX12Shader> shaders;
        RCArray<DX12ShaderHitGroup> hitGroups;
        ID3D12RootSignature* emptyGlobalRootSignature;
        ID3D12RootSignature* emptyLocalRootSignature;
        ID3D12StateObject* pipelineState;
        int maxRecursionDepth;

    public:
        DX12Pipeline();
        ~DX12Pipeline();
        void addShader(DX12Shader* shader);
        void addHitGroup(DX12ShaderHitGroup* hitGroup);
        void setMaxRecursionDepth(int maxRecursionDepth);
        void createPipelineState(ID3D12Device5* zDevice,
            PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature);
        ID3D12StateObject* zPipelineState() const;
        DX12ShaderBindingTable* createShaderBindingTable();
        DX12GlobalDescriptorHeap* createGlobalDescriptorHeap();
        const RCArray<DX12Shader>& getShaders() const;
    }; // namespace Framework

    struct DX12ShaderRegisterInput
    {
        DX12ShaderRegister registerType;
        ReferenceCounter*
            inputResource; // Can be Texture*, DXBuffer*, or DX12TLAS*
    };

    class DX12GlobalDescriptorHeap : public ReferenceCounter
    {
    private:
        DX12Pipeline* pipeline;
        ID3D12DescriptorHeap* descriptorHeap;
        Array<DX12ShaderRegisterInput*> registerInputs;
        int lastDescriptorHeapSize;

    public:
        DX12GlobalDescriptorHeap(DX12Pipeline* pipeline);
        ~DX12GlobalDescriptorHeap();

    private:
        void addInput(DX12ShaderRegister type, ReferenceCounter* inputResource);

    public:
        void addTextureInput(DX12ShaderRegister type, Texture* zTexture);
        void addBufferInput(DX12ShaderRegister type, DXBuffer* zBuffer);
        void addTLASInput(DX12ShaderRegister type, DX12TLAS* zTLAS);
        void updateDescriptorHeap(ID3D12Device5* zDevice);
        DX12Pipeline* zPipeline() const;
        ID3D12DescriptorHeap* zDescriptorHeap() const;
    };

    class DX12ShaderBindingTable : public ReferenceCounter
    {
    private:
        DX12Pipeline* pipeline;
        DX12Buffer* shaderBindingTableBuffer;
        DX12GlobalDescriptorHeap* globalDescriptorHeap;
        int rayGenRecordSize;
        int rayGenCount;
        int missRecordSize;
        int missCount;
        int hitGroupRecordSize;
        int hitGroupCount;
        bool changed;

    public:
        DX12ShaderBindingTable(DX12Pipeline* pipeline);
        ~DX12ShaderBindingTable();
        void setGlobalDescriptorHeap(
            DX12GlobalDescriptorHeap* zGlobalDescriptorHeap);
        void startUpdate();
        void setShaderInputs(DX12ShaderFunction* zFunction,
            std::initializer_list<unsigned __int64> gpuAddresses);
        void setHitGroupShaderInputs(int instanceIndex,
            DX12ShaderHitGroup* zHitGroup,
            std::initializer_list<unsigned __int64> gpuAddresses);
        void endUpdate(ID3D12Device5* zDevice);
        void fillDispatchRaysDesc(D3D12_DISPATCH_RAYS_DESC* dispatchRaysDesc);
        DX12Pipeline* zPipeline() const;
    };
} // namespace Framework