ATOM-15939 Add support to capture attachment for ParentPass (#1887)

* ATOM-15939 Add support to capture attachment for ParentPass
- Moved the attachment read back support to Pass class so it supports both ParentPass and RenderPass.
- Added support to output input or output state of an InputOutput attachment.
- Enabled showing ParentPass attachments in PassTree tool.

Signed-off-by: Tao <qingtao@amazon.com>

Code/Editor/TrackView/AtomOutputFrameCapture.cpp

@@ -73,7 +73,7 @@ namespace TrackView
         bool startedCapture = false;
         AZ::Render::FrameCaptureRequestBus::BroadcastResult(
             startedCapture, &AZ::Render::FrameCaptureRequestBus::Events::CapturePassAttachmentWithCallback, m_passHierarchy,
-            AZStd::string("Output"), attachmentReadbackCallback);
+            AZStd::string("Output"), attachmentReadbackCallback, AZ::RPI::PassAttachmentReadbackOption::Output);
 
         return startedCapture;
     }

Gems/Atom/Feature/Common/Code/Include/Atom/Feature/Utils/FrameCaptureBus.h

@@ -46,14 +46,16 @@ namespace AZ
             //! Save a buffer attachment or a image attachment binded to a pass's slot to a data file.
             //! @param passHierarchy For finding the pass by using PassHierarchyFilter
             //! @param slotName Name of the pass's slot. The attachment bound to this slot will be captured.
+            //! @param option Only valid for an InputOutput attachment. Use PassAttachmentReadbackOption::Input to capture the input state
+            //!               and use PassAttachmentReadbackOption::Output to capture the output state
             //! @param outputFilename The output file path. 
             virtual bool CapturePassAttachment(const AZStd::vector<AZStd::string>& passHierarchy, const AZStd::string& slotName
-                , const AZStd::string& outputFilePath) = 0;
+                , const AZStd::string& outputFilePath, RPI::PassAttachmentReadbackOption option) = 0;
 
             //! Similar to CapturePassAttachment. But instead of saving the read back result to a file, it will call the callback function provide
             //! in the input when callback is finished
             virtual bool CapturePassAttachmentWithCallback(const AZStd::vector<AZStd::string>& passHierarchy, const AZStd::string& slotName
-                , RPI::AttachmentReadback::CallbackFunction callback) = 0;
+                , RPI::AttachmentReadback::CallbackFunction callback, RPI::PassAttachmentReadbackOption option) = 0;
 
         };
         using FrameCaptureRequestBus = EBus<FrameCaptureRequests>;

Gems/Atom/Feature/Common/Code/Source/FrameCaptureSystemComponent.cpp

@@ -343,7 +343,7 @@ namespace AZ
         }
 
         bool FrameCaptureSystemComponent::CapturePassAttachment(const AZStd::vector<AZStd::string>& passHierarchy, const AZStd::string& slot,
-            const AZStd::string& outputFilePath)
+            const AZStd::string& outputFilePath, RPI::PassAttachmentReadbackOption option)
         {
             InitReadback();
 
@@ -376,40 +376,22 @@ namespace AZ
                 return false;
             }
 
-            AZ::RPI::RenderPass* renderPass = azrtti_cast<AZ::RPI::RenderPass*>(foundPasses[0]);
+            AZ::RPI::Pass* pass = foundPasses[0];
-            if (renderPass)
+            if (pass->ReadbackAttachment(m_readback, Name(slot), option))
             {
-                Name slotName = Name(slot);
+                m_state = State::Pending;
-                AZ::RPI::PassAttachment* attachment = nullptr;
+                m_result = FrameCaptureResult::None;
-                for (auto& binding : renderPass->GetAttachmentBindings())
+                SystemTickBus::Handler::BusConnect();
-                {
+                return true;
-                    if (binding.m_name == slotName)
-                    {
-                        attachment = binding.m_attachment.get();
-                        break;
-                    }
-                }
-                if (attachment)
-                {
-                    m_state = State::Pending;
-                    m_result = FrameCaptureResult::None;
-                    SystemTickBus::Handler::BusConnect();
-                    renderPass->ReadbackAttachment(m_readback, attachment);
-                }
-                else
-                {
-                    AZ_Warning("FrameCaptureSystemComponent", false, "Failed to find attachment bound to pass [%s] slot [%s]",
-                        renderPass->GetName().GetCStr(), slotName.GetCStr());
-                    return false;
-                }
             }
