KSG-Framework/DX12Shader.cpp

#include "DX12Shader.h"

#include "DX12Texture.h"
#include "DX12TLAS.h"
#include "Logging.h"

using namespace Framework;

Framework::DX12ShaderSignature::DX12ShaderSignature()
    : ReferenceCounter(),
      signature(0),
      changed(1)
{}

Framework::DX12ShaderSignature::~DX12ShaderSignature()
{
    if (signature)
    {
        signature->Release();
    }
}

void Framework::DX12ShaderSignature::addRegisterUsageLinkedToShaderBindingTable(
    DX12ShaderRegister registerType, int registerIndex, int spaceIndex)
{
    addRegisterUsageLinkedToDescriptorHeap(
        -1, registerType, registerIndex, spaceIndex);
}

void Framework::DX12ShaderSignature::addRegisterUsageLinkedToDescriptorHeap(
    int descriptorHeapIndex,
    DX12ShaderRegister registerType,
    int registerIndex,
    int spaceIndex)
{
    // register usages sould be sorted by registerType -> spaceIndex ->
    // registerIndex
    ArrayIterator<DX12ShaderRegisterUsage> it = registerUsages.begin();
    bool found = 0;
    while (it)
    {
        const DX12ShaderRegisterUsage& usage = it.val();
        if (usage.registerType > registerType)
        {
            found = 1;
            break;
        }
        if (usage.registerType == registerType)
        {
            if (usage.spaceIndex > spaceIndex)
            {
                found = 1;
                break;
            }
            if (usage.spaceIndex == spaceIndex)
            {
                if (usage.registerIndex >= registerIndex)
                {
                    found = 1;
                    break;
                }
            }
        }
        ++it;
    }
    if (found)
    {
        if (it.val().registerIndex == registerIndex)
        {
            Logging::error()
                << "Duplicate register usage in root signature: "
                << registerType << " " << registerIndex << " " << spaceIndex;
            return;
        }
        it.addBefore(
            {registerType, registerIndex, spaceIndex, descriptorHeapIndex});
    }
    else
    {
        registerUsages.add(
            {registerType, registerIndex, spaceIndex, descriptorHeapIndex});
    }
    changed = 1;
}

void Framework::DX12ShaderSignature::createSignature(ID3D12Device5* zDevice,
    PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature)
{
    if (!changed)
    {
        return;
    }
    changed = 0;
    if (signature)
    {
        signature->Release();
        signature = 0;
    }
    D3D12_ROOT_PARAMETER* descriptorTable
        = new D3D12_ROOT_PARAMETER[registerUsages.getEntryCount()];
    D3D12_DESCRIPTOR_RANGE* descriptorRanges
        = new D3D12_DESCRIPTOR_RANGE[registerUsages.getEntryCount()];
    ArrayIterator<DX12ShaderRegisterUsage> it = registerUsages.begin();
    int index = 0;
    while (it)
    {
        const auto& usage = it.val();
        if (usage.descriptorHeapIndex >= 0)
        {
            descriptorTable[index].ParameterType
                = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE;
            descriptorTable[index].ShaderVisibility
                = D3D12_SHADER_VISIBILITY_ALL;
            descriptorTable[index].DescriptorTable.NumDescriptorRanges
                = registerUsages.getEntryCount();
            D3D12_DESCRIPTOR_RANGE* range = &descriptorRanges[index];
            switch (usage.registerType)
            {
            case DX12_SHADER_REGISTER_B_CONST_BUFFER:
                range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_CBV;
                break;
            case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
                range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
                break;
            case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
                range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_UAV;
                break;
            default:
                Logging::error() << "Unknown register type for root signature: "
                                 << usage.registerType;
                range->RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV;
            }
            ArrayIterator<DX12ShaderRegisterUsage> next = it.next();
            int size = 1;
            while (next && next.val().registerType == usage.registerType
                   && next.val().spaceIndex == usage.spaceIndex
                   && next.val().registerIndex == usage.registerIndex + size
                   && next.val().descriptorHeapIndex
                          == usage.descriptorHeapIndex + size)
            {
                ++size;
                it = next;
                ++next;
            }
            range->NumDescriptors = size;
            range->BaseShaderRegister = usage.registerIndex;
            range->RegisterSpace = usage.spaceIndex;
            range->OffsetInDescriptorsFromTableStart
                = usage.descriptorHeapIndex;
            descriptorTable[index].DescriptorTable.pDescriptorRanges = range;
        }
        else
        {
            switch (usage.registerType)
            {
            case DX12_SHADER_REGISTER_B_CONST_BUFFER:
                descriptorTable[index].ParameterType
                    = D3D12_ROOT_PARAMETER_TYPE_CBV;
                break;
            case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
                descriptorTable[index].ParameterType
                    = D3D12_ROOT_PARAMETER_TYPE_SRV;
                break;
            case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
                descriptorTable[index].ParameterType
                    = D3D12_ROOT_PARAMETER_TYPE_UAV;
                break;
            }
            descriptorTable[index].ShaderVisibility
                = D3D12_SHADER_VISIBILITY_ALL;
            descriptorTable[index].Descriptor.ShaderRegister
                = usage.registerIndex;
            descriptorTable[index].Descriptor.RegisterSpace = usage.spaceIndex;
        }
        ++it;
        ++index;
    }
    D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
    rootDesc.NumParameters = index;
    rootDesc.pParameters = descriptorTable;
    rootDesc.Flags = D3D12_ROOT_SIGNATURE_FLAG_LOCAL_ROOT_SIGNATURE;
    ID3DBlob* pSigBlob = 0;
    ID3DBlob* pErrorBlob = 0;
    HRESULT hr = pfnD3D12SerializeRootSignature(
        &rootDesc, D3D_ROOT_SIGNATURE_VERSION_1_0, &pSigBlob, &pErrorBlob);
    if (pSigBlob)
    {
        zDevice->CreateRootSignature(0,
            pSigBlob->GetBufferPointer(),
            pSigBlob->GetBufferSize(),
            __uuidof(ID3D12RootSignature),
            (void**)&signature);
        pSigBlob->Release();
    }
    if (pErrorBlob)
    {
        std::string errorMessage(
            static_cast<const char*>(pErrorBlob->GetBufferPointer()),
            pErrorBlob->GetBufferSize());
        Logging::error() << "Failed to serialize root signature: "
                         << errorMessage;
        pErrorBlob->Release();
    }
    delete[] descriptorRanges;
}

