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;
struct ID3D12StateObjectProperties;

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;
        int bindingTableIndex;
    };

    class DX12ShaderHeap;

    class DX12ShaderSignature : public ReferenceCounter
    {
    private:
        ID3D12RootSignature* signature;
        Array<DX12ShaderRegisterUsage*> descriptorHeapBindings;
        Array<DX12ShaderRegisterUsage*> bindingTableBindings;
        bool changed;
        int shaderBindingTableParamCount;

    public:
        DLLEXPORT DX12ShaderSignature();
        DLLEXPORT ~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)
         * \return pointer to the offset in the shader binding table where this
         * parameter needs to be placed. This pointer is only valid after the
         * call to createSignature() and should be used to fill the shader
         * binding table with the correct GPU addresses of the resources.
         */
        DLLEXPORT int* 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.
         * For optimal performance the registers that are linked to the
         * descriptor heaps should be consecutive and use consecutive descriptor
         * heap spaces. Example: register(b3), register(b4), register(b5) linked
         * to descriptor heap index 2, 3, 4.
         *
         * \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)
         */
        DLLEXPORT void addRegisterUsageLinkedToDescriptorHeap(
            int descriptorHeapIndex,
            DX12ShaderRegister registerType,
            int registerIndex,
            int spaceIndex = 0);
        /**
         * Creates the root signature.
         */
        DLLEXPORT void createSignature(ID3D12Device5* zDevice,
            PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature);
        DLLEXPORT ID3D12RootSignature* zSignature() const;
        DLLEXPORT const Array<DX12ShaderRegisterUsage*>&
        getDescriptorHeapBindings() const;
        DLLEXPORT int gerShaderBindingTableParamCount() const;
    };

    enum DX12ShaderFunctionType
    {
        DX12_SHADER_FUNCTION_TYPE_RAY_GEN,
        DX12_SHADER_FUNCTION_TYPE_INTERSECTION,
        DX12_SHADER_FUNCTION_TYPE_ANY_HIT,
        DX12_SHADER_FUNCTION_TYPE_CLOSEST_HIT,
        DX12_SHADER_FUNCTION_TYPE_MISS,
        DX12_SHADER_FUNCTION_TYPE_CALLABLE,
    };

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

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

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

    public:
        DLLEXPORT DX12Shader(
            const unsigned char* shaderBytes, int shaderBytesSize);
        DLLEXPORT ~DX12Shader();
        DLLEXPORT void addFunction(DX12ShaderFunction* function);
        DLLEXPORT int getShaderBytesSize() const;
        DLLEXPORT const unsigned char* getShaderBytes() const;
        DLLEXPORT const RCArray<DX12ShaderFunction>& getFunctions() const;
        DLLEXPORT 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:
        DLLEXPORT DX12ShaderHitGroup(const Text name);
        DLLEXPORT ~DX12ShaderHitGroup();
        DLLEXPORT void setClosestHitShaderFunction(
            DX12ShaderFunction* zClosestHitShaderFunction);
        DLLEXPORT void setAnyHitShaderFunction(
            DX12ShaderFunction* zAnyHitShaderFunction);
        DLLEXPORT void setIntersectionShaderFunction(
            DX12ShaderFunction* zIntersectionShaderFunction);
        DLLEXPORT void setPayloadSize(int payloadSize);
        DLLEXPORT void setAttributeSize(int attributeSize);
        DLLEXPORT const Text& getName() const;
        DLLEXPORT DX12ShaderFunction* zClosestHitShaderFunction() const;
        DLLEXPORT DX12ShaderFunction* zAnyHitShaderFunction() const;
        DLLEXPORT DX12ShaderFunction* zIntersectionShaderFunction() const;
        DLLEXPORT int getPayloadSize() const;
        DLLEXPORT int getAttributeSize() const;
        DLLEXPORT D3D12_HIT_GROUP_DESC* zHitGroupDesc() const;
        DLLEXPORT DX12ShaderSignature* zSignature() const;
    };

    class DX12ShaderBindingTable;
    class DX12GlobalDescriptorHeap;

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

    public:
        DLLEXPORT DX12Pipeline();
        DLLEXPORT ~DX12Pipeline();
        DLLEXPORT void addShader(DX12Shader* shader);
        DLLEXPORT void addHitGroup(DX12ShaderHitGroup* hitGroup);
        DLLEXPORT void setMaxRecursionDepth(int maxRecursionDepth);
        DLLEXPORT void createPipelineState(ID3D12Device5* zDevice,
            PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature);
        DLLEXPORT ID3D12StateObject* zPipelineState() const;
        DLLEXPORT DX12ShaderBindingTable* createShaderBindingTable();
        DLLEXPORT DX12GlobalDescriptorHeap* createGlobalDescriptorHeap();
        DLLEXPORT const RCArray<DX12Shader>& getShaders() const;
        DLLEXPORT const Array<const DX12ShaderFunction*>&
        getFunctionsWithoutHitGroups() const;
        DLLEXPORT const RCArray<DX12ShaderHitGroup>& getHitGroups() 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;
        ID3D12Device5* zDevice;

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

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

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

    class DX12ShaderBindingTable : public ReferenceCounter
    {
    private:
        DX12Pipeline* pipeline;
        DX12Buffer* shaderBindingTableBuffer;
        DX12GlobalDescriptorHeap* globalDescriptorHeap;
        int rayGenRecordSize;
        int rayGenCount;
        int missRecordSize;
        int missCount;
        int callableRecordSize;
        int callableCount;
        int hitGroupRecordSize;
        int hitGroupCount;
        char* tableBuffer;
        int tableBufferSize;
        Array<char*> tempBuffers;
        int nextHitGroupOffset;
        ID3D12StateObjectProperties* stateObjectProperties;

    public:
        DLLEXPORT DX12ShaderBindingTable(DX12Pipeline* pipeline);
        DLLEXPORT ~DX12ShaderBindingTable();
        DLLEXPORT void setGlobalDescriptorHeap(
            DX12GlobalDescriptorHeap* zGlobalDescriptorHeap);
        DLLEXPORT void startUpdate();

    private:
        DLLEXPORT void set(int index, void* data, int size);

    public:
        /**
         * sets inputs for shader outside of any hit groups for the current
         * update of the shader binding table.
         *
         * \param zFunction the sample shader function for which the inputs are
         * set. The function must be part of the pipeline.
         * \param offsetPointer pointer for the offset in the shader binding
         * table witch was obtained by adding the register usage to the
         * segnature of the shader function by calling
         * DX12ShaderSignature::addRegisterUsageLinkedToShaderBindingTable.
         * \param gpuAddress the parameter to be placed in the shader binding
         * table for this input. This is usually the GPU address of a buffer or
         * texture.
         */
        DLLEXPORT void setShaderInput(DX12ShaderFunction* zFunction,
            int* offsetPointer,
            __int64 gpuAddress);
        /**
         * adds a hit group to the current update of the shader binding table.
         * each hit group mus be readded for each update of the shader binding
         * table.
         *
         * \param zHitGroup the hit group to be added. The hit group must be
         * part of the pipeline.
         * \return the index of the hit group in the shader binding table. This
         * index can be used to set the inputs for this hit group by calling
         * setHitGroupShaderInput.
         */
        DLLEXPORT int addHitGroup(DX12ShaderHitGroup* zHitGroup);
        /**
         * sets a specific input for a previously added hitgroup.
         *
         * \param hitGroupOffset the hit group offset returned by addHitGroup
         * for the hit group for which the input is set.
         * \param offsetPointer pointer for the offset in the shader binding
         * table witch was obtained by adding the register usage to the
         * segnature of the shader function by calling
         * DX12ShaderSignature::addRegisterUsageLinkedToShaderBindingTable.
         * \param gpuAddress the parameter to be placed in the shader binding
         * table for this input. This is usually the GPU address of a buffer or
         * texture.
         */
        DLLEXPORT void setHitGroupShaderInput(
            int hitGroupOffset, int* offsetPointer, __int64 gpuAddress);
        DLLEXPORT void endUpdate(
            ID3D12Device5* zDevice, DX12CommandQueue* zQueue);
        DLLEXPORT void fillDispatchRaysDesc(
            D3D12_DISPATCH_RAYS_DESC* dispatchRaysDesc);
        DLLEXPORT DX12Pipeline* zPipeline() const;
    };
} // namespace Framework