-            return true;
+            AZ_Warning("FrameCaptureSystemComponent", false, "Failed to readback the attachment bound to pass [%s] slot [%s]", pass->GetName().GetCStr(), slot.c_str());
+            return false;
         }
 
         bool FrameCaptureSystemComponent::CapturePassAttachmentWithCallback(const AZStd::vector<AZStd::string>& passHierarchy, const AZStd::string& slotName
-            , RPI::AttachmentReadback::CallbackFunction callback)
+            , RPI::AttachmentReadback::CallbackFunction callback, RPI::PassAttachmentReadbackOption option)
         {
-            bool result = CapturePassAttachment(passHierarchy, slotName, "");
+            bool result = CapturePassAttachment(passHierarchy, slotName, "", option);
 
             // Append state change to user provided call back
             AZ::RPI::AttachmentReadback::CallbackFunction callbackSetState = [&, callback](const AZ::RPI::AttachmentReadback::ReadbackResult& result)

Gems/Atom/Feature/Common/Code/Source/FrameCaptureSystemComponent.h

@@ -36,9 +36,10 @@ namespace AZ
             bool CaptureScreenshot(const AZStd::string& filePath) override;
             bool CaptureScreenshotForWindow(const AZStd::string& filePath, AzFramework::NativeWindowHandle windowHandle) override;
             bool CaptureScreenshotWithPreview(const AZStd::string& outputFilePath) override;
-            bool CapturePassAttachment(const AZStd::vector<AZStd::string>& passHierarchy, const AZStd::string& slotName, const AZStd::string& outputFilePath) override;
+            bool CapturePassAttachment(const AZStd::vector<AZStd::string>& passHierarchy, const AZStd::string& slotName, const AZStd::string& outputFilePath,
+                RPI::PassAttachmentReadbackOption option) override;
             bool CapturePassAttachmentWithCallback(const AZStd::vector<AZStd::string>& passHierarchy, const AZStd::string& slotName
-                , RPI::AttachmentReadback::CallbackFunction callback) override;
+                , RPI::AttachmentReadback::CallbackFunction callback, RPI::PassAttachmentReadbackOption option) override;
 
         private:
             void CaptureAttachmentCallback(const AZ::RPI::AttachmentReadback::ReadbackResult& readbackResult);

Gems/Atom/Feature/Common/Code/Source/LuxCore/LuxCoreTexturePass.cpp

@@ -64,23 +64,9 @@ namespace AZ
                 // Set up read back attachment before children prepare
                 if (m_readback->IsReady())
                 {
-                    AZ::RPI::RenderPass* renderPass = azrtti_cast<AZ::RPI::RenderPass*>(m_renderTargetPass.get());
+                    if (m_renderTargetPass)
-                    if (renderPass)
                     {
-                        RPI::PassAttachment* attachment = nullptr;
+                        m_attachmentReadbackComplete = m_renderTargetPass->ReadbackAttachment(m_readback, AZ::Name("RenderTargetOutput"));                           
-                        for (auto& binding : renderPass->GetAttachmentBindings())
-                        {
-                            if (binding.m_slotType == RPI::PassSlotType::Output)
-                            {
-                                attachment = binding.m_attachment.get();
-                                break;
-                            }
-                        }
-                        if (attachment)
-                        {
-                            renderPass->ReadbackAttachment(m_readback, attachment);
-                            m_attachmentReadbackComplete = true;
-                        }
                     }
                 }
             }

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

@@ -57,6 +57,7 @@ namespace AZ
         class PassTemplate;
         struct PassRequest;
         struct PassValidationResults;
+        class AttachmentReadback;
 
         using SortedPipelineViewTags = AZStd::set<PipelineViewTag, AZNameSortAscending>;
         using PassesByDrawList = AZStd::map<RHI::DrawListTag, const Pass*>;
@@ -65,7 +66,12 @@ namespace AZ
         const uint32_t PassInputBindingCountMax = 16;
         const uint32_t PassInputOutputBindingCountMax = PassInputBindingCountMax;
         const uint32_t PassOutputBindingCountMax = PassInputBindingCountMax;
-
+                
+        enum class PassAttachmentReadbackOption : uint8_t
+        {
+            Input = 0,
+            Output
+        };
 
         //! Atom's base pass class (every pass class in Atom must derive from this class).
         //! 
