Various new metal fixes (#1465)

* - Fix the second window related tabbing issue
- Merge calls to UseResoources acrooaa all the SRGs
- Move SamplerCache to the device to ensure only one cache to reduce duplication
- Fixes to compute threading numbers getting reset to 0,0,0
- Cleanup withing BufferPoolResolver
- Argument buffers are now queued to be cleaned up upon shutdown

Code/Framework/AzFramework/Platform/Mac/AzFramework/Windowing/NativeWindow_Mac.mm

@@ -75,6 +75,8 @@ namespace AzFramework
         // Add a fullscreen button in the upper right of the title bar.
         [m_nativeWindow setCollectionBehavior:NSWindowCollectionBehaviorFullScreenPrimary];
 
+        m_nativeWindow.tabbingMode = NSWindowTabbingModeDisallowed;
+        
         // Make the window active
         [m_nativeWindow makeKeyAndOrderFront:nil];
         m_nativeWindow.title = m_windowTitle;

Gems/Atom/Feature/Common/Assets/Passes/ReflectionCopyFrameBuffer.pass

@@ -21,7 +21,7 @@
                     "SlotType": "Output",
                     "ScopeAttachmentUsage": "RenderTarget",
                     "LoadStoreAction": {
-                        "LoadAction": "Load"
+                        "LoadAction": "DontCare"
                     }
                 }
             ],

Gems/Atom/RHI/Metal/Code/Source/RHI/ArgumentBuffer.cpp

@@ -35,7 +35,6 @@ namespace AZ
             {
                 m_device = device;
                 m_srgLayout = srgLayout;
-                m_srgPool = srgPool;
                                                 
                 m_constantBufferSize = srgLayout->GetConstantDataSize();
                 if (m_constantBufferSize)
@@ -93,9 +92,6 @@ namespace AZ
                     
                     //Attach the constant buffer
                     AttachConstantBuffer();
-                    
-                    m_samplerCache = [[NSCache alloc]init];
-                    [m_samplerCache setName:@"SamplerCache"];
                 }
             }
         }
@@ -211,8 +207,8 @@ namespace AZ
                 }
                 else
                 {
-                    RHI::Ptr<Memory> nullMtlBufferMemPtr = m_device->GetNullDescriptorManager().GetNullImage(shaderInputImage.m_type).GetMemory();
-                    mtlTextures[imageArrayLen] = nullMtlBufferMemPtr->GetGpuAddress<id<MTLTexture>>();
+                    RHI::Ptr<Memory> nullMtlImagePtr = m_device->GetNullDescriptorManager().GetNullImage(shaderInputImage.m_type).GetMemory();
+                    mtlTextures[imageArrayLen] = nullMtlImagePtr->GetGpuAddress<id<MTLTexture>>();
                 }
                 imageArrayLen++;
             }
@@ -345,15 +341,20 @@ namespace AZ
                 m_device->GetArgumentBufferAllocator().DeAllocate(m_argumentBuffer);
             }
 #endif
-            m_argumentBuffer = {};
-            m_constantBuffer = {};
             
-            [m_samplerCache removeAllObjects];
-            [m_samplerCache release];
-            m_samplerCache = nil;
+            if(m_argumentBuffer.IsValid())
+            {
+                m_device->QueueForRelease(m_argumentBuffer);
+            }
             
+            if(m_constantBuffer.IsValid())
+            {
+                m_device->QueueForRelease(m_constantBuffer);
+            }
+
             [m_argumentEncoder release];
             m_argumentEncoder = nil;
+             
             Base::Shutdown();
         }
         
@@ -374,23 +375,22 @@ namespace AZ
     
         id<MTLSamplerState> ArgumentBuffer::GetMtlSampler(MTLSamplerDescriptor* samplerDesc)
         {
-            id<MTLSamplerState> mtlSamplerState = [m_samplerCache objectForKey:samplerDesc];
+            const NSCache* samplerCache = m_device->GetSamplerCache();
+            id<MTLSamplerState> mtlSamplerState = [samplerCache objectForKey:samplerDesc];
             if(mtlSamplerState == nil)
             {
                 mtlSamplerState = [m_device->GetMtlDevice() newSamplerStateWithDescriptor:samplerDesc];
-                [m_samplerCache setObject:mtlSamplerState forKey:samplerDesc];
+                [samplerCache setObject:mtlSamplerState forKey:samplerDesc];
             }
             
             return mtlSamplerState;
         }
     
