Fixes to multi-device resources (o3de#17638)

Signed-off-by: Martin Winter <[email protected]>
Co-authored-by: Joerg H. Mueller <[email protected]>

Author: martinwinter-huawei

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

@@ -20,7 +20,7 @@ namespace AZ::RHI
 
         IndirectArgumentsTemplate(
             uint32_t maxSequenceCount,
-            const IndirectBufferView& indirectBuffer,
+            const IndirectBufferViewClass& indirectBuffer,
             uint64_t indirectBufferByteOffset)
             : IndirectArgumentsTemplate(
                 maxSequenceCount,
@@ -32,9 +32,9 @@ namespace AZ::RHI
 
         IndirectArgumentsTemplate(
             uint32_t maxSequenceCount,
-            const IndirectBufferView& indirectBuffer,
+            const IndirectBufferViewClass& indirectBuffer,
             uint64_t indirectBufferByteOffset,
-            const Buffer* countBuffer,
+            const BufferClass* countBuffer,
             uint64_t countBufferByteOffset)
             : m_maxSequenceCount(maxSequenceCount)
             , m_indirectBufferView(&indirectBuffer)

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

@@ -20,7 +20,7 @@ namespace AZ::RHI
     struct MultiDeviceBufferMapResponse
     {
         //! Will hold the mapped data for each device selected in the MultiDeviceBuffer
-        AZStd::vector<void*> m_data;
+        AZStd::unordered_map<int, void*> m_data;
     };
 
     using MultiDeviceBufferInitRequest = BufferInitRequestTemplate<MultiDeviceBuffer>;

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

@@ -53,7 +53,7 @@ namespace AZ::RHI
             case DispatchType::Direct:
                 return DispatchArguments(m_direct);
             case DispatchType::Indirect:
-                return DispatchArguments(DispatchIndirect{m_mdIndirect.m_maxSequenceCount, m_mdIndirect.m_indirectBufferView->GetDeviceIndirectBufferView(deviceIndex), m_mdIndirect.m_indirectBufferByteOffset, m_mdIndirect.m_countBuffer->GetDeviceBuffer(deviceIndex).get(), m_mdIndirect.m_countBufferByteOffset});
+                return DispatchArguments(DispatchIndirect{m_mdIndirect.m_maxSequenceCount, m_mdIndirect.m_indirectBufferView->GetDeviceIndirectBufferView(deviceIndex), m_mdIndirect.m_indirectBufferByteOffset, m_mdIndirect.m_countBuffer ? m_mdIndirect.m_countBuffer->GetDeviceBuffer(deviceIndex).get() : nullptr, m_mdIndirect.m_countBufferByteOffset});
             default:
                 return DispatchArguments();
             }

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

@@ -32,18 +32,55 @@ namespace AZ::RHI
             uint32_t byteCount,
             uint32_t byteStride);
 
+        //! The mutex stops the default generation
+        MultiDeviceIndirectBufferView(const MultiDeviceIndirectBufferView& other)
+            : m_hash{ other.m_hash }
+            , m_mdBuffer{ other.m_mdBuffer }
+            , m_mdSignature{ other.m_mdSignature }
+            , m_byteOffset{ other.m_byteOffset }
+            , m_byteCount{ other.m_byteCount }
+            , m_byteStride{ other.m_byteStride }
+        {
+        }
+
+        MultiDeviceIndirectBufferView& operator=(const MultiDeviceIndirectBufferView& other)
+        {
+            this->m_hash = other.m_hash;
+            this->m_mdBuffer = other.m_mdBuffer;
+            this->m_mdSignature = other.m_mdSignature;
+            this->m_byteOffset = other.m_byteOffset;
+            this->m_byteCount = other.m_byteCount;
+            this->m_byteStride = other.m_byteStride;
+
+            return *this;
+        }
+
         //! Returns the device-specific IndirectBufferView for the given index
