Updating feature processors and components to remove direct dependency on material assignment data types. Instead, the mesh feature processor will manage its own custom materials. This will limit the amount of data that gets copied into the feature processor and carried around with every mesh handle.

Signed-off-by: gadams3 <[email protected]>

Gems/Atom/Feature/Common/Code/Include/Atom/Feature/Mesh/MeshFeatureProcessor.h

@@ -23,7 +23,6 @@
 #include <Atom/RPI.Public/Shader/ShaderSystemInterface.h>
 
 #include <Atom/Feature/Mesh/MeshFeatureProcessorInterface.h>
-#include <Atom/Feature/Material/MaterialAssignment.h>
 #include <Atom/Feature/Material/MaterialAssignmentBus.h>
 #include <Atom/Feature/TransformService/TransformServiceFeatureProcessor.h>
 #include <Atom/Feature/Mesh/ModelReloaderSystemInterface.h>
@@ -71,10 +70,10 @@ namespace AZ
                 void OnCatalogAssetAdded(const AZ::Data::AssetId& assetId) override;
 
                 void OnModelReloaded(Data::Asset<Data::AssetData> asset);
-                ModelReloadedEvent::Handler m_modelReloadedEventHandler { [&](Data::Asset<RPI::ModelAsset> modelAsset)
-                                                                  {
+                ModelReloadedEvent::Handler m_modelReloadedEventHandler{ [&](Data::Asset<RPI::ModelAsset> modelAsset)
+                                                                         {
                                                                              OnModelReloaded(modelAsset);
-                                                                  } };
+                                                                         } };
                 MeshFeatureProcessorInterface::ModelChangedEvent m_modelChangedEvent;
                 Data::Asset<RPI::ModelAsset> m_modelAsset;
                 ModelDataInstance* m_parent = nullptr;
@@ -88,8 +87,10 @@ namespace AZ
                 RayTracingFeatureProcessor* rayTracingFeatureProcessor,
                 TransformServiceFeatureProcessor* transformServiceFeatureProcessor);
             void RemoveRayTracingData(RayTracingFeatureProcessor* rayTracingFeatureProcessor);
-            void SetIrradianceData(RayTracingFeatureProcessor::SubMesh& subMesh,
-                    const Data::Instance<RPI::Material> material, const Data::Instance<RPI::Image> baseColorImage);
+            void SetIrradianceData(
+                RayTracingFeatureProcessor::SubMesh& subMesh,
+                const Data::Instance<RPI::Material> material,
+                const Data::Instance<RPI::Image> baseColorImage);
             void SetSortKey(RHI::DrawItemSortKey sortKey);
             RHI::DrawItemSortKey GetSortKey() const;
             void SetMeshLodConfiguration(RPI::Cullable::LodConfiguration meshLodConfig);
@@ -108,10 +109,12 @@ namespace AZ
             // MaterialAssignmentNotificationBus overrides
             void OnRebuildMaterialInstance() override;
 
+            CustomMaterialInfo GetCustomMaterialWithFallback(const CustomMaterialId& id) const;
+
             RPI::MeshDrawPacketLods m_drawPacketListsByLod;
 
             RPI::Cullable m_cullable;
-            MaterialAssignmentMap m_materialAssignments;
+            CustomMaterialMap m_customMaterials;
 
             typedef AZStd::unordered_map<Data::Instance<RPI::Material>, RPI::Material::ChangeId> MaterialChangeIdMap;
             MaterialChangeIdMap m_materialChangeIds;
@@ -176,12 +179,8 @@ namespace AZ
             void OnEndPrepareRender() override;
 
             TransformServiceFeatureProcessorInterface::ObjectId GetObjectId(const MeshHandle& meshHandle) const override;
-            MeshHandle AcquireMesh(
-                const MeshHandleDescriptor& descriptor,
-                const MaterialAssignmentMap& materials = {}) override;
-            MeshHandle AcquireMesh(
-                const MeshHandleDescriptor& descriptor,
-                const Data::Instance<RPI::Material>& material) override;
+            MeshHandle AcquireMesh(const MeshHandleDescriptor& descriptor, const CustomMaterialMap& materials = {}) override;
+            MeshHandle AcquireMesh(const MeshHandleDescriptor& descriptor, const Data::Instance<RPI::Material>& material) override;
             bool ReleaseMesh(MeshHandle& meshHandle) override;
             MeshHandle CloneMesh(const MeshHandle& meshHandle) override;
 
@@ -190,9 +189,9 @@ namespace AZ
             const RPI::MeshDrawPacketLods& GetDrawPackets(const MeshHandle& meshHandle) const override;
             const AZStd::vector<Data::Instance<RPI::ShaderResourceGroup>>& GetObjectSrgs(const MeshHandle& meshHandle) const override;
             void QueueObjectSrgForCompile(const MeshHandle& meshHandle) const override;
