resourceRegistry.h

//
// Copyright 2016 Pixar
//
// Licensed under the terms set forth in the LICENSE.txt file available at
// https://openusd.org/license.
//
#ifndef PXR_IMAGING_HD_ST_RESOURCE_REGISTRY_H
#define PXR_IMAGING_HD_ST_RESOURCE_REGISTRY_H
 
#include "pxr/pxr.h"
#include "pxr/base/vt/dictionary.h"
 
#include "pxr/imaging/hdSt/api.h"
#include "pxr/imaging/hdSt/bufferArrayRegistry.h"
#include "pxr/imaging/hdSt/enums.h"
 
#include "pxr/imaging/hgi/hgi.h"
 
#include "pxr/imaging/hd/bufferArrayRange.h"
#include "pxr/imaging/hd/bufferSource.h"
#include "pxr/imaging/hd/bufferSpec.h"
#include "pxr/imaging/hd/instanceRegistry.h"
#include "pxr/imaging/hd/resourceRegistry.h"
 
#include <tbb/concurrent_vector.h>
 
#include <atomic>
#include <map>
#include <memory>
 
#ifdef PXR_MATERIALX_SUPPORT_ENABLED
#include <MaterialXCore/Library.h>
MATERIALX_NAMESPACE_BEGIN
    using ShaderPtr = std::shared_ptr<class Shader>;
MATERIALX_NAMESPACE_END
#endif
 
PXR_NAMESPACE_OPEN_SCOPE
 
using HdStComputationSharedPtr = std::shared_ptr<class HdStComputation>;
using HdStDispatchBufferSharedPtr = std::shared_ptr<class HdStDispatchBuffer>;
using HdStGLSLProgramSharedPtr = std::shared_ptr<class HdStGLSLProgram>;
using HioGlslfxSharedPtr = std::shared_ptr<class HioGlslfx>;
 
using HdSt_BasisCurvesTopologySharedPtr =
    std::shared_ptr<class HdSt_BasisCurvesTopology>;
 
using HdStShaderCodePtr =
    std::weak_ptr<class HdStShaderCode>;
using HdSt_GeometricShaderSharedPtr =
    std::shared_ptr<class HdSt_GeometricShader>;
 
using HdStTextureHandleSharedPtr =
    std::shared_ptr<class HdStTextureHandle>;
using HdStTextureObjectSharedPtr =
    std::shared_ptr<class HdStTextureObject>;
using HdStBufferResourceSharedPtr = 
    std::shared_ptr<class HdStBufferResource>;
using HdStResourceRegistrySharedPtr = 
    std::shared_ptr<class HdStResourceRegistry>;
using HdSt_VertexAdjacencyBuilderSharedPtr = 
    std::shared_ptr<class HdSt_VertexAdjacencyBuilder>;
using HdSt_MeshTopologySharedPtr = 
    std::shared_ptr<class HdSt_MeshTopology>;
using HgiResourceBindingsSharedPtr = 
    std::shared_ptr<HgiResourceBindingsHandle>;
using HgiGraphicsPipelineSharedPtr = 
    std::shared_ptr<HgiGraphicsPipelineHandle>;
using HgiComputePipelineSharedPtr = 
    std::shared_ptr<HgiComputePipelineHandle>;
 
class HdStTextureIdentifier;
class HdSamplerParameters;
class HdStStagingBuffer;
 
enum HdStComputeQueue {
    HdStComputeQueueZero=0,
    HdStComputeQueueOne,
    HdStComputeQueueTwo,
    HdStComputeQueueThree,
    HdStComputeQueueCount};
 
using HdStComputationComputeQueuePairVector = 
    std::vector<std::pair<HdStComputationSharedPtr, HdStComputeQueue>>;
 
 
class HdStResourceRegistry final : public HdResourceRegistry 
{
public:
    HF_MALLOC_TAG_NEW("new HdStResourceRegistry");
 
    HDST_API
    explicit HdStResourceRegistry(Hgi* hgi);
 
    HDST_API
    ~HdStResourceRegistry() override;
 
    HDST_API
    void InvalidateShaderRegistry() override;
 
    HDST_API
    void ReloadResource(TfToken const& resourceType,
                        std::string const& path) override;
 
