Add memory heaps for tracking to some of Atom assets (#18424)

* Add Memory Heaps to some of Atom assets

- Add allocator to ImageMipChainAsset, ShaderStageFunction, StreamingImageAsset and BufferAsset
- Add conversion for old type
- Add ImgGuiHeapMemory profiler to visualize all heap allocators

Signed-off-by: Akio Gaule <[email protected]>

Code/Framework/AzCore/AzCore/Memory/AllocatorManager.cpp

@@ -784,7 +784,13 @@ namespace AZ
 
             if (outStats)
             {
-                outStats->emplace(outStats->end(), allocator->GetName(), allocator->NumAllocatedBytes(), allocator->Capacity());
+                const char* parentName = "";
+                if (auto childAllocatorSchema = azrtti_cast<AZ::ChildAllocatorSchemaBase*>(allocator); childAllocatorSchema != nullptr)
+                {
+                    auto parentAllocator = childAllocatorSchema->GetParentAllocator();
+                    parentName = parentAllocator != nullptr ? parentAllocator->GetName() : "";
+                }
+                outStats->emplace(outStats->end(), allocator->GetName(), parentName, allocator->NumAllocatedBytes(), allocator->Capacity());
             }
         }
     }

Code/Framework/AzCore/AzCore/Memory/AllocatorManager.h

@@ -107,13 +107,15 @@ namespace AZ
 
         struct AllocatorStats
         {
-            AllocatorStats(const char* name, size_t allocatedBytes, size_t capacityBytes)
+            AllocatorStats(const char* name, const char* parentName, size_t allocatedBytes, size_t capacityBytes)
                 : m_name(name)
+                , m_parentName(parentName)
                 , m_allocatedBytes(allocatedBytes)
                 , m_capacityBytes(capacityBytes)
             {}
 
             AZStd::string m_name;
+            AZStd::string m_parentName;
             size_t m_allocatedBytes;
             size_t m_capacityBytes;
         };

Gems/Atom/Asset/ImageProcessingAtom/Code/Source/ImageBuilderComponent.cpp