-            void SetMaterialAssignmentMap(const MeshHandle& meshHandle, const Data::Instance<RPI::Material>& material) override;
-            void SetMaterialAssignmentMap(const MeshHandle& meshHandle, const MaterialAssignmentMap& materials) override;
-            const MaterialAssignmentMap& GetMaterialAssignmentMap(const MeshHandle& meshHandle) const override;
+            void SetCustomMaterials(const MeshHandle& meshHandle, const Data::Instance<RPI::Material>& material) override;
+            void SetCustomMaterials(const MeshHandle& meshHandle, const CustomMaterialMap& materials) override;
+            const CustomMaterialMap& GetCustomMaterials(const MeshHandle& meshHandle) const override;
             void ConnectModelChangeEventHandler(const MeshHandle& meshHandle, ModelChangedEvent::Handler& handler) override;
 
             void SetTransform(const MeshHandle& meshHandle, const AZ::Transform& transform,

Gems/Atom/Feature/Common/Code/Include/Atom/Feature/Mesh/MeshFeatureProcessorInterface.h

@@ -5,18 +5,19 @@
  * SPDX-License-Identifier: Apache-2.0 OR MIT
  *
  */
+
 #pragma once
 
-#include <AzCore/EBus/Event.h>
-#include <AzCore/Outcome/Outcome.h>
-#include <AzCore/std/functional.h>
-#include <Atom/Feature/Material/MaterialAssignment.h>
+#include <Atom/Feature/TransformService/TransformServiceFeatureProcessorInterface.h>
 #include <Atom/RPI.Public/Culling.h>
 #include <Atom/RPI.Public/FeatureProcessor.h>
 #include <Atom/RPI.Public/MeshDrawPacket.h>
+#include <Atom/RPI.Public/Model/Model.h>
 #include <Atom/RPI.Reflect/Model/ModelAsset.h>
 #include <Atom/Utils/StableDynamicArray.h>
-#include <Atom/Feature/TransformService/TransformServiceFeatureProcessorInterface.h>
+#include <AzCore/EBus/Event.h>
+#include <AzCore/Outcome/Outcome.h>
+#include <AzCore/std/functional.h>
 
 namespace AZ
 {
@@ -45,9 +46,21 @@ namespace AZ
             bool m_excludeFromReflectionCubeMaps = false;
         };
 
-        //! MeshFeatureProcessorInterface provides an interface to acquire and release a MeshHandle from the underlying MeshFeatureProcessor
-        class MeshFeatureProcessorInterface
-            : public RPI::FeatureProcessor
+        using CustomMaterialLodIndex = AZ::u64;
+        using CustomMaterialId = AZStd::pair<CustomMaterialLodIndex, uint32_t>;
+        struct CustomMaterialInfo
+        {
+            Data::Instance<RPI::Material> m_material;
+            RPI::MaterialModelUvOverrideMap m_uvMapping;
+        };
+        using CustomMaterialMap = AZStd::unordered_map<CustomMaterialId, CustomMaterialInfo>;
+        static const CustomMaterialLodIndex DefaultCustomMaterialLodIndex = AZStd::numeric_limits<CustomMaterialLodIndex>::max();
+        static const CustomMaterialId DefaultCustomMaterialId = CustomMaterialId(DefaultCustomMaterialLodIndex, 0);
+        static const CustomMaterialMap DefaultCustomMaterialMap = CustomMaterialMap();
+
+        //! MeshFeatureProcessorInterface provides an interface to acquire and release a MeshHandle from the underlying
+        //! MeshFeatureProcessor
+        class MeshFeatureProcessorInterface : public RPI::FeatureProcessor
         {
         public:
             AZ_RTTI(AZ::Render::MeshFeatureProcessorInterface, "{975D7F0C-2E7E-4819-94D0-D3C4E2024721}", AZ::RPI::FeatureProcessor);
@@ -58,15 +71,11 @@ namespace AZ
             //! Returns the object id for a mesh handle.
             virtual TransformServiceFeatureProcessorInterface::ObjectId GetObjectId(const MeshHandle& meshHandle) const = 0;
 
-            //! Acquires a model with an optional collection of material assignments.
+            //! Acquires a model with an optional collection of custom materials.
             //! @param requiresCloneCallback The callback indicates whether cloning is required for a given model asset.
-            virtual MeshHandle AcquireMesh(
-                const MeshHandleDescriptor& descriptor,
-                const MaterialAssignmentMap& materials = {}) = 0;
+            virtual MeshHandle AcquireMesh(const MeshHandleDescriptor& descriptor, const CustomMaterialMap& materials = {}) = 0;
             //! Acquires a model with a single material applied to all its meshes.
-            virtual MeshHandle AcquireMesh(
-                const MeshHandleDescriptor& descriptor,
-                const Data::Instance<RPI::Material>& material) = 0;
+            virtual MeshHandle AcquireMesh(const MeshHandleDescriptor& descriptor, const Data::Instance<RPI::Material>& material) = 0;
             //! Releases the mesh handle
             virtual bool ReleaseMesh(MeshHandle& meshHandle) = 0;
             //! Creates a new instance and handle of a mesh using an existing MeshId. Currently, this will reset the new mesh to default materials.
@@ -90,13 +99,13 @@ namespace AZ
             virtual const AZStd::vector<Data::Instance<RPI::ShaderResourceGroup>>& GetObjectSrgs(const MeshHandle& meshHandle) const = 0;
             //! Queues the object srg for compile.
             virtual void QueueObjectSrgForCompile(const MeshHandle& meshHandle) const = 0;
-            //! Sets the MaterialAssignmentMap for a meshHandle, using just a single material for the DefaultMaterialAssignmentId.
-            //! Note if there is already a material assignment map, this will replace the entire map with just a single material.
-            virtual void SetMaterialAssignmentMap(const MeshHandle& meshHandle, const Data::Instance<RPI::Material>& material) = 0;
-            //! Sets the MaterialAssignmentMap for a meshHandle.
-            virtual void SetMaterialAssignmentMap(const MeshHandle& meshHandle, const MaterialAssignmentMap& materials) = 0;
-            //! Gets the MaterialAssignmentMap for a meshHandle.
-            virtual const MaterialAssignmentMap& GetMaterialAssignmentMap(const MeshHandle& meshHandle) const = 0;
+            //! Sets the CustomMaterialMap for a meshHandle, using just a single material for the DefaultCustomMaterialId.
+            //! Note if there is already a CustomMaterialMap, this will replace the entire map with just a single material.
+            virtual void SetCustomMaterials(const MeshHandle& meshHandle, const Data::Instance<RPI::Material>& material) = 0;
+            //! Sets the CustomMaterialMap for a meshHandle.
+            virtual void SetCustomMaterials(const MeshHandle& meshHandle, const CustomMaterialMap& materials) = 0;
+            //! Gets the CustomMaterialMap for a meshHandle.
+            virtual const CustomMaterialMap& GetCustomMaterials(const MeshHandle& meshHandle) const = 0;
             //! Connects a handler to any changes to an RPI::Model. Changes include loading and reloading.
             virtual void ConnectModelChangeEventHandler(const MeshHandle& meshHandle, ModelChangedEvent::Handler& handler) = 0;
 

Gems/Atom/Feature/Common/Code/Mocks/MockMeshFeatureProcessor.h

@@ -23,14 +23,14 @@ namespace UnitTest
         MOCK_CONST_METHOD1(GetDrawPackets, const AZ::RPI::MeshDrawPacketLods&(const MeshHandle&));
         MOCK_CONST_METHOD1(GetObjectSrgs, const AZStd::vector<AZStd::intrusive_ptr<AZ::RPI::ShaderResourceGroup>>&(const MeshHandle&));
         MOCK_CONST_METHOD1(QueueObjectSrgForCompile, void(const MeshHandle&));
-        MOCK_CONST_METHOD1(GetMaterialAssignmentMap, const AZ::Render::MaterialAssignmentMap&(const MeshHandle&));
+        MOCK_CONST_METHOD1(GetCustomMaterials, const AZ::Render::CustomMaterialMap&(const MeshHandle&));
         MOCK_METHOD2(ConnectModelChangeEventHandler, void(const MeshHandle&, ModelChangedEvent::Handler&));
         MOCK_METHOD3(SetTransform, void(const MeshHandle&, const AZ::Transform&, const AZ::Vector3&));
         MOCK_METHOD2(SetExcludeFromReflectionCubeMaps, void(const MeshHandle&, bool));
         MOCK_METHOD2(SetIsAlwaysDynamic, void(const MeshHandle&, bool));
         MOCK_CONST_METHOD1(GetIsAlwaysDynamic, bool(const MeshHandle&));
-        MOCK_METHOD2(SetMaterialAssignmentMap, void(const MeshHandle&, const AZ::Data::Instance<AZ::RPI::Material>&));
-        MOCK_METHOD2(SetMaterialAssignmentMap, void(const MeshHandle&, const AZ::Render::MaterialAssignmentMap&));
+        MOCK_METHOD2(SetCustomMaterials, void(const MeshHandle&, const AZ::Data::Instance<AZ::RPI::Material>&));
+        MOCK_METHOD2(SetCustomMaterials, void(const MeshHandle&, const AZ::Render::CustomMaterialMap&));
         MOCK_METHOD1(GetTransform, AZ::Transform(const MeshHandle&));
         MOCK_METHOD1(GetNonUniformScale, AZ::Vector3(const MeshHandle&));
         MOCK_METHOD2(SetLocalAabb, void(const MeshHandle&, const AZ::Aabb&));
@@ -39,7 +39,7 @@ namespace UnitTest
         MOCK_CONST_METHOD1(GetSortKey, AZ::RHI::DrawItemSortKey(const MeshHandle&));
         MOCK_METHOD2(SetMeshLodConfiguration, void(const MeshHandle&, const AZ::RPI::Cullable::LodConfiguration&));
         MOCK_CONST_METHOD1(GetMeshLodConfiguration, AZ::RPI::Cullable::LodConfiguration(const MeshHandle&));
-        MOCK_METHOD2(AcquireMesh, MeshHandle (const AZ::Render::MeshHandleDescriptor&, const AZ::Render::MaterialAssignmentMap&));
+        MOCK_METHOD2(AcquireMesh, MeshHandle (const AZ::Render::MeshHandleDescriptor&, const AZ::Render::CustomMaterialMap&));
         MOCK_METHOD2(AcquireMesh, MeshHandle (const AZ::Render::MeshHandleDescriptor&, const AZ::Data::Instance<AZ::RPI::Material>&));
         MOCK_METHOD2(SetRayTracingEnabled, void (const MeshHandle&, bool));
         MOCK_CONST_METHOD1(GetRayTracingEnabled, bool(const MeshHandle&));

Gems/Atom/Feature/Common/Code/Source/Mesh/MeshFeatureProcessor.cpp

@@ -307,17 +307,17 @@ namespace AZ
         }
 
         MeshFeatureProcessor::MeshHandle MeshFeatureProcessor::AcquireMesh(
-            const MeshHandleDescriptor& descriptor,
-            const MaterialAssignmentMap& materials)
+            const MeshHandleDescriptor& descriptor, const CustomMaterialMap& materials)
         {
             AZ_PROFILE_SCOPE(AzRender, "MeshFeatureProcessor: AcquireMesh");
 
-            // don't need to check the concurrency during emplace() because the StableDynamicArray won't move the other elements during insertion
+            // don't need to check the concurrency during emplace() because the StableDynamicArray won't move the other elements during
+            // insertion
             MeshHandle meshDataHandle = m_modelData.emplace();
 
             meshDataHandle->m_descriptor = descriptor;
             meshDataHandle->m_scene = GetParentScene();
-            meshDataHandle->m_materialAssignments = materials;
+            meshDataHandle->m_customMaterials = materials;
             meshDataHandle->m_objectId = m_transformService->ReserveObjectId();
             meshDataHandle->m_rayTracingUuid = AZ::Uuid::CreateRandom();
             meshDataHandle->m_originalModelAsset = descriptor.m_modelAsset;
@@ -335,12 +335,13 @@ namespace AZ
         }
 
         MeshFeatureProcessor::MeshHandle MeshFeatureProcessor::AcquireMesh(
-            const MeshHandleDescriptor& descriptor,
-            const Data::Instance<RPI::Material>& material)
+            const MeshHandleDescriptor& descriptor, const Data::Instance<RPI::Material>& material)
         {
-            Render::MaterialAssignmentMap materials;
-            Render::MaterialAssignment& defaultMaterial = materials[AZ::Render::DefaultMaterialAssignmentId];
-            defaultMaterial.m_materialInstance = material;
+            Render::CustomMaterialMap materials;
+            if (material)
+            {
+                materials[AZ::Render::DefaultCustomMaterialId] = { material };
+            }
 
             return AcquireMesh(descriptor, materials);
         }