@@ -222,6 +228,14 @@ namespace AZ
             //! Enables/Disables PipelineStatistics queries for this pass
             virtual void SetPipelineStatisticsQueryEnabled(bool enable);
 
+            //! Readback an attachment attached to the specified slot name
+            //! @param readback The AttachmentReadback object which is used for readback. Its callback function will be called when readback is finished.
+            //! @param slotName The attachment bind to the slot with this slotName is to be readback
+            //! @param option The option is used for choosing input or output state when readback an InputOutput attachment.
+            //!               It's ignored if the attachment isn't an InputOutput attachment.
+            //! Return true if the readback request was successful. User may expect the AttachmentReadback's callback function would be called. 
+            bool ReadbackAttachment(AZStd::shared_ptr<AttachmentReadback> readback, const Name& slotName, PassAttachmentReadbackOption option = PassAttachmentReadbackOption::Output);
+
             //! Returns whether the Timestamp queries is enabled/disabled for this pass
             bool IsTimestampQueryEnabled() const;
 
@@ -266,7 +280,6 @@ namespace AZ
             // Update output bindings on this pass that are connected to bindings on other passes
             void UpdateConnectedOutputBindings();
 
-
         protected:
             explicit Pass(const PassDescriptor& descriptor);
 
@@ -349,6 +362,7 @@ namespace AZ
             void FrameEnd();
             virtual void FrameEndInternal() { }
 
+            void UpdateReadbackAttachment(FramePrepareParams params, bool beforeAddScopes);
 
             // --- Protected Members ---
 
@@ -442,7 +456,10 @@ namespace AZ
             // Sort type to be used by the default sort implementation. Passes can also provide
             // fully custom sort implementations by overriding the SortDrawList() function.
             RHI::DrawListSortType m_drawListSortType = RHI::DrawListSortType::KeyThenDepth;
-
+            
+            // For read back attachment
+            AZStd::shared_ptr<AttachmentReadback> m_attachmentReadback;
+            PassAttachmentReadbackOption m_readbackOption;
 
         private:
             // Return the Timestamp result of this pass

Gems/Atom/RPI/Code/Include/Atom/RPI.Public/Pass/RenderPass.h

@@ -55,10 +55,7 @@ namespace AZ
 
             //! Get MultisampleState of this pass from its output attachments
             RHI::MultisampleState GetMultisampleState() const;
-
+            
-            //! Capture pass's output and input/output attachments in following frames
-            void ReadbackAttachment(AZStd::shared_ptr<AttachmentReadback> readback, const PassAttachment* attachment);
-
             //! Returns a pointer to the Pass ShaderResourceGroup
             Data::Instance<ShaderResourceGroup> GetShaderResourceGroup();
 
@@ -145,9 +142,6 @@ namespace AZ
             // Readback the results from the ScopeQueries
             void ReadbackScopeQueryResults();
 
-            // For read back attachments
-            AZStd::shared_ptr<AttachmentReadback> m_attachmentReadback;
-
             AZStd::weak_ptr<ImageAttachmentCopy> m_attachmentCopy;
 
             // Readback results from the Timestamp queries

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

@@ -166,9 +166,6 @@ namespace AZ
             AZ::Vector2 m_position = AZ::Vector2(0, 0.6f);
             AZ::Vector2 m_size = AZ::Vector2(0.4f, 0.4f);
             bool m_keepAspectRatio = true;
-
-            // For readback the output image attachment 
-            AZStd::shared_ptr<AttachmentReadback> m_readback;
         };
     }   // namespace RPI
 }   // namespace AZ

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

@@ -60,9 +60,6 @@ namespace AZ
             // Name of the template used to create the child pass. Needed for Recreate()
             Name m_childTemplateName;
 
-            // For read back output
-            AZStd::shared_ptr<AttachmentReadback> m_readback;
-
             Ptr<PassAttachment> m_outputAttachment;
 
             // saved settings for this pass

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

@@ -15,7 +15,6 @@
 #include <Atom/RHI.Reflect/Format.h>
 #include <Atom/RHI.Reflect/SwapChainDescriptor.h>
 
-#include <Atom/RPI.Public/Pass/AttachmentReadback.h>
 #include <Atom/RPI.Public/Pass/ParentPass.h>
 
 #include <AzFramework/Windowing/WindowBus.h>

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

