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::addRegisterUsage(
    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)
    {
        it.addBefore({registerType, registerIndex, spaceIndex});
    }
    else
    {
        registerUsages.add({registerType, registerIndex, spaceIndex});
    }
    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;
    descriptorTable.ParameterType = D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE;
    descriptorTable.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
    descriptorTable.DescriptorTable.NumDescriptorRanges
        = registerUsages.getEntryCount();
    D3D12_DESCRIPTOR_RANGE* descriptorRanges
        = new D3D12_DESCRIPTOR_RANGE[registerUsages.getEntryCount()];
    int index = 0;
    ArrayIterator<DX12ShaderRegisterUsage> it = registerUsages.begin();
    while (it)
    {
        const auto& usage = it.val();
        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)
        {
            ++size;
            it = next;
            ++next;
        }
        range->NumDescriptors = size;
        range->BaseShaderRegister = usage.registerIndex;
        range->RegisterSpace = usage.spaceIndex;
        range->OffsetInDescriptorsFromTableStart = index;
        ++it;
        index++;
    }
    descriptorTable.DescriptorTable.pDescriptorRanges = descriptorRanges;
    D3D12_ROOT_SIGNATURE_DESC rootDesc = {};
    rootDesc.NumParameters = 1;
    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::getRegisterUsagesOrder() const
{
    return registerUsages;
}