    HDST_API
    VtDictionary GetResourceAllocation() const override;
 
    HDST_API
    Hgi* GetHgi();
 
 
    HDST_API
    HdStTextureHandleSharedPtr AllocateTextureHandle(
        const HdStTextureIdentifier &textureId,
        HdStTextureType textureType,
        const HdSamplerParameters &samplerParams,
        size_t memoryRequest,
        HdStShaderCodePtr const &shaderCode);
 
    HDST_API
    HdStTextureObjectSharedPtr AllocateTextureObject(
        const HdStTextureIdentifier &textureId,
        HdStTextureType textureType);
 
    HDST_API
    void SetMemoryRequestForTextureType(
        HdStTextureType textureType,
        size_t memoryRequest);
 
 
    HDST_API
    HdBufferArrayRangeSharedPtr AllocateNonUniformBufferArrayRange(
        TfToken const &role,
        HdBufferSpecVector const &bufferSpecs,
        HdBufferArrayUsageHint usageHint);
 
    HDST_API
    HdBufferArrayRangeSharedPtr AllocateNonUniformImmutableBufferArrayRange(
        TfToken const &role,
        HdBufferSpecVector const &bufferSpecs,
        HdBufferArrayUsageHint usageHint);
 
    HDST_API
    HdBufferArrayRangeSharedPtr AllocateUniformBufferArrayRange(
        TfToken const &role,
        HdBufferSpecVector const &bufferSpecs,
        HdBufferArrayUsageHint usageHint);
 
    HDST_API
    HdBufferArrayRangeSharedPtr AllocateShaderStorageBufferArrayRange(
        TfToken const &role,
        HdBufferSpecVector const &bufferSpecs,
        HdBufferArrayUsageHint usageHint);
 
    HDST_API
    HdBufferArrayRangeSharedPtr AllocateSingleBufferArrayRange(
        TfToken const &role,
        HdBufferSpecVector const &bufferSpecs,
        HdBufferArrayUsageHint usageHint);
 
    
    HDST_API
    HdBufferArrayRangeSharedPtr UpdateNonUniformBufferArrayRange(
        TfToken const &role,
        HdBufferArrayRangeSharedPtr const& curRange,
        HdBufferSpecVector const &updatedOrAddedSpecs,
        HdBufferSpecVector const& removedSpecs,
        HdBufferArrayUsageHint usageHint);
 
    HDST_API
    HdBufferArrayRangeSharedPtr UpdateNonUniformImmutableBufferArrayRange(
        TfToken const &role,
        HdBufferArrayRangeSharedPtr const& curRange,
        HdBufferSpecVector const &updatedOrAddedSpecs,
        HdBufferSpecVector const& removedSpecs,
        HdBufferArrayUsageHint usageHint);
 
    HDST_API
    HdBufferArrayRangeSharedPtr UpdateUniformBufferArrayRange(
        TfToken const &role,
        HdBufferArrayRangeSharedPtr const& curRange,
        HdBufferSpecVector const &updatedOrAddedSpecs,
        HdBufferSpecVector const& removedSpecs,
        HdBufferArrayUsageHint usageHint);
 
    HDST_API
    HdBufferArrayRangeSharedPtr UpdateShaderStorageBufferArrayRange(
        TfToken const &role,
        HdBufferArrayRangeSharedPtr const& curRange,
        HdBufferSpecVector const &updatedOrAddedSpecs,
        HdBufferSpecVector const& removedSpecs,
        HdBufferArrayUsageHint usageHint);
 
 
    HDST_API
    void AddSources(HdBufferArrayRangeSharedPtr const &range,
                    HdBufferSourceSharedPtrVector &&sources);
 
    HDST_API
    void AddSource(HdBufferArrayRangeSharedPtr const &range,
                   HdBufferSourceSharedPtr const &source);
 
    HDST_API
    void AddSource(HdBufferSourceSharedPtr const &source);
 
    HDST_API
    void AddComputation(HdBufferArrayRangeSharedPtr const &range,
                        HdStComputationSharedPtr const &computation,
                        HdStComputeQueue const queue);
 
 
    HDST_API
    HdStDispatchBufferSharedPtr RegisterDispatchBuffer(
        TfToken const &role, int count, int commandNumUints);
 