@@ -15,7 +15,6 @@
 #include <Atom/RHI.Reflect/Format.h>
 #include <Atom/RHI.Reflect/SwapChainDescriptor.h>
 
-#include <Atom/RPI.Public/Pass/AttachmentReadback.h>
 #include <Atom/RPI.Public/Pass/ParentPass.h>
 
 #include <AzFramework/Windowing/WindowBus.h>
@@ -88,9 +87,6 @@ namespace AZ
 
             // Name of the template used to create the child pass. Needed for Recreate()
             Name m_childTemplateName;
-
-            // For read back swap chain
-            AZStd::shared_ptr<AttachmentReadback> m_swapChainReadback;
         };
 
     }   // namespace RPI

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

@@ -19,6 +19,7 @@
 #include <Atom/RPI.Public/Buffer/Buffer.h>
 #include <Atom/RPI.Public/Image/AttachmentImage.h>
 #include <Atom/RPI.Reflect/Image/Image.h>
+#include <Atom/RPI.Public/Pass/AttachmentReadback.h>
 #include <Atom/RPI.Public/Pass/ParentPass.h>
 #include <Atom/RPI.Public/Pass/Pass.h>
 #include <Atom/RPI.Public/Pass/PassLibrary.h>
@@ -36,8 +37,7 @@
 namespace AZ
 {
     namespace RPI
-    {
+    {             
-
         // --- Constructors ---
 
         Pass::Pass(const PassDescriptor& descriptor)
@@ -1285,9 +1285,15 @@ namespace AZ
             CreateTransientAttachments(params.m_frameGraphBuilder->GetAttachmentDatabase());
             ImportAttachments(params.m_frameGraphBuilder->GetAttachmentDatabase());
 
+            // readback attachment with input state
+            UpdateReadbackAttachment(params, true);
+
             // FrameBeginInternal needs to be the last function be called in FrameBegin because its implementation expects 
             // all the attachments are imported to database (for example, ImageAttachmentPreview)
             FrameBeginInternal(params);
+            
+            // readback attachment with output state
+            UpdateReadbackAttachment(params, false);
 
             UpdateConnectedOutputBindings();
         }
@@ -1426,6 +1432,57 @@ namespace AZ
             m_flags.m_pipelineStatisticsQueryEnabled = enable;
         }
 
+        bool Pass::ReadbackAttachment(AZStd::shared_ptr<AttachmentReadback> readback, const Name& slotName, PassAttachmentReadbackOption option)
+        {
+            // Return false if it's already readback
+            if (m_attachmentReadback)
+            {
+                AZ_Warning("Pass", false, "ReadbackAttachment: skip readback pass [%s] slot [%s]because there is an another active readback", m_name.GetCStr(), slotName.GetCStr());
+                return false;
+            }
+            uint32_t bindingIndex = 0;
+            for (auto& binding : m_attachmentBindings)
+            {
+                if (slotName == binding.m_name)
+                {
+                    RHI::AttachmentType type = binding.m_attachment->GetAttachmentType();
+                    if (type == RHI::AttachmentType::Buffer || type == RHI::AttachmentType::Image)
+                    {
+                        RHI::AttachmentId attachmentId = binding.m_attachment->GetAttachmentId();
+
+                        // Append slot index and pass name so the read back's name won't be same as the attachment used in other passes.
+                        AZStd::string readbackName = AZStd::string::format("%s_%d_%s", attachmentId.GetCStr(),
+                            bindingIndex, GetName().GetCStr());
+                        if (readback->ReadPassAttachment(binding.m_attachment.get(), AZ::Name(readbackName)))
+                        {
+                            m_readbackOption = PassAttachmentReadbackOption::Output;
+                            // The m_readbackOption is only meaningful if the attachment is used for InputOutput.
+                            if (binding.m_slotType == PassSlotType::InputOutput)
+                            {
+                                m_readbackOption = option;
+                            }
+                            m_attachmentReadback = readback;
+                            return true;
+                        }
+                        return false;
+                    }
+                }
+                bindingIndex++;
+            }
+            AZ_Warning("Pass", false, "ReadbackAttachment: failed to find slot [%s] from pass [%s]", slotName.GetCStr(), m_name.GetCStr());
+            return false;
+        }
+
+        void Pass::UpdateReadbackAttachment(FramePrepareParams params, bool beforeAddScopes)
+        {
+            if (beforeAddScopes == (m_readbackOption == PassAttachmentReadbackOption::Input) && m_attachmentReadback)
+            {
+                // Read the attachment for one frame. The reference can be released afterwards
+                m_attachmentReadback->FrameBegin(params);
+                m_attachmentReadback = nullptr;
+            }
+        }
+
         bool Pass::IsTimestampQueryEnabled() const
         {
             return m_flags.m_timestampQueryEnabled;

Gems/Atom/RPI/Code/Source/RPI.Public/Pass/RenderPass.cpp

@@ -188,14 +188,8 @@ namespace AZ
             {
                 SetScopeId(RHI::ScopeId(GetPathName()));
             }
-            params.m_frameGraphBuilder->ImportScopeProducer(*this);
 
-            // Read the attachment for one frame. The reference can be released afterwards
+            params.m_frameGraphBuilder->ImportScopeProducer(*this);
-            if (m_attachmentReadback)
-            {
-                m_attachmentReadback->FrameBegin(params);
-                m_attachmentReadback = nullptr;
-            }
 
             // Read back the ScopeQueries submitted from previous frames
             ReadbackScopeQueryResults();
@@ -414,31 +408,6 @@ namespace AZ
             m_flags.m_hasPipelineViewTag = !viewTag.IsEmpty();
         }
 
-        void RenderPass::ReadbackAttachment(AZStd::shared_ptr<AttachmentReadback> readback, const PassAttachment* attachment)
-        {
-            m_attachmentReadback = readback;
-
-            uint32_t bindingIndex = 0;
-            for (auto& binding : m_attachmentBindings)
-            {
-                if (attachment == binding.m_attachment)
-                {
-                    RHI::AttachmentType type = binding.m_attachment->GetAttachmentType();
-                    if (type == RHI::AttachmentType::Buffer || type == RHI::AttachmentType::Image)
-                    {
-                        RHI::AttachmentId attachmentId = binding.m_attachment->GetAttachmentId();
-
-                        // Append slot index and pass name so the read back's name won't be same as the attachment used in other passes.
-                        AZStd::string readbackName = AZStd::string::format("%s_%d_%s", attachmentId.GetCStr(),
-                            bindingIndex, GetName().GetCStr());
-                        m_attachmentReadback->ReadPassAttachment(binding.m_attachment.get(), AZ::Name(readbackName));
-                        return;
-                    }
-                }
-                bindingIndex++;
-            }
-        }
-
         TimestampResult RenderPass::GetTimestampResultInternal() const
         {
             return m_timestampResult;

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

@@ -334,22 +334,16 @@ namespace AZ
                 m_passSrg->Compile();
                 m_passSrgChanged = false;
             }
-
-            // Read preview output
-            if (m_readback)
-            {
-                m_readback->FrameBegin(params);
-                m_readback = nullptr;
-            }
         }
 
         bool ImageAttachmentPreviewPass::ReadbackOutput(AZStd::shared_ptr<AttachmentReadback> readback)
         {
             if (m_outputColorAttachment)
             {
-                m_readback = readback;
+                m_readbackOption = PassAttachmentReadbackOption::Output;
+                m_attachmentReadback = readback;
                 AZStd::string readbackName = AZStd::string::format("%s_%s", m_outputColorAttachment->GetAttachmentId().GetCStr(), GetName().GetCStr());
-                return m_readback->ReadPassAttachment(m_outputColorAttachment.get(), AZ::Name(readbackName));
+                return m_attachmentReadback->ReadPassAttachment(m_outputColorAttachment.get(), AZ::Name(readbackName));
             }
             return false;
         }

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