ID3D12RootSignature* Framework::DX12ShaderSignature::zSignature() const
{
    return signature;
}

const Array<DX12ShaderRegisterUsage>&
Framework::DX12ShaderSignature::getRegisterUsages() const
{
    return registerUsages;
}

Framework::DX12ShaderFunction::DX12ShaderFunction(
    const Text& functionName, DX12ShaderSignature* signature)
    : ReferenceCounter(),
      functionName(functionName),
      signature(signature),
      exportDesc(new D3D12_EXPORT_DESC())
{
    wchar_t* wc = new wchar_t[functionName.getLength() + 1];
    mbtowc(wc, functionName.getText(), functionName.getLength() + 1);
    wc[functionName.getLength()] = 0;
    exportDesc->Name = wc;
    exportDesc->ExportToRename = 0;
    exportDesc->Flags = D3D12_EXPORT_FLAG_NONE;
}

Framework::DX12ShaderFunction::~DX12ShaderFunction()
{
    signature->release();
    delete[] exportDesc->Name;
    delete exportDesc;
}

const Text& Framework::DX12ShaderFunction::getFunctionName() const
{
    return functionName;
}

DX12ShaderSignature* Framework::DX12ShaderFunction::zSignature() const
{
    return signature;
}

D3D12_EXPORT_DESC* Framework::DX12ShaderFunction::zExportDesc() const
{
    return exportDesc;
}

Framework::DX12Shader::DX12Shader(
    const unsigned char* shaderBytes, int shaderBytesSize)
    : ReferenceCounter(),
      shaderBytes(shaderBytes),
      shaderBytesSize(shaderBytesSize),
      libraryDesc(new D3D12_DXIL_LIBRARY_DESC())
{
    libraryDesc->DXILLibrary.pShaderBytecode = shaderBytes;
    libraryDesc->DXILLibrary.BytecodeLength = shaderBytesSize;
    libraryDesc->NumExports = 0;
    libraryDesc->pExports = 0;
}

Framework::DX12Shader::~DX12Shader()
{
    delete[] libraryDesc->pExports;
    delete libraryDesc;
}

void Framework::DX12Shader::addFunction(DX12ShaderFunction* function)
{
    functions.add(function);
}

int Framework::DX12Shader::getShaderBytesSize() const
{
    return shaderBytesSize;
}

const unsigned char* Framework::DX12Shader::getShaderBytes() const
{
    return shaderBytes;
}

const RCArray<DX12ShaderFunction>& Framework::DX12Shader::getFunctions() const
{
    return functions;
}

D3D12_DXIL_LIBRARY_DESC* Framework::DX12Shader::zLibraryDesc() const
{
    if (libraryDesc->NumExports != functions.getEntryCount())
    {
        delete[] libraryDesc->pExports;
        libraryDesc->NumExports = functions.getEntryCount();
        D3D12_EXPORT_DESC* pExports
            = new D3D12_EXPORT_DESC[functions.getEntryCount()];
        int index = 0;
        for (const auto& function : functions)
        {
            memcpy(pExports + index,
                function->zExportDesc(),
                sizeof(D3D12_EXPORT_DESC));
            index++;
        }
        libraryDesc->pExports = pExports;
    }
    return libraryDesc;
}