ShaderHeap* Framework::DX12ShaderSignature::createShaderHeap()
{
    return new ShaderHeap(dynamic_cast<DX12ShaderSignature*>(getThis()));
}

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 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 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* intersectionShaderFunction)
{
    if (closestHitShaderFunction
        && closestHitShaderFunction->zSignature()
               != intersectionShaderFunction->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 (anyHitShaderFunction
        && anyHitShaderFunction->zSignature()
               != intersectionShaderFunction->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->intersectionShaderFunction)
    {
        this->intersectionShaderFunction->release();
    }
    this->intersectionShaderFunction = intersectionShaderFunction;
    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;
}

Framework::ShaderHeap::ShaderHeap(DX12ShaderSignature* signature)
    : ReferenceCounter(),
      signature(signature),
      descriptorHeap(0),
      lastDescriptorHeapSize(0)
{}

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

void Framework::ShaderHeap::setRegisterInput(DX12ShaderRegister registerType,
    int registerIndex,
    int spaceIndex,
    ReferenceCounter* inputResource)
{
    for (DX12ShaderRegisterInput* input : registerInputs)
    {
        if (input->registerIndex == registerIndex
            && input->spaceIndex == spaceIndex
            && input->registerType == registerType)
        {
            input->inputResource->release();
            input->inputResource = inputResource->getThis();
            return;
        }
    }
    bool found = 0;
    for (const DX12ShaderRegisterUsage& usage :
        signature->getRegisterUsagesOrder())
    {
        if (usage.registerIndex == registerIndex
            && usage.spaceIndex == spaceIndex
            && usage.registerType == registerType)
        {
            found = 1;
            break;
        }
    }
    if (found)
    {
        registerInputs.add(new DX12ShaderRegisterInput{
            registerType, registerIndex, spaceIndex, inputResource->getThis()});
    }
    else
    {
        Logging::error() << "Register type " << registerType
                         << ", register index " << registerIndex
                         << ", space index " << spaceIndex
                         << " is not used in the shader signature. The given "
                            "input will be ignored.";
    }
}

void Framework::ShaderHeap::setRegisterInput(
    Texture* zTexture, int registerIndex, int spaceIndex)
{
    setRegisterInput(DX12_SHADER_REGISTER_U_UNORDERED_ACCESS,
        registerIndex,
        spaceIndex,
        zTexture);
}

void Framework::ShaderHeap::setRegisterInput(
    DXBuffer* zBuffer, int registerIndex, int spaceIndex)
{
    setRegisterInput(DX12_SHADER_REGISTER_B_CONST_BUFFER,
        registerIndex,
        spaceIndex,
        zBuffer);
}

void Framework::ShaderHeap::setRegisterInput(
    DX12TLAS* zTLAS, int registerIndex, int spaceIndex)
{
    setRegisterInput(DX12_SHADER_REGISTER_T_SHADER_RESOURCE,
        registerIndex,
        spaceIndex,
        zTLAS);
}

void Framework::ShaderHeap::updateDescriptorHeap(ID3D12Device5* zDevice)
{
    const Array<DX12ShaderRegisterUsage>& registerUsages
        = signature->getRegisterUsagesOrder();
    if (!descriptorHeap
        || lastDescriptorHeapSize
               != signature->getRegisterUsagesOrder().getEntryCount())
    {
        if (descriptorHeap)
        {
            descriptorHeap->Release();
        }
        D3D12_DESCRIPTOR_HEAP_DESC desc = {};
        desc.NumDescriptors = registerUsages.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 = registerUsages.getEntryCount();
    }

    D3D12_CPU_DESCRIPTOR_HANDLE descriptorHeapHandle
        = descriptorHeap->GetCPUDescriptorHandleForHeapStart();
    ArrayIterator<DX12ShaderRegisterUsage> it = registerUsages.begin();
    while (it)
    {
        const DX12ShaderRegisterUsage& usage = it.val();
        ArrayIterator<DX12ShaderRegisterInput*> inputIt
            = registerInputs.begin();
        bool found = 0;
        while (inputIt)
        {
            if (inputIt->registerType == usage.registerType
                && inputIt->registerIndex == usage.registerIndex
                && inputIt->spaceIndex == usage.spaceIndex)
            {
                switch (inputIt->registerType)
                {
                case DX12_SHADER_REGISTER_B_CONST_BUFFER:
                    D3D12_CONSTANT_BUFFER_VIEW_DESC cbvDesc = {};
                    DX12Buffer* buffer
                        = dynamic_cast<DX12Buffer*>(inputIt->inputResource);
                    cbvDesc.BufferLocation
                        = buffer->zBuffer()->GetGPUVirtualAddress();
                    cbvDesc.SizeInBytes = buffer->getElementCount()
                                        * buffer->getElementLength();
                    zDevice->CreateConstantBufferView(
                        &cbvDesc, descriptorHeapHandle);
                    break;
                case DX12_SHADER_REGISTER_T_SHADER_RESOURCE:
                    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc;
                    srvDesc.Format = DXGI_FORMAT_UNKNOWN;
                    srvDesc.ViewDimension
                        = D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE;
                    srvDesc.Shader4ComponentMapping
                        = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
                    srvDesc.RaytracingAccelerationStructure.Location
                        = dynamic_cast<DX12TLAS*>(inputIt->inputResource)
                              ->zResultBuffer()
                              ->zBuffer()
                              ->GetGPUVirtualAddress();
                    zDevice->CreateShaderResourceView(
                        0, &srvDesc, descriptorHeapHandle);
                    break;
                case DX12_SHADER_REGISTER_U_UNORDERED_ACCESS:
                    D3D12_UNORDERED_ACCESS_VIEW_DESC uavDesc = {};
                    uavDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D;
                    uavDesc.Format = DXGI_FORMAT_UNKNOWN;
                    uavDesc.Texture2D.MipSlice = 0;
                    uavDesc.Texture2D.PlaneSlice = 0;
                    zDevice->CreateUnorderedAccessView(
                        dynamic_cast<DX12Texture*>(inputIt->inputResource)
                            ->zResource(),
                        0,
                        &uavDesc,
                        descriptorHeapHandle);
                    break;
                default:
                    Logging::error()
                        << "Unknown register type for descriptor heap: "
                        << inputIt->registerType;
                }
                found = 1;
                break;
            }
            ++inputIt;
        }
        if (!found)
        {
            Logging::error()
                << "No input resource found for register type "
                << usage.registerType << ", register index "
                << usage.registerIndex << ", space index " << usage.spaceIndex
                << ". This will result in an uninitialized descriptor in the "
                   "descriptor heap. Access to the register in the shader "
                   "might lead to undefined behaviour.";
        }
        descriptorHeapHandle.ptr += zDevice->GetDescriptorHandleIncrementSize(
            D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV);
        ++it;
    }
}

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