@@ -8,6 +8,7 @@
 #include <Atom/RHI/FrameScheduler.h>
 #include <Atom/RHI.Reflect/Format.h>
 
+#include <Atom/RPI.Public/Pass/AttachmentReadback.h>
 #include <Atom/RPI.Public/Pass/PassSystemInterface.h>
 #include <Atom/RPI.Public/Pass/PassUtils.h>
 #include <Atom/RPI.Public/Pass/Specific/RenderToTexturePass.h>
@@ -77,17 +78,6 @@ namespace AZ
             params.m_viewportState = m_viewport;
 
             Base::FrameBeginInternal(params);
-
-            // for read back output
-            if (m_readback)
-            {
-                m_readback->FrameBegin(params);
-                if (m_readback->IsFinished())
-                {
-                    // Done reading. Remove the reference
-                    m_readback = nullptr;
-                }
-            }
         }
 
         void RenderToTexturePass::ResizeOutput(uint32_t width, uint32_t height)
@@ -118,9 +108,10 @@ namespace AZ
         {
             if (m_outputAttachment)
             {
-                m_readback = readback;
+                m_readbackOption = PassAttachmentReadbackOption::Output;
+                m_attachmentReadback = readback;
                 AZStd::string readbackName = AZStd::string::format("%s_%s", m_outputAttachment->GetAttachmentId().GetCStr(), GetName().GetCStr());
-                m_readback->ReadPassAttachment(m_outputAttachment.get(), AZ::Name(readbackName));
+                m_attachmentReadback->ReadPassAttachment(m_outputAttachment.get(), AZ::Name(readbackName));
             }
         }
 

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

