KSG-Framework/DX12BLAS.cpp

#include "DX12BLAS.h"

#include "DX12Buffer.h"
#include "DX12CommandQueue.h"

Framework::DX12BLAS::DX12BLAS(
    ID3D12Device5* zDevice, DX12DirectCommandQueue* zDirectQueue)
    : ReferenceCounter(),
      scratchBuffer(new DX12Buffer(1,
          zDevice,
          dynamic_cast<DX12CommandQueue*>(zDirectQueue->getThis()),
          D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS)),
      resultBuffer(new DX12Buffer(1,
          zDevice,
          dynamic_cast<DX12CommandQueue*>(zDirectQueue->getThis()),
          D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS)),
      vertexDataBuffer(0),
      indexBuffer(0),
      polygonSizeBuffer(0),
      zDevice(zDevice),
      zDirectQueue(zDirectQueue),
      geometryDesc(new D3D12_RAYTRACING_GEOMETRY_DESC())
{}

Framework::DX12BLAS::~DX12BLAS()
{
    scratchBuffer->release();
    resultBuffer->release();
    if (vertexDataBuffer)
    {
        vertexDataBuffer->release();
    }
    if (indexBuffer)
    {
        indexBuffer->release();
    }
    if (polygonSizeBuffer)
    {
        polygonSizeBuffer->release();
    }
    delete geometryDesc;
}

void Framework::DX12BLAS::build(DX12Buffer* zVertexBuffer,
    DX12Buffer* vertexDataBuffer,
    DX12Buffer* indexBuffer,
    DX12Buffer* polygonSizeBuffer)
{
    if (this->indexBuffer)
    {
        this->indexBuffer->release();
    }
    if (this->vertexDataBuffer)
    {
        this->vertexDataBuffer->release();
    }
    if (this->polygonSizeBuffer)
    {
        this->polygonSizeBuffer->release();
    }
    this->vertexDataBuffer = vertexDataBuffer;
    this->indexBuffer = indexBuffer;
    this->polygonSizeBuffer = polygonSizeBuffer;

    memset(geometryDesc, 0, sizeof(D3D12_RAYTRACING_GEOMETRY_DESC));
    geometryDesc->Type = D3D12_RAYTRACING_GEOMETRY_TYPE_TRIANGLES;
    geometryDesc->Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_NONE;
    geometryDesc->Triangles.VertexBuffer.StartAddress
        = zVertexBuffer->zBuffer()->GetGPUVirtualAddress();
    geometryDesc->Triangles.VertexBuffer.StrideInBytes
        = zVertexBuffer->getElementLength();
    geometryDesc->Triangles.VertexCount
        = (unsigned)zVertexBuffer->getElementCount();
    geometryDesc->Triangles.VertexFormat = DXGI_FORMAT_R32G32B32_FLOAT;
    geometryDesc->Triangles.IndexBuffer
        = indexBuffer->zBuffer()->GetGPUVirtualAddress();
    geometryDesc->Triangles.IndexCount
        = (unsigned)indexBuffer->getElementCount();
    geometryDesc->Triangles.IndexFormat = DXGI_FORMAT_R32_UINT;

    D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS prebuildDesc;
    prebuildDesc.Type
        = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL;
    prebuildDesc.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
    prebuildDesc.NumDescs = 1;
    prebuildDesc.pGeometryDescs = geometryDesc;
    prebuildDesc.Flags
        = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_NONE;

    D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO info = {};
    zDevice->GetRaytracingAccelerationStructurePrebuildInfo(
        &prebuildDesc, &info);

    scratchBuffer->setLength(
        ROUND_UP_POWER_OF_2((int)info.ScratchDataSizeInBytes, 256));
    scratchBuffer->createBufferWithoutData(D3D12_RESOURCE_STATE_COMMON);
    resultBuffer->setLength(
        ROUND_UP_POWER_OF_2((int)info.ResultDataMaxSizeInBytes, 256));
    resultBuffer->createBufferWithoutData(
        D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE);

    D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_DESC buildDesc;
    buildDesc.Inputs.Type
        = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL;
    buildDesc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY;
    buildDesc.Inputs.NumDescs = 1;
    buildDesc.Inputs.pGeometryDescs = geometryDesc;
    buildDesc.DestAccelerationStructureData
        = {resultBuffer->zBuffer()->GetGPUVirtualAddress()};
    buildDesc.ScratchAccelerationStructureData
        = {scratchBuffer->zBuffer()->GetGPUVirtualAddress()};
    buildDesc.SourceAccelerationStructureData = 0;
    buildDesc.Inputs.Flags
        = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_NONE;

    // Build the AS
    zDirectQueue->zCommandList()->BuildRaytracingAccelerationStructure(
        &buildDesc, 0, nullptr);
}

Framework::DX12Buffer* Framework::DX12BLAS::zScratchBuffer() const
{
    return scratchBuffer;
}

Framework::DX12Buffer* Framework::DX12BLAS::zResultBuffer() const
{
    return resultBuffer;
}

Framework::DX12Buffer* Framework::DX12BLAS::zVertexDataBuffer() const
{
    return vertexDataBuffer;
}

Framework::DX12Buffer* Framework::DX12BLAS::zIndexBuffer() const
{
    return indexBuffer;
}

Framework::DX12Buffer* Framework::DX12BLAS::zPolygonSizeBuffer() const
{
    return polygonSizeBuffer;
}