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 "GR_Uniforms.h"
#include "GR_UtilsVK.h"
#include "GR_Light.h"

#include <GT/GT_PrimPointMesh.h>
#include <UT/UT_RadixSort.h>

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

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
    DispatchIndex  // for compute workgroup dispatch
};

enum GSplatProjectionModeHint
{
    GSPLAT_PROJECTION_MODE_PERSPECTIVE = 0,
    GSPLAT_PROJECTION_MODE_TANGENTIAL = 1
};

enum GSplatSortingModeHint
{
    GSPLAT_SORTING_MODE_Z_DEPTH = 0,
    GSPLAT_SORTING_MODE_CAMERA_DISTANCE = 1
};

struct GSplatAttribStream
{
    GT_Storage cpuFormat = GT_STORE_INVALID;
    RV_GPUType gpuFormat = RV_GPU_UNKNOWN;
    uint component = 0;
    exint size = 0;
    std::variant<UT_Array<fpreal16>, UT_Array<fpreal32>> data;
    UT_UniquePtr<RV_ShaderBlock> block;
};

struct GSplatSHAttribs
{
    GSplatAttribStream r;
    GSplatAttribStream g;
    GSplatAttribStream b;
};

struct GSplatAttribs
{
    GSplatSHAttribs    sh;
};

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                getVkDecorConnectGroup(
                            RV_PrimType render_prim,
                            bool is_uv_view,
                            int &out_connect_start,
                            int &out_connect_num) override;

    exint               getNumGLPrimitives() const override;
    bool                hasTransparency() const override { return false; }

    void                retireGeometry() override;
    
    static bool         initShaders(RV_Instance* inst);
    static void         cleanupShaders();

    // True while any gsplat renderer is live. DM_RenderVulkan uses this to build
    // the secondary "gsplat-excluded" shadow map only for scenes with gsplats.
    static bool         sceneHasGsplats();

private:
    void                drawPrivate(RV_Render *rv, GR_Uniforms* u,
                                    GR_RenderMode mode, GR_RenderFlags flags,
                                    GR_DrawParms dp);

    void                drawUtility(RV_Render *rv, GR_Uniforms* u,
                                    GR_RenderMode mode, GR_RenderFlags flags,
                                    GR_DrawParms dp);

    void                refreshViewDependentState(RV_Render *rv, GR_Uniforms *uniforms);
    void                frustumCullingGPU(RV_Render *rv, GR_Uniforms *uniforms);
    void                refreshPreprocess(RV_Render *rv, GR_Uniforms *uniforms);

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

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

    void                drawDepth(RV_Render* rv, GR_Uniforms* uniforms,
                                GR_RenderMode mode, GR_RenderFlags flags,
                                GR_DrawParms dp); 
    
    void                drawShadow(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                initIndirectCmd(RV_Render *rv);
    
    void                setupAllSets(RV_Render *rv);
    void                uploadAllStaticBuffer(RV_Render *rv);

    void                prepareGaussiansData(const GT_PrimPointMesh& pmesh);
    
    void                setupDrawSet(RV_Render *rv);
    void                setupDrawPointSet(RV_Render *rv);
    void                setupKeysSets(RV_Render *rv);
    void                setupDepthSet(RV_Render *rv);
    bool                bindDepthShadowLight(RV_Render *rv,
                                            GR_Light *light,
                                            GR_DrawParms dp,
                                            RV_ShaderProgram *shader);
    void                setupPreprocessBuffer(RV_Render *rv);

    void                updateDrawBuffer(RV_Render *rv);
    void                updateKeysBuffer(RV_Render *rv, GR_Uniforms *uniforms);

    void                uploadHintBuffer(RV_Render *rv);
    void                uploadShadowDrawIndirectBuffer(RV_Render *rv);
    void                uploadStaticPreprocessBuffer(RV_Render *rv);
    void                uploadGSplatAttribStream(RV_Render *rv,
                                            const GSplatAttribStream &stream);
    
    bool                computeKeys(RV_Render *rv, GR_Uniforms *uniforms);
    bool                computePreprocess(RV_Render *rv, GR_Uniforms *uniforms);

    exint               getGSplatAttribScalarCount(const GSplatAttribStream &stream);
    void                bindGSplatAttribFormat(RV_ShaderBlock *block, const char *name,
                                            const GSplatAttribStream &stream);
    void                bindGSplatGeometryAttribFormat(
                                            RV_ShaderBlock *block, 
                                            const char *name,
                                            const char *attrib_name);

    GSplatProjectionModeHint parseProjectionModeHint(const GT_DataArrayHandle &attrib);
    GSplatSortingModeHint    parseSortingModeHint(const GT_DataArrayHandle &attrib);
    

    static RV_ShaderProgram     *theGSplatDrawShader;
    static RV_ShaderProgram     *theDepthDrawShader;
    static RV_ShaderProgram     *theDepthShaderWithDistantShadow;
    static RV_ShaderProgram     *theDepthShaderWithPointShadow;
    static RV_ShaderProgram     *thePointDrawShader;
    static RV_ShaderProgram     *thePointPickShader;

    static RV_ShaderProgram     *theShadowCubeShader;
    static RV_ShaderProgram     *theShadowCascadeShader;

    static RV_ShaderCompute     *theKeysComputeShader;
    static RV_ShaderCompute     *thePreprocessComputeShader;

    // Options
    bool                               myUseSPH;
    bool                               myObjectSelected;
    bool                               myIsWireframe;
    fpreal                             myGSplatsAlphaCulling;
    GSplatProjectionModeHint           myProjectionModeHint;
    GSplatSortingModeHint              mySortingModeHint;

    exint                              myNumInstances; 
    exint                              myNumVisibles;
    UT_Matrix4DArray                   myInstanceTransforms;

    UT_UniquePtr<RV_Geometry>          myFullScreenGeoVK;

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

    VkDrawIndirectCommand              myDrawIndirectCmd;

    UT_UniquePtr<RV_ShaderVariableSet> myKeysSet;
    UT_UniquePtr<RV_ShaderVariableSet> myDrawSet;
    UT_UniquePtr<RV_ShaderVariableSet> myDepthSet;
    UT_UniquePtr<RV_ShaderVariableSet> myPointDrawSet;
    UT_UniquePtr<RV_ShaderVariableSet> myPointPickSet;
    UT_UniquePtr<RV_ShaderVariableSet> myDrawShadowSet;
    UT_UniquePtr<RV_ShaderVariableSet> myPreprocessSet;
    UT_UniquePtr<RV_ShaderVariableSet> myPreprocessShadowSet;

    UT_UniquePtr<RV_ShaderBlock> myKeysBlock;
    UT_UniquePtr<RV_ShaderBlock> myHintsBlock;
    UT_UniquePtr<RV_ShaderBlock> myIndicesBlock;
    UT_UniquePtr<RV_ShaderBlock> myPreprocessBlock;
    UT_UniquePtr<RV_ShaderBlock> myFullIndicesBlock;
    UT_UniquePtr<RV_ShaderBlock> myDrawIndirectBlock;
    UT_UniquePtr<RV_VKBuffer>    myShadowDrawIndirectBuffer;
    UT_UniquePtr<RV_ShaderBlock> myShadowAttribFormatsBlock;
    UT_UniquePtr<RV_ShaderBlock> myPreprocessAttribFormatsBlock;

    UT_UniquePtr<GR_Utils::RadixSortGPU> myRadixSortGPU;

    GSplatAttribs               myGSplatAttribs;
};

#endif