@@ -8,6 +8,7 @@
 #include <Atom/RHI/FrameScheduler.h>
 
 #include <Atom/RPI.Public/WindowContext.h>
+#include <Atom/RPI.Public/Pass/AttachmentReadback.h>
 #include <Atom/RPI.Public/Pass/PassSystemInterface.h>
 #include <Atom/RPI.Public/Pass/Specific/SwapChainPass.h>
 #include <Atom/RHI/RHISystemInterface.h>
@@ -134,13 +135,6 @@ namespace AZ
             attachmentDatabase.ImportSwapChain(m_windowContext->GetSwapChainAttachmentId(), m_windowContext->GetSwapChain());
 
             ParentPass::FrameBeginInternal(params);
-
-            // Read swap chain for one frame. The reference can be released afterwards
-            if (m_swapChainReadback)
-            {
-                m_swapChainReadback->FrameBegin(params);
-                m_swapChainReadback = nullptr;
-            }
         }
         
         void SwapChainPass::OnWindowResized([[maybe_unused]] uint32_t width, [[maybe_unused]] uint32_t height)
@@ -152,9 +146,11 @@ namespace AZ
         {
             if (m_swapChainAttachment)
             {
-                m_swapChainReadback = readback;
+                m_readbackOption = PassAttachmentReadbackOption::Output;
+                m_attachmentReadback = readback;
+
                 AZStd::string readbackName = AZStd::string::format("%s_%s", m_swapChainAttachment->GetAttachmentId().GetCStr(), GetName().GetCStr());
-                m_swapChainReadback->ReadPassAttachment(m_swapChainAttachment.get(), AZ::Name(readbackName));
+                m_attachmentReadback->ReadPassAttachment(m_swapChainAttachment.get(), AZ::Name(readbackName));
             }
         }
 

Gems/Atom/Utils/Code/Include/Atom/Utils/ImGuiPassTree.h

@@ -32,11 +32,14 @@ namespace AZ
 
             void ReadbackCallback(const AZ::RPI::AttachmentReadback::ReadbackResult& readbackResult);
 
+            void DrawPassAttachments(AZ::RPI::Pass* pass);
+
             bool m_previewAttachment = false;
             bool m_showAttachments = false;
 
             AZ::RPI::Pass* m_selectedPass = nullptr;
             AZ::RHI::AttachmentId m_attachmentId;
+            AZ::Name m_slotName;
             bool m_selectedChanged = false;
 
             AZStd::shared_ptr<AZ::RPI::AttachmentReadback> m_readback;

Gems/Atom/Utils/Code/Include/Atom/Utils/ImGuiPassTree.inl

@@ -76,6 +76,7 @@ namespace AZ::Render
                 {
                     m_selectedChanged = true;
                     m_attachmentId = AZ::RHI::AttachmentId{};
+                    m_slotName = AZ::Name{};
                 }
             }
 
@@ -88,13 +89,12 @@ namespace AZ::Render
                     m_readback->SetCallback(AZStd::bind(&ImGuiPassTree::ReadbackCallback, this, AZStd::placeholders::_1));
                 }
 