-        void ArgumentBuffer::AddUntrackedResourcesToEncoder(id<MTLCommandEncoder> commandEncoder, const ShaderResourceGroupVisibility& srgResourcesVisInfo) const
+        void ArgumentBuffer::CollectUntrackedResources(id<MTLCommandEncoder> commandEncoder,
+                                                            const ShaderResourceGroupVisibility& srgResourcesVisInfo,
+                                                            ComputeResourcesToMakeResidentMap& resourcesToMakeResidentCompute,
+                                                            GraphicsResourcesToMakeResidentMap& resourcesToMakeResidentGraphics) const
         {
-            //Map to cache all the resources based on the usage as we can batch all the resources for a given usage
-            ComputeResourcesToMakeResidentMap resourcesToMakeResidentCompute;
-            //Map to cache all the resources based on the usage and shader stage as we can batch all the resources for a given usage/shader usage
-            GraphicsResourcesToMakeResidentMap resourcesToMakeResidentGraphics;
-            
             //Cache the constant buffer associated with a srg
             if (m_constantBufferSize)
             {
@@ -434,25 +434,6 @@ namespace AZ
                     }
                 }
             }
-            
-            //Call UseResource on all resources for Compute stage
-            for (const auto& key : resourcesToMakeResidentCompute)
-            {
-                AZStd::vector<id <MTLResource>> resourcesToProcessVec(key.second.begin(), key.second.end());
-                [static_cast<id<MTLComputeCommandEncoder>>(commandEncoder) useResources: &resourcesToProcessVec[0]
-                                                                                  count: resourcesToProcessVec.size()
-                                                                                  usage: key.first];
-            }
-            
-            //Call UseResource on all resources for Vertex and Fragment stages
-            for (const auto& key : resourcesToMakeResidentGraphics)
-            {
-                AZStd::vector<id <MTLResource>> resourcesToProcessVec(key.second.begin(), key.second.end());
-                [static_cast<id<MTLRenderCommandEncoder>>(commandEncoder) useResources: &resourcesToProcessVec[0]
-                                                                                 count: resourcesToProcessVec.size()
-                                                                                 usage: key.first.first
-                                                                                stages: key.first.second];
-            }
         }
 
         void ArgumentBuffer::CollectResourcesForCompute(id<MTLCommandEncoder> encoder,

Gems/Atom/RHI/Metal/Code/Source/RHI/ArgumentBuffer.h

@@ -97,7 +97,15 @@ namespace AZ
             id<MTLBuffer> GetArgEncoderBuffer() const;
             size_t GetOffset() const;
             
-            void AddUntrackedResourcesToEncoder(id<MTLCommandEncoder> commandEncoder, const ShaderResourceGroupVisibility& srgResourcesVisInfo) const;
+            //Map to cache all the resources based on the usage as we can batch all the resources for a given usage.
+            using ComputeResourcesToMakeResidentMap = AZStd::unordered_map<MTLResourceUsage, AZStd::unordered_set<id <MTLResource>>>;
+            //Map to cache all the resources based on the usage and shader stage as we can batch all the resources for a given usage/shader usage.
+            using GraphicsResourcesToMakeResidentMap = AZStd::unordered_map<AZStd::pair<MTLResourceUsage,MTLRenderStages>, AZStd::unordered_set<id <MTLResource>>>;
+
+            void CollectUntrackedResources(id<MTLCommandEncoder> commandEncoder,
+                                           const ShaderResourceGroupVisibility& srgResourcesVisInfo,
+                                           ComputeResourcesToMakeResidentMap& resourcesToMakeResidentCompute,
+                                           GraphicsResourcesToMakeResidentMap& resourcesToMakeResidentGraphics) const;
             
             void ClearResourceTracking();
             
@@ -120,11 +128,7 @@ namespace AZ
             ResourceBindingsMap m_resourceBindings;
             
             static const int MaxEntriesInArgTable = 31;
-            //Map to cache all the resources based on the usage as we can batch all the resources for a given usage.
-            using ComputeResourcesToMakeResidentMap = AZStd::unordered_map<MTLResourceUsage, AZStd::unordered_set<id <MTLResource>>>;
-            //Map to cache all the resources based on the usage and shader stage as we can batch all the resources for a given usage/shader usage.
-            using GraphicsResourcesToMakeResidentMap = AZStd::unordered_map<AZStd::pair<MTLResourceUsage,MTLRenderStages>, AZStd::unordered_set<id <MTLResource>>>;
-            
+                        
             void CollectResourcesForCompute(id<MTLCommandEncoder> encoder,
                                             const ResourceBindingsSet& resourceBindingData,
                                             ComputeResourcesToMakeResidentMap& resourcesToMakeResidentMap) const;
@@ -153,9 +157,6 @@ namespace AZ
             MemoryView m_argumentBuffer;
             MemoryView m_constantBuffer;
 #endif
-            
-            ShaderResourceGroupPool* m_srgPool = nullptr;
-            NSCache* m_samplerCache;
         };
     }
 }