@@ -365,8 +366,7 @@ namespace AZ
         {
             if (meshHandle.IsValid())
             {
-                MeshHandle clone = AcquireMesh(meshHandle->m_descriptor, meshHandle->m_materialAssignments);
-                return clone;
+                return AcquireMesh(meshHandle->m_descriptor, meshHandle->m_customMaterials);
             }
             return MeshFeatureProcessor::MeshHandle();
         }
@@ -405,16 +405,14 @@ namespace AZ
             }
         }
 
-        void MeshFeatureProcessor::SetMaterialAssignmentMap(const MeshHandle& meshHandle, const Data::Instance<RPI::Material>& material)
+        void MeshFeatureProcessor::SetCustomMaterials(const MeshHandle& meshHandle, const Data::Instance<RPI::Material>& material)
         {
-            Render::MaterialAssignmentMap materials;
-            Render::MaterialAssignment& defaultMaterial = materials[AZ::Render::DefaultMaterialAssignmentId];
-            defaultMaterial.m_materialInstance = material;
-
-            return SetMaterialAssignmentMap(meshHandle, materials);
+            Render::CustomMaterialMap materials;
+            materials[AZ::Render::DefaultCustomMaterialId] = { material };
+            return SetCustomMaterials(meshHandle, materials);
         }
 