-                if (m_selectedPass && !m_attachmentId.IsEmpty())
+                if (m_selectedPass && !m_slotName.IsEmpty())
                 {
-                    AZ::RPI::RenderPass* renderPass = azrtti_cast<AZ::RPI::RenderPass*>(m_selectedPass);
+                    bool readbackResult = m_selectedPass->ReadbackAttachment(m_readback, m_slotName);
-                    if (renderPass)
+                    if (!readbackResult)
                     {
-                        AZ::RPI::PassAttachment* attachment = FindPassAttachment(renderPass, m_attachmentId);
+                        AZ_Error("ImGuiPassTree", false, "Failed to readback attachment from pass [%s] slot [%s]", m_selectedPass->GetName().GetCStr(), m_slotName.GetCStr());
-                        renderPass->ReadbackAttachment(m_readback, attachment);
                     }
                 }
             }
@@ -144,6 +144,77 @@ namespace AZ::Render
         }
         ImGui::End();
     }
+        
+    inline void ImGuiPassTree::DrawPassAttachments(AZ::RPI::Pass* pass)
+    {
+        for (const auto& binding : pass->GetAttachmentBindings())
+        {
+            // Binding info: [slot type] [slot name]
+            AZStd::string label = AZStd::string::format("[%s] [%s]", AZ::RPI::ToString(binding.m_slotType),
+                binding.m_name.GetCStr());
+
+            // Append attachment info if the attachment exists
+            if (binding.m_attachment)
+            {
+                AZ::RHI::AttachmentType type = binding.m_attachment->GetAttachmentType();
+
+                // Append attachment info: [attachment type] attachment name
+                label += AZStd::string::format(" [%s] %s",
+                    AZ::RHI::ToString(type),
+                    binding.m_attachment->m_name.GetCStr());
+
+                if (type == AZ::RHI::AttachmentType::Image)
+                {
+                    // Append image info: [format] [size] [msaa]
+                    AZ::RHI::ImageDescriptor descriptor;
+                    if (binding.m_attachment->m_importedResource)
+                    {
+                        AZ::RPI::Image* image = static_cast<AZ::RPI::Image*>(binding.m_attachment->m_importedResource.get());
+                        descriptor = image->GetRHIImage()->GetDescriptor();
+                    }
+                    else
+                    {
+                        descriptor = binding.m_attachment->m_descriptor.m_image;
+                    }
+                    auto format = descriptor.m_format;
+                    auto size = descriptor.m_size;
+                    label += AZStd::string::format(" [%s] [%dx%d]", AZ::RHI::ToString(format), size.m_width, size.m_height);
+
+                    if (descriptor.m_multisampleState.m_samples > 1)
+                    {
+                        if (descriptor.m_multisampleState.m_customPositionsCount > 0)
+                        {
+                            label += AZStd::string::format(" [MSAA_Custom_%dx]", descriptor.m_multisampleState.m_samples);
+                        }
+                        else
+                        {
+                            label += AZStd::string::format(" [MSAA_%dx]", descriptor.m_multisampleState.m_samples);
+                        }
+                    }
+                }
+                else if (type == AZ::RHI::AttachmentType::Buffer)
+                {
+                    // Append buffer info: [size]
+                    auto size = binding.m_attachment->m_descriptor.m_buffer.m_byteCount;
+                    label += AZStd::string::format(" [%llu]", size);
+                }
+
+                if (Scriptable_ImGui::Selectable(label.c_str(), m_attachmentId == binding.m_attachment->GetAttachmentId()))
+                {
+                    m_selectedPass = pass;
+                    m_attachmentId = binding.m_attachment->GetAttachmentId();
+                    m_slotName = binding.m_name;
+                    m_selectedChanged = true;
+                }
+            }
+            else
+            {
+                // Only draw text (not selectable) if there is no attachment binded to the slot.
+                ImGui::Text(label.c_str());
+            }
+        }
+
+    }
 
     inline void ImGuiPassTree::DrawTreeView(AZ::RPI::Pass* pass)
     {
@@ -164,6 +235,7 @@ namespace AZ::Render
                 {
                     m_selectedPass = pass;
                     m_attachmentId = AZ::RHI::AttachmentId{};
+                    m_slotName = AZ::Name{};
                     m_selectedChanged = true;
                 }
             }
