Move swapchain and image recreation to the end of the frame (#5388)

Signed-off-by: jiaweig <jiaweig@amazon.com>

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

@@ -75,6 +75,9 @@ namespace AZ
             //! Return True if the swap chain prefers exclusive full screen mode and a transition happened, false otherwise.
             virtual bool SetExclusiveFullScreenState([[maybe_unused]]bool fullScreenState) { return false; }
 
+            //! Recreate the swapchain if it becomes invalid during presenting. This should happen at the end of the frame
+            //! due to images being used as attachments in the frame graph.
+            virtual void ProcessRecreation() {};
         protected:
             SwapChain();
 
@@ -98,6 +101,14 @@ namespace AZ
 
             //////////////////////////////////////////////////////////////////////////
 
+            //! Shutdown and clear all the images.
+            void ShutdownImages();
+
+            //! Initialized all the images.
+            ResultCode InitImages();
+
+            //! Flag indicating if swapchain recreation is needed at the end of the frame.
+            bool m_pendingRecreation = false;
         private:
 
             bool ValidateDescriptor(const SwapChainDescriptor& descriptor) const;

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

@@ -134,7 +134,6 @@ namespace AZ
             m_scopeAttachmentLookup.clear();
             m_imageAttachments.clear();
             m_bufferAttachments.clear();
-            m_swapChainAttachments.clear();
             m_importedImageAttachments.clear();
             m_importedBufferAttachments.clear();
             m_transientImageAttachments.clear();
@@ -153,6 +152,13 @@ namespace AZ
                 delete attachment;
             }
             m_attachments.clear();
+
+            for (auto swapchainAttachment : m_swapChainAttachments)
+            {
+                swapchainAttachment->GetSwapChain()->ProcessRecreation();
+            }
+
+            m_swapChainAttachments.clear();
         }
 
         ImageDescriptor FrameGraphAttachmentDatabase::GetImageDescriptor(const AttachmentId& attachmentId) const

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

@@ -58,9 +58,32 @@ namespace AZ
                 // Overwrite descriptor dimensions with the native ones (the ones assigned by the platform) returned by InitInternal.
                 m_descriptor.m_dimensions = nativeDimensions;
 
+                resultCode = InitImages();
+            }
+
+            return resultCode;
+        }
+
+        void SwapChain::ShutdownImages()
+        {
+            // Shutdown existing set of images.
+            uint32_t imageSize = aznumeric_cast<uint32_t>(m_images.size());
+            for (uint32_t imageIdx = 0; imageIdx < imageSize; ++imageIdx)
+            {
+                m_images[imageIdx]->Shutdown();
+            }
+
+            m_images.clear();
+        }
+
+        ResultCode SwapChain::InitImages()
+        {
+            ResultCode resultCode = ResultCode::Success;
+
             m_images.reserve(m_descriptor.m_dimensions.m_imageCount);
 
-                for (uint32_t imageIdx = 0; imageIdx < m_descriptor.m_dimensions.m_imageCount; ++imageIdx)
+            // If the new display mode has more buffers, add them.
+            for (uint32_t i = 0; i < m_descriptor.m_dimensions.m_imageCount; ++i)
             {
                 m_images.emplace_back(RHI::Factory::Get().CreateImage());
             }
@@ -80,8 +103,7 @@ namespace AZ
                 request.m_imageIndex = imageIdx;
 
                 resultCode = ImagePoolBase::InitImage(
-                        request.m_image,
+                    request.m_image, imageDescriptor,
-                        imageDescriptor,
                     [this, &request]()
                     {
                         return InitImageInternal(request);
@@ -89,11 +111,14 @@ namespace AZ
 
                 if (resultCode != ResultCode::Success)
                 {
+                    AZ_Error("Swapchain", false, "Failed to initialize images.");
                     Shutdown();
                     break;
                 }
             }
-            }
+
+            // Reset the current index back to 0 so we match the platform swap chain.
+            m_currentImageIndex = 0;
 
             return resultCode;
         }