-        void MeshFeatureProcessor::SetMaterialAssignmentMap(const MeshHandle& meshHandle, const MaterialAssignmentMap& materials)
+        void MeshFeatureProcessor::SetCustomMaterials(const MeshHandle& meshHandle, const CustomMaterialMap& materials)
         {
             if (meshHandle.IsValid())
             {
@@ -422,12 +420,12 @@ namespace AZ
                 {
                     Data::Instance<RPI::Model> model = meshHandle->m_model;
                     meshHandle->DeInit(m_rayTracingFeatureProcessor);
-                    meshHandle->m_materialAssignments = materials;
+                    meshHandle->m_customMaterials = materials;
                     meshHandle->QueueInit(model);
                 }
                 else
                 {
-                    meshHandle->m_materialAssignments = materials;
+                    meshHandle->m_customMaterials = materials;
                 }
 
                 meshHandle->UpdateMaterialChangeIds();
@@ -436,9 +434,9 @@ namespace AZ
             }
         }
 
-        const MaterialAssignmentMap& MeshFeatureProcessor::GetMaterialAssignmentMap(const MeshHandle& meshHandle) const
+        const CustomMaterialMap& MeshFeatureProcessor::GetCustomMaterials(const MeshHandle& meshHandle) const
         {
-            return meshHandle.IsValid() ? meshHandle->m_materialAssignments : DefaultMaterialAssignmentMap;
+            return meshHandle.IsValid() ? meshHandle->m_customMaterials : DefaultCustomMaterialMap;
         }
 
         void MeshFeatureProcessor::ConnectModelChangeEventHandler(const MeshHandle& meshHandle, ModelChangedEvent::Handler& handler)
