RV_VKAccelerationStructure.h

Go to the documentation of this file.

/*
 * PROPRIETARY INFORMATION.  This software is proprietary to
 * Side Effects Software Inc., and is not to be reproduced,
 * transmitted, or disclosed in any way without written permission.
 *
 * NAME:        RV_VKAccelerationStructure.h ( RV Library, C++)
 *
 * COMMENTS:
 *      Class to manage Vulkan Acceleration Structure object
 */

#ifndef RV_VKAccelerationStructure_h
#define RV_VKAccelerationStructure_h

#include "RV_API.h"
#include "RV_TypePtrs.h"
#include "RV_VKQuery.h"

#include <VE/VE_VK.h>

#include <UT/UT_Array.h>
#include <UT/UT_Matrix4.h>
#include <UT/UT_UniquePtr.h>

class RV_Instance;
class RV_VKBuffer;


struct RV_AccelStructShaderInfo
{
    // Defines the ray hit behavior.
    int tag = RV_ACCEL_STRUCT_OPAQUE_OBJECT;
    // Light ID > -1 if the accel struct is a light.
    int lightId = -1;
};

// -----------------
// RV_VKAccelerationStructureGeometry
//
// Specifies a single geometry within an acceleration structure
class RV_API RV_VKAccelerationStructureGeometry
{
public:
    void setOpaque(bool opaque)
        { myIsOpaque = opaque; }
    bool isOpaque() const
        { return myIsOpaque; }
    void setSingleHit(bool singleHit)
        { myIsSingleHit = singleHit; }
    bool isSingleHit() const
        { return myIsSingleHit; }

    void setTriangles(
            VkFormat vertex_format,
            VkDeviceAddress vertex_buffer,
            exint vertex_stride,
            uint32_t num_vertices,
            VkDeviceAddress index_buffer,
            uint32_t num_indices);

    void setAABBs(
            VkDeviceAddress aabb_buffer,
            exint aabb_stride,
            uint32_t num_boxes);

    void setInstances(
            VkDeviceAddress instance_buffer,
            uint32_t num_instances);

    void fillGeometryInfo();

protected:
    VkAccelerationStructureGeometryKHR myVkGeometry = {
        VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
        nullptr,
        VK_GEOMETRY_TYPE_TRIANGLES_KHR, /* geometryType */
        {},                             /* geometry */
        0                               /* flags */
    };

    VkAccelerationStructureGeometryDataKHR myGeometryData = {};
    VkAccelerationStructureBuildRangeInfoKHR myBuildRangeInfo = {};
    VkGeometryTypeKHR myGeometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;

    bool myIsOpaque = true;
    bool myIsSingleHit = false;

    friend class RV_VKAccelerationStructureBuildInfo;
};

// -----------------
// RV_VKAccelerationStructureBuildInfo
//
// Specifies the info for creating and building an acceleration structure
class RV_API RV_VKAccelerationStructureBuildInfo
{
public:
    void setUpdate(bool update)
        { myIsUpdate = update; }
    bool isUpdate() const
        { return myIsUpdate; }
    void setTopLevel(bool top_level)
        { myIsTopLevel = top_level; }
    bool isTopLevel() const
        { return myIsTopLevel; }
    void setPreference(RV_AccelStructBuildPref preference)
        { myPreference = preference; }
    RV_AccelStructBuildPref getPreference() const
        { return myPreference; }

    exint getSize() const
        { return myBuildSizeInfo.accelerationStructureSize; }
    exint getScratchSize() const
        { return !myIsUpdate ? myBuildSizeInfo.buildScratchSize :
                               myBuildSizeInfo.updateScratchSize; }
    const VkAccelerationStructureBuildRangeInfoKHR* getBuildRangeInfos()
        { return myBuildRangeInfos.data(); }

    void addGeometry(RV_VKAccelerationStructureGeometry& geometry);

    RV_VKAccelerationStructureBuildInfo() = default;
    UT_UniquePtr<RV_VKAccelerationStructureBuildInfo> clone() const
    {
        return UT_UniquePtr<RV_VKAccelerationStructureBuildInfo>(
                    new RV_VKAccelerationStructureBuildInfo(*this));
    }

protected:
    VkAccelerationStructureCreateInfoKHR myVkCreateInfo = {
        VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR,
        nullptr,
        0,                                               /* createFlags */
        VK_NULL_HANDLE,                                  /* buffer */
        0,                                               /* offset */
        0,                                               /* size */
        VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR, /* type */
        0                                                /* deviceAddress */
    };