@@ -106,62 +131,14 @@ namespace AZ
 
         ResultCode SwapChain::Resize(const RHI::SwapChainDimensions& dimensions)
         {
-            // Shutdown existing set of images.
+            ShutdownImages();
-            for (uint32_t imageIdx = 0; imageIdx < GetImageCount(); ++imageIdx)
-            {
-                m_images[imageIdx]->Shutdown();
-            }
 
             SwapChainDimensions nativeDimensions = dimensions;
             ResultCode resultCode = ResizeInternal(dimensions, &nativeDimensions);
             if (resultCode == ResultCode::Success)
             {
                 m_descriptor.m_dimensions = nativeDimensions;
-                m_images.reserve(m_descriptor.m_dimensions.m_imageCount);
+                resultCode = InitImages();
-
-                // If the new display mode has more buffers, add them.
-                while (m_images.size() < static_cast<size_t>(m_descriptor.m_dimensions.m_imageCount))
-                {
-                    m_images.emplace_back(RHI::Factory::Get().CreateImage());
-                }
-
-                // If it has fewer, trim down.
-                while (m_images.size() > static_cast<size_t>(m_descriptor.m_dimensions.m_imageCount))
-                {
-                    m_images.pop_back();
-                }
-
-                InitImageRequest request;
-
-                RHI::ImageDescriptor& imageDescriptor = request.m_descriptor;
-                imageDescriptor.m_dimension = RHI::ImageDimension::Image2D;
-                imageDescriptor.m_bindFlags = RHI::ImageBindFlags::Color;
-                imageDescriptor.m_size.m_width = m_descriptor.m_dimensions.m_imageWidth;
-                imageDescriptor.m_size.m_height = m_descriptor.m_dimensions.m_imageHeight;
-                imageDescriptor.m_format = m_descriptor.m_dimensions.m_imageFormat;
-
-                for (uint32_t imageIdx = 0; imageIdx < GetImageCount(); ++imageIdx)
-                {
-                    request.m_image = m_images[imageIdx].get();
-                    request.m_imageIndex = imageIdx;
-
-                    resultCode = ImagePoolBase::InitImage(
-                        request.m_image,
-                        imageDescriptor,
-                        [this, &request]()
-                    {
-                        return InitImageInternal(request);
-                    });
-
-                    if (resultCode != ResultCode::Success)
-                    {
-                        Shutdown();
-                        break;
-                    }
-                }
-
-                // Reset the current index back to 0 so we match the platform swap chain.
-                m_currentImageIndex = 0;
             }
 
             return resultCode;
@@ -188,7 +165,7 @@ namespace AZ
 
         uint32_t SwapChain::GetImageCount() const
         {
-            return static_cast<uint32_t>(m_images.size());
+            return aznumeric_cast<uint32_t>(m_images.size());
         }
 
         uint32_t SwapChain::GetCurrentImageIndex() const
@@ -209,8 +186,18 @@ namespace AZ
         void SwapChain::Present()
         {
             AZ_TRACE_METHOD();
+            // Due to swapchain recreation, the images are refreshed.
+            // There is no need to present swapchain for this frame.
+            const uint32_t imageCount = aznumeric_cast<uint32_t>(m_images.size());
+            if (imageCount == 0)
+            {
+                return;
+            }
+            else
+            {
                 m_currentImageIndex = PresentInternal();
-            AZ_Assert(m_currentImageIndex < m_images.size(), "Invalid image index");
+                AZ_Assert(m_currentImageIndex < imageCount, "Invalid image index");
+            }
         }
     }
 }

Gems/Atom/RHI/Vulkan/Code/Source/RHI/SwapChain.cpp

@@ -59,6 +59,19 @@ namespace AZ
             m_swapChainBarrier.m_isValid = true;
         }
 
+        void SwapChain::ProcessRecreation()
+        {
+            if (m_pendingRecreation)
+            {
+                ShutdownImages();
+                InvalidateNativeSwapChain();
+                CreateSwapchain();
+                InitImages();
+
+                m_pendingRecreation = false;
+            }
+        }
+
         void SwapChain::SetVerticalSyncIntervalInternal(uint32_t previousVsyncInterval)
         {
             if (GetDescriptor().m_verticalSyncInterval == 0 || previousVsyncInterval == 0)
@@ -231,8 +244,7 @@ namespace AZ
                 // VK_SUBOPTIMAL_KHR is treated as success, but we better update the surface info as well.
                 if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR)
                 {
-                    InvalidateNativeSwapChain();
+                    m_pendingRecreation = true;
-                    CreateSwapchain();
                 }
                 else
                 {
@@ -246,18 +258,16 @@ namespace AZ
                 }
             };
 
-            m_presentationQueue->QueueCommand(AZStd::move(presentCommand));
-
             uint32_t acquiredImageIndex = GetCurrentImageIndex();
             RHI::ResultCode result = AcquireNewImage(&acquiredImageIndex);
             if (result == RHI::ResultCode::Fail)
             {
-                InvalidateNativeSwapChain();
+                m_pendingRecreation = true;
-                CreateSwapchain();
                 return 0;
             }
             else
             {
+                m_presentationQueue->QueueCommand(AZStd::move(presentCommand));
                 return acquiredImageIndex;
             }
         }

Gems/Atom/RHI/Vulkan/Code/Source/RHI/SwapChain.h

@@ -51,6 +51,7 @@ namespace AZ
 
             void QueueBarrier(const VkPipelineStageFlags src, const VkPipelineStageFlags dst, const VkImageMemoryBarrier& imageBarrier);
 
+            void ProcessRecreation() override;
         private:
             SwapChain() = default;