@@ -179,76 +251,13 @@ namespace AZ::Render
                 {
                     m_selectedPass = pass;
                     m_attachmentId = AZ::RHI::AttachmentId{};
+                    m_slotName = AZ::Name{};
                     m_selectedChanged = true;
                 }
 
                 if (nodeOpen)
                 {
-                    for (const auto& binding : pass->GetAttachmentBindings())
+                    DrawPassAttachments(pass);
-                    {
-                        // Binding info: [slot type] [slot name]
-                        AZStd::string label = AZStd::string::format("[%s] [%s]", AZ::RPI::ToString(binding.m_slotType),
-                            binding.m_name.GetCStr());
-
-                        // Append attachment info if the attachment exists
-                        if (binding.m_attachment)
-                        {
-                            AZ::RHI::AttachmentType type = binding.m_attachment->GetAttachmentType();
-
-                            // Append attachment info: [attachment type] attachment name
-                            label += AZStd::string::format(" [%s] %s",
-                                AZ::RHI::ToString(type),
-                                binding.m_attachment->m_name.GetCStr());
-
-                            if (type == AZ::RHI::AttachmentType::Image)
-                            {
-                                // Append image info: [format] [size] [msaa]
-                                AZ::RHI::ImageDescriptor descriptor;
-                                if (binding.m_attachment->m_importedResource)
-                                {
-                                    AZ::RPI::Image* image = static_cast<AZ::RPI::Image*>(binding.m_attachment->m_importedResource.get());
-                                    descriptor = image->GetRHIImage()->GetDescriptor();
-                                }
-                                else
-                                {
-                                    descriptor = binding.m_attachment->m_descriptor.m_image;
-                                }
-                                auto format = descriptor.m_format;
-                                auto size = descriptor.m_size;
-                                label += AZStd::string::format(" [%s] [%dx%d]", AZ::RHI::ToString(format), size.m_width, size.m_height);
-
-                                if (descriptor.m_multisampleState.m_samples > 1)
-                                {
-                                    if (descriptor.m_multisampleState.m_customPositionsCount > 0)
-                                    {
-                                        label += AZStd::string::format(" [MSAA_Custom_%dx]", descriptor.m_multisampleState.m_samples);
-                                    }
-                                    else
-                                    {
-                                        label += AZStd::string::format(" [MSAA_%dx]", descriptor.m_multisampleState.m_samples);
-                                    }
-                                }
-                            }
-                            else if (type == AZ::RHI::AttachmentType::Buffer)
-                            {
-                                // Append buffer info: [size]
-                                auto size = binding.m_attachment->m_descriptor.m_buffer.m_byteCount;
-                                label += AZStd::string::format(" [%llu]", size);
-                            }
-
-                            if (Scriptable_ImGui::Selectable(label.c_str(), m_attachmentId == binding.m_attachment->GetAttachmentId()))
-                            {
-                                m_selectedPass = pass;
-                                m_attachmentId = binding.m_attachment->GetAttachmentId();
-                                m_selectedChanged = true;
-                            }
-                        }
-                        else
-                        {
-                            // Only draw text (not selectable) if there is no attachment binded to the slot.
-                            ImGui::Text(label.c_str());
-                        }
-                    }
 
                     Scriptable_ImGui::TreePop();
                 }
@@ -266,11 +275,13 @@ namespace AZ::Render
             {
                 m_selectedPass = pass;
                 m_attachmentId = AZ::RHI::AttachmentId{};
+                m_slotName = AZ::Name{};
                 m_selectedChanged = true;
             }
 
             if (nodeOpen)
             {
+                DrawPassAttachments(pass);
                 for (const auto& child : asParent->GetChildren())
                 {
                     DrawTreeView(child.get());
@@ -354,6 +365,7 @@ namespace AZ::Render
 
         m_selectedPass = nullptr;
         m_attachmentId = AZ::RHI::AttachmentId{};
+         m_slotName = AZ::Name{};
         m_selectedChanged = false;
         m_readback = nullptr;
         m_previewPass = nullptr;

Gems/AtomLyIntegration/CommonFeatures/Code/Source/Thumbnails/Rendering/ThumbnailRendererSteps/CaptureStep.cpp

@@ -114,7 +114,7 @@ namespace AZ
                     Render::FrameCaptureRequestBus::BroadcastResult(
                         startedCapture,
                         &Render::FrameCaptureRequestBus::Events::CapturePassAttachmentWithCallback,
-                        m_context->GetData()->m_passHierarchy, AZStd::string("Output"), readbackCallback);
+                        m_context->GetData()->m_passHierarchy, AZStd::string("Output"), readbackCallback, RPI::PassAttachmentReadbackOption::Output);
                     // Reset the capture flag if the capture request was successful. Otherwise try capture it again next tick.
                     if (startedCapture)
                     {