Fix multi-threaded creation of PipelineStates in PipelineStateCache (o3de#18091)

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

Author: martinwinter-huawei

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

@@ -44,6 +44,9 @@ namespace AZ::RHI
             const PipelineStateDescriptor& descriptor,
             PipelineLibrary* pipelineLibrary = nullptr);
 
+        //! Preinitializes a pipeline state to allow for safe usage with multi-threaded DrawPacket creation
+        void PreInitialize(MultiDevice::DeviceMask deviceMask);
+
         PipelineStateType GetType() const;
 
     private:

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

@@ -29,6 +29,19 @@ namespace AZ::RHI
         return true;
     }
 
+    void PipelineState::PreInitialize(MultiDevice::DeviceMask deviceMask)
+    {
+        const int deviceCount = RHI::RHISystemInterface::Get()->GetDeviceCount();
+
+        for (int deviceIndex = 0; deviceIndex < deviceCount; ++deviceIndex)
+        {
+            if ((AZStd::to_underlying(deviceMask) >> deviceIndex) & 1)
+            {
+                m_deviceObjects[deviceIndex] = Factory::Get().CreatePipelineState();
+            }
+        }
+    }
+
     ResultCode PipelineState::Init(
         MultiDevice::DeviceMask deviceMask, const PipelineStateDescriptor& descriptor, PipelineLibrary* pipelineLibrary)
     {
@@ -51,7 +64,11 @@ namespace AZ::RHI
             {
                 auto* device = RHISystemInterface::Get()->GetDevice(deviceIndex);
 
-                m_deviceObjects[deviceIndex] = Factory::Get().CreatePipelineState();
+                if(!m_deviceObjects.contains(deviceIndex))
+                {
+                    m_deviceObjects[deviceIndex] = Factory::Get().CreatePipelineState();
+                }
+
                 switch (descriptor.GetType())
                 {
                 case PipelineStateType::Draw:
@@ -106,8 +123,8 @@ namespace AZ::RHI
 
         if (resultCode != ResultCode::Success)
         {
-            // Reset already initialized device-specific PipelineStates and set deviceMask to 0
-            m_deviceObjects.clear();
+            // Only reset the device mask but the the device-specific PipelineStates, as other
+            // threads might be using them already
             MultiDeviceObject::Init(static_cast<MultiDevice::DeviceMask>(0u));
         }
 
@@ -116,9 +133,9 @@ namespace AZ::RHI
 
     void PipelineState::Shutdown()
     {
+        m_deviceObjects.clear();
         if (IsInitialized())
         {
-            m_deviceObjects.clear();
             MultiDeviceObject::Shutdown();
         }
     }

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

@@ -376,6 +376,8 @@ namespace AZ::RHI
             // but don't initialize it yet. We can safely allocate the 'empty' instance and cache it.
             pipelineState = aznew PipelineState;
 
+            pipelineState->PreInitialize(m_deviceMask);
+
             [[maybe_unused]] bool success =
                 InsertPipelineState(pendingCache, PipelineStateEntry(pipelineStateHash, pipelineState, descriptor));
             AZ_Assert(success, "PipelineStateEntry already exists in the pending cache.");