    HDST_API
    HdStBufferResourceSharedPtr RegisterBufferResource(
        TfToken const &role, 
        HdTupleType tupleType,
        HgiBufferUsage bufferUsage);
 
    HDST_API
    void GarbageCollectDispatchBuffers();
 
    HDST_API
    void GarbageCollectBufferResources();
 
    
 
    HDST_API
    HdInstance<HdSt_MeshTopologySharedPtr>
    RegisterMeshTopology(HdInstance<HdSt_MeshTopologySharedPtr>::ID id);
 
    HDST_API
    HdInstance<HdSt_BasisCurvesTopologySharedPtr>
    RegisterBasisCurvesTopology(
        HdInstance<HdSt_BasisCurvesTopologySharedPtr>::ID id);
 
    HDST_API
    HdInstance<HdSt_VertexAdjacencyBuilderSharedPtr>
    RegisterVertexAdjacencyBuilder(
        HdInstance<HdSt_VertexAdjacencyBuilderSharedPtr>::ID id);
 
    HDST_API
    HdInstance<HdBufferArrayRangeSharedPtr>
    RegisterMeshIndexRange(
        HdInstance<HdBufferArrayRangeSharedPtr>::ID id, TfToken const &name);
 
    HDST_API
    HdInstance<HdBufferArrayRangeSharedPtr>
    RegisterBasisCurvesIndexRange(
       HdInstance<HdBufferArrayRangeSharedPtr>::ID id, TfToken const &name);
 
    HDST_API
    HdInstance<HdBufferArrayRangeSharedPtr>
    RegisterPrimvarRange(
        HdInstance<HdBufferArrayRangeSharedPtr>::ID id);
 
    HDST_API
    HdInstance<HdBufferArrayRangeSharedPtr>
    RegisterExtComputationDataRange(
        HdInstance<HdBufferArrayRangeSharedPtr>::ID id);
 
    HDST_API
    HdInstance<HdSt_GeometricShaderSharedPtr>
    RegisterGeometricShader(HdInstance<HdSt_GeometricShaderSharedPtr>::ID id);
 
    HDST_API
    HdInstance<HdStGLSLProgramSharedPtr>
    RegisterGLSLProgram(HdInstance<HdStGLSLProgramSharedPtr>::ID id);
 
    HDST_API
    HdInstance<HioGlslfxSharedPtr>
    RegisterGLSLFXFile(HdInstance<HioGlslfxSharedPtr>::ID id);
 
#ifdef PXR_MATERIALX_SUPPORT_ENABLED
    HDST_API
    HdInstance<MaterialX::ShaderPtr>
    RegisterMaterialXShader(HdInstance<MaterialX::ShaderPtr>::ID id);
#endif
 
    HDST_API
    HdInstance<HgiResourceBindingsSharedPtr>
    RegisterResourceBindings(HdInstance<HgiResourceBindingsSharedPtr>::ID id);
 
    HDST_API
    HdInstance<HgiGraphicsPipelineSharedPtr>
    RegisterGraphicsPipeline(HdInstance<HgiGraphicsPipelineSharedPtr>::ID id);
 
    HDST_API
    HdInstance<HgiComputePipelineSharedPtr>
    RegisterComputePipeline(HdInstance<HgiComputePipelineSharedPtr>::ID id);
 
    HDST_API
    HdResourceRegistry*
    FindOrCreateSubResourceRegistry(
        const std::string& identifier,
        const std::function<std::unique_ptr<HdResourceRegistry>()>& factory);
 
    HDST_API
    HgiBlitCmds* GetGlobalBlitCmds();
 
    HDST_API
    HgiComputeCmds* GetGlobalComputeCmds(
        HgiComputeDispatch dispatchMethod = HgiComputeDispatchSerial);
 
    HDST_API
    void SubmitBlitWork(HgiSubmitWaitType wait = HgiSubmitWaitTypeNoWait);
 
    HDST_API
    void SubmitComputeWork(HgiSubmitWaitType wait = HgiSubmitWaitTypeNoWait);
 
    HDST_API
    HdStStagingBuffer* GetStagingBuffer();
 
public:
    //
    // Unit test API
    //
 
