Fixes to multi-device *Items and DrawPacketBuilder (o3de#17535)

Fixes to multi-device *Items and DrawPacketBuilder

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

Author: martinwinter-huawei

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

@@ -128,11 +128,11 @@ namespace AZ::RHI
 
         //! Array of shader resource groups to bind (count must match m_shaderResourceGroupCount).
         void SetShaderResourceGroups(
-            const AZStd::span<const MultiDeviceShaderResourceGroup*> shaderResourceGroups, uint8_t shaderResourceGroupCount)
+            const AZStd::span<const MultiDeviceShaderResourceGroup*> shaderResourceGroups)
         {
             for (auto& [deviceIndex, dispatchItem] : m_deviceDispatchItems)
             {
-                dispatchItem.m_shaderResourceGroupCount = shaderResourceGroupCount;
+                dispatchItem.m_shaderResourceGroupCount = static_cast<uint8_t>(shaderResourceGroups.size());
                 for (int i = 0; i < dispatchItem.m_shaderResourceGroupCount; ++i)
                 {
                     dispatchItem.m_shaderResourceGroups[i] = shaderResourceGroups[i]->GetDeviceShaderResourceGroup(deviceIndex).get();

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

@@ -55,13 +55,13 @@ namespace AZ::RHI
         }
 
         //! The number of rays to cast
-        void SetDimensions(uint32_t width, uint32_t height, uint32_t depth)
+        void SetDimensionsDirect(uint32_t width, uint32_t height, uint32_t depth)
         {
             for (auto& [deviceIndex, dispatchRaysItem] : m_deviceDispatchRaysItems)
             {
-                dispatchRaysItem.m_width = width;
-                dispatchRaysItem.m_height = height;
-                dispatchRaysItem.m_depth = depth;
+                dispatchRaysItem.m_arguments.m_direct.m_width = width;
+                dispatchRaysItem.m_arguments.m_direct.m_height = height;
+                dispatchRaysItem.m_arguments.m_direct.m_depth = depth;
             }
         }
 
@@ -70,7 +70,7 @@ namespace AZ::RHI
         {
             for (auto& [deviceIndex, dispatchRaysItem] : m_deviceDispatchRaysItems)
             {
-                dispatchRaysItem.m_rayTracingPipelineState = rayTracingPipelineState->GetDevicePipelineState(deviceIndex).get();
+                dispatchRaysItem.m_rayTracingPipelineState = rayTracingPipelineState->GetDeviceRayTracingPipelineState(deviceIndex).get();
             }
         }
 

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

@@ -76,6 +76,10 @@ namespace AZ::RHI
             }
         }
 
+        MultiDeviceDrawPacketBuilder(const MultiDeviceDrawPacketBuilder& other);
+        MultiDeviceDrawPacketBuilder& operator=(const MultiDeviceDrawPacketBuilder& other);
+        AZ_DISABLE_MOVE(MultiDeviceDrawPacketBuilder)
+
         // NOTE: This is configurable; just used to control the amount of memory held by the builder.
         static const size_t DrawItemCountMax = 16;
 

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

@@ -39,6 +39,38 @@ namespace AZ::RHI
                 m_drawFilterMask };
     }
 
+    MultiDeviceDrawPacketBuilder::MultiDeviceDrawPacketBuilder(const MultiDeviceDrawPacketBuilder& other)
+    {
+        m_deviceMask = other.m_deviceMask;
+
+        m_drawRequests = other.m_drawRequests;
+
+        m_drawPacketInFlight = aznew MultiDeviceDrawPacket;
+        if (other.m_drawPacketInFlight)
+        {
+            m_drawPacketInFlight->m_drawListMask = other.m_drawPacketInFlight->m_drawListMask;
+        }
+
+        m_deviceDrawPacketBuilders = other.m_deviceDrawPacketBuilders;
+    }
+
+    MultiDeviceDrawPacketBuilder& MultiDeviceDrawPacketBuilder::operator=(const MultiDeviceDrawPacketBuilder& other)
+    {
+        m_deviceMask = other.m_deviceMask;
+
+        m_drawRequests = other.m_drawRequests;
+
+        m_drawPacketInFlight = aznew MultiDeviceDrawPacket;
+        if (other.m_drawPacketInFlight)
+        {
+            m_drawPacketInFlight->m_drawListMask = other.m_drawPacketInFlight->m_drawListMask;
+        }
+
+        m_deviceDrawPacketBuilders = other.m_deviceDrawPacketBuilders;
+
+        return *this;
+    }
+
     void MultiDeviceDrawPacketBuilder::Begin(IAllocator* allocator)
     {
         AZ_Error(
@@ -189,6 +221,7 @@ namespace AZ::RHI
         m_drawPacketInFlight = aznew MultiDeviceDrawPacket;
 
         auto drawRequestCount{ original->m_drawListTags.size() };
+        m_drawPacketInFlight->m_drawListMask = original->m_drawListMask;
         m_drawPacketInFlight->m_drawListTags.resize_no_construct(drawRequestCount);
         m_drawPacketInFlight->m_drawFilterMasks.resize_no_construct(drawRequestCount);
         m_drawPacketInFlight->m_drawItemSortKeys.resize_no_construct(drawRequestCount);