GR_GaussianSplatVK.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:        GR_GaussianSplatVK ( GR Library, C++)
 *
 * COMMENTS:
 *      3D Gaussian Splat Renderer.
 */

#ifndef GR_GaussianSplatVK_h
#define GR_GaussianSplatVK_h

#include "GR_API.h"
#include "GR_GeoRenderVK.h"

#include <GT/GT_PrimPointMesh.h>

#include <UT/UT_RadixSort.h>

class RV_Geometry;
class RV_Render;
class RV_ShaderVariableSet;
class RV_ShaderBlock;

const uint K_WORKGROUP_SIZE = 256;
const uint K_RADIX_SORT_BINS = 256; 
const uint K_RADIX_PASS = 4;

enum  ConnectIndex : uint8_t 
{
    TriangleIndex, // for gsplat and fullscreen drawing (quad)
    PointIndex     // for point marker decoration
};

// Unified structure for indirect draw and compute dispatch commands.
// This allows reuse of a single buffer across different pipeline stages
struct IndirectCmd
{
    // Indirect draw
    uint    vertexCount;
    uint    instanceCount;
    uint    firstVertex;
    uint    firstInstance;

    // Indirect compute
    uint    dispatchX;
    uint    dispatchY;
    uint    dispatchZ;
    uint    dispatchXX;
    uint    dispatchYY;
    uint    dispatchZZ;
};

struct MVPUboCPU
{
    UT_Matrix4F proj;
    UT_Matrix4F view;
    UT_Matrix4F obj;
};

struct DrawSettingsCPU
{
    bool useSph;
};

struct GaussianCPU
{
    UT_Vector4F mean;
    UT_Vector4F cdAlpha;

    UT_Vector4F rotate;
    UT_Vector4F scale;
};

struct GaussianSphCPU
{
    UT_Matrix4F sphR;
    UT_Matrix4F sphG;
    UT_Matrix4F sphB;
};

class GR_API GR_GaussianSplatVK : public GR_GeoRenderVK
{
public:
             GR_GaussianSplatVK(GR_Primitive *prim);
            ~GR_GaussianSplatVK() override;

    const char         *className() const override { return "GR_GaussianSplatVK"; }
    
    void                update(RE_RenderContext rc,
                               const GT_PrimitiveHandle &primh,
                               const GR_UpdateParms     &parms,
                               const char               *cache_name) override;
    
    void                viewUpdate(RE_RenderContext r,
                                const GR_DisplayOption &opts,
                                const char *cache_name,
                                RE_CacheVersion geo_version,
                                RE_CacheVersion select_version) override;

    void                draw(RE_RenderContext rc,
                             GR_RenderMode mode,
                             GR_RenderFlags flags,
                             GR_DrawParms dp) override;

    void                getPickConnectGroup(unsigned int pick_mode,
                                        GR_PickGeometry pick_geo,
                                        int &out_connect_start,
                                        int &out_connect_num) override;

    void                getVkDecorGeometry(
                            GR_DecorationRender::PrimitiveType prim_type,
                            RV_PrimType render_prim,
                            RV_Geometry &geo,
                            int &out_connect_start,
                            int &out_connect_num) override;

    exint               getNumGLPrimitives() const override;
    bool                hasTransparency() const override { return true; }
    bool                supportsInstances() const override { return myUseCPUSort; }
    
    static bool         initShaders(RV_Instance* inst);
    static void         cleanupShaders();

private:
    void                sortGaussians(RV_Render *rv, const GR_Uniforms &uniforms);

    void                drawGaussians(RV_Render* rv, GR_Uniforms* uniforms,
                                    GR_RenderMode mode, GR_RenderFlags flags,
                                    GR_DrawParms dp);

    void                createGSplatGeom(RV_Render *rv);
    void                createFullScreenGeom(RV_Render *rv);

    void                fillPointAttrib(RV_Render *rv, 
                                            const GT_PrimitiveHandle &primh, 
                                            const GT_PrimPointMesh *pmesh,
                                            const GR_UpdateParms &parms);

    void                createGSplatFramebuffer(RV_Render *rv, 
                                            const GR_Uniforms &uniforms);
    void                transitionToSampling(RV_Render* rv);