Gems/Atom/RHI/Metal/Code/Source/RHI/BufferPoolResolver.cpp

@@ -40,7 +40,7 @@ namespace AZ
                 buffer->m_pendingResolves++;
 
                 uploadRequest.m_attachmentBuffer = buffer;
-                uploadRequest.m_byteOffset = buffer->GetMemoryView().GetOffset() + request.m_byteOffset;
+                uploadRequest.m_byteOffset = request.m_byteOffset;
                 uploadRequest.m_stagingBuffer = stagingBuffer;
 
                 return stagingBuffer->GetMemoryView().GetCpuAddress();
@@ -51,6 +51,12 @@ namespace AZ
 
         void BufferPoolResolver::Compile()
         {
+            for (BufferUploadPacket& packet : m_uploadPackets)
+            {
+                Buffer* stagingBuffer = packet.m_stagingBuffer.get();
+                //Inform the GPU that the CPU has modified the staging buffer.
+                Platform::SynchronizeBufferOnCPU(stagingBuffer->GetMemoryView().GetGpuAddress<id<MTLBuffer>>(), stagingBuffer->GetMemoryView().GetOffset(), stagingBuffer->GetMemoryView().GetSize());
+            }
         }
 
         void BufferPoolResolver::Resolve(CommandList& commandList) const
@@ -62,15 +68,12 @@ namespace AZ
                 Buffer* destBuffer = packet.m_attachmentBuffer;
                 AZ_Assert(stagingBuffer, "Staging Buffer is null.");
                 AZ_Assert(destBuffer, "Attachment Buffer is null.");
-
-                //Inform the GPU that the CPU has modified the staging buffer.
-                Platform::SynchronizeBufferOnCPU(stagingBuffer->GetMemoryView().GetGpuAddress<id<MTLBuffer>>(), stagingBuffer->GetMemoryView().GetOffset(), stagingBuffer->GetMemoryView().GetSize());
                 
                 RHI::CopyBufferDescriptor copyDescriptor;
                 copyDescriptor.m_sourceBuffer = stagingBuffer;
                 copyDescriptor.m_sourceOffset = stagingBuffer->GetMemoryView().GetOffset();
                 copyDescriptor.m_destinationBuffer = destBuffer;
-                copyDescriptor.m_destinationOffset = static_cast<uint32_t>(packet.m_byteOffset);
+                copyDescriptor.m_destinationOffset = destBuffer->GetMemoryView().GetOffset() + static_cast<uint32_t>(packet.m_byteOffset);
                 copyDescriptor.m_size = stagingBuffer->GetMemoryView().GetSize();
 
                 commandList.Submit(RHI::CopyItem(copyDescriptor));