    VkAccelerationStructureBuildGeometryInfoKHR myVkBuildInfo = {
        VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
        nullptr,
        VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR, /* type */
        VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_COMPACTION_BIT_KHR, /* flags */
        VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR,  /* mode */
        VK_NULL_HANDLE,                                  /* srcAccelerationStructure */
        VK_NULL_HANDLE,                                  /* dstAccelerationStructure */
        0,                                               /* geometryCount */
        nullptr,                                         /* pGeometries */
        nullptr,                                         /* ppGeometries */
        {}                                               /* scratchData */
    };

    VkAccelerationStructureBuildSizesInfoKHR myBuildSizeInfo = {
        VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR,
        nullptr,
        0,
        0,
        0
    };

    RV_VKAccelerationStructureBuildInfo(const RV_VKAccelerationStructureBuildInfo&) = default;
    RV_VKAccelerationStructureBuildInfo &operator=(const RV_VKAccelerationStructureBuildInfo&) = delete;

    void querySizeInfo(RV_Instance* inst);
    void fillCreateInfo(RV_VKBuffer* buffer);
    void fillBuildInfo(
        RV_VKBuffer* scratch_buffer,
        VkAccelerationStructureKHR accel_struct);

    friend class RV_VKAccelerationStructure;

private:
    UT_Array<VkAccelerationStructureGeometryKHR> myVkGeometries;
    UT_Array<VkAccelerationStructureBuildRangeInfoKHR> myBuildRangeInfos;

    bool myIsUpdate = false;
    bool myIsTopLevel = false;
    RV_AccelStructBuildPref myPreference = RV_ACCEL_STRUCT_FAST_BUILD;
};

class RV_API RV_VKAccelerationStructureInstance
{
public:
    void setTransform(UT_Matrix4F transform)
        { myTransform = transform; }
    void setInstCustomIndex(exint index)
        { myVkAccelStructInst.instanceCustomIndex = index; }
    void setVisibilityMask(exint mask)
        { myVkAccelStructInst.mask = mask; }
    void setSBTOffset(exint offset)
        { myVkAccelStructInst.instanceShaderBindingTableRecordOffset = offset; }
    void setFaceCullingEnabled(bool enable)
        { myFaceCullingEnabled = enable; }
    
    void fillInstanceInfo();

protected:
    VkAccelerationStructureInstanceKHR myVkAccelStructInst = {};

    void setAccelStructReference(VkDeviceAddress ref)
        { myVkAccelStructInst.accelerationStructureReference = ref; }
    
    friend class RV_VKAccelerationStructure;

private:
    UT_Matrix4F myTransform = UT_Matrix4F::getIdentityMatrix();
    bool myFaceCullingEnabled = true;
};

// -----------------
// RV_VKAccelerationStructure
//
// Class holding VkAcceleratrionStructureKHR handle and bound
/// memory allocation
class RV_API RV_VKAccelerationStructure
{
public:
    // Build the compact version of an acceleration structure if the
    // its compact size is less than %90 of its worst-case size
    static constexpr float theCompactSizeThreshold = 0.9f;

    static RV_VKAccelerationStructurePtr create(
            RV_Instance* inst,
            RV_VKAccelerationStructureBuildInfo* as_info,
            const UT_StringHolder &name = UT_StringHolder());

    RV_VKAccelerationStructure(
            RV_Instance* inst,
            RV_VKAccelerationStructureBuildInfo* as_info,
            VkAccelerationStructureKHR vk_as,
            RV_VKBufferPtr buffer);

    ~RV_VKAccelerationStructure();