@@ -74,7 +74,7 @@ namespace ImageProcessingAtom
         builderDescriptor.m_busId = azrtti_typeid<ImageBuilderWorker>();
         builderDescriptor.m_createJobFunction = AZStd::bind(&ImageBuilderWorker::CreateJobs, &m_imageBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
         builderDescriptor.m_processJobFunction = AZStd::bind(&ImageBuilderWorker::ProcessJob, &m_imageBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
-        builderDescriptor.m_version = 34;   // change ImageMipchainAsset reference to NoLoad
+        builderDescriptor.m_version = 35;   // Added MipmapChain and StreamingImage allocator
         builderDescriptor.m_analysisFingerprint = ImageProcessingAtom::BuilderSettingManager::Instance()->GetAnalysisFingerprint();
         m_imageBuilder.BusConnect(builderDescriptor.m_busId);
         AssetBuilderSDK::AssetBuilderBus::Broadcast(&AssetBuilderSDK::AssetBuilderBusTraits::RegisterBuilderInformation, builderDescriptor);

Gems/Atom/Asset/Shader/Code/Source/Editor/AzslShaderBuilderSystemComponent.cpp

@@ -83,7 +83,7 @@ namespace AZ
             // Register Shader Asset Builder
             AssetBuilderSDK::AssetBuilderDesc shaderAssetBuilderDescriptor;
             shaderAssetBuilderDescriptor.m_name = "Shader Asset Builder";
-            shaderAssetBuilderDescriptor.m_version = 125; // Add subpass input support
+            shaderAssetBuilderDescriptor.m_version = 126; // ShaderStageFunction allocator
             shaderAssetBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern( AZStd::string::format("*.%s", RPI::ShaderSourceData::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
             shaderAssetBuilderDescriptor.m_busId = azrtti_typeid<ShaderAssetBuilder>();
             shaderAssetBuilderDescriptor.m_createJobFunction = AZStd::bind(&ShaderAssetBuilder::CreateJobs, &m_shaderAssetBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
@@ -108,7 +108,7 @@ namespace AZ
                 shaderVariantAssetBuilderDescriptor.m_name = "Shader Variant Asset Builder";
                 // Both "Shader Variant Asset Builder" and "Shader Asset Builder" produce ShaderVariantAsset products. If you update
                 // ShaderVariantAsset you will need to update BOTH version numbers, not just "Shader Variant Asset Builder".
-                shaderVariantAssetBuilderDescriptor.m_version = 42; // Add batching to shader variant assets
+                shaderVariantAssetBuilderDescriptor.m_version = 43; // ShaderStageFunction allocator
                 shaderVariantAssetBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", HashedVariantListSourceData::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
                 shaderVariantAssetBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", HashedVariantInfoSourceData::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
                 shaderVariantAssetBuilderDescriptor.m_busId = azrtti_typeid<ShaderVariantAssetBuilder>();
@@ -121,7 +121,7 @@ namespace AZ
                 // Register Shader Variant List Builder
                 AssetBuilderSDK::AssetBuilderDesc shaderVariantListBuilderDescriptor;
                 shaderVariantListBuilderDescriptor.m_name = "Shader Variant List Builder";
-                shaderVariantListBuilderDescriptor.m_version = 3; // Add batching to shader variant assets
+                shaderVariantListBuilderDescriptor.m_version = 4; // ShaderStageFunction allocator
                 shaderVariantListBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", RPI::ShaderVariantListSourceData::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
                 shaderVariantListBuilderDescriptor.m_busId = azrtti_typeid<ShaderVariantListBuilder>();
                 shaderVariantListBuilderDescriptor.m_createJobFunction = AZStd::bind(&ShaderVariantListBuilder::CreateJobs, &m_shaderVariantListBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
@@ -135,7 +135,7 @@ namespace AZ
             // Register Precompiled Shader Builder
             AssetBuilderSDK::AssetBuilderDesc precompiledShaderBuilderDescriptor;
             precompiledShaderBuilderDescriptor.m_name = "Precompiled Shader Builder";
-            precompiledShaderBuilderDescriptor.m_version = 14; // Fixing warnings from having both source and job dependencies
+            precompiledShaderBuilderDescriptor.m_version = 15; // ShaderStageFunction allocator
             precompiledShaderBuilderDescriptor.m_patterns.push_back(AssetBuilderSDK::AssetBuilderPattern(AZStd::string::format("*.%s", AZ::PrecompiledShaderBuilder::Extension), AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
             precompiledShaderBuilderDescriptor.m_busId = azrtti_typeid<PrecompiledShaderBuilder>();
             precompiledShaderBuilderDescriptor.m_createJobFunction = AZStd::bind(&PrecompiledShaderBuilder::CreateJobs, &m_precompiledShaderBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);

Gems/Atom/RHI/Code/Include/Atom/RHI.Reflect/Allocators.h

@@ -0,0 +1,18 @@
+/*
+ * Copyright (c) Contributors to the Open 3D Engine Project.
+ * For complete copyright and license terms please see the LICENSE at the root of this distribution.
+ *
+ * SPDX-License-Identifier: Apache-2.0 OR MIT
+ *
+ */
+
+#pragma once
+
+#include <AzCore/Memory/SystemAllocator.h>
+#include <AzCore/Memory/ChildAllocatorSchema.h>
+
+namespace AZ::RHI
+{
+    AZ_CHILD_ALLOCATOR_WITH_NAME(
+        ShaderStageFunctionAllocator, "Atom::ShaderStageFunctionAllocator", "{15F285F1-74D5-4FAE-8CE4-B7D235A92F23}", AZ::SystemAllocator);
+} // namespace RHI

Gems/Atom/RHI/Code/Include/Atom/RHI.Reflect/ShaderStageFunction.h

@@ -11,6 +11,7 @@
 #include <AzCore/Name/Name.h>
 #include <AzCore/Utils/TypeHash.h>
 #include <Atom/RHI.Reflect/ShaderStages.h>
+#include <Atom/RHI.Reflect/Allocators.h>
 
 namespace AZ
 {
@@ -31,6 +32,8 @@ namespace AZ::RHI
         : public AZStd::intrusive_base
     {
     public:
+        using Allocator = ShaderStageFunctionAllocator_for_std_t;
+
         AZ_RTTI(ShaderStageFunction, "{D7D800D1-4A14-4C3F-A7E4-B2645137FAC6}");
         static void Reflect(AZ::ReflectContext* context);
 

Gems/Atom/RHI/Code/Source/RHI.Private/FactoryManagerSystemComponent.cpp

@@ -20,7 +20,7 @@
 
 #include <AzFramework/API/ApplicationAPI.h>
 #include <AzFramework/IO/LocalFileIO.h>
-
+#include <AzFramework/StringFunc/StringFunc.h>
 
 namespace AZ::RHI
 {
@@ -38,20 +38,27 @@ namespace AZ::RHI
                 ec->Class<FactoryManagerSystemComponent>("Atom RHI Manager", "Atom Renderer")
                     ->ClassElement(AZ::Edit::ClassElements::EditorData, "")
                     ->Attribute(AZ::Edit::Attributes::AutoExpand, true)
-                    ->DataElement(AZ::Edit::UIHandlers::Default, &FactoryManagerSystemComponent::m_factoriesPriority, "RHI Priority list", "Priorities for RHI Implementations")
-                    ->DataElement(AZ::Edit::UIHandlers::ComboBox, &FactoryManagerSystemComponent::m_validationMode, "Validation Layer Mode", "Set the validation mode for the RHI. It only applies for non release builds")
-                    ->Attribute(AZ::Edit::Attributes::EnumValues,
-                        AZStd::vector<AZ::Edit::EnumConstant<RHI::ValidationMode>>
-                {
-                    AZ::Edit::EnumConstant<RHI::ValidationMode>(RHI::ValidationMode::Disabled,
-                        "Disable - Disables any validation."),
-                        AZ::Edit::EnumConstant<RHI::ValidationMode>(RHI::ValidationMode::Enabled,
-                            "Enable - Enables warnings and errors validation messages."),
-                        AZ::Edit::EnumConstant<RHI::ValidationMode>(RHI::ValidationMode::Verbose,
-                            "Verbose - Enables warnings, error and information messages."),
-                        AZ::Edit::EnumConstant<RHI::ValidationMode>(RHI::ValidationMode::GPU,
-                            "GPU - Enables based validation."),
-                });
+                    ->DataElement(
+                        AZ::Edit::UIHandlers::Default,
+                        &FactoryManagerSystemComponent::m_factoriesPriority,
+                        "RHI Priority list",
+                        "Priorities for RHI Implementations")
+                    ->DataElement(
+                        AZ::Edit::UIHandlers::ComboBox,
+                        &FactoryManagerSystemComponent::m_validationMode,
+                        "Validation Layer Mode",
+                        "Set the validation mode for the RHI. It only applies for non release builds")
+                    ->Attribute(
+                        AZ::Edit::Attributes::EnumValues,
+                        AZStd::vector<AZ::Edit::EnumConstant<RHI::ValidationMode>>{
+                            AZ::Edit::EnumConstant<RHI::ValidationMode>(
+                                RHI::ValidationMode::Disabled, "Disable - Disables any validation."),
+                            AZ::Edit::EnumConstant<RHI::ValidationMode>(
+                                RHI::ValidationMode::Enabled, "Enable - Enables warnings and errors validation messages."),
+                            AZ::Edit::EnumConstant<RHI::ValidationMode>(
+                                RHI::ValidationMode::Verbose, "Verbose - Enables warnings, error and information messages."),
+                            AZ::Edit::EnumConstant<RHI::ValidationMode>(RHI::ValidationMode::GPU, "GPU - Enables based validation."),
+                        });
             }
                 
         }

Gems/Atom/RHI/Code/atom_rhi_reflect_files.cmake

@@ -131,4 +131,5 @@ set(FILES
     Include/Atom/RHI.Reflect/ShaderInputNameIndex.h
     Source/RHI.Reflect/ShaderInputNameIndex.cpp
     Include/Atom/RHI.Reflect/VariableRateShadingEnums.h
+    Include/Atom/RHI.Reflect/Allocators.h
 )

Gems/Atom/RHI/DX12/Code/Include/Atom/RHI.Reflect/DX12/ShaderStageFunction.h

@@ -30,7 +30,7 @@ namespace AZ
 
     namespace DX12
     {
-        using ShaderByteCode = AZStd::vector<uint8_t>;
+        using ShaderByteCode = AZStd::vector<uint8_t, RHI::ShaderStageFunction::Allocator>;
         using ShaderByteCodeView = AZStd::span<const uint8_t>;
 
         //! Sentinel value used when patching shaders for specialization constants
@@ -55,13 +55,13 @@ namespace AZ
         {
         public:
             AZ_RTTI(ShaderStageFunction, "{1BAEE536-96CA-4AEB-BA73-D5D72EE35B45}", RHI::ShaderStageFunction);
-            AZ_CLASS_ALLOCATOR(ShaderStageFunction, AZ::SystemAllocator);
+            AZ_CLASS_ALLOCATOR_DECL
             static void Reflect(AZ::ReflectContext* context);
 
             static RHI::Ptr<ShaderStageFunction> Create(RHI::ShaderStage shaderStage);
 
             /// Assigns byte code to the function.
-            void SetByteCode(uint32_t subStageIndex, const ShaderByteCode& byteCode);
+            void SetByteCode(uint32_t subStageIndex, const AZStd::vector<uint8_t>& byteCode);
 
             /// Returns the assigned byte code.
             ShaderByteCodeView GetByteCode(uint32_t subStageIndex = 0) const;

Gems/Atom/RHI/DX12/Code/Source/RHI.Builders/ShaderPlatformInterface.cpp

@@ -47,7 +47,7 @@ namespace AZ
         {
             RHI::Ptr<ShaderStageFunction> newShaderStageFunction =  ShaderStageFunction::Create(RHI::ToRHIShaderStage(stageDescriptor.m_stageType));
 
-            const DX12::ShaderByteCode& byteCode = stageDescriptor.m_byteCode;
+            const auto& byteCode = stageDescriptor.m_byteCode;
             const int byteCodeIndex = 0;
             newShaderStageFunction->SetByteCode(byteCodeIndex, byteCode);
 

Gems/Atom/RHI/DX12/Code/Source/RHI.Reflect/ShaderStageFunction.cpp

@@ -14,12 +14,47 @@ namespace AZ
 {
     namespace DX12
     {
+        AZ_CLASS_ALLOCATOR_IMPL(ShaderStageFunction, RHI::ShaderStageFunctionAllocator)
+
+        static bool ConvertOldVersions(AZ::SerializeContext& context, AZ::SerializeContext::DataElementNode& classElement)
+        {
+            if (classElement.GetVersion() < 3)
+            {
+                // We need to convert the shader byte code because we added a custom allocator
+                // and the serialization system doens't automatically convert between two different classes
+                auto crc32 = AZ::Crc32("m_byteCodes");
+                auto* arrayElement = classElement.FindSubElement(crc32);
+                if (arrayElement)
+                {
+                    // When we convert all sub elements are removed, so we first need to
+                    // get the child data, so we can re-insert it after we convert the element
+                    AZStd::array<AZStd::vector<uint8_t>, ShaderSubStageCountMax> oldData;
+                    if (classElement.GetChildData(crc32, oldData))
+                    {
+                        // Convert the array with the new allocator
+                        arrayElement->Convert(context, AZ::AzTypeInfo<AZStd::array<ShaderByteCode, ShaderSubStageCountMax>>::Uuid());
+                        for (const auto& element : oldData)
+                        {
+                            // Convert each vector and re add it. During convertion all sub elements were removed.
+                            ShaderByteCode newData(element.size());
+                            ::memcpy(newData.data(), element.data(), element.size());
+                            arrayElement->AddElementWithData<ShaderByteCode>(context, "element", newData);
+                        }
+                    }
+                }
+            }
+            return true;
+        }
+
         void ShaderStageFunction::Reflect(ReflectContext* context)
         {
             if (SerializeContext* serializeContext = azrtti_cast<SerializeContext*>(context))
             {
+                // Need to register the old type with the Serializer so we can read it in order to convert it
+                serializeContext->RegisterGenericType<AZStd::array<AZStd::vector<uint8_t>, ShaderSubStageCountMax>>();
+
                 serializeContext->Class<ShaderStageFunction, RHI::ShaderStageFunction>()
-                    ->Version(2)
+                    ->Version(3, &ConvertOldVersions)
                     ->Field("m_byteCodes", &ShaderStageFunction::m_byteCodes)
                     ->Field("m_specializationOffsets", &ShaderStageFunction::m_specializationOffsets);
             }
@@ -35,9 +70,10 @@ namespace AZ
             return aznew ShaderStageFunction(shaderStage);
         }
 
-        void ShaderStageFunction::SetByteCode(uint32_t subStageIndex, const ShaderByteCode& byteCode)
+        void ShaderStageFunction::SetByteCode(uint32_t subStageIndex, const AZStd::vector<uint8_t>& byteCode)
         {
-            m_byteCodes[subStageIndex] = byteCode;
+            m_byteCodes[subStageIndex].resize(byteCode.size());
+            ::memcpy(m_byteCodes[subStageIndex].data(), byteCode.data(), byteCode.size());
         }
 
         ShaderByteCodeView ShaderStageFunction::GetByteCode(uint32_t subStageIndex) const

Gems/Atom/RHI/Metal/Code/Include/Atom/RHI.Reflect/Metal/ShaderStageFunction.h

@@ -32,14 +32,14 @@ namespace AZ
     namespace Metal
     {
         using ShaderSourceCode = AZStd::vector<char>;
-        using ShaderByteCode = AZStd::vector<uint8_t>;
+        using ShaderByteCode = AZStd::vector<uint8_t, RHI::ShaderStageFunction::Allocator>;
 
         class ShaderStageFunction
             : public RHI::ShaderStageFunction
         {
         public:
             AZ_RTTI(ShaderStageFunction, "{44E51B8E-CFEE-4A63-8DC2-65CDCA0E373B}", RHI::ShaderStageFunction);
-            AZ_CLASS_ALLOCATOR(ShaderStageFunction, AZ::SystemAllocator);
+            AZ_CLASS_ALLOCATOR_DECL
 
             static void Reflect(AZ::ReflectContext* context);
 
@@ -53,7 +53,7 @@ namespace AZ
             const AZStd::string& GetSourceCode() const;
 
             /// Assigns byte code and byte code length.
-            void SetByteCode(const ShaderByteCode& byteCode);
+            void SetByteCode(const AZStd::vector<uint8_t>& byteCode);
 
             /// Assigns entry function name.
             void SetEntryFunctionName(AZStd::string_view entryFunctionName);

Gems/Atom/RHI/Metal/Code/Source/RHI.Builders/ShaderPlatformInterface.cpp

@@ -126,7 +126,7 @@ namespace AZ
                 newShaderStageFunction->SetSourceCode(sourceCode);
             }
             
-            const Metal::ShaderByteCode& byteCode = stageDescriptor.m_byteCode;
+            const auto& byteCode = stageDescriptor.m_byteCode;
             const AZStd::string& entryFunctionName = stageDescriptor.m_entryFunctionName;
             newShaderStageFunction->SetByteCode(byteCode);
             newShaderStageFunction->SetEntryFunctionName(entryFunctionName);

Gems/Atom/RHI/Metal/Code/Source/RHI.Reflect/ShaderStageFunction.cpp

@@ -14,12 +14,42 @@ namespace AZ
 {
     namespace Metal
     {
+        AZ_CLASS_ALLOCATOR_IMPL(ShaderStageFunction, RHI::ShaderStageFunctionAllocator)
+
+        static bool ConvertOldVersions(AZ::SerializeContext& context, AZ::SerializeContext::DataElementNode& classElement)
+        {
+            if (classElement.GetVersion() < 4)
+            {
+                auto crc32 = AZ::Crc32("m_byteCode");
+                auto* vectorElement = classElement.FindSubElement(crc32);
+                if (vectorElement)
+                {
+                    // Get the old data
+                    AZStd::vector<uint8_t> oldData;
+                    if (vectorElement->GetData(oldData))
+                    {
+                        // Convert the vector with the new allocator
+                        vectorElement->Convert(context, AZ::AzTypeInfo<ShaderByteCode>::Uuid());
+                        // Copy old data to new data
+                        ShaderByteCode newData(oldData.size());
+                        ::memcpy(newData.data(), oldData.data(), newData.size());
+                        // Set the new data
+                        vectorElement->SetData(context, newData);
+                    }
+                }
+            }
+            return true;
+        }
+
         void ShaderStageFunction::Reflect(ReflectContext* context)
         {
             if (SerializeContext* serializeContext = azrtti_cast<SerializeContext*>(context))
             {
+                // Need to register the old type with the Serializer so we can read it in order to convert it
+                serializeContext->RegisterGenericType<AZStd::vector<uint8_t>>();
+
                 serializeContext->Class<ShaderStageFunction, RHI::ShaderStageFunction>()
-                    ->Version(3)
+                    ->Version(4, &ConvertOldVersions)
                     ->Field("m_sourceCode", &ShaderStageFunction::m_sourceCode)
                     ->Field("m_byteCode", &ShaderStageFunction::m_byteCode)
                     ->Field("m_byteCodeLength", &ShaderStageFunction::m_byteCodeLength)
@@ -51,9 +81,10 @@ namespace AZ
             return m_sourceCode;
         }
 
-        void ShaderStageFunction::SetByteCode(const ShaderByteCode& byteCode)
+        void ShaderStageFunction::SetByteCode(const AZStd::vector<uint8_t>& byteCode)
         {
-            m_byteCode = byteCode;
+            m_byteCode.resize(byteCode.size());
+            ::memcpy(m_byteCode.data(), byteCode.data(), m_byteCode.size());
             m_byteCodeLength = aznumeric_cast<uint32_t>(byteCode.size());
         }
 

Gems/Atom/RHI/Vulkan/Code/Include/Atom/RHI.Reflect/Vulkan/ShaderStageFunction.h

@@ -32,7 +32,7 @@ namespace AZ
 
     namespace Vulkan
     {
-        using ShaderByteCode = AZStd::vector<uint8_t>;
+        using ShaderByteCode = AZStd::vector<uint8_t, RHI::ShaderStageFunction::Allocator>;
         using ShaderByteCodeView = AZStd::span<const uint8_t>;
 
         /**
@@ -53,14 +53,14 @@ namespace AZ
 
         public:
             AZ_RTTI(ShaderStageFunction, "{A606478A-97E9-402D-A776-88EE72DAC6F9}", RHI::ShaderStageFunction);
-            AZ_CLASS_ALLOCATOR(ShaderStageFunction, AZ::SystemAllocator);
+            AZ_CLASS_ALLOCATOR_DECL
 
             static void Reflect(AZ::ReflectContext* context);
 
             static RHI::Ptr<ShaderStageFunction> Create(RHI::ShaderStage shaderStage);
 
             /// Assigns byte code to the function.
-            void SetByteCode(uint32_t subStageIndex, const ShaderByteCode& byteCode, const AZStd::string_view& entryFunctionName);
+            void SetByteCode(uint32_t subStageIndex, const AZStd::vector<uint8_t>& byteCode, const AZStd::string_view& entryFunctionName);
 
             /// Returns the assigned bytecode.
             ShaderByteCodeView GetByteCode(uint32_t subStageIndex) const;

Gems/Atom/RHI/Vulkan/Code/Source/RHI.Builders/ShaderPlatformInterface.cpp

@@ -40,7 +40,7 @@ namespace AZ
         {
             RHI::Ptr<ShaderStageFunction> newShaderStageFunction = ShaderStageFunction::Create(RHI::ToRHIShaderStage(stageDescriptor.m_stageType));
 
-            const Vulkan::ShaderByteCode& byteCode = stageDescriptor.m_byteCode;
+            const auto& byteCode = stageDescriptor.m_byteCode;
             const AZStd::string& entryFunctionName = stageDescriptor.m_entryFunctionName;
             const int byteCodeIndex = 0;
             newShaderStageFunction->SetByteCode(byteCodeIndex, byteCode, entryFunctionName);

Gems/Atom/RHI/Vulkan/Code/Source/RHI.Reflect/ShaderStageFunction.cpp

@@ -15,12 +15,47 @@ namespace AZ
 {
     namespace Vulkan
     {
+        AZ_CLASS_ALLOCATOR_IMPL(ShaderStageFunction, RHI::ShaderStageFunctionAllocator)
+
+        static bool ConvertOldVersions(AZ::SerializeContext& context, AZ::SerializeContext::DataElementNode& classElement)
+        {
+            if (classElement.GetVersion() < 3)
+            {
+                // We need to convert the shader byte code because we added a custom allocator
+                // and the serialization system doens't automatically convert between two different classes
+                auto crc32 = AZ::Crc32("m_byteCodes");
+                auto* arrayElement = classElement.FindSubElement(crc32);
+                if (arrayElement)
+                {
+                    // When we convert all sub elements are removed, so we first need to
+                    // get the child data, so we can re-insert it after we convert the element
+                    AZStd::array<AZStd::vector<uint8_t>, ShaderSubStageCountMax> oldData;
+                    if (classElement.GetChildData(crc32, oldData))
+                    {
+                        // Convert the array with the new allocator
+                        arrayElement->Convert(context, AZ::AzTypeInfo<AZStd::array<ShaderByteCode, ShaderSubStageCountMax>>::Uuid());
+                        for (const auto& element : oldData)
+                        {
+                            // Convert each vector and re add it. During convertion all sub elements were removed.
+                            ShaderByteCode newData(element.size());
+                            ::memcpy(newData.data(), element.data(), element.size());
+                            arrayElement->AddElementWithData<ShaderByteCode>(context, "element", newData);
+                        }
+                    }
+                }
+            }
+            return true;
+        }
+
         void ShaderStageFunction::Reflect(ReflectContext* context)
         {
             if (SerializeContext* serializeContext = azrtti_cast<SerializeContext*>(context))
             {
+                // Need to register the old type with the Serializer so we can read it in order to convert it
+                serializeContext->RegisterGenericType<AZStd::array<AZStd::vector<uint8_t>, ShaderSubStageCountMax>>();
+
                 serializeContext->Class<ShaderStageFunction, RHI::ShaderStageFunction>()
-                    ->Version(2)
+                    ->Version(3, &ConvertOldVersions)
                     ->Field("m_byteCodes", &ShaderStageFunction::m_byteCodes)
                     ->Field("m_entryFunctionName", &ShaderStageFunction::m_entryFunctionNames);
             }
@@ -36,10 +71,12 @@ namespace AZ
             return aznew ShaderStageFunction(shaderStage);
         }
 
-        void ShaderStageFunction::SetByteCode(uint32_t subStageIndex, const ShaderByteCode& byteCode, const AZStd::string_view& entryFunctionName)
+        void ShaderStageFunction::SetByteCode(
+            uint32_t subStageIndex, const AZStd::vector<uint8_t>& byteCode, const AZStd::string_view& entryFunctionName)
         {
             AZ_Assert(subStageIndex < ShaderSubStageCountMax, "SubStage index is out of bound.");
-            m_byteCodes[subStageIndex] = byteCode;
+            m_byteCodes[subStageIndex].resize(byteCode.size());
+            ::memcpy(m_byteCodes[subStageIndex].data(), byteCode.data(), byteCode.size());
             m_entryFunctionNames[subStageIndex] = entryFunctionName;
         }
 

Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Allocators.h

@@ -0,0 +1,22 @@
+/*
+ * Copyright (c) Contributors to the Open 3D Engine Project.
+ * For complete copyright and license terms please see the LICENSE at the root of this distribution.
+ *
+ * SPDX-License-Identifier: Apache-2.0 OR MIT
+ *
+ */
+
+#pragma once
+
+#include <AzCore/Memory/SystemAllocator.h>
+#include <AzCore/Memory/ChildAllocatorSchema.h>
+
+namespace AZ::RPI
+{
+    AZ_CHILD_ALLOCATOR_WITH_NAME(
+        ImageMipChainAssetAllocator, "Atom::ImageMipChainAssetAllocator", "{E3F40786-8675-461A-BF96-70812B4CFAF5}", AZ::SystemAllocator);
+    AZ_CHILD_ALLOCATOR_WITH_NAME(
+        StreamingImageAssetAllocator, "Atom::StreamingImageAssetAllocator", "{FE311FD6-105B-484E-A603-C773BFF25BA6}", AZ::SystemAllocator);
+    AZ_CHILD_ALLOCATOR_WITH_NAME(
+        BufferAssetAllocator, "Atom::BufferAssetAllocator", "{3C9BE3B7-65C8-48E0-BF03-6BE228C1842F}", AZ::SystemAllocator);
+} // namespace RPI

Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Buffer/BufferAsset.h

@@ -13,6 +13,7 @@
 
 #include <Atom/RPI.Public/Buffer/BufferSystemInterface.h>
 
+#include <Atom/RPI.Reflect/Allocators.h>
 #include <Atom/RPI.Reflect/Asset/AssetHandler.h>
 #include <Atom/RPI.Reflect/ResourcePoolAsset.h>
 
@@ -39,7 +40,7 @@ namespace AZ
             static constexpr const char* Extension{ "azbuffer" };
 
             AZ_RTTI(BufferAsset, "{F6C5EA8A-1DB3-456E-B970-B6E2AB262AED}", Data::AssetData);
-            AZ_CLASS_ALLOCATOR(BufferAsset, AZ::SystemAllocator);
+            AZ_CLASS_ALLOCATOR_DECL;
 
             static void Reflect(AZ::ReflectContext* context);
 
@@ -60,6 +61,8 @@ namespace AZ
             const AZStd::string& GetName() const;
 
         private:
+            using Allocator = BufferAssetAllocator_for_std_t;
+
             // AssetData overrides...
             bool HandleAutoReload() override
             {
@@ -76,7 +79,7 @@ namespace AZ
 
             AZStd::string m_name;
 
-            AZStd::vector<uint8_t> m_buffer;
+            AZStd::vector<uint8_t, Allocator> m_buffer;
 
             RHI::BufferDescriptor m_bufferDescriptor;
 

Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Image/ImageMipChainAsset.h

@@ -9,7 +9,7 @@
 #pragma once
 
 #include <Atom/RPI.Reflect/Asset/AssetHandler.h>
-
+#include <Atom/RPI.Reflect/Allocators.h>
 #include <Atom/RHI.Reflect/ImageSubresource.h>
 
 #include <Atom/RHI/StreamingImagePool.h>
@@ -48,8 +48,10 @@ namespace AZ
             static constexpr const char* Group{ "Image" };
             static constexpr const char* Extension{ "imagemipchain" };
 
+            using Allocator = ImageMipChainAssetAllocator_for_std_t;
+
             AZ_RTTI(ImageMipChainAsset, "{CB403C8A-6982-4C9F-8090-78C9C36FBEDB}", Data::AssetData);
-            AZ_CLASS_ALLOCATOR(ImageMipChainAsset, AZ::SystemAllocator);
+            AZ_CLASS_ALLOCATOR_DECL
 
             static void Reflect(AZ::ReflectContext* context);
 
@@ -86,7 +88,6 @@ namespace AZ
             bool HandleAutoReload() override { return false; }
 
         private:
-
             // Copy content from another ImageMipChainAsset
             void CopyFrom(const ImageMipChainAsset& source);
 
@@ -100,7 +101,7 @@ namespace AZ
             MipSliceList m_mipSlices;
 
             // The list of subresource data, fixed up from serialization.
-            AZStd::vector<RHI::StreamingImageSubresourceData> m_subImageDatas;
+            AZStd::vector<RHI::StreamingImageSubresourceData, Allocator> m_subImageDatas;
 
             // [Serialized] Topology of sub-images in the mip group.
             uint16_t m_mipLevels = 0;
@@ -113,10 +114,10 @@ namespace AZ
             AZStd::array<RHI::DeviceImageSubresourceLayout, RHI::Limits::Image::MipCountMax> m_subImageLayouts;
 
             // [Serialized] Contains a flat list of sub-images which reference the flat data blob.
-            AZStd::vector<AZ::u64> m_subImageDataOffsets;
+            AZStd::vector<AZ::u64, Allocator> m_subImageDataOffsets;
 
             // [Serialized] Flat image data interpreted by m_subImages.
-            AZStd::vector<uint8_t> m_imageData;
+            AZStd::vector<uint8_t, Allocator> m_imageData;
         };
 
         class ImageMipChainAssetHandler final

Gems/Atom/RPI/Code/Include/Atom/RPI.Reflect/Image/StreamingImageAsset.h

@@ -57,8 +57,10 @@ namespace AZ
             static constexpr const char* Group{ "Image" };
             static constexpr const char* Extension{ "streamingimage" };
 
+            using Allocator = StreamingImageAssetAllocator_for_std_t;
+
             AZ_RTTI(StreamingImageAsset, "{3C96A826-9099-4308-A604-7B19ADBF8761}", ImageAsset);
-            AZ_CLASS_ALLOCATOR(StreamingImageAsset , AZ::SystemAllocator)
+            AZ_CLASS_ALLOCATOR_DECL
 
             static void Reflect(ReflectContext* context);
 
@@ -107,7 +109,8 @@ namespace AZ
             bool HasFullMipChainAssets() const;
 
             //! Returns the image tags
-            const AZStd::vector<AZ::Name>& GetTags() const;
+            using TagList = AZStd::vector<AZ::Name, Allocator>;
+            const TagList& GetTags() const;
 
             //! Removes up to `mipChainLevel` mipchains, reducing quality (used by the image tag system).
             //! The last mipchain won't be removed
@@ -153,7 +156,7 @@ namespace AZ
             //! that resides in the tail mip chain.
             const ImageMipChainAsset* GetImageMipChainAsset(AZ::u32 mipLevel) const;
 
-            AZStd::vector<AZ::Name> m_tags;
+            TagList m_tags;
         };
     }
 }

Gems/Atom/RPI/Code/Source/RPI.Builders/Pass/PassBuilder.cpp

@@ -55,7 +55,7 @@ namespace AZ
         {
             AssetBuilderSDK::AssetBuilderDesc builder;
             builder.m_name = PassBuilderJobKey;
-            builder.m_version = 17; // Add RHI::ScopeAttachmentUsage to PassTemplate
+            builder.m_version = 18; // Add Allocator to ShaderStageFunction
             builder.m_busId = azrtti_typeid<PassBuilder>();
             builder.m_createJobFunction = AZStd::bind(&PassBuilder::CreateJobs, this, AZStd::placeholders::_1, AZStd::placeholders::_2);
             builder.m_processJobFunction = AZStd::bind(&PassBuilder::ProcessJob, this, AZStd::placeholders::_1, AZStd::placeholders::_2);

Gems/Atom/RPI/Code/Source/RPI.Reflect/Buffer/BufferAsset.cpp

@@ -18,12 +18,14 @@ namespace AZ
 
     namespace RPI
     {
+        AZ_CLASS_ALLOCATOR_IMPL(BufferAsset, BufferAssetAllocator)
+
         void BufferAsset::Reflect(ReflectContext* context)
         {
             if (auto* serializeContext = azrtti_cast<SerializeContext*>(context))
             {
                 serializeContext->Class<BufferAsset>()
-                    ->Version(2)
+                    ->Version(3)
                     ->Field("Name", &BufferAsset::m_name)
                     ->Field("Buffer", &BufferAsset::m_buffer)
                     ->Field("BufferDescriptor", &BufferAsset::m_bufferDescriptor)

Gems/Atom/RPI/Code/Source/RPI.Reflect/Image/ImageMipChainAsset.cpp

@@ -7,6 +7,7 @@
  */
 
 #include <Atom/RPI.Reflect/Image/ImageMipChainAsset.h>
+#include <Atom/RPI.Reflect/Allocators.h>
 
 #include <AzCore/Serialization/SerializeContext.h>
 
@@ -14,12 +15,66 @@ namespace AZ
 {
     namespace RPI
     {
+
+        AZ_CLASS_ALLOCATOR_IMPL(ImageMipChainAsset, ImageMipChainAssetAllocator)
+
+        static bool ConvertOldVersions(AZ::SerializeContext& context, AZ::SerializeContext::DataElementNode& classElement)
+        {
+            if (classElement.GetVersion() < 1)
+            {
+                // We need to convert the vectors because we added a custom allocator
+                // and the serialization system doens't automatically convert between two different classes
+                {
+                    auto crc32 = AZ::Crc32("m_imageData");
+                    auto* vectorElement = classElement.FindSubElement(crc32);
+                    if (vectorElement)
+                    {
+                        // Get the old data
+                        AZStd::vector<uint8_t> oldData;
+                        if (vectorElement->GetData(oldData))
+                        {
+                            // Convert the vector with the new allocator
+                            vectorElement->Convert(context, AZ::AzTypeInfo<AZStd::vector<uint8_t, ImageMipChainAsset::Allocator>>::Uuid());
+                            // Copy old data to new data
+                            AZStd::vector<uint8_t, ImageMipChainAsset::Allocator> newData(oldData.size());
+                            ::memcpy(newData.data(), oldData.data(), newData.size());
+                            // Set the new data
+                            vectorElement->SetData(context, newData);
+                        }
+                    }
+                }
+                {
+                    auto crc32 = AZ::Crc32("m_subImageDataOffsets");
+                    auto* vectorElement = classElement.FindSubElement(crc32);
+                    if (vectorElement)
+                    {
+                        AZStd::vector<AZ::u64> oldData;
+                        if (classElement.GetChildData(crc32, oldData))
+                        {
+                            // Convert the vector with the new allocator
+                            vectorElement->Convert(context, AZ::AzTypeInfo<AZStd::vector<AZ::u64, ImageMipChainAsset::Allocator>>::Uuid());
+                            for (const auto& element : oldData)
+                            {
+                                // Convert each vector and re add it. During convertion all sub elements were removed.
+                                vectorElement->AddElementWithData<AZ::u64>(context, "element", element);
+                            }
+                        }
+                    }
+                }
+            }
+            return true;
+        }
+
         void ImageMipChainAsset::Reflect(ReflectContext* context)
         {
             if (auto* serializeContext = azrtti_cast<SerializeContext*>(context))
             {
+                // Need to register the old type with the Serializer so we can read it in order to convert it
+                serializeContext->RegisterGenericType<AZStd::vector<uint8_t>>();
+                serializeContext->RegisterGenericType<AZStd::vector<u64>>();
+
                 serializeContext->Class<ImageMipChainAsset, Data::AssetData>()
-                    ->Version(0)
+                    ->Version(1, &ConvertOldVersions)
                     ->Field("m_mipLevels", &ImageMipChainAsset::m_mipLevels)
                     ->Field("m_arraySize", &ImageMipChainAsset::m_arraySize)
                     ->Field("m_mipToSubImageOffset", &ImageMipChainAsset::m_mipToSubImageOffset)

Gems/Atom/RPI/Code/Source/RPI.Reflect/Image/StreamingImageAsset.cpp

@@ -7,6 +7,7 @@
  */
 
 #include <Atom/RPI.Reflect/Image/StreamingImageAsset.h>
+#include <Atom/RPI.Reflect/Allocators.h>
 
 #include <AzCore/Asset/AssetSerializer.h>
 #include <AzCore/Serialization/SerializeContext.h>
@@ -15,10 +16,41 @@ namespace AZ
 {
     namespace RPI
     {
+
+        AZ_CLASS_ALLOCATOR_IMPL(StreamingImageAsset, StreamingImageAssetAllocator)
+
+        static bool ConvertOldVersions(AZ::SerializeContext& context, AZ::SerializeContext::DataElementNode& classElement)
+        {
+            if (classElement.GetVersion() < 3)
+            {
+                auto crc32 = AZ::Crc32("m_tags");
+                auto* vectorElement = classElement.FindSubElement(crc32);
+                if (vectorElement)
+                {
+                    // Get the old data
+                    AZStd::vector<AZ::Name> oldData;
+                    if (classElement.GetChildData(crc32, oldData))
+                    {
+                        // Convert the vector with the new allocator
+                        vectorElement->Convert(context, AZ::AzTypeInfo<StreamingImageAsset::TagList>::Uuid());
+                        for (const auto& element : oldData)
+                        {
+                            // Re add the elements
+                            vectorElement->AddElementWithData<AZ::Name>(context, "element", element);
+                        }
+                    }
+                }
+            }
+            return true;
+        }
+
         void StreamingImageAsset::Reflect(ReflectContext* context)
         {
             if (auto* serializeContext = azrtti_cast<SerializeContext*>(context))
             {
+                // Need to register the old type with the Serializer so we can read it in order to convert it
+                serializeContext->RegisterGenericType<AZStd::vector<AZ::Name>>();
+
                 serializeContext->Class<MipChain>()
                     ->Field("m_mipOffset", &MipChain::m_mipOffset)
                     ->Field("m_mipCount", &MipChain::m_mipCount)
@@ -26,7 +58,7 @@ namespace AZ
                     ;
 
                 serializeContext->Class<StreamingImageAsset, ImageAsset>()
-                    ->Version(2) // Added m_averageColor field
+                    ->Version(3, &ConvertOldVersions) // Added m_averageColor field
                     ->Field("m_mipLevelToChainIndex", &StreamingImageAsset::m_mipLevelToChainIndex)
                     ->Field("m_mipChains", &StreamingImageAsset::m_mipChains)
                     ->Field("m_flags", &StreamingImageAsset::m_flags)
@@ -130,7 +162,7 @@ namespace AZ
             return imageDescriptor;
         }
 
-        const AZStd::vector<AZ::Name>& StreamingImageAsset::GetTags() const
+        const StreamingImageAsset::TagList& StreamingImageAsset::GetTags() const
         {
             return m_tags;
         }

Gems/Atom/RPI/Code/atom_rpi_reflect_files.cmake

@@ -7,6 +7,7 @@
 #
 
 set(FILES
+    Include/Atom/RPI.Reflect/Allocators.h
     Include/Atom/RPI.Reflect/AssetCreator.h
     Include/Atom/RPI.Reflect/Base.h
     Include/Atom/RPI.Reflect/FeatureProcessorDescriptor.h

Gems/AtomLyIntegration/AtomImGuiTools/Code/Source/AtomImGuiToolsSystemComponent.cpp

@@ -104,7 +104,6 @@ namespace AtomImGuiTools
                 m_showTransientAttachmentProfiler = m_imguiTransientAttachmentProfiler.Draw(transientStats);
             }
         }
-
         m_showMaterialDetails = m_imguiMaterialDetails.Tick(m_materialDetailsController.GetMeshDrawPackets(), m_materialDetailsController.GetSelectionName().c_str());
     }
 

Gems/AtomLyIntegration/AtomViewportDisplayInfo/Code/Source/AtomViewportDisplayInfoSystemComponent.cpp

@@ -335,10 +335,24 @@ namespace AZ::Render
         size_t deviceResident = 0;
         size_t deviceReserved = 0;
 
-        for (const auto& pool : stats->m_pools)
+        for (const auto& heap : stats->m_heaps)
         {
-            deviceReserved += pool.m_memoryUsage.GetHeapMemoryUsage(RHI::HeapMemoryLevel::Device).m_totalResidentInBytes;
-            deviceResident += pool.m_memoryUsage.GetHeapMemoryUsage(RHI::HeapMemoryLevel::Device).m_usedResidentInBytes;
+            if (heap.m_heapMemoryType == RHI::HeapMemoryLevel::Device)
+            {
+                deviceResident += heap.m_memoryUsage.m_usedResidentInBytes;                
+                deviceReserved += heap.m_memoryUsage.m_totalResidentInBytes;
+            }
+        }
+
+        // Fallback if the information from the heaps is not available, we use the pools instead.
+        if (deviceResident == 0 || deviceReserved == 0)
+        {
+            deviceResident = deviceReserved = 0;
+            for (const auto& pool : stats->m_pools)
+            {
+                deviceReserved += pool.m_memoryUsage.GetHeapMemoryUsage(RHI::HeapMemoryLevel::Device).m_totalResidentInBytes;
+                deviceResident += pool.m_memoryUsage.GetHeapMemoryUsage(RHI::HeapMemoryLevel::Device).m_usedResidentInBytes;
+            }
         }
 
         // Query for available device memory

Gems/Profiler/Code/Source/ImGuiHeapMemoryProfiler.cpp

@@ -0,0 +1,105 @@
+/*
+ * Copyright (c) Contributors to the Open 3D Engine Project.
+ * For complete copyright and license terms please see the LICENSE at the root of this distribution.
+ *
+ * SPDX-License-Identifier: Apache-2.0 OR MIT
+ *
+ */
+
+#if defined(IMGUI_ENABLED)
+
+#include <ImGuiHeapMemoryProfiler.h>
+#include <AzCore/std/sort.h>
+
+namespace Profiler
+{
+    enum HeapProfilerColumnID
+    {
+        HeapProfilerColumnID_Name = 0,
+        HeapProfilerColumnID_ParentName,
+        HeapProfilerColumnID_AllocatedMem,
+        HeapProfilerColumnID_CapacityMem
+    };
+
+    void ImGuiHeapMemoryProfiler::Draw(bool& draw)
+    {
+        using namespace AZ;
+        static constexpr size_t KB = 1u << 10;
+
+        ImGui::SetNextWindowPos(ImVec2(300, 60), ImGuiCond_FirstUseEver);
+        ImGui::SetNextWindowSize(ImVec2(800, 700), ImGuiCond_FirstUseEver);
+        if (ImGui::Begin("Heap Profiler", &draw, ImGuiWindowFlags_None))
+        {
+            ImGui::Text("Allocator Name Filter (inc, -exc)");
+            ImGui::SameLine();
+            m_filter.Draw("");            
+            ImGuiTableFlags flags = ImGuiTableFlags_Borders | ImGuiTableFlags_RowBg | ImGuiTableFlags_Sortable;
+            constexpr int NumColumns = 4;
+            if (ImGui::BeginTable("table", NumColumns, flags))
+            {
+                ImGui::TableSetupColumn(
+                    "Allocator Name", ImGuiTableColumnFlags_WidthFixed | ImGuiTableColumnFlags_DefaultSort, 0.f, HeapProfilerColumnID_Name);
+                ImGui::TableSetupColumn("Allocated Memory (kB)", 0, 0.f, HeapProfilerColumnID_AllocatedMem);
+                ImGui::TableSetupColumn("Capacity Memory (kB)", 0, 0.f, HeapProfilerColumnID_CapacityMem);
+                ImGui::TableSetupColumn("Parent Name", 0, 0.f, HeapProfilerColumnID_ParentName);
+                ImGui::TableHeadersRow();
+
+                size_t requestedBytes;
+                size_t capacityBytes;
+                AZStd::vector<AZ::AllocatorManager::AllocatorStats> stats;
+                AZ::AllocatorManager::Instance().GetAllocatorStats(requestedBytes, capacityBytes, &stats);
+
+                if (ImGuiTableSortSpecs* sorts_specs = ImGui::TableGetSortSpecs())
+                {
+                    AZStd::sort(
+                        stats.begin(),
+                        stats.end(),
+                        [&](const AZ::AllocatorManager::AllocatorStats& lhs, const AZ::AllocatorManager::AllocatorStats& rhs)
+                        {
+                            const AZ::AllocatorManager::AllocatorStats* left = &lhs;
+                            const AZ::AllocatorManager::AllocatorStats* right = &rhs;
+                            if (sorts_specs->Specs->SortDirection == ImGuiSortDirection_Descending)
+                            {
+                                AZStd::swap(left, right);
+                            }
+
+                            switch (sorts_specs->Specs->ColumnUserID)
+                            {
+                            case HeapProfilerColumnID_Name:
+                                return left->m_name < right->m_name;
+                            case HeapProfilerColumnID_AllocatedMem:
+                                return left->m_allocatedBytes < right->m_allocatedBytes;
+                            case HeapProfilerColumnID_CapacityMem:
+                                return left->m_capacityBytes < right->m_capacityBytes;
+                            case HeapProfilerColumnID_ParentName:
+                                return left->m_parentName < right->m_parentName;
+                            default:
+                                return false;
+                            }
+                        });
+                        sorts_specs->SpecsDirty = false;
+                }
+
+                for (const auto& stat : stats)
+                {
+                    if (m_filter.PassFilter(stat.m_name.c_str()))
+                    {
+                        ImGui::TableNextRow();
+                        ImGui::TableNextColumn();
+                        ImGui::TextUnformatted(stat.m_name.c_str());
+                        ImGui::TableNextColumn();
+                        ImGui::Text("%.1f", static_cast<float>(stat.m_allocatedBytes) / KB);
+                        ImGui::TableNextColumn();
+                        ImGui::Text("%.1f", static_cast<float>(stat.m_capacityBytes) / KB);
+                        ImGui::TableNextColumn();
+                        ImGui::TextUnformatted(stat.m_parentName.c_str());
+                    }
+                }
+
+                ImGui::EndTable();
+            }
+        }
+        ImGui::End();       
+    }
+}
+#endif

Gems/Profiler/Code/Source/ImGuiHeapMemoryProfiler.h

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) Contributors to the Open 3D Engine Project.
+ * For complete copyright and license terms please see the LICENSE at the root of this distribution.
+ *
+ * SPDX-License-Identifier: Apache-2.0 OR MIT
+ *
+ */
+
+#pragma once
+
+#if defined(IMGUI_ENABLED)
+#include <imgui/imgui.h>
+#include <AzCore/Memory/AllocatorManager.h>
+
+namespace Profiler
+{
+    //! Profiler for displaying information about Memory Heaps
+    class ImGuiHeapMemoryProfiler
+    {
+    public:
+        ImGuiHeapMemoryProfiler() = default;
+        ~ImGuiHeapMemoryProfiler() = default;
+
+        void Draw(bool& draw);
+
+    private:
+        ImGuiTextFilter m_filter;
+    };
+}
+#endif

Gems/Profiler/Code/Source/ProfilerImGuiSystemComponent.cpp

@@ -111,6 +111,10 @@ namespace Profiler
         {
             ShowCpuProfilerWindow(m_showCpuProfiler);
         }
+        if (m_showHeapMemoryProfiler)
+        {
+            m_imguiHeapMemoryProfiler.Draw(m_showHeapMemoryProfiler);
+        }
     }
 
     void ProfilerImGuiSystemComponent::OnImGuiMainMenuUpdate()
@@ -121,6 +125,7 @@ namespace Profiler
             {
                 AZ::Debug::ProfilerSystemInterface::Get()->SetActive(m_showCpuProfiler);
             }
+            ImGui::MenuItem("Heap Memory", "", &m_showHeapMemoryProfiler);
             ImGui::EndMenu();
         }
     }

Gems/Profiler/Code/Source/ProfilerImGuiSystemComponent.h

@@ -11,6 +11,7 @@
 #include <Profiler/ProfilerImGuiBus.h>
 
 #include <ImGuiCpuProfiler.h>
+#include <ImGuiHeapMemoryProfiler.h>
 
 #include "ImGuiTreemapImpl.h"
 
@@ -61,7 +62,9 @@ namespace Profiler
 #if defined(IMGUI_ENABLED)
         ImGuiTreemapFactoryImpl m_imguiTreemapFactory;
         ImGuiCpuProfiler m_imguiCpuProfiler;
+        ImGuiHeapMemoryProfiler m_imguiHeapMemoryProfiler;
         bool m_showCpuProfiler{ false };
+        bool m_showHeapMemoryProfiler{ false };
 #endif // defined(IMGUI_ENABLED)
     };
 

Gems/Profiler/Code/profiler_imgui_shared_files.cmake

@@ -16,4 +16,6 @@ set(FILES
     Source/ProfilerImGuiModule.cpp
     Source/ProfilerImGuiSystemComponent.cpp
     Source/ProfilerImGuiSystemComponent.h
+    Source/ImGuiHeapMemoryProfiler.cpp
+    Source/ImGuiHeapMemoryProfiler.h
 )

Gems/SceneProcessing/Code/Source/SceneBuilder/SceneBuilderComponent.cpp

@@ -44,7 +44,7 @@ namespace SceneBuilder
         builderDescriptor.m_createJobFunction = AZStd::bind(&SceneBuilderWorker::CreateJobs, &m_sceneBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
         builderDescriptor.m_processJobFunction = AZStd::bind(&SceneBuilderWorker::ProcessJob, &m_sceneBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);
 
-        builderDescriptor.m_version = 12; // bump this to rebuild everything.
+        builderDescriptor.m_version = 13; // Add BufferAssetAllocator
         builderDescriptor.m_analysisFingerprint = m_sceneBuilder.GetFingerprint(); // bump this to at least re-analyze everything.
 
         m_sceneBuilder.BusConnect(builderDescriptor.m_busId);

Gems/TextureAtlas/Code/Source/Editor/AtlasBuilderComponent.cpp

@@ -36,7 +36,7 @@ namespace TextureAtlasBuilder
     {
         AssetBuilderSDK::AssetBuilderDesc builderDescriptor;
         builderDescriptor.m_name = "Atlas Worker Builder";
-        builderDescriptor.m_version = 1;
+        builderDescriptor.m_version = 2; // Add MipImageAsset allocator
         builderDescriptor.m_patterns.emplace_back(AssetBuilderSDK::AssetBuilderPattern("*.texatlas", AssetBuilderSDK::AssetBuilderPattern::PatternType::Wildcard));
         builderDescriptor.m_busId = azrtti_typeid<AtlasBuilderWorker>();
         builderDescriptor.m_createJobFunction = AZStd::bind(&AtlasBuilderWorker::CreateJobs, &m_atlasBuilder, AZStd::placeholders::_1, AZStd::placeholders::_2);