    void                initIndirectBuffer(RV_Render *rv);
    // draw
    void                prepareGaussiansData(const GT_PrimPointMesh& pmesh);
    void                setupDrawSet(RV_Render *rv);
    void                uploadStaticDrawBuffer(RV_Render *rv);
    void                updateDrawBuffer(RV_Render *rv);
    // keys
    void                setupKeysSets(RV_Render *rv);
    void                uploadStaticKeysBuffer(RV_Render *rv);
    void                updateKeysBuffer(RV_Render *rv, const GR_Uniforms& uniforms);
    bool                computeKeys(RV_Render *rv);
    // update dispatch params
    void                setupIndirectParamsSet(RV_Render *rv);
    bool                computeIndirectParams(RV_Render *rv);
    // hist
    void                setupHistSets(RV_Render *rv);
    bool                computeHist(RV_Render *rv);
    // prefix sum
    void                setupPrefixSumSets(RV_Render *rv);
    bool                computePrefixSum(RV_Render *rv);
    // sort
    void                setupSortSets(RV_Render *rv);
    void                clearSortBuffer(RV_Render *rv);
    void                updateSortBuffer(RV_Render *rv, uint pass);
    bool                computeSort(RV_Render *rv);
    void                pingPongBuffer(RV_Render *rv, uint pass);

    // full screen with color image and depth reconstruction
    void                setupDepthSet(RV_Render *rv);
    void                drawDepth(RV_Render* rv, GR_Uniforms* uniforms,
                                GR_RenderMode mode, GR_RenderFlags flags,
                                GR_DrawParms dp); 

    void                sortOnCPU(RV_Render *rv,
                                const GR_Uniforms &uniforms);

    static RV_ShaderProgram     *theGSplatDrawShader;
    static RV_ShaderProgram     *theDepthDrawShader;

    static RV_ShaderCompute     *theKeysComputeShader;
    static RV_ShaderCompute     *theIndirectParamsComputeShader;
    static RV_ShaderCompute     *theHistComputeShader;
    static RV_ShaderCompute     *thePrefixSumComputeShader;
    static RV_ShaderCompute     *theSortComputeShader;

    // Options
    bool                               myUseCPUSort;
    bool                               myUseSPH;
    bool                               myObjectSelected;

    std::mutex                         myMutex;

    exint                              myNumInstances; 
    exint                              myNumVisibles;

    UT_UniquePtr<RV_Geometry>          myFullScreenGeoVK;

    UT_UniquePtr<RV_Framebuffer>       myFramebuffer;
    UT_UniquePtr<RV_VKImage>           myColorImage;
    UT_UniquePtr<RV_VKImage>           myDepthDataImage;

    // indirect draw
    IndirectCmd                        myIndirectCmd;
    UT_UniquePtr<RV_ShaderBlock>       myIndirectBlock;

    // gaussians data on CPU
    MVPUboCPU                          myMVPUboCPU;
    DrawSettingsCPU                    myDrawSettingsCPU;
    UT_Array<GaussianSphCPU>           mySphArray;
    UT_Array<GaussianCPU>              myGaussianArray;

    // depth sorting data on CPU
    UT_Array<UT_Vector4F>              myMeanArray;
    UT_Array<uint>                     myKeysArray;
    UT_Array<uint>                     myIndicesArray;

    // draw shader
    UT_UniquePtr<RV_ShaderVariableSet> myDrawSet;
    UT_UniquePtr<RV_ShaderBlock>       myDrawBlock;
    UT_UniquePtr<RV_ShaderBlock>       myDrawSphBlock;
    UT_UniquePtr<RV_ShaderBlock>       myDrawSettingsBlock;

    // keys generation compute shader 
    UT_UniquePtr<RV_ShaderVariableSet> myKeysSet;
    UT_UniquePtr<RV_ShaderVariableSet> myKeysUboSet;

    UT_UniquePtr<RV_ShaderBlock>       myMVPBlock;
    UT_UniquePtr<RV_ShaderBlock>       myMeanBlock;
    UT_UniquePtr<RV_ShaderBlock>       myKeysBlock;
    UT_UniquePtr<RV_ShaderBlock>       myIndicesBlock;
    UT_UniquePtr<RV_ShaderBlock>       myPushConstBlock;
    UT_UniquePtr<RV_ShaderBlock>       myIndirectParamsBlock;

    UT_UniquePtr<RV_ShaderVariableSet> myIndirectParamsSet;

    // depth-sorting compute shader
    UT_UniquePtr<RV_ShaderVariableSet> myHistSet;

    UT_UniquePtr<RV_ShaderVariableSet> myPrefixSumSet;

    UT_UniquePtr<RV_ShaderVariableSet> mySortSet;
    UT_UniquePtr<RV_ShaderVariableSet> mySortUboSet;
    UT_UniquePtr<RV_ShaderVariableSet> mySortPingPongSet;

    UT_UniquePtr<RV_ShaderBlock>       myHistogramBlock;
    UT_UniquePtr<RV_ShaderBlock>       myGlobalHistBlock;
    UT_UniquePtr<RV_ShaderBlock>       mySortedKeysBlock;
    UT_UniquePtr<RV_ShaderBlock>       mySortedIndicesBlock;
    UT_UniquePtr<RV_ShaderBlock>       mySortPushConstBlock;

    UT_UniquePtr<RV_ShaderVariableSet> myDepthSet;
};

#endif