@@ -913,7 +911,7 @@ namespace AZ
             RemoveRayTracingData(rayTracingFeatureProcessor);
 
             m_drawPacketListsByLod.clear();
-            m_materialAssignments.clear();
+            m_customMaterials.clear();
             m_materialChangeIds.clear();
             m_objectSrgList = {};
             m_model = {};
@@ -943,10 +941,9 @@ namespace AZ
                 objectIdIndex.AssertValid();
             }
 
-            for (const auto& materialAssignment : m_materialAssignments)
+            for (const auto& customMaterial : m_customMaterials)
             {
-                const AZ::Data::Instance<RPI::Material>& materialInstance = materialAssignment.second.m_materialInstance;
-                if (materialInstance.get())
+                if (const AZ::Data::Instance<RPI::Material>& materialInstance = customMaterial.second.m_material)
                 {
                     MaterialAssignmentNotificationBus::MultiHandler::BusConnect(materialInstance->GetAssetId());
                 }
@@ -978,15 +975,10 @@ namespace AZ
             {
                 const RPI::ModelLod::Mesh& mesh = modelLod.GetMeshes()[meshIndex];
 
-                Data::Instance<RPI::Material> material = mesh.m_material;
-
-                // Determine if there is a material override specified for this sub mesh
-                const MaterialAssignmentId materialAssignmentId(modelLodIndex, mesh.m_materialSlotStableId);
-                const MaterialAssignment& materialAssignment = GetMaterialAssignmentFromMapWithFallback(m_materialAssignments, materialAssignmentId);
-                if (materialAssignment.m_materialInstance.get())
-                {
-                    material = materialAssignment.m_materialInstance;
-                }
+                // Determine if there is a custom material specified for this submission
+                const CustomMaterialId customMaterialId(aznumeric_cast<AZ::u64>(modelLodIndex), mesh.m_materialSlotStableId);
+                const auto& customMaterialInfo = GetCustomMaterialWithFallback(customMaterialId);
+                const auto& material = customMaterialInfo.m_material ? customMaterialInfo.m_material : mesh.m_material;
 
                 if (!material)
                 {
@@ -1027,7 +1019,7 @@ namespace AZ
                 }
 
                 // setup the mesh draw packet
-                RPI::MeshDrawPacket drawPacket(modelLod, meshIndex, material, meshObjectSrg, materialAssignment.m_matModUvOverrides);
+                RPI::MeshDrawPacket drawPacket(modelLod, meshIndex, material, meshObjectSrg, customMaterialInfo.m_uvMapping);
 
                 // set the shader option to select forward pass IBL specular if necessary
                 if (!drawPacket.SetShaderOption(s_o_meshUseForwardPassIBLSpecular_Name, AZ::RPI::ShaderOptionValue{ m_descriptor.m_useForwardPassIblSpecular }))
@@ -1133,14 +1125,9 @@ namespace AZ
                 const RPI::ModelLod::Mesh& mesh = modelLod->GetMeshes()[meshIndex];
 
                 // retrieve the material
-                Data::Instance<RPI::Material> material = mesh.m_material;
-
-                const MaterialAssignmentId materialAssignmentId(rayTracingLod, mesh.m_materialSlotStableId);
-                const MaterialAssignment& materialAssignment = GetMaterialAssignmentFromMapWithFallback(m_materialAssignments, materialAssignmentId);
-                if (materialAssignment.m_materialInstance.get())
-                {
-                    material = materialAssignment.m_materialInstance;
-                }
+                const CustomMaterialId customMaterialId(rayTracingLod, mesh.m_materialSlotStableId);
+                const auto& customMaterialInfo = GetCustomMaterialWithFallback(customMaterialId);
+                const auto& material = customMaterialInfo.m_material ? customMaterialInfo.m_material : mesh.m_material;
 
                 if (!material)
                 {
@@ -1156,7 +1143,7 @@ namespace AZ
                     nullptr,
                     shaderInputContract,
                     meshIndex,
-                    materialAssignment.m_matModUvOverrides,
+                    customMaterialInfo.m_uvMapping,
                     material->GetAsset()->GetMaterialTypeAsset()->GetUvNameMap());
                 AZ_Assert(result, "Failed to retrieve mesh stream buffer views");
 
@@ -1741,10 +1728,9 @@ namespace AZ
             // update the material changeId list with the current material assignments
             m_materialChangeIds.clear();
 
-            for (const auto& materialAssignment : m_materialAssignments)
+            for (const auto& customMaterial : m_customMaterials)
             {
-                const AZ::Data::Instance<RPI::Material>& materialInstance = materialAssignment.second.m_materialInstance;
-                if (materialInstance.get())
+                if (const AZ::Data::Instance<RPI::Material>& materialInstance = customMaterial.second.m_material)
                 {
                     m_materialChangeIds[materialInstance] = materialInstance->GetCurrentChangeId();
                 }
@@ -1754,20 +1740,21 @@ namespace AZ
         bool ModelDataInstance::CheckForMaterialChanges() const
         {
             // check for the same number of materials
-            if (m_materialChangeIds.size() != m_materialAssignments.size())
+            if (m_materialChangeIds.size() != m_customMaterials.size())
             {
                 return true;
             }
 
             // check for material changes using the changeId
-            for (const auto& materialAssignment : m_materialAssignments)
+            for (const auto& customMaterial : m_customMaterials)
             {
-                const AZ::Data::Instance<RPI::Material>& materialInstance = materialAssignment.second.m_materialInstance;
-
-                MaterialChangeIdMap::const_iterator it = m_materialChangeIds.find(materialInstance);
-                if (it == m_materialChangeIds.end() || it->second != materialInstance->GetCurrentChangeId())
+                if (const AZ::Data::Instance<RPI::Material>& materialInstance = customMaterial.second.m_material)
                 {
-                    return true;
+                    MaterialChangeIdMap::const_iterator it = m_materialChangeIds.find(materialInstance);
+                    if (it == m_materialChangeIds.end() || it->second != materialInstance->GetCurrentChangeId())
+                    {
+                        return true;
+                    }
                 }
             }
 
@@ -1787,5 +1774,26 @@ namespace AZ
                 }
             }
         }
+
+        CustomMaterialInfo ModelDataInstance::GetCustomMaterialWithFallback(const CustomMaterialId& id) const
+        {
+            if (auto itr = m_customMaterials.find(id); itr != m_customMaterials.end() && itr->second.m_material)
+            {
+                return itr->second;
+            }
+
+            const CustomMaterialId ignoreLodId(DefaultCustomMaterialLodIndex, id.second);
+            if (auto itr = m_customMaterials.find(ignoreLodId); itr != m_customMaterials.end() && itr->second.m_material)
+            {
+                return itr->second;
+            }
+
+            if (auto itr = m_customMaterials.find(DefaultCustomMaterialId); itr != m_customMaterials.end() && itr->second.m_material)
+            {
+                return itr->second;
+            }
+
+            return CustomMaterialInfo{};
+        }
     } // namespace Render
 } // namespace AZ

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

