From a1be8f832e4218db213fa2a4b131a8a426332c3e Mon Sep 17 00:00:00 2001
From: dmcdiarmid-ly <63674186+dmcdiarmid-ly@users.noreply.github.com>
Date: Wed, 15 Dec 2021 19:38:39 -0700
Subject: [PATCH] Build BLAS objects from the BLAS instance list instead of the
 mesh list

Signed-off-by: dmcdiarmid-ly <63674186+dmcdiarmid-ly@users.noreply.github.com>
---
 .../RayTracingAccelerationStructurePass.cpp   | 13 ++++----
 .../RayTracing/RayTracingFeatureProcessor.cpp | 19 ++++-------
 .../RayTracing/RayTracingFeatureProcessor.h   | 32 ++++++++++---------
 3 files changed, 30 insertions(+), 34 deletions(-)

diff --git a/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingAccelerationStructurePass.cpp b/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingAccelerationStructurePass.cpp
index abdf59a915..2073c0b8c0 100644
--- a/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingAccelerationStructurePass.cpp
+++ b/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingAccelerationStructurePass.cpp
@@ -150,18 +150,17 @@ namespace AZ
             }
 
             // build newly added BLAS objects
-            // [GFX TODO][ATOM-14159] Add changelist for meshes in the RayTracingFeatureProcessor
-            RayTracingFeatureProcessor::MeshMap& rayTracingMeshes = rayTracingFeatureProcessor->GetMeshes();
-            for (auto& rayTracingMesh : rayTracingMeshes)
+            RayTracingFeatureProcessor::BlasInstanceMap& blasInstances = rayTracingFeatureProcessor->GetBlasInstances();
+            for (auto& blasInstance : blasInstances)
             {
-                if (rayTracingMesh.second.m_blasBuilt == false)
+                if (blasInstance.second.m_blasBuilt == false)
                 {
-                    for (auto& rayTracingSubMesh : rayTracingMesh.second.m_subMeshes)
+                    for (auto& blasInstanceSubMesh : blasInstance.second.m_subMeshes)
                     {
-                        context.GetCommandList()->BuildBottomLevelAccelerationStructure(*rayTracingSubMesh.m_blas);
+                        context.GetCommandList()->BuildBottomLevelAccelerationStructure(*blasInstanceSubMesh.m_blas);
                     }
 
-                    rayTracingMesh.second.m_blasBuilt = true;
+                    blasInstance.second.m_blasBuilt = true;
                 }
             }
 
diff --git a/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.cpp b/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.cpp
index bef95ced17..0a9782980f 100644
--- a/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.cpp
+++ b/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.cpp
@@ -149,23 +149,18 @@ namespace AZ
                 {
                     AZ_Assert(blasInstanceFound == false, "Partial set of RayTracingBlas objects found for mesh");
 
-                    // create the BLAS object
-                    subMesh.m_blas = AZ::RHI::RayTracingBlas::CreateRHIRayTracingBlas();
+                    // create the BLAS object and store it in the BLAS list
+                    RHI::Ptr<RHI::RayTracingBlas> rayTracingBlas = AZ::RHI::RayTracingBlas::CreateRHIRayTracingBlas();
+                    itMeshBlasInstance->second.m_subMeshes.push_back({ rayTracingBlas });
 
-                    // create the buffers from the descriptor
-                    subMesh.m_blas->CreateBuffers(*device, &blasDescriptor, *m_bufferPools);
+                    // create the buffers from the BLAS descriptor
+                    rayTracingBlas->CreateBuffers(*device, &blasDescriptor, *m_bufferPools);
 
-                    // store the BLAS in the side list
-                    itMeshBlasInstance->second.m_subMeshes.push_back({ subMesh.m_blas });
+                    // store the BLAS in the mesh
+                    subMesh.m_blas = rayTracingBlas;
                 }
             }
 
-            if (blasInstanceFound)
-            {
-                // set the mesh BLAS flag so we don't try to rebuild it in the RayTracingAccelerationStructurePass
-                mesh.m_blasBuilt = true;
-            }
-
             // set initial transform
             mesh.m_transform = m_transformServiceFeatureProcessor->GetTransformForId(objectId);
             mesh.m_nonUniformScale = m_transformServiceFeatureProcessor->GetNonUniformScaleForId(objectId);
diff --git a/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.h b/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.h
index fbed3b80cf..d098fca35a 100644
--- a/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.h
+++ b/Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.h
@@ -127,9 +127,6 @@ namespace AZ
 
                 // mesh non-uniform scale
                 AZ::Vector3 m_nonUniformScale = AZ::Vector3::CreateOne();
-
-                // flag indicating if the Blas objects in the sub-meshes are built
-                bool m_blasBuilt = false;
             };
 
             using MeshMap = AZStd::map<uint32_t, Mesh>;
@@ -184,6 +181,23 @@ namespace AZ
             //! Updates the RayTracingSceneSrg and RayTracingMaterialSrg, called after the TLAS is allocated in the RayTracingAccelerationStructurePass
             void UpdateRayTracingSrgs();
 
+            struct SubMeshBlasInstance
+            {
+                RHI::Ptr<RHI::RayTracingBlas> m_blas;
+            };
+
+            struct MeshBlasInstance
+            {
+                uint32_t m_count = 0;
+                AZStd::vector<SubMeshBlasInstance> m_subMeshes;
+
+                // flag indicating if the Blas objects in the sub-mesh list are built
+                bool m_blasBuilt = false;
+            };
+
+            using BlasInstanceMap = AZStd::unordered_map<AZ::Data::AssetId, MeshBlasInstance>;
+            BlasInstanceMap& GetBlasInstances() { return m_blasInstanceMap; }
+
         private:
 
             AZ_DISABLE_COPY_MOVE(RayTracingFeatureProcessor);
@@ -266,18 +280,6 @@ namespace AZ
             bool m_materialInfoBufferNeedsUpdate = false;
 
             // side list for looking up existing BLAS objects so they can be re-used when the same mesh is added multiple times
-            struct SubMeshBlasInstance
-            {
-                RHI::Ptr<RHI::RayTracingBlas> m_blas;
-            };
-
-            struct MeshBlasInstance
-            {
-                uint32_t m_count = 0;
-                AZStd::vector<SubMeshBlasInstance> m_subMeshes;
-            };
-
-            using BlasInstanceMap = AZStd::unordered_map<AZ::Data::AssetId, MeshBlasInstance>;
             BlasInstanceMap m_blasInstanceMap;
 
             // Cache view pointers so we dont need to update them if none changed from frame to frame.