Framework::DX12ShaderHitGroup::DX12ShaderHitGroup(const Text name)
    : ReferenceCounter(),
      name(name),
      closestHitShaderFunction(0),
      anyHitShaderFunction(0),
      intersectionShaderFunction(0),
      payloadSize(0),
      attributeSize(0),
      hitGroupDesc(new D3D12_HIT_GROUP_DESC())
{
    wchar_t* wc = new wchar_t[name.getLength() + 1];
    mbtowc(wc, name.getText(), name.getLength() + 1);
    wc[name.getLength()] = 0;
    hitGroupDesc->HitGroupExport = wc;
    hitGroupDesc->IntersectionShaderImport = 0;
    hitGroupDesc->AnyHitShaderImport = 0;
    hitGroupDesc->ClosestHitShaderImport = 0;
}

Framework::DX12ShaderHitGroup::~DX12ShaderHitGroup()
{
    delete[] hitGroupDesc->HitGroupExport;
    delete hitGroupDesc;
    if (closestHitShaderFunction)
    {
        closestHitShaderFunction->release();
    }
    if (anyHitShaderFunction)
    {
        anyHitShaderFunction->release();
    }
    if (intersectionShaderFunction)
    {
        intersectionShaderFunction->release();
    }
}

void Framework::DX12ShaderHitGroup::setClosestHitShaderFunction(
    DX12ShaderFunction* closestHitShaderFunction)
{
    if (anyHitShaderFunction
        && anyHitShaderFunction->zSignature()
               != closestHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Any-hit shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Any Hit Shader Function: '"
            << anyHitShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (intersectionShaderFunction
        && intersectionShaderFunction->zSignature()
               != closestHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << intersectionShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

    if (this->closestHitShaderFunction)
    {
        this->closestHitShaderFunction->release();
    }
    this->closestHitShaderFunction = closestHitShaderFunction;
    hitGroupDesc->ClosestHitShaderImport = 0;
    if (closestHitShaderFunction)
    {
        hitGroupDesc->ClosestHitShaderImport
            = closestHitShaderFunction->zExportDesc()->Name;
    }
}

void Framework::DX12ShaderHitGroup::setAnyHitShaderFunction(
    DX12ShaderFunction* anyHitShaderFunction)
{
    if (closestHitShaderFunction
        && closestHitShaderFunction->zSignature()
               != anyHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Any-hit shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Any Hit Shader Function: '"
            << anyHitShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (intersectionShaderFunction
        && intersectionShaderFunction->zSignature()
               != anyHitShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and any-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << intersectionShaderFunction->getFunctionName().getText()
            << "' Any Hit Shader Function: '"
            << anyHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

    if (this->anyHitShaderFunction)
    {
        this->anyHitShaderFunction->release();
    }
    this->anyHitShaderFunction = anyHitShaderFunction;
    hitGroupDesc->AnyHitShaderImport = 0;
    if (anyHitShaderFunction)
    {
        hitGroupDesc->AnyHitShaderImport
            = anyHitShaderFunction->zExportDesc()->Name;
    }
}

void Framework::DX12ShaderHitGroup::setIntersectionShaderFunction(
    DX12ShaderFunction* zIntersectionShaderFunction)
{
    if (closestHitShaderFunction
        && closestHitShaderFunction->zSignature()
               != zIntersectionShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and closest-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << zIntersectionShaderFunction->getFunctionName().getText()
            << "' Closest Hit Shader Function: '"
            << closestHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }
    if (anyHitShaderFunction
        && anyHitShaderFunction->zSignature()
               != zIntersectionShaderFunction->zSignature())
    {
        Logging::error()
            << "Intersection shader function and any-hit shader "
               "function must have the same root signature when they are "
               "combined in the same hit group. HitGroup Name: '"
            << name << "' Intersection Shader Function: '"
            << zIntersectionShaderFunction->getFunctionName().getText()
            << "' Any Hit Shader Function: '"
            << anyHitShaderFunction->getFunctionName().getText() << "'";
        throw std::runtime_error("Incompatible root signatures in hit group");
    }

    if (intersectionShaderFunction)
    {
        intersectionShaderFunction->release();
    }
    intersectionShaderFunction = dynamic_cast<DX12ShaderFunction*>(
        zIntersectionShaderFunction->getThis());
    hitGroupDesc->IntersectionShaderImport = 0;
    if (intersectionShaderFunction)
    {
        hitGroupDesc->IntersectionShaderImport
            = intersectionShaderFunction->zExportDesc()->Name;
    }
}

void Framework::DX12ShaderHitGroup::setPayloadSize(int payloadSize)
{
    this->payloadSize = payloadSize;
}

void Framework::DX12ShaderHitGroup::setAttributeSize(int attributeSize)
{
    this->attributeSize = attributeSize;
}

const Text& Framework::DX12ShaderHitGroup::getName() const
{
    return name;
}

DX12ShaderFunction*
Framework::DX12ShaderHitGroup::zClosestHitShaderFunction() const
{
    return closestHitShaderFunction;
}

DX12ShaderFunction* Framework::DX12ShaderHitGroup::zAnyHitShaderFunction() const
{
    return anyHitShaderFunction;
}

DX12ShaderFunction*
Framework::DX12ShaderHitGroup::zIntersectionShaderFunction() const
{
    return intersectionShaderFunction;
}

int Framework::DX12ShaderHitGroup::getPayloadSize() const
{
    return payloadSize;
}

int Framework::DX12ShaderHitGroup::getAttributeSize() const
{
    return attributeSize;
}

D3D12_HIT_GROUP_DESC* Framework::DX12ShaderHitGroup::zHitGroupDesc() const
{
    if (hitGroupDesc->IntersectionShaderImport)
    {
        hitGroupDesc->Type = D3D12_HIT_GROUP_TYPE_PROCEDURAL_PRIMITIVE;
    }
    else
    {
        hitGroupDesc->Type = D3D12_HIT_GROUP_TYPE_TRIANGLES;
    }
    return hitGroupDesc;
}

Framework::DX12Pipeline::DX12Pipeline()
    : ReferenceCounter(),
      emptyGlobalRootSignature(0),
      emptyLocalRootSignature(0),
      pipelineState(0),
      maxRecursionDepth(0)
{}

Framework::DX12Pipeline::~DX12Pipeline()
{
    if (emptyGlobalRootSignature)
    {
        emptyGlobalRootSignature->Release();
    }
    if (emptyLocalRootSignature)
    {
        emptyLocalRootSignature->Release();
    }
    if (pipelineState)
    {
        pipelineState->Release();
    }
}

void Framework::DX12Pipeline::addShader(DX12Shader* shader)
{
    shaders.add(shader);
}

void Framework::DX12Pipeline::addHitGroup(DX12ShaderHitGroup* hitGroup)
{
    hitGroups.add(hitGroup);
}

void Framework::DX12Pipeline::setMaxRecursionDepth(int maxRecursionDepth)
{
    this->maxRecursionDepth = maxRecursionDepth;
}

void Framework::DX12Pipeline::createPipelineState(ID3D12Device5* zDevice,
    PFN_D3D12_SERIALIZE_ROOT_SIGNATURE pfnD3D12SerializeRootSignature)
{
    if (!emptyGlobalRootSignature)
    {
        D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
        rootDesc.NumParameters = 0;
        rootDesc.pParameters = 0;
        rootDesc.Flags = D3D12_ROOT_SIGNATURE_FLAG_NONE;
        ID3DBlob* pSigBlob = 0;
        ID3DBlob* pErrorBlob = 0;
        HRESULT hr = pfnD3D12SerializeRootSignature(
            &rootDesc, D3D_ROOT_SIGNATURE_VERSION_1_0, &pSigBlob, &pErrorBlob);
        if (pSigBlob)
        {
            zDevice->CreateRootSignature(0,
                pSigBlob->GetBufferPointer(),
                pSigBlob->GetBufferSize(),
                __uuidof(ID3D12RootSignature),
                (void**)&emptyGlobalRootSignature);
            pSigBlob->Release();
        }
        if (pErrorBlob)
        {
            std::string errorMessage(
                static_cast<const char*>(pErrorBlob->GetBufferPointer()),
                pErrorBlob->GetBufferSize());
            Logging::error()
                << "Failed to serialize empty root signature: " << errorMessage;
            pErrorBlob->Release();
        }
    }
    if (!emptyLocalRootSignature)
    {
        D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
        rootDesc.NumParameters = 0;
        rootDesc.pParameters = 0;
        rootDesc.Flags = D3D12_ROOT_SIGNATURE_FLAG_LOCAL_ROOT_SIGNATURE;
        ID3DBlob* pSigBlob = 0;
        ID3DBlob* pErrorBlob = 0;
        HRESULT hr = pfnD3D12SerializeRootSignature(
            &rootDesc, D3D_ROOT_SIGNATURE_VERSION_1_0, &pSigBlob, &pErrorBlob);
        if (pSigBlob)
        {
            zDevice->CreateRootSignature(0,
                pSigBlob->GetBufferPointer(),
                pSigBlob->GetBufferSize(),
                __uuidof(ID3D12RootSignature),
                (void**)&emptyGlobalRootSignature);
            pSigBlob->Release();
        }
        if (pErrorBlob)
        {
            std::string errorMessage(
                static_cast<const char*>(pErrorBlob->GetBufferPointer()),
                pErrorBlob->GetBufferSize());
            Logging::error()
                << "Failed to serialize empty root signature: " << errorMessage;
            pErrorBlob->Release();
        }
    }
    unsigned int subobjectCount
        = shaders.getEntryCount() + hitGroups.getEntryCount() + 5;
    Array<DX12ShaderSignature*> distinctSignatures;
    for (const DX12Shader* shader : shaders)
    {
        for (const DX12ShaderFunction* function : shader->getFunctions())
        {
            DX12ShaderSignature* signature = function->zSignature();
            if (distinctSignatures.getValueIndex(signature) < 0)
            {
                distinctSignatures.add(signature);
            }
        }
    }
    subobjectCount += distinctSignatures.getEntryCount()
                    * 2; // Local root signatures for each distinct signature
    D3D12_STATE_SUBOBJECT* subobjects
        = new D3D12_STATE_SUBOBJECT[subobjectCount];
    int index = 0;
    for (const DX12Shader* shader : shaders)
    {
        subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_DXIL_LIBRARY;
        subobjects[index].pDesc = shader->zLibraryDesc();
        index++;
    }
    int maxPayloadSize = 0;
    int maxAttributeSize = 0;
    for (const DX12ShaderHitGroup* hitGroup : hitGroups)
    {
        subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_HIT_GROUP;
        subobjects[index].pDesc = hitGroup->zHitGroupDesc();
        if (hitGroup->getPayloadSize() > maxPayloadSize)
        {
            maxPayloadSize = hitGroup->getPayloadSize();
        }
        if (hitGroup->getAttributeSize() > maxAttributeSize)
        {
            maxAttributeSize = hitGroup->getAttributeSize();
        }
        index++;
    }
    D3D12_RAYTRACING_SHADER_CONFIG shaderDesc = {};
    shaderDesc.MaxPayloadSizeInBytes = maxPayloadSize;
    shaderDesc.MaxAttributeSizeInBytes = maxAttributeSize;

    subobjects[index].Type
        = D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_SHADER_CONFIG;
    subobjects[index].pDesc = &shaderDesc;
    index++;
    // create export list for hole pipeline with all hitgroups and all functions
    // that are not part of a hit group
    Array<const DX12ShaderFunction*> functionsWithoutHitGroups;
    for (const DX12Shader* shader : shaders)
    {
        for (const DX12ShaderFunction* function : shader->getFunctions())
        {
            bool found = 0;
            for (const DX12ShaderHitGroup* hitGroup : hitGroups)
            {
                if (hitGroup->zClosestHitShaderFunction() == function
                    || hitGroup->zAnyHitShaderFunction() == function
                    || hitGroup->zIntersectionShaderFunction() == function)
                {
                    found = 1;
                    break;
                }
            }
            if (!found)
            {
                functionsWithoutHitGroups.add(function);
            }
        }
    }
    const wchar_t** functionAndHitGroupNames
        = new const wchar_t*[functionsWithoutHitGroups.getEntryCount()
                             + hitGroups.getEntryCount()];
    int nameIndex = 0;
    for (const DX12ShaderFunction* function : functionsWithoutHitGroups)
    {
        functionAndHitGroupNames[nameIndex]
            = const_cast<wchar_t*>(function->zExportDesc()->Name);
        nameIndex++;
    }
    for (const DX12ShaderHitGroup* hitGroup : hitGroups)
    {
        functionAndHitGroupNames[nameIndex]
            = const_cast<wchar_t*>(hitGroup->zHitGroupDesc()->HitGroupExport);
        nameIndex++;
    }
    D3D12_SUBOBJECT_TO_EXPORTS_ASSOCIATION shaderPayloadAssociation = {};
    shaderPayloadAssociation.NumExports
        = functionsWithoutHitGroups.getEntryCount() + hitGroups.getEntryCount();
    shaderPayloadAssociation.pExports = functionAndHitGroupNames;
    shaderPayloadAssociation.pSubobjectToAssociate
        = &subobjects[index - 1]; // shader config
    subobjects[index].Type
        = D3D12_STATE_SUBOBJECT_TYPE_SUBOBJECT_TO_EXPORTS_ASSOCIATION;
    subobjects[index].pDesc = &shaderPayloadAssociation;
    index++;
    D3D12_LOCAL_ROOT_SIGNATURE* localRootSignatures
        = new D3D12_LOCAL_ROOT_SIGNATURE[distinctSignatures.getEntryCount()];
    int rootSignatureIndex = 0;
    const wchar_t*** rootSignatureExports
        = new const wchar_t**[distinctSignatures.getEntryCount()];
    D3D12_SUBOBJECT_TO_EXPORTS_ASSOCIATION* localRootAssociations
        = new D3D12_SUBOBJECT_TO_EXPORTS_ASSOCIATION[distinctSignatures
                .getEntryCount()];
    for (DX12ShaderSignature* signature : distinctSignatures)
    {
        signature->createSignature(zDevice, pfnD3D12SerializeRootSignature);
        subobjects[index].Type
            = D3D12_STATE_SUBOBJECT_TYPE_LOCAL_ROOT_SIGNATURE;
        localRootSignatures[rootSignatureIndex].pLocalRootSignature
            = signature->zSignature();
        subobjects[index].pDesc = &localRootSignatures[rootSignatureIndex];
        index++;

        Array<const DX12ShaderFunction*> functionsWithThisSignature;
        for (const DX12Shader* shader : shaders)
        {
            for (const DX12ShaderFunction* function : shader->getFunctions())
            {
                if (function->zSignature() == signature)
                {
                    bool found = 0;
                    for (DX12ShaderHitGroup* hitGroup : hitGroups)
                    {
                        if (hitGroup->zClosestHitShaderFunction() == function
                            || hitGroup->zAnyHitShaderFunction() == function
                            || hitGroup->zIntersectionShaderFunction()
                                   == function)
                        {
                            found = 1;
                            break;
                        }
                    }
                    if (!found)
                    {
                        functionsWithThisSignature.add(function);
                    }
                }
            }
        }
        Array<const DX12ShaderHitGroup*> hitGroupsWithThisSignature;
        for (const DX12ShaderHitGroup* hitGroup : hitGroups)
        {
            if (hitGroup->zClosestHitShaderFunction()
                && hitGroup->zClosestHitShaderFunction()->zSignature()
                       == signature)
            {
                hitGroupsWithThisSignature.add(hitGroup);
            }
            else if (hitGroup->zAnyHitShaderFunction()
                     && hitGroup->zAnyHitShaderFunction()->zSignature()
                            == signature)
            {
                hitGroupsWithThisSignature.add(hitGroup);
            }
            else if (hitGroup->zIntersectionShaderFunction()
                     && hitGroup->zIntersectionShaderFunction()->zSignature()
                            == signature)
            {
                hitGroupsWithThisSignature.add(hitGroup);
            }
        }
        rootSignatureExports[rootSignatureIndex]
            = new const wchar_t*[functionsWithThisSignature.getEntryCount()
                                 + hitGroupsWithThisSignature.getEntryCount()];
        int nameIndex = 0;
        for (const DX12ShaderFunction* function : functionsWithThisSignature)
        {
            rootSignatureExports[rootSignatureIndex][nameIndex]
                = function->zExportDesc()->Name;
            nameIndex++;
        }
        for (const DX12ShaderHitGroup* hitGroup : hitGroupsWithThisSignature)
        {
            rootSignatureExports[rootSignatureIndex][nameIndex]
                = hitGroup->zHitGroupDesc()->HitGroupExport;
            nameIndex++;
        }
        localRootAssociations[rootSignatureIndex].NumExports
            = functionsWithThisSignature.getEntryCount()
            + hitGroupsWithThisSignature.getEntryCount();
        localRootAssociations[rootSignatureIndex].pExports
            = rootSignatureExports[rootSignatureIndex];
        localRootAssociations[rootSignatureIndex].pSubobjectToAssociate
            = &subobjects[index - 1]; // local root signature
        subobjects[index].Type
            = D3D12_STATE_SUBOBJECT_TYPE_SUBOBJECT_TO_EXPORTS_ASSOCIATION;
        subobjects[index].pDesc = &localRootAssociations[rootSignatureIndex];
        index++;

        rootSignatureIndex++;
    }

    subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_GLOBAL_ROOT_SIGNATURE;
    subobjects[index].pDesc = &emptyGlobalRootSignature;
    index++;

    subobjects[index].Type = D3D12_STATE_SUBOBJECT_TYPE_LOCAL_ROOT_SIGNATURE;
    subobjects[index].pDesc = &emptyLocalRootSignature;
    index++;

    D3D12_RAYTRACING_PIPELINE_CONFIG pipelineConfig = {};
    pipelineConfig.MaxTraceRecursionDepth = maxRecursionDepth;
    subobjects[index].Type
        = D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_PIPELINE_CONFIG;
    subobjects[index].pDesc = &pipelineConfig;
    index++;

    D3D12_STATE_OBJECT_DESC pipelineDesc = {};
    pipelineDesc.Type = D3D12_STATE_OBJECT_TYPE_RAYTRACING_PIPELINE;
    pipelineDesc.NumSubobjects = subobjectCount;
    pipelineDesc.pSubobjects = subobjects;
    HRESULT hr = zDevice->CreateStateObject(
        &pipelineDesc, __uuidof(ID3D12StateObject), (void**)&pipelineState);
    if (FAILED(hr))
    {
        Logging::error()
            << "Failed to create raytracing pipeline state object: " << std::hex
            << hr;
        throw std::logic_error("Could not create the raytracing state object");
    }

    delete[] functionAndHitGroupNames;
    delete[] localRootSignatures;
    for (int i = 0; i < distinctSignatures.getEntryCount(); i++)
    {
        delete[] rootSignatureExports[i];
    }
    delete[] rootSignatureExports;
    delete[] localRootAssociations;
}

ID3D12StateObject* Framework::DX12Pipeline::zPipelineState() const
{
    return pipelineState;
}

DX12ShaderBindingTable* Framework::DX12Pipeline::createShaderBindingTable()
{
    return new DX12ShaderBindingTable(dynamic_cast<DX12Pipeline*>(getThis()));
}

DX12GlobalDescriptorHeap* Framework::DX12Pipeline::createGlobalDescriptorHeap()
{
    return new DX12GlobalDescriptorHeap(dynamic_cast<DX12Pipeline*>(getThis()));
}

Framework::DX12GlobalDescriptorHeap::DX12GlobalDescriptorHeap(
    DX12Pipeline* pipeline)
    : ReferenceCounter(),
      descriptorHeap(0),
      lastDescriptorHeapSize(0)
{}

Framework::DX12GlobalDescriptorHeap::~DX12GlobalDescriptorHeap()
{
    pipeline->release();
    if (descriptorHeap)
    {
        descriptorHeap->Release();
    }
    for (const DX12ShaderRegisterInput* input : registerInputs)
    {
        input->inputResource->release();
        delete input;
    }
}

void Framework::DX12GlobalDescriptorHeap::addInput(
    DX12ShaderRegister type, ReferenceCounter* inputResource)
{
    bool found = 0;
    for (DX12Shader* shader : pipeline->getShaders())
    {
        for (DX12ShaderFunction* function : shader->getFunctions())
        {
            for (const DX12ShaderRegisterUsage& usage :
                function->zSignature()->getRegisterUsages())
            {
                if (usage.descriptorHeapIndex == registerInputs.getEntryCount())
                {
                    if (usage.registerType != type)
                    {
                        Logging::error()
                            << "Register type mismatch for register index "
                            << usage.registerIndex << ", space index "
                            << usage.spaceIndex << ". Expected register type: "
                            << usage.registerType
                            << ", given register type: " << type
                            << ". The register type is specified in the "
                               "signature of shader function '"
                            << function->getFunctionName() << "'";
                        throw std::logic_error(
                            "Register type mismatch for shader input");
                    }
                    else
                    {
                        found = 1;
                        break;
                    }
                }
            }
            if (found)
            {
                break;
            }
        }
        if (found)
        {
            break;
        }
    }
    registerInputs.add(
        new DX12ShaderRegisterInput{type, inputResource->getThis()});
}

void Framework::DX12GlobalDescriptorHeap::addTextureInput(
    DX12ShaderRegister type, Texture* zTexture)
{
    addInput(type, zTexture);
}

void Framework::DX12GlobalDescriptorHeap::addBufferInput(
    DX12ShaderRegister type, DXBuffer* zBuffer)
{
    addInput(type, zBuffer);
}

void Framework::DX12GlobalDescriptorHeap::addTLASInput(
    DX12ShaderRegister type, DX12TLAS* zTLAS)
{
    addInput(type, zTLAS);
}

void Framework::DX12GlobalDescriptorHeap::updateDescriptorHeap(
    ID3D12Device5* zDevice)
{
    if (!descriptorHeap
        || lastDescriptorHeapSize != registerInputs.getEntryCount())
    {
        if (descriptorHeap)
        {
            descriptorHeap->Release();
        }
        D3D12_DESCRIPTOR_HEAP_DESC desc = {};
        desc.NumDescriptors = registerInputs.getEntryCount();
        desc.Type = D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV;
        desc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
        desc.NodeMask = 0;

        HRESULT r = zDevice->CreateDescriptorHeap(
            &desc, __uuidof(ID3D12DescriptorHeap), (void**)&descriptorHeap);
        lastDescriptorHeapSize = registerInputs.getEntryCount();
    }

    D3D12_CPU_DESCRIPTOR_HANDLE descriptorHeapHandle
        = descriptorHeap->GetCPUDescriptorHandleForHeapStart();
    for (const DX12ShaderRegisterInput* input : registerInputs)
    {
        DX12TLAS* zTLAS = dynamic_cast<DX12TLAS*>(input->inputResource);
        DX12Texture* zTexture
            = dynamic_cast<DX12Texture*>(input->inputResource);
        DX12Buffer* zBuffer = dynamic_cast<DX12Buffer*>(input->inputResource);
        switch (input->registerType)
        {
        case DX12_SHADER_REGISTER_B_CONST_BUFFER:
            {
                D3D12_CONSTANT_BUFFER_VIEW_DESC cbvDesc = {};
                if (!zBuffer)
                {
                    Logging::error()
                        << "Expected a buffer resource for register type "
                        << input->registerType;
                    throw std::logic_error(
                        "Expected a buffer resource for register type "
                        + std::to_string(input->registerType));
                }
                cbvDesc.BufferLocation
                    = zBuffer->zBuffer()->GetGPUVirtualAddress();
                cbvDesc.SizeInBytes = (unsigned)zBuffer->getElementCount()
                                    * zBuffer->getElementLength();
                zDevice->CreateConstantBufferView(
                    &cbvDesc, descriptorHeapHandle);
                break;
            }
        case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
            {
                D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc;
                srvDesc.Format = DXGI_FORMAT_UNKNOWN;
                if (zTLAS)
                {
                    srvDesc.ViewDimension
                        = D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE;
                    srvDesc.RaytracingAccelerationStructure.Location
                        = zTLAS->zResultBuffer()
                              ->zBuffer()
                              ->GetGPUVirtualAddress();
                }
                else if (zTexture)
                {
                    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
                    srvDesc.Texture2D.MipLevels = 0;
                    srvDesc.Texture2D.MostDetailedMip = 0;
                    srvDesc.Texture2D.PlaneSlice = 0;
                    srvDesc.Texture2D.ResourceMinLODClamp = 0.0f;
                }
                else if (zBuffer)
                {
                    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_BUFFER;
                    srvDesc.Buffer.FirstElement = 0;
                    srvDesc.Buffer.NumElements
                        = (unsigned)zBuffer->getElementCount();
                    srvDesc.Buffer.StructureByteStride
                        = zBuffer->getElementLength();
                    srvDesc.Buffer.Flags = D3D12_BUFFER_SRV_FLAG_NONE;
                }
                srvDesc.Shader4ComponentMapping
                    = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
                zDevice->CreateShaderResourceView(
                    zTexture ? zTexture->zResource()
                             : (zBuffer ? zBuffer->zBuffer() : 0),
                    &srvDesc,
                    descriptorHeapHandle);
                break;
            }
        case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
            {
                D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};
                if (zTexture)
                {
                    uavDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D;
                    uavDesc.Format = DXGI_FORMAT_UNKNOWN;
                    uavDesc.Texture2D.MipSlice = 0;
                    uavDesc.Texture2D.PlaneSlice = 0;
                }
                else if (zBuffer)
                {
                    uavDesc.ViewDimension = D3D12_UAV_DIMENSION_BUFFER;
                    uavDesc.Buffer.FirstElement = 0;
                    uavDesc.Buffer.NumElements
                        = (unsigned)zBuffer->getElementCount();
                    uavDesc.Buffer.StructureByteStride
                        = zBuffer->getElementLength();
                    uavDesc.Buffer.CounterOffsetInBytes = 0;
                    uavDesc.Buffer.Flags = D3D12_BUFFER_UAV_FLAG_NONE;
                }
                else
                {
                    Logging::error()
                        << "Expected a texture or buffer resource for "
                           "register type "
                        << input->registerType;
                    throw std::logic_error(
                        "Expected a texture or buffer resource for register "
                        "type "
                        + std::to_string(input->registerType));
                }
                zDevice->CreateUnorderedAccessView(
                    zTexture ? zTexture->zResource() : zBuffer->zBuffer(),
                    0,
                    &uavDesc,
                    descriptorHeapHandle);
                break;
            }
        }
        descriptorHeapHandle.ptr += zDevice->GetDescriptorHandleIncrementSize(
            D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
    }
}

DX12Pipeline* Framework::DX12GlobalDescriptorHeap::zPipeline() const
{
    return pipeline;
}

ID3D12DescriptorHeap*
Framework::DX12GlobalDescriptorHeap::zDescriptorHeap() const
{
    return descriptorHeap;
}

Framework::DX12ShaderBindingTable::DX12ShaderBindingTable(
    DX12Pipeline* pipeline)
    : ReferenceCounter(),
      pipeline(pipeline)
{}