Gems/Atom/RHI/Metal/Code/Source/RHI/CommandList.cpp

@@ -85,6 +85,7 @@ namespace AZ
                               destinationOffset:descriptor.m_destinationOffset
                                            size:descriptor.m_size];
                     
+                    Platform::SynchronizeBufferOnGPU(blitEncoder, destinationBuffer->GetMemoryView().GetGpuAddress<id<MTLBuffer>>());
                     break;
                 }
                 case RHI::CopyItemType::Image:
@@ -114,6 +115,8 @@ namespace AZ
                                        destinationSlice: descriptor.m_destinationSubresource.m_arraySlice
                                        destinationLevel: descriptor.m_destinationSubresource.m_mipSlice
                                       destinationOrigin: destinationOrigin];
+                    
+                    Platform::SynchronizeTextureOnGPU(blitEncoder, destinationImage->GetMemoryView().GetGpuAddress<id<MTLTexture>>());
                     break;
                 }
                 case RHI::CopyItemType::BufferToImage:
@@ -266,6 +269,11 @@ namespace AZ
             mtlVertexArgBufferOffsets.fill(0);
             mtlFragmentOrComputeArgBufferOffsets.fill(0);
             
+            //Map to cache all the resources based on the usage as we can batch all the resources for a given usage
+            ArgumentBuffer::ComputeResourcesToMakeResidentMap resourcesToMakeResidentCompute;
+            //Map to cache all the resources based on the usage and shader stage as we can batch all the resources for a given usage/shader usage
+            ArgumentBuffer::GraphicsResourcesToMakeResidentMap resourcesToMakeResidentGraphics;
+            
             for (uint32_t slot = 0; slot < RHI::Limits::Pipeline::ShaderResourceGroupCountMax; ++slot)
             {
                 const ShaderResourceGroup* shaderResourceGroup = bindings.m_srgsBySlot[slot];
@@ -291,7 +299,6 @@ namespace AZ
                         //For graphics and compute shader stages, cache all the argument buffers, offsets and track the min/max indices
                         if(m_commandEncoderType == CommandEncoderType::Render)
                         {
-                            id<MTLRenderCommandEncoder> renderEncoder = GetEncoder<id<MTLRenderCommandEncoder>>();
                             uint8_t numBitsSet = RHI::CountBitsSet(static_cast<uint64_t>(srgVisInfo));
                             if( numBitsSet > 1 || srgVisInfo == RHI::ShaderStageMask::Vertex)
                             {
@@ -334,11 +341,11 @@ namespace AZ
                         //format compatible with the appropriate metal function.
                         if(m_commandEncoderType == CommandEncoderType::Render)
                         {
-                            shaderResourceGroup->AddUntrackedResourcesToEncoder(m_encoder, srgResourcesVisInfo);
+                            shaderResourceGroup->CollectUntrackedResources(m_encoder, srgResourcesVisInfo, resourcesToMakeResidentCompute, resourcesToMakeResidentGraphics);
                         }
                         else if(m_commandEncoderType == CommandEncoderType::Compute)
                         {
-                            shaderResourceGroup->AddUntrackedResourcesToEncoder(m_encoder, srgResourcesVisInfo);
+                            shaderResourceGroup->CollectUntrackedResources(m_encoder, srgResourcesVisInfo, resourcesToMakeResidentCompute, resourcesToMakeResidentGraphics);
                         }
                     }
                 }
@@ -368,6 +375,32 @@ namespace AZ
                                     mtlFragmentOrComputeArgBufferOffsets);
             }
             
