ATOM-16249 Adding draw srg caching to DynamicDrawContext since creating SRG can be expensive for some backend (#4491)

* ATOM-16249 Adding draw srg caching to DynamicDrawContext since creating SRG can be expensive for some backend

Signed-off-by: Qing Tao <qingtao@amazon.com>

Gems/Atom/RHI/Code/Include/Atom/RHI/ShaderResourceGroupData.h

@@ -166,6 +166,10 @@ namespace AZ
             AZStd::array_view<ConstPtr<ImageView>> GetImageGroup() const;
             AZStd::array_view<ConstPtr<BufferView>> GetBufferGroup() const;
             AZStd::array_view<SamplerState> GetSamplerGroup() const;
+            
+            //! Reset image and buffer views setup for this ShaderResourceGroupData
+            //! So it won't hold references for any RHI resources
+            void ResetViews();
 
             //! Returns the opaque constant data populated by calls to SetConstant and SetConstantData.
             //! 

Gems/Atom/RHI/Code/Source/RHI/ShaderResourceGroupData.cpp

@@ -330,6 +330,14 @@ namespace AZ
             return m_samplers;
         }
 
+        void ShaderResourceGroupData::ResetViews()
+        {
+            m_imageViews.assign(m_imageViews.size(), nullptr);
+            m_bufferViews.assign(m_bufferViews.size(), nullptr);
+            m_imageViewsUnboundedArray.assign(m_imageViewsUnboundedArray.size(), nullptr);
+            m_bufferViewsUnboundedArray.assign(m_bufferViewsUnboundedArray.size(), nullptr);
+        }
+
         AZStd::array_view<uint8_t> ShaderResourceGroupData::GetConstantData() const
         {
             return m_constantsData.GetConstantData();

Gems/Atom/RPI/Code/Include/Atom/RPI.Public/DynamicDraw/DynamicDrawContext.h

@@ -268,6 +268,8 @@ namespace AZ
             AZStd::vector<RHI::StreamBufferView> m_cachedStreamBufferViews;
             AZStd::vector<RHI::IndexBufferView> m_cachedIndexBufferViews;
             AZStd::vector<Data::Instance<ShaderResourceGroup>> m_cachedDrawSrg;
+
+            uint32_t m_nextDrawSrgIdx = 0;
             
             // structure includes DrawItem and stream and index buffer index
             using BufferViewIndexType = uint32_t;

Gems/Atom/RPI/Code/Include/Atom/RPI.Public/Shader/ShaderResourceGroup.h

@@ -133,6 +133,9 @@ namespace AZ
             /// Returns an array of RPI buffers associated with the buffer shader input index.
             AZStd::array_view<Data::Instance<Buffer>> GetBufferArray(RHI::ShaderInputNameIndex& inputIndex) const;
             AZStd::array_view<Data::Instance<Buffer>> GetBufferArray(RHI::ShaderInputBufferIndex inputIndex) const;
+            
+            //! Reset image and buffer views so that it won't hold references for any RHI resources
+            void ResetViews();
 
             //////////////////////////////////////////////////////////////////////////
             // Methods for assignment / access of RHI Image types.

Gems/Atom/RPI/Code/Source/RPI.Public/DynamicDraw/DynamicDrawContext.cpp

@@ -518,7 +518,6 @@ namespace AZ
             if (drawSrg)
             {
                 drawItem.m_uniqueShaderResourceGroup = drawSrg->GetRHIShaderResourceGroup();
-                m_cachedDrawSrg.push_back(drawSrg);
             }
 
             // Set scissor per draw if scissor is enabled.
@@ -608,7 +607,6 @@ namespace AZ
             if (drawSrg)
             {
                 drawItem.m_uniqueShaderResourceGroup = drawSrg->GetRHIShaderResourceGroup();
-                m_cachedDrawSrg.push_back(drawSrg);
             }
 
             // Set scissor per draw if scissor is enabled.
@@ -635,7 +633,22 @@ namespace AZ
             {
                 return nullptr;
             }
-            auto drawSrg = AZ::RPI::ShaderResourceGroup::Create(m_shader->GetAsset(), m_shader->GetSupervariantIndex(), m_drawSrgLayout->GetName());
+
+            Data::Instance<ShaderResourceGroup> drawSrg;
+            if (m_nextDrawSrgIdx == m_cachedDrawSrg.size())
+            {
+                drawSrg = AZ::RPI::ShaderResourceGroup::Create(m_shader->GetAsset(), m_shader->GetSupervariantIndex(), m_drawSrgLayout->GetName());
+                m_cachedDrawSrg.push_back(drawSrg);
+            }
+            else if (m_nextDrawSrgIdx < m_cachedDrawSrg.size())
+            {
+                drawSrg = m_cachedDrawSrg[m_nextDrawSrgIdx];
+            }
+            else
+            {
+                AZ_Assert(false, "Unexpected next draw srg index");
+            }
+            m_nextDrawSrgIdx++;
 
             // Set fallback value for shader variant if draw srg contains constant for shader variant fallback 
             if (m_hasShaderVariantKeyFallbackEntry)
@@ -727,7 +740,7 @@ namespace AZ
             }
 
             for (auto& drawItemProperties : m_cachedDrawList)
-            {                
+            {
                 view->AddDrawItem(m_drawListTag, drawItemProperties);
             }
         }
@@ -743,9 +756,14 @@ namespace AZ
             m_cachedDrawItems.clear();
             m_cachedStreamBufferViews.clear();
             m_cachedIndexBufferViews.clear();
-            m_cachedDrawSrg.clear();
             m_cachedDrawList.clear();
+            m_nextDrawSrgIdx = 0;
             m_drawFinalized = false;
+
+            for (auto srg:m_cachedDrawSrg)
+            {
+                srg->ResetViews();
+            }
         }
 
         const RHI::PipelineState* DynamicDrawContext::GetCurrentPipelineState()

Gems/Atom/RPI/Code/Source/RPI.Public/Shader/ShaderResourceGroup.cpp

@@ -580,6 +580,11 @@ namespace AZ
             return {};
         }
 
+        void ShaderResourceGroup::ResetViews()
+        {
+            m_data.ResetViews();
+        }
+
         const RHI::SamplerState& ShaderResourceGroup::GetSampler(RHI::ShaderInputNameIndex& inputIndex, uint32_t arrayIndex) const
         {
             inputIndex.ValidateOrFindSamplerIndex(GetLayout());