    void SetNonUniformAggregationStrategy(
                std::unique_ptr<HdStAggregationStrategy> &&strategy) {
        _nonUniformAggregationStrategy = std::move(strategy);
    }
 
    void SetNonUniformImmutableAggregationStrategy(
                std::unique_ptr<HdStAggregationStrategy> &&strategy) {
        _nonUniformImmutableAggregationStrategy = std::move(strategy);
    }
 
    void SetUniformAggregationStrategy(
                std::unique_ptr<HdStAggregationStrategy> &&strategy) {
        _uniformUboAggregationStrategy = std::move(strategy);
    }
 
    void SetShaderStorageAggregationStrategy(
                std::unique_ptr<HdStAggregationStrategy> &&strategy) {
        _uniformSsboAggregationStrategy = std::move(strategy);
    }
 
    void SetSingleStorageAggregationStrategy(
                std::unique_ptr<HdStAggregationStrategy> &&strategy) {
        _singleAggregationStrategy = std::move(strategy);
    }
 
    HDST_API
    friend std::ostream &operator <<(
        std::ostream &out,
        const HdStResourceRegistry& self);
 
protected:
    void _Commit() override;
    void _GarbageCollect() override;
 
private:
    void _CommitTextures();
    // Wrapper function for BAR allocation
    HdBufferArrayRangeSharedPtr _AllocateBufferArrayRange(
        HdStAggregationStrategy *strategy,
        HdStBufferArrayRegistry &bufferArrayRegistry,
        TfToken const &role,
        HdBufferSpecVector const &bufferSpecs,
        HdBufferArrayUsageHint usageHint);
    
    HdBufferArrayRangeSharedPtr _UpdateBufferArrayRange(
        HdStAggregationStrategy *strategy,
        HdStBufferArrayRegistry &bufferArrayRegistry,
        TfToken const &role,
        HdBufferArrayRangeSharedPtr const& curRange,
        HdBufferSpecVector const &updatedOrAddedSpecs,
        HdBufferSpecVector const& removedSpecs,
        HdBufferArrayUsageHint usageHint);
 
    // Tally resources by key into the given dictionary. Any additions should
    // be cumulative with the existing key values. 
    void _TallyResourceAllocation(VtDictionary *result) const;
 
    // TODO: this is a transient structure. we'll revisit the BufferSource
    // interface later.
    struct _PendingSource {
        _PendingSource(HdBufferArrayRangeSharedPtr const &range)
            : range(range)
            , sources() 
        {
        }
 
        _PendingSource(HdBufferArrayRangeSharedPtr const &range,
                       HdBufferSourceSharedPtr     const &source)
            : range(range)
            , sources(1, source)
        {
        }
 
        _PendingSource(HdBufferArrayRangeSharedPtr const &range,
                       HdBufferSourceSharedPtrVector     && sources)
            : range(range)
            , sources(std::move(sources))
        {
        }
 
        HdBufferArrayRangeSharedPtr range;
        HdBufferSourceSharedPtrVector sources;
    };
 
    Hgi* _hgi;
 
    typedef tbb::concurrent_vector<_PendingSource> _PendingSourceList;
    _PendingSourceList    _pendingSources;
    std::atomic_size_t    _numBufferSourcesToResolve;
    
    struct _PendingComputation{
        _PendingComputation(HdBufferArrayRangeSharedPtr const &range,
                            HdStComputationSharedPtr const &computation)
            : range(range), computation(computation) { }
        HdBufferArrayRangeSharedPtr range;
        HdStComputationSharedPtr computation;
    };
 
    // If we need more 'compute queues' we can increase HdStComputeQueueCount.
    // We could also consider tbb::concurrent_priority_queue if we want this
    // to be dynamically scalable.
    typedef tbb::concurrent_vector<_PendingComputation> _PendingComputationList;
    _PendingComputationList  _pendingComputations[HdStComputeQueueCount];
 
    // aggregated buffer array
    HdStBufferArrayRegistry _nonUniformBufferArrayRegistry;
    HdStBufferArrayRegistry _nonUniformImmutableBufferArrayRegistry;
    HdStBufferArrayRegistry _uniformUboBufferArrayRegistry;
    HdStBufferArrayRegistry _uniformSsboBufferArrayRegistry;
    HdStBufferArrayRegistry _singleBufferArrayRegistry;
 