+            id<MTLRenderCommandEncoder> renderEncoder = GetEncoder<id<MTLRenderCommandEncoder>>();
+            id<MTLComputeCommandEncoder> computeEncoder = GetEncoder<id<MTLComputeCommandEncoder>>();
+            
+            //Call UseResource on all resources for Compute stage
+            for (const auto& key : resourcesToMakeResidentCompute)
+            {
+                AZStd::vector<id <MTLResource>> resourcesToProcessVec(key.second.begin(), key.second.end());
+                
+                [computeEncoder useResources: &resourcesToProcessVec[0]
+                                       count: resourcesToProcessVec.size()
+                                       usage: key.first];
+                 
+            }
+            
+            //Call UseResource on all resources for Vertex and Fragment stages
+            for (const auto& key : resourcesToMakeResidentGraphics)
+            {
+                
+                AZStd::vector<id <MTLResource>> resourcesToProcessVec(key.second.begin(), key.second.end());
+                
+                [renderEncoder useResources: &resourcesToProcessVec[0]
+                                      count: resourcesToProcessVec.size()
+                                      usage: key.first.first
+                                     stages: key.first.second];
+            }
+            
             return true;
         }
     

Gems/Atom/RHI/Metal/Code/Source/RHI/CommandListBase.cpp

@@ -83,7 +83,7 @@ namespace AZ
                     for (id<MTLHeap> residentHeap : *m_residentHeaps)
                     {
                         [renderEncoder useHeap : residentHeap
-                                       stages  : MTLRenderStageFragment];
+                                       stages  : MTLRenderStageVertex | MTLRenderStageFragment];
                     }
                     break;
                 }

Gems/Atom/RHI/Metal/Code/Source/RHI/Device.cpp

@@ -80,6 +80,9 @@ namespace AZ
 
             m_nullDescriptorManager.Init(*this);
             
+            m_samplerCache = [[NSCache alloc]init];
+            [m_samplerCache setName:@"SamplerCache"];
+            
             return RHI::ResultCode::Success;
         }
     
@@ -101,6 +104,10 @@ namespace AZ
             m_releaseQueue.Shutdown();
             m_pipelineLayoutCache.Shutdown();
             
+            [m_samplerCache removeAllObjects];
+            [m_samplerCache release];
+            m_samplerCache = nil;
+            
             for (AZ::u32 i = 0; i < CommandEncoderTypeCount; ++i)
             {
                 m_commandListPools[i].Shutdown();

Gems/Atom/RHI/Metal/Code/Source/RHI/Device.h

@@ -144,6 +144,11 @@ namespace AZ
                 return m_asyncUploadQueue;
             }
             
+            const NSCache* GetSamplerCache() const
+            {
+                return m_samplerCache;
+            }
+            
             BufferMemoryAllocator& GetArgBufferConstantBufferAllocator() { return m_argumentBufferConstantsAllocator;}
             BufferMemoryAllocator& GetArgumentBufferAllocator() { return m_argumentBufferAllocator;}
             
@@ -194,6 +199,7 @@ namespace AZ
             RHI::HeapMemoryUsage m_argumentBufferAllocatorMemoryUsage;
             
             NullDescriptorManager m_nullDescriptorManager;
+            NSCache* m_samplerCache;
         };
     }
 }

Gems/Atom/RHI/Metal/Code/Source/RHI/ShaderResourceGroup.cpp

@@ -11,6 +11,7 @@
 */
 #include "Atom_RHI_Metal_precompiled.h"
 
+#include <RHI/ArgumentBuffer.h>
 #include <RHI/ImageView.h>
 #include <RHI/ShaderResourceGroup.h>
 
@@ -33,10 +34,12 @@ namespace AZ
             return *m_compiledArgBuffers[m_compiledDataIndex];
         }
     
-        void ShaderResourceGroup::AddUntrackedResourcesToEncoder(id<MTLCommandEncoder> commandEncoder,
-                                            const ShaderResourceGroupVisibility& srgResourcesVisInfo) const
+        void ShaderResourceGroup::CollectUntrackedResources(id<MTLCommandEncoder> commandEncoder,
+                                                            const ShaderResourceGroupVisibility& srgResourcesVisInfo,
+                                                            ArgumentBuffer::ComputeResourcesToMakeResidentMap& resourcesToMakeResidentCompute,
+                                                            ArgumentBuffer::GraphicsResourcesToMakeResidentMap& resourcesToMakeResidentGraphics) const
         {
-            GetCompiledArgumentBuffer().AddUntrackedResourcesToEncoder(commandEncoder, srgResourcesVisInfo);
+            GetCompiledArgumentBuffer().CollectUntrackedResources(commandEncoder, srgResourcesVisInfo, resourcesToMakeResidentCompute, resourcesToMakeResidentGraphics);
         }
     }
 }