-        IndirectBufferView GetDeviceIndirectBufferView(int deviceIndex) const
+        const IndirectBufferView& GetDeviceIndirectBufferView(int deviceIndex) const
         {
             AZ_Error("MultiDeviceIndirectBufferView", m_mdSignature, "No MultiDeviceIndirectBufferSignature available\n");
             AZ_Error("MultiDeviceIndirectBufferView", m_mdBuffer, "No MultiDeviceBuffer available\n");
 
-            return IndirectBufferView(
-                *m_mdBuffer->GetDeviceBuffer(deviceIndex),
-                *m_mdSignature->GetDeviceIndirectBufferSignature(deviceIndex),
-                m_byteOffset,
-                m_byteCount,
-                m_byteStride);
+            AZStd::lock_guard lock(m_bufferViewMutex);
+            auto iterator{ m_cache.find(deviceIndex) };
+            if (iterator == m_cache.end())
+            {
+                //! Buffer view is not yet in the cache
+                auto [new_iterator, inserted]{ m_cache.insert(AZStd::make_pair(
+                    deviceIndex,
+                    IndirectBufferView(
+                        *m_mdBuffer->GetDeviceBuffer(deviceIndex),
+                        *m_mdSignature->GetDeviceIndirectBufferSignature(deviceIndex),
+                        m_byteOffset,
+                        m_byteCount,
+                        m_byteStride))) };
+                if (inserted)
+                {
+                    return new_iterator->second;
+                }
+            }
+
+            return iterator->second;
         }
 
         //! Returns the hash of the view. This hash is precomputed at creation time.
@@ -72,5 +109,9 @@ namespace AZ::RHI
         uint32_t m_byteOffset = 0;
         uint32_t m_byteCount = 0;
         uint32_t m_byteStride = 0;
+
+        //! Safe-guard access to IndirectBufferView cache during parallel access
+        mutable AZStd::mutex m_bufferViewMutex{};
+        mutable AZStd::unordered_map<int, IndirectBufferView> m_cache;
     };
 } // namespace AZ::RHI

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

@@ -54,7 +54,7 @@ namespace AZ::RHI
         ShaderInputConstantIndex FindShaderInputConstantIndex(const Name& name) const;
 
         //! Sets one image view for the given shader input index.
-        bool SetImageView(ShaderInputImageIndex inputIndex, const MultiDeviceImageView* imageView, uint32_t arrayIndex);
+        bool SetImageView(ShaderInputImageIndex inputIndex, const MultiDeviceImageView* imageView, uint32_t arrayIndex = 0);
 
         //! Sets an array of image view for the given shader input index.
         bool SetImageViewArray(

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

@@ -40,6 +40,8 @@ namespace AZ::RHI
         //! it is explicitly initialized with just a deviceIndex
         ResultCode Init(int deviceIndex, const SwapChainDescriptor& descriptor);
 
+        Ptr<SwapChain> GetDeviceSwapChain() const;
+
         //! Presents the swap chain to the display, and rotates the images.
         void Present();
 

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

@@ -230,7 +230,7 @@ namespace AZ::RHI
             }
             else
             {
-                response.m_data.push_back(deviceMapResponse.m_data);
+                response.m_data[deviceIndex] = deviceMapResponse.m_data;
             }
 
             return resultCode;

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

@@ -16,8 +16,6 @@ namespace AZ::RHI
     IndirectBufferSignatureDescriptor MultiDeviceIndirectBufferSignatureDescriptor::GetDeviceIndirectBufferSignatureDescriptor(
         int deviceIndex) const
     {
-        AZ_Assert(m_pipelineState, "No MultiDevicePipelineState available\n");
-
         IndirectBufferSignatureDescriptor descriptor{ m_layout };
 
         if (m_pipelineState)

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

@@ -75,7 +75,7 @@ namespace AZ::RHI
     }
 
     bool MultiDeviceShaderResourceGroupData::SetImageView(
-        ShaderInputImageIndex inputIndex, const MultiDeviceImageView* imageView, uint32_t arrayIndex = 0)
+        ShaderInputImageIndex inputIndex, const MultiDeviceImageView* imageView, uint32_t arrayIndex)
     {
         AZStd::array<const MultiDeviceImageView* const, 1> imageViews = { { imageView } };
         return SetImageViewArray(inputIndex, imageViews, arrayIndex);

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

@@ -85,6 +85,12 @@ namespace AZ::RHI
         return resultCode;
     }
 
+    Ptr<SwapChain> MultiDeviceSwapChain::GetDeviceSwapChain() const
+    {
+        // As MultiDeviceSwapChain is always initialized for one single device, the method returns this single item by accessing map.begin()
+        return AZStd::static_pointer_cast<SwapChain>(m_deviceObjects.begin()->second);
+    }
+
     void MultiDeviceSwapChain::ShutdownImages()
     {
         // Shutdown existing set of images.