    void setShaderInfoIndex(uint32_t index)
        { myShaderInfoIndex = index; }
    uint32_t getShaderInfoIndex() const
        { return myShaderInfoIndex; }
    void setTag(int tag)
        { myShaderInfo.tag = tag; }
    void setLightId(int light_id)
        { myShaderInfo.lightId = light_id; }
    RV_AccelStructShaderInfo getShaderInfo() const
        { return myShaderInfo; }
    VkAccelerationStructureKHR getVkAccelStruct()
        { return myVkAccelStruct; }
    exint getScratchSize() const
        { return myInfo->getScratchSize(); }
    bool isBuilt() const
        { return myIsBuilt; }
    const RV_ResourceID &getID() const
        { return myId; }
    const UT_StringHolder &getName() const
        { return myName; }
    exint getSize() const;
    VkDeviceAddress getDeviceAddress();

    void addInstance(RV_VKAccelerationStructureInstance as_inst);
    void clearInstances();
    void getVkInstances(UT_Array<VkAccelerationStructureInstanceKHR> &as_insts);

    // Only call this when the latest instances have been included
    // in a top-level acceleration structure.
    void clearInstDirtyFlag()
        { myIsInstDirty = false; }
    bool isInstDirty() const
        { return myIsInstDirty; }

    void build(RV_VKCommandBuffer* cb, RV_VKBuffer* scratch_buffer);
    void buildCompact(
        RV_VKCommandBuffer* cb,
        exint compact_size,
        VkAccelerationStructureKHR& compact_vk_as,
        RV_VKBufferPtr& compact_buffer);
    void reset();
    void update(
        VkAccelerationStructureKHR compact_vk_as,
        RV_VKBufferPtr compact_buffer);

    // return the size of the struct in GPU memory bytes
    exint gpuMemoryUsage() const;

    // Debug print out of acceleration structure properties
    void print() const;

private:
    RV_Instance* myInst;
    RV_ResourceID myId;
    UT_StringHolder myName;
    UT_UniquePtr<RV_VKAccelerationStructureBuildInfo> myInfo;
    VkAccelerationStructureKHR myVkAccelStruct;
    RV_VKBufferPtr myBuffer;
    UT_Array<RV_VKAccelerationStructureInstance> myInstances;

    bool myIsBuilt = false;
    bool myIsInstDirty = true;
    uint32_t myShaderInfoIndex = 0;

    RV_AccelStructShaderInfo myShaderInfo;

    exint getVisibilityMask() const;
};

// --------------------------------------------------
// RV_BuildAccelerationStructures
//
// Helper class to build acceleration structures
class RV_API RV_BuildAccelerationStructures
{
public:
    RV_BuildAccelerationStructures(RV_Instance* inst);

    void addAccelerationStructure(RV_VKAccelerationStructure* accel_struct);

    void build(RV_VKCommandBuffer* cb);
    void buildAndQueryCompactSize(RV_VKCommandBuffer* cb);
    void buildCompact(RV_VKCommandBuffer* cb);
    void updateCompact();

    bool isBuilt() const
        { return myIsBuilt; }
    bool isBuiltCompact() const
        { return myIsBuiltCompact; }
    
    exint getTotalAccelStructSize() const
        { return myTotalAccelStructSize; }
    exint getTotalScratchSize() const
        { return myTotalScratchSize; }
    exint getTotalCompactSize();

private:
    RV_Instance* myInst;
    UT_UniquePtr<RV_VKQuery> myQuery;
    UT_Array<RV_VKAccelerationStructure*> myAccelStructs;
    UT_Array<RV_VKBufferPtr> myScratchBuffers;

    UT_Array<VkAccelerationStructureKHR> myCompactVkAccelStructs;
    UT_Array<RV_VKBufferPtr> myCompactBuffers;

    bool myIsBuilt = false;
    bool myIsBuiltCompact = false;

    exint myTotalAccelStructSize = 0;
    exint myTotalScratchSize = 0;
    exint myTotalCompactSize = 0;
};

class RV_ResetAccelStructTask
{
public:
    RV_VKAccelerationStructure* myAccelStruct;

    void operator()(RV_Instance*)
    {
        myAccelStruct->reset();
    }

    RV_ResetAccelStructTask(
        RV_VKAccelerationStructure* accel_struct)
    : myAccelStruct(accel_struct)
    {}
};

RV_VKCommandBuffer::Callback RVmakeResetAccelStructTask(
    RV_VKAccelerationStructure* accel_struct);
#endif