@@ -67,7 +67,7 @@ namespace AZ
                 Data::AssetBus::MultiHandler::BusDisconnect(asset.GetId());
 
                 m_visualizationMaterial = AZ::RPI::Material::FindOrCreate(m_visualizationMaterialAsset);
-                m_meshFeatureProcessor->SetMaterialAssignmentMap(m_visualizationMeshHandle, m_visualizationMaterial);
+                m_meshFeatureProcessor->SetCustomMaterials(m_visualizationMeshHandle, m_visualizationMaterial);
             }
         }
 

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

@@ -35,7 +35,7 @@ namespace AZ
                 m_visualizationMaterialAsset = asset;
                 Data::AssetBus::MultiHandler::BusDisconnect(asset.GetId());
 
-                m_meshFeatureProcessor->SetMaterialAssignmentMap(m_visualizationMeshHandle, AZ::RPI::Material::FindOrCreate(m_visualizationMaterialAsset));
+                m_meshFeatureProcessor->SetCustomMaterials(m_visualizationMeshHandle, AZ::RPI::Material::FindOrCreate(m_visualizationMaterialAsset));
             }
         }
 

Gems/AtomLyIntegration/AtomImGuiTools/Code/Include/AtomLyIntegration/AtomImGuiTools/AtomImGuiToolsBus.h

@@ -11,7 +11,6 @@
 #include <AzCore/Component/ComponentBus.h>
 #include <AzCore/Component/EntityId.h>
 #include <Atom/RPI.Public/MeshDrawPacket.h>
