resourceRegistry.h

//
// Copyright 2016 Pixar
//
// Licensed under the Apache License, Version 2.0 (the "Apache License")
// with the following modification; you may not use this file except in
// compliance with the Apache License and the following modification to it:
// Section 6. Trademarks. is deleted and replaced with:
//
// 6. Trademarks. This License does not grant permission to use the trade
//    names, trademarks, service marks, or product names of the Licensor
//    and its affiliates, except as required to comply with Section 4(c) of
//    the License and to reproduce the content of the NOTICE file.
//
// You may obtain a copy of the Apache License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the Apache License with the above modification is
// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the Apache License for the specific
// language governing permissions and limitations under the Apache 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/hgi/hgi.h"

#include "pxr/imaging/hd/bufferArrayRange.h"
#include "pxr/imaging/hd/bufferArrayRegistry.h"
#include "pxr/imaging/hd/bufferSource.h"
#include "pxr/imaging/hd/bufferSpec.h"
#include "pxr/imaging/hd/enums.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 HdComputationSharedPtr = std::shared_ptr<class HdComputation>;
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 Hd_VertexAdjacencySharedPtr = 
    std::shared_ptr<class Hd_VertexAdjacency>;
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 HdStComputationSharedPtrVector = 
    std::vector<std::pair<HdComputationSharedPtr, 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,
        HdTextureType textureType,
        const HdSamplerParameters &samplerParams,
        size_t memoryRequest,
        HdStShaderCodePtr const &shaderCode);

    HDST_API
    HdStTextureObjectSharedPtr AllocateTextureObject(
        const HdStTextureIdentifier &textureId,
        HdTextureType textureType);

    HDST_API
    void SetMemoryRequestForTextureType(
        HdTextureType 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,
                        HdComputationSharedPtr 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);

    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<Hd_VertexAdjacencySharedPtr>
    RegisterVertexAdjacency(HdInstance<Hd_VertexAdjacencySharedPtr>::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
    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<HdAggregationStrategy> &&strategy) {
        _nonUniformAggregationStrategy = std::move(strategy);
    }

    void SetNonUniformImmutableAggregationStrategy(
                std::unique_ptr<HdAggregationStrategy> &&strategy) {
        _nonUniformImmutableAggregationStrategy = std::move(strategy);
    }

    void SetUniformAggregationStrategy(
                std::unique_ptr<HdAggregationStrategy> &&strategy) {
        _uniformUboAggregationStrategy = std::move(strategy);
    }

    void SetShaderStorageAggregationStrategy(
                std::unique_ptr<HdAggregationStrategy> &&strategy) {
        _uniformSsboAggregationStrategy = std::move(strategy);
    }

    void SetSingleStorageAggregationStrategy(
                std::unique_ptr<HdAggregationStrategy> &&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(
        HdAggregationStrategy *strategy,
        HdBufferArrayRegistry &bufferArrayRegistry,
        TfToken const &role,
        HdBufferSpecVector const &bufferSpecs,
        HdBufferArrayUsageHint usageHint);
    
    HdBufferArrayRangeSharedPtr _UpdateBufferArrayRange(
        HdAggregationStrategy *strategy,
        HdBufferArrayRegistry &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,
                            HdComputationSharedPtr const &computation)
            : range(range), computation(computation) { }
        HdBufferArrayRangeSharedPtr range;
        HdComputationSharedPtr 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
    HdBufferArrayRegistry _nonUniformBufferArrayRegistry;
    HdBufferArrayRegistry _nonUniformImmutableBufferArrayRegistry;
    HdBufferArrayRegistry _uniformUboBufferArrayRegistry;
    HdBufferArrayRegistry _uniformSsboBufferArrayRegistry;
    HdBufferArrayRegistry _singleBufferArrayRegistry;

    // current aggregation strategies
    std::unique_ptr<HdAggregationStrategy> _nonUniformAggregationStrategy;
    std::unique_ptr<HdAggregationStrategy>
                                _nonUniformImmutableAggregationStrategy;
    std::unique_ptr<HdAggregationStrategy> _uniformUboAggregationStrategy;
    std::unique_ptr<HdAggregationStrategy> _uniformSsboAggregationStrategy;
    std::unique_ptr<HdAggregationStrategy> _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<Hd_VertexAdjacencySharedPtr>
        _vertexAdjacencyRegistry;

    // 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;

    HgiBlitCmdsUniquePtr _blitCmds;
    HgiComputeCmdsUniquePtr _computeCmds;

    std::unique_ptr<HdStStagingBuffer> _stagingBuffer;
};


PXR_NAMESPACE_CLOSE_SCOPE

#endif //PXR_IMAGING_HD_ST_RESOURCE_REGISTRY_H