Gems/Atom/RHI/Metal/Code/Source/RHI/ShaderResourceGroup.h

@@ -47,7 +47,10 @@ namespace AZ
             const ImageView* GetImageView(const int index) const;
             void UpdateCompiledDataIndex();            
             const ArgumentBuffer& GetCompiledArgumentBuffer() const;            
-            void AddUntrackedResourcesToEncoder(id<MTLCommandEncoder> commandEncoder, const ShaderResourceGroupVisibility& srgResourcesVisInfo) const;
+            void CollectUntrackedResources(id<MTLCommandEncoder> commandEncoder,
+                                           const ShaderResourceGroupVisibility& srgResourcesVisInfo,
+                                            ArgumentBuffer::ComputeResourcesToMakeResidentMap& resourcesToMakeResidentCompute,
+                                            ArgumentBuffer::GraphicsResourcesToMakeResidentMap& resourcesToMakeResidentGraphics) const;
 
         private:
             ShaderResourceGroup() = default;

Gems/Atom/RPI/Code/Include/Atom/RPI.Public/Pass/Specific/DownsampleMipChainPass.h

@@ -30,6 +30,7 @@ namespace AZ
         //! It does this by recursively creating Compute Passes to write to each mip using the Compute Shader.
         class DownsampleMipChainPass
             : public ParentPass
+            , private ShaderReloadNotificationBus::Handler
         {
             AZ_RPI_PASS(DownsampleMipChainPass);
 
@@ -39,6 +40,7 @@ namespace AZ
           
             //! Creates a new pass without a PassTemplate
             static Ptr<DownsampleMipChainPass> Create(const PassDescriptor& descriptor);
+            virtual ~DownsampleMipChainPass();
             
         protected:
             explicit DownsampleMipChainPass(const PassDescriptor& descriptor);
@@ -49,6 +51,11 @@ namespace AZ
             void BuildInternal() override;
             void FrameBeginInternal(FramePrepareParams params) override;
 
+            // ShaderReloadNotificationBus::Handler overrides...
+            void OnShaderReinitialized(const Shader& shader) override;
+            void OnShaderAssetReinitialized(const Data::Asset<ShaderAsset>& shaderAsset) override;
+            void OnShaderVariantReinitialized(const ShaderVariant& shaderVariant) override;
+            
         private:
 
             // Gets target height, width and mip levels from the input/output image attachment

Gems/Atom/RPI/Code/Source/RPI.Public/Pass/Specific/DownsampleMipChainPass.cpp

@@ -54,8 +54,14 @@ namespace AZ
             }
 
             m_passData = *passData;
+            ShaderReloadNotificationBus::Handler::BusConnect(passData->m_shaderReference.m_assetId);
         }
 
+        DownsampleMipChainPass::~DownsampleMipChainPass()
+        {
+            ShaderReloadNotificationBus::Handler::BusDisconnect();
+        }
+    
         void DownsampleMipChainPass::ResetInternal()
         {
             RemoveChildren();
@@ -206,5 +212,19 @@ namespace AZ
             ParentPass::FrameBeginInternal(params);
         }
 
+        void DownsampleMipChainPass::OnShaderReinitialized([[maybe_unused]] const Shader& shader)
+        {
+            m_needToUpdateChildren = true;
+        }
+    
+        void DownsampleMipChainPass::OnShaderAssetReinitialized([[maybe_unused]] const Data::Asset<ShaderAsset>& shaderAsset)
+        {
+            m_needToUpdateChildren = true;
+        }
+    
+        void DownsampleMipChainPass::OnShaderVariantReinitialized([[maybe_unused]] const ShaderVariant& shaderVariant)
+        {
+            m_needToUpdateChildren = true;
+        }
     }   // namespace RPI
 }   // namespace AZ