-#include <Atom/Feature/Material/MaterialAssignmentId.h>
 
 namespace AZ
 {

Gems/AtomLyIntegration/CommonFeatures/Code/Include/AtomLyIntegration/CommonFeatures/PostProcess/DisplayMapper/DisplayMapperComponentBus.h

@@ -8,7 +8,6 @@
 #pragma once
 
 #include <AzCore/Component/ComponentBus.h>
-#include <Atom/Feature/Material/MaterialAssignment.h>
 #include <ACES/Aces.h>
 
 namespace AZ

Gems/AtomLyIntegration/CommonFeatures/Code/Source/Mesh/MeshComponentController.cpp

@@ -359,7 +359,14 @@ namespace AZ
         {
             if (m_meshFeatureProcessor)
             {
-                m_meshFeatureProcessor->SetMaterialAssignmentMap(m_meshHandle, materials);
+                AZ::Render::CustomMaterialMap customMaterials;
+                for (const auto& materialAssignment : materials)
+                {
+                    customMaterials.emplace(
+                        AZ::Render::CustomMaterialId{ materialAssignment.first.m_lodIndex, materialAssignment.first.m_materialSlotStableId },
+                        AZ::Render::CustomMaterialInfo{ materialAssignment.second.m_materialInstance, materialAssignment.second.m_matModUvOverrides });
+                }
+                m_meshFeatureProcessor->SetCustomMaterials(m_meshHandle, customMaterials);
             }
         }
 
@@ -409,6 +416,14 @@ namespace AZ
                 MaterialAssignmentMap materials;
                 MaterialComponentRequestBus::EventResult(materials, entityId, &MaterialComponentRequests::GetMaterialMap);
 
+                AZ::Render::CustomMaterialMap customMaterials;
+                for (const auto& materialAssignment : materials)
+                {
+                    customMaterials.emplace(
+                        AZ::Render::CustomMaterialId{ materialAssignment.first.m_lodIndex, materialAssignment.first.m_materialSlotStableId },
+                        AZ::Render::CustomMaterialInfo{ materialAssignment.second.m_materialInstance, materialAssignment.second.m_matModUvOverrides });
+                }
+
                 m_meshFeatureProcessor->ReleaseMesh(m_meshHandle);
                 MeshHandleDescriptor meshDescriptor;
                 meshDescriptor.m_modelAsset = m_configuration.m_modelAsset;
@@ -417,7 +432,7 @@ namespace AZ
                 meshDescriptor.m_isRayTracingEnabled = m_configuration.m_isRayTracingEnabled;
                 meshDescriptor.m_excludeFromReflectionCubeMaps = m_configuration.m_excludeFromReflectionCubeMaps;
                 meshDescriptor.m_isAlwaysDynamic = m_configuration.m_isAlwaysDynamic;
-                m_meshHandle = m_meshFeatureProcessor->AcquireMesh(meshDescriptor, materials);
+                m_meshHandle = m_meshFeatureProcessor->AcquireMesh(meshDescriptor, customMaterials);
                 m_meshFeatureProcessor->ConnectModelChangeEventHandler(m_meshHandle, m_changeEventHandler);
 
                 const AZ::Transform& transform =

Gems/AtomLyIntegration/EMotionFXAtom/Code/Source/AtomActorInstance.cpp

@@ -239,7 +239,14 @@ namespace AZ::Render
     {
         if (m_meshFeatureProcessor)
         {
-            m_meshFeatureProcessor->SetMaterialAssignmentMap(*m_meshHandle, materials);
+            AZ::Render::CustomMaterialMap customMaterials;
+            for (const auto& materialAssignment : materials)
+            {
+                customMaterials.emplace(
+                    AZ::Render::CustomMaterialId{ materialAssignment.first.m_lodIndex, materialAssignment.first.m_materialSlotStableId },
+                    AZ::Render::CustomMaterialInfo{ materialAssignment.second.m_materialInstance, materialAssignment.second.m_matModUvOverrides });
+            }
+            m_meshFeatureProcessor->SetCustomMaterials(*m_meshHandle, customMaterials);
         }
     }
 
@@ -670,8 +677,15 @@ namespace AZ::Render
             meshDescriptor.m_isAlwaysDynamic = true;
             meshDescriptor.m_excludeFromReflectionCubeMaps = true;
 
+            AZ::Render::CustomMaterialMap customMaterials;
+            for (const auto& materialAssignment : materials)
+            {
+                customMaterials.emplace(
+                    AZ::Render::CustomMaterialId{ materialAssignment.first.m_lodIndex, materialAssignment.first.m_materialSlotStableId },
+                    AZ::Render::CustomMaterialInfo{ materialAssignment.second.m_materialInstance, materialAssignment.second.m_matModUvOverrides });
+            }
             m_meshHandle = AZStd::make_shared<MeshFeatureProcessorInterface::MeshHandle>(
-                m_meshFeatureProcessor->AcquireMesh(meshDescriptor, materials));
+                m_meshFeatureProcessor->AcquireMesh(meshDescriptor, customMaterials));
         }
 
         // If render proxies already exist, they will be auto-freed

Gems/Terrain/Code/Source/TerrainRenderer/Components/TerrainSurfaceMaterialsListComponent.h

@@ -13,7 +13,6 @@
 #include <AzCore/Component/TickBus.h>
 #include <AzCore/Serialization/EditContext.h>
 
-#include <Atom/Feature/Material/MaterialAssignment.h>
 #include <Atom/RPI.Reflect/Material/MaterialAsset.h>
 
 #include <LmbrCentral/Shape/ShapeComponentBus.h>

Gems/WhiteBox/Code/Source/Rendering/Atom/WhiteBoxAtomRenderMesh.cpp

