Merge pull request #4200 from aws-lumberyard-dev/Atom/dmcdiar/ATOM-16437

Instance the ray tracing BLAS objects

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

@@ -951,7 +951,7 @@ namespace AZ
                 subMeshes.push_back(subMesh);
             }
 
-            rayTracingFeatureProcessor->SetMesh(m_objectId, subMeshes);
+            rayTracingFeatureProcessor->SetMesh(m_objectId, m_model->GetModelAsset()->GetId(), subMeshes);
         }
 
         void MeshDataInstance::SetSortKey(RHI::DrawItemSortKey sortKey)

Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.cpp

@@ -79,7 +79,7 @@ namespace AZ
             AZ_Assert(m_rayTracingMaterialSrg, "Failed to create RayTracingMaterialSrg");
         }
 
-        void RayTracingFeatureProcessor::SetMesh(const ObjectId objectId, const SubMeshVector& subMeshes)
+        void RayTracingFeatureProcessor::SetMesh(const ObjectId objectId, const AZ::Data::AssetId& assetId, const SubMeshVector& subMeshes)
         {
             if (!m_rayTracingEnabled)
             {
@@ -89,10 +89,13 @@ namespace AZ
             RHI::Ptr<RHI::Device> device = RHI::RHISystemInterface::Get()->GetDevice();
             uint32_t objectIndex = objectId.GetIndex();
 
+            // lock the mutex to protect the mesh and BLAS lists
+            AZStd::unique_lock<AZStd::mutex> lock(m_mutex);
+
             MeshMap::iterator itMesh = m_meshes.find(objectIndex);
             if (itMesh == m_meshes.end())
             {
-                m_meshes.insert(AZStd::make_pair(objectIndex, Mesh{ subMeshes }));
+                m_meshes.insert(AZStd::make_pair(objectIndex, Mesh{ assetId, subMeshes }));
             }
             else
             {
@@ -102,9 +105,12 @@ namespace AZ
                 m_meshes[objectIndex].m_subMeshes = subMeshes;
             }
 
-            // create the BLAS buffers for each sub-mesh
+            // create the BLAS buffers for each sub-mesh, or re-use existing BLAS objects if they were already created.
+            // Note: all sub-meshes must either create new BLAS objects or re-use existing ones, otherwise it's an error (it's the same model in both cases)
             // Note: the buffer is just reserved here, the BLAS is built in the RayTracingAccelerationStructurePass
             Mesh& mesh = m_meshes[objectIndex];
+            bool blasInstanceFound = false;
+
             for (auto& subMesh : mesh.m_subMeshes)
             {
                 RHI::RayTracingBlasDescriptor blasDescriptor;
@@ -115,11 +121,37 @@ namespace AZ
                         ->IndexBuffer(subMesh.m_indexBufferView)
                 ;
 
-                // create the BLAS object
-                subMesh.m_blas = AZ::RHI::RayTracingBlas::CreateRHIRayTracingBlas();
+                // search for an existing BLAS object for this model
+                RayTracingBlasMap::iterator itBlas = m_blasMap.find(assetId);
+                if (itBlas != m_blasMap.end())
+                {
+                    // re-use existing BLAS
+                    subMesh.m_blas = itBlas->second.m_blas;
+                    itBlas->second.m_count++;
 
-                // create the buffers from the descriptor
-                subMesh.m_blas->CreateBuffers(*device, &blasDescriptor, *m_bufferPools);
+                    // keep track of the fact that we re-used a BLAS
+                    blasInstanceFound = true;
+                }
+                else
+                {
+                    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 buffers from the descriptor
+                    subMesh.m_blas->CreateBuffers(*device, &blasDescriptor, *m_bufferPools);
+
+                    // store the BLAS in the side list
+                    RayTracingBlasInstance blasInstance = { subMesh.m_blas, 1 };
+                    m_blasMap.insert({ assetId, blasInstance });
+                }
+            }
+
+            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
@@ -140,12 +172,26 @@ namespace AZ
                 return;
             }
 
+            // lock the mutex to protect the mesh and BLAS lists
+            AZStd::unique_lock<AZStd::mutex> lock(m_mutex);
+
             MeshMap::iterator itMesh = m_meshes.find(objectId.GetIndex());
             if (itMesh != m_meshes.end())
             {
                 m_subMeshCount -= aznumeric_cast<uint32_t>(itMesh->second.m_subMeshes.size());
                 m_meshes.erase(itMesh);
                 m_revision++;
+
+                // decrement the count from the BLAS instance, and check to see if we can remove it
+                RayTracingBlasMap::iterator itBlas = m_blasMap.find(itMesh->second.m_assetId);
+                if (itBlas != m_blasMap.end())
+                {
+                    itBlas->second.m_count--;
+                    if (itBlas->second.m_count == 0)
+                    {
+                        m_blasMap.erase(itBlas);
+                    }
+                }                
             }
 
             m_meshInfoBufferNeedsUpdate = true;

Gems/Atom/Feature/Common/Code/Source/RayTracing/RayTracingFeatureProcessor.h

@@ -116,6 +116,9 @@ namespace AZ
             //! Contains data for the top level mesh, including the list of sub-meshes
             struct Mesh
             {
+                // assetId of the model
+                AZ::Data::AssetId m_assetId = AZ::Data::AssetId{};
+
                 // sub-mesh list
                 SubMeshVector m_subMeshes;
 
@@ -134,7 +137,7 @@ namespace AZ
 
             //! Sets ray tracing data for a mesh.
             //! This will cause an update to the RayTracing acceleration structure on the next frame
-            void SetMesh(const ObjectId objectId, const SubMeshVector& subMeshes);
+            void SetMesh(const ObjectId objectId, const AZ::Data::AssetId& assetId, const SubMeshVector& subMeshes);
 
             //! Removes ray tracing data for a mesh.
             //! This will cause an update to the RayTracing acceleration structure on the next frame
@@ -220,6 +223,9 @@ namespace AZ
             // cached TransformServiceFeatureProcessor
             TransformServiceFeatureProcessor* m_transformServiceFeatureProcessor = nullptr;
 
+            // mutex for the mesh and BLAS lists
+            AZStd::mutex m_mutex;
+
             // structure for data in the m_meshInfoBuffer, shaders that use the buffer must match this type
             struct MeshInfo
             {
@@ -260,6 +266,16 @@ namespace AZ
 
             // flag indicating we need to update the materialInfo buffer
             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 RayTracingBlasInstance
+            {
+                RHI::Ptr<RHI::RayTracingBlas> m_blas;
+                uint32_t m_count = 0;
+            };
+
+            using RayTracingBlasMap = AZStd::unordered_map<AZ::Data::AssetId, RayTracingBlasInstance>;
+            RayTracingBlasMap m_blasMap;
         };
     }
 }