Fixes to parallel cache access and to GetTlas*Buffer (o3de#17546)

Signed-off-by: Martin Winter <[email protected]>

Author: martinwinter-huawei

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

@@ -104,6 +104,8 @@ namespace AZ::RHI
         }
 
     private:
+        //! Safe-guard access to BufferView cache during parallel access
+        mutable AZStd::mutex m_bufferViewMutex;
         //! A raw pointer to a multi-device buffer
         ConstPtr<RHI::MultiDeviceBuffer> m_buffer;
         //! The corresponding BufferViewDescriptor for this view.

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

@@ -150,6 +150,8 @@ namespace AZ::RHI
         }
 
     private:
+        //! Safe-guard access to ImageView cache during parallel access
+        mutable AZStd::mutex m_imageViewMutex;
         //! A raw pointer to a multi-device image
         ConstPtr<RHI::MultiDeviceImage> m_image;
         //! The corresponding ImageViewDescriptor for this view.

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

@@ -227,5 +227,7 @@ namespace AZ::RHI
 
     private:
         MultiDeviceRayTracingTlasDescriptor m_mdDescriptor;
+        mutable RHI::Ptr<RHI::MultiDeviceBuffer> m_tlasBuffer;
+        mutable RHI::Ptr<RHI::MultiDeviceBuffer> m_tlasInstancesBuffer;
     };
 } // namespace AZ::RHI

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

@@ -66,6 +66,7 @@ namespace AZ::RHI
     //! Given a device index, return the corresponding BufferView for the selected device
     const RHI::Ptr<RHI::BufferView> MultiDeviceBufferView::GetDeviceBufferView(int deviceIndex) const
     {
+        AZStd::lock_guard lock(m_bufferViewMutex);
         auto iterator{ m_cache.find(deviceIndex) };
         if (iterator == m_cache.end())
         {

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

@@ -105,6 +105,7 @@ namespace AZ::RHI
     //! Given a device index, return the corresponding BufferView for the selected device
     const RHI::Ptr<RHI::ImageView> MultiDeviceImageView::GetDeviceImageView(int deviceIndex) const
     {
+        AZStd::lock_guard lock(m_imageViewMutex);
         auto iterator{ m_cache.find(deviceIndex) };
         if (iterator == m_cache.end())
         {

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

@@ -267,25 +267,30 @@ namespace AZ::RHI
             return nullptr;
         }
 
-        auto tlasBuffer = aznew RHI::MultiDeviceBuffer;
-        tlasBuffer->Init(GetDeviceMask());
+        if (m_tlasBuffer)
+        {
+            return m_tlasBuffer;
+        }
+
+        m_tlasBuffer = aznew RHI::MultiDeviceBuffer;
+        m_tlasBuffer->Init(GetDeviceMask());
 
         IterateObjects<RayTracingTlas>(
-            [&tlasBuffer](int deviceIndex, auto deviceRayTracingTlas)
+            [this](int deviceIndex, auto deviceRayTracingTlas)
             {
-                tlasBuffer->m_deviceObjects[deviceIndex] = deviceRayTracingTlas->GetTlasBuffer();
+                m_tlasBuffer->m_deviceObjects[deviceIndex] = deviceRayTracingTlas->GetTlasBuffer();
 
-                if (!tlasBuffer->m_deviceObjects[deviceIndex])
+                if (!m_tlasBuffer->m_deviceObjects[deviceIndex])
                 {
-                    tlasBuffer = nullptr;
+                    m_tlasBuffer = nullptr;
                     return ResultCode::Fail;
                 }
 
-                tlasBuffer->SetDescriptor(tlasBuffer->GetDeviceBuffer(deviceIndex)->GetDescriptor());
+                m_tlasBuffer->SetDescriptor(m_tlasBuffer->GetDeviceBuffer(deviceIndex)->GetDescriptor());
                 return ResultCode::Success;
             });
 
-        return tlasBuffer;
+        return m_tlasBuffer;
     }
 
     const RHI::Ptr<RHI::MultiDeviceBuffer> MultiDeviceRayTracingTlas::GetTlasInstancesBuffer() const
@@ -295,23 +300,28 @@ namespace AZ::RHI
             return nullptr;
         }
 
-        auto tlasInstancesBuffer = aznew RHI::MultiDeviceBuffer;
-        tlasInstancesBuffer->Init(GetDeviceMask());
+        if (m_tlasInstancesBuffer)
+        {
+            return m_tlasInstancesBuffer;
+        }
+
+        m_tlasInstancesBuffer = aznew RHI::MultiDeviceBuffer;
+        m_tlasInstancesBuffer->Init(GetDeviceMask());
 
         IterateObjects<RayTracingTlas>(
-            [&tlasInstancesBuffer](int deviceIndex, auto deviceRayTracingTlas)
+            [this](int deviceIndex, auto deviceRayTracingTlas)
             {
-                tlasInstancesBuffer->m_deviceObjects[deviceIndex] = deviceRayTracingTlas->GetTlasBuffer();
+                m_tlasInstancesBuffer->m_deviceObjects[deviceIndex] = deviceRayTracingTlas->GetTlasBuffer();
 
-                if (!tlasInstancesBuffer->m_deviceObjects[deviceIndex])
+                if (!m_tlasInstancesBuffer->m_deviceObjects[deviceIndex])
                 {
-                    tlasInstancesBuffer = nullptr;
+                    m_tlasInstancesBuffer = nullptr;
                     return false;
                 }
 
-                tlasInstancesBuffer->SetDescriptor(tlasInstancesBuffer->GetDeviceBuffer(deviceIndex)->GetDescriptor());
+                m_tlasInstancesBuffer->SetDescriptor(m_tlasInstancesBuffer->GetDeviceBuffer(deviceIndex)->GetDescriptor());
                 return true;
             });
-        return tlasInstancesBuffer;
+        return m_tlasInstancesBuffer;
     }
 } // namespace AZ::RHI