@@ -32,6 +32,7 @@ namespace WhiteBox
 
     AtomRenderMesh::~AtomRenderMesh()
     {
+        m_materialInstance = {};
         if (m_meshHandle.IsValid() && m_meshFeatureProcessor)
         {
             m_meshFeatureProcessor->ReleaseMesh(m_meshHandle);
@@ -167,10 +168,7 @@ namespace WhiteBox
         
         if (auto materialAsset = AZ::RPI::AssetUtils::LoadAssetByProductPath<AZ::RPI::MaterialAsset>(TexturedMaterialPath.data()))
         {
-            auto materialOverrideInstance = AZ::RPI::Material::FindOrCreate(materialAsset);
-            auto& materialAssignment = m_materialMap[AZ::Render::DefaultMaterialAssignmentId];
-            materialAssignment.m_materialAsset = materialAsset;
-            materialAssignment.m_materialInstance = materialOverrideInstance;
+            m_materialInstance = AZ::RPI::Material::FindOrCreate(materialAsset);
 
             AZ::RPI::ModelMaterialSlot materialSlot;
             materialSlot.m_stableId = OneMaterialSlotId;
@@ -201,7 +199,7 @@ namespace WhiteBox
         }
 
         m_meshFeatureProcessor->ReleaseMesh(m_meshHandle);
-        m_meshHandle = m_meshFeatureProcessor->AcquireMesh(AZ::Render::MeshHandleDescriptor{ m_modelAsset });
+        m_meshHandle = m_meshFeatureProcessor->AcquireMesh(AZ::Render::MeshHandleDescriptor{ m_modelAsset }, m_materialInstance);
         AZ::Render::MeshHandleStateNotificationBus::Event(m_entityId, &AZ::Render::MeshHandleStateNotificationBus::Events::OnMeshHandleSet, &m_meshHandle);
 
         return true;
@@ -271,20 +269,24 @@ namespace WhiteBox
 
     void AtomRenderMesh::UpdateMaterial(const WhiteBoxMaterial& material)
     {
-        if (m_meshFeatureProcessor)
-        {
-            auto& materialAssignment = m_materialMap[AZ::Render::DefaultMaterialAssignmentId];
-            materialAssignment.m_propertyOverrides[AZ::Name("baseColor.color")] = AZ::Color(material.m_tint);
-            materialAssignment.m_propertyOverrides[AZ::Name("baseColor.useTexture")] = material.m_useTexture;
-            // if ApplyProperties fails, defer updating the material assignment map 
-            // on the next tick, and try applying properties again
-            if (materialAssignment.ApplyProperties())
+        if (m_meshFeatureProcessor && m_materialInstance)
+        {            
+            if (const auto& materialPropertyIndex = m_materialInstance->FindPropertyIndex(AZ::Name("baseColor.color"));
+                materialPropertyIndex.IsValid())
             {
-                if (AZ::TickBus::Handler::BusIsConnected())
-                {
-                    AZ::TickBus::Handler::BusDisconnect();
-                }
-                m_meshFeatureProcessor->SetMaterialAssignmentMap(m_meshHandle, m_materialMap);
+                m_materialInstance->SetPropertyValue(materialPropertyIndex, AZ::Color(material.m_tint));
+            }
+
+            if (const auto& materialPropertyIndex = m_materialInstance->FindPropertyIndex(AZ::Name("baseColor.useTexture"));
+                materialPropertyIndex.IsValid())
+            {
+                m_materialInstance->SetPropertyValue(materialPropertyIndex, material.m_useTexture);
+            }
+
+            // If the material changes were successfully applied then disconnect from the tick bus. Otherwise, make another attempt on the next tick.
+            if (!m_materialInstance->NeedsCompile() || m_materialInstance->Compile())
+            {
+                AZ::TickBus::Handler::BusDisconnect();
             }
             else if (!AZ::TickBus::Handler::BusIsConnected())
             {
@@ -295,10 +297,8 @@ namespace WhiteBox
 
     void AtomRenderMesh::OnTick([[maybe_unused]] float deltaTime, [[maybe_unused]] AZ::ScriptTimePoint time)
     {
-        auto& materialAssignment = m_materialMap[AZ::Render::DefaultMaterialAssignmentId];
-        if (materialAssignment.ApplyProperties())
+        if (!m_materialInstance || !m_materialInstance->NeedsCompile() || m_materialInstance->Compile())
         {
-            m_meshFeatureProcessor->SetMaterialAssignmentMap(m_meshHandle, m_materialMap);
             AZ::TickBus::Handler::BusDisconnect();
         }
     }

Gems/WhiteBox/Code/Source/Rendering/Atom/WhiteBoxAtomRenderMesh.h

@@ -93,7 +93,7 @@ namespace WhiteBox
         AZ::Data::Instance<AZ::RPI::Model> m_model;
         AZ::Render::MeshFeatureProcessorInterface* m_meshFeatureProcessor = nullptr;
         AZ::Render::MeshFeatureProcessorInterface::MeshHandle m_meshHandle;
-        AZ::Render::MaterialAssignmentMap m_materialMap;
+        AZ::Data::Instance<AZ::RPI::Material> m_materialInstance;
         uint32_t m_vertexCount = 0;
         AZStd::unique_ptr<IndexBuffer> m_indexBuffer;
         AZStd::array<