    // current aggregation strategies
    std::unique_ptr<HdStAggregationStrategy> _nonUniformAggregationStrategy;
    std::unique_ptr<HdStAggregationStrategy>
                                _nonUniformImmutableAggregationStrategy;
    std::unique_ptr<HdStAggregationStrategy> _uniformUboAggregationStrategy;
    std::unique_ptr<HdStAggregationStrategy> _uniformSsboAggregationStrategy;
    std::unique_ptr<HdStAggregationStrategy> _singleAggregationStrategy;
 
    typedef std::vector<HdStDispatchBufferSharedPtr>
        _DispatchBufferRegistry;
    _DispatchBufferRegistry _dispatchBufferRegistry;
 
    typedef std::vector<HdStBufferResourceSharedPtr>
        _BufferResourceRegistry;
    _BufferResourceRegistry _bufferResourceRegistry;
 
    // Register mesh topology.
    HdInstanceRegistry<HdSt_MeshTopologySharedPtr>
        _meshTopologyRegistry;
 
    // Register basisCurves topology.
    HdInstanceRegistry<HdSt_BasisCurvesTopologySharedPtr>
        _basisCurvesTopologyRegistry;
 
    // Register vertex adjacency.
    HdInstanceRegistry<HdSt_VertexAdjacencyBuilderSharedPtr>
        _vertexAdjacencyBuilderRegistry;
 
    // Register topology index buffers.
    typedef HdInstanceRegistry<HdBufferArrayRangeSharedPtr>
        _TopologyIndexRangeInstanceRegistry;
    typedef tbb::concurrent_unordered_map< TfToken,
                                           _TopologyIndexRangeInstanceRegistry,
                                           TfToken::HashFunctor >
        _TopologyIndexRangeInstanceRegMap;
 
    _TopologyIndexRangeInstanceRegMap _meshTopologyIndexRangeRegistry;
    _TopologyIndexRangeInstanceRegMap _basisCurvesTopologyIndexRangeRegistry;
 
    // Register shared primvar buffers.
    HdInstanceRegistry<HdBufferArrayRangeSharedPtr>
        _primvarRangeRegistry;
 
    // Register ext computation resource.
    HdInstanceRegistry<HdBufferArrayRangeSharedPtr>
        _extComputationDataRangeRegistry;
 
    // geometric shader registry
    HdInstanceRegistry<HdSt_GeometricShaderSharedPtr>
        _geometricShaderRegistry;
 
    // glsl shader program registry
    HdInstanceRegistry<HdStGLSLProgramSharedPtr>
        _glslProgramRegistry;
 
    // glslfx file registry
    HdInstanceRegistry<HioGlslfxSharedPtr>
        _glslfxFileRegistry;
 
#ifdef PXR_MATERIALX_SUPPORT_ENABLED
    // MaterialX glslfx shader registry
    HdInstanceRegistry<MaterialX::ShaderPtr> _materialXShaderRegistry;
#endif
 
    // texture handle registry
    std::unique_ptr<class HdSt_TextureHandleRegistry> _textureHandleRegistry;
 
    // Hgi resource bindings registry
    HdInstanceRegistry<HgiResourceBindingsSharedPtr>
        _resourceBindingsRegistry;
 
    // Hgi graphics pipeline registry
    HdInstanceRegistry<HgiGraphicsPipelineSharedPtr>
        _graphicsPipelineRegistry;
 
    // Hgi compute pipeline registry
    HdInstanceRegistry<HgiComputePipelineSharedPtr>
        _computePipelineRegistry;
 
    using _SubResourceRegistryMap =
        tbb::concurrent_unordered_map<std::string,
                                    std::unique_ptr<HdResourceRegistry>>;
    _SubResourceRegistryMap _subResourceRegistries;
 
    HgiBlitCmdsUniquePtr _blitCmds;
    HgiComputeCmdsUniquePtr _computeCmds;
 
    std::unique_ptr<HdStStagingBuffer> _stagingBuffer;
};
 
 
PXR_NAMESPACE_CLOSE_SCOPE
 
#endif //PXR_IMAGING_HD_ST_RESOURCE_REGISTRY_H