- added support for tlas update

Author: GBDixonAlex

plugins/ecs_examples/src/raytraced_shadows.rs

@@ -16,7 +16,8 @@ pub struct BLAS {
 
 #[derive(Component)]
 pub struct TLAS {
-    pub tlas: Option<gfx_platform::RaytracingTLAS>
+    pub tlas: Option<gfx_platform::RaytracingTLAS>,
+    pub instance_buffer: Option<gfx_platform::Buffer>
 }
 
 /// Setup multiple draw calls with draw indexed and per draw call push constants for transformation matrix etc.
@@ -28,9 +29,10 @@ pub fn raytraced_shadows(client: &mut Client<gfx_platform::Device, os_platform::
             "setup_raytraced_shadows_scene"
         ],
         update: systems![
-            "setup_raytraced_shadows_tlas",
+            "animate_meshes",
             "animate_lights",
-            "batch_lights"
+            "batch_lights",
+            "update_tlas"
         ],
         render_graph: "mesh_lit_rt_shadow",
         ..Default::default()
@@ -64,8 +66,9 @@ pub fn setup_raytraced_shadows_scene(
     mut commands: Commands) -> Result<(), hotline_rs::Error> {
 
     let cube_mesh = hotline_rs::primitives::create_cube_mesh(&mut device.0);
+
     let tourus_mesh = hotline_rs::primitives::create_tourus_mesh(&mut device.0, 32);
-    let helix_mesh = hotline_rs::primitives::create_helix_mesh(&mut device.0, 32, 4);
+    let teapot_mesh = hotline_rs::primitives::create_teapot_mesh(&mut device.0, 32);
     let tube_mesh = hotline_rs::primitives::create_tube_prism_mesh(&mut device.0, 32, 0, 32, true, true, 1.0, 0.66, 1.0);
     let triangle_mesh = hotline_rs::primitives::create_tube_prism_mesh(&mut device.0, 3, 0, 3, false, true, 0.33, 0.66, 1.0);
 
@@ -115,9 +118,9 @@ pub fn setup_raytraced_shadows_scene(
         Position(vec3f(shape_bounds * -0.3, shape_bounds * -0.6, shape_bounds * 0.8)),
         Scale(splat3f(tourus_size * 2.0)),
         Rotation(Quatf::identity()),
-        MeshComponent(helix_mesh.clone()),
+        MeshComponent(teapot_mesh.clone()),
         WorldMatrix(Mat34f::identity()),
-        blas_from_mesh(&mut device, &helix_mesh)?
+        blas_from_mesh(&mut device, &teapot_mesh)?
     ));
 
     // tube
@@ -190,15 +193,16 @@ pub fn setup_raytraced_shadows_scene(
 
     commands.spawn(
         TLAS {
-            tlas: None
+            tlas: None,
+            instance_buffer: None
         }
     );
 
     Ok(())
 }
 
 #[export_update_fn]
-pub fn setup_raytraced_shadows_tlas(
+pub fn update_tlas(
     mut device: ResMut<DeviceRes>,
     mut pmfx: ResMut<PmfxRes>,
     mut entities_query: Query<(&mut Position, &mut Scale, &mut Rotation, &BLAS)>,
@@ -207,31 +211,31 @@ pub fn setup_raytraced_shadows_tlas(
 
     // ..
     for mut t in &mut tlas_query {
+        let mut instances = Vec::new();
+        for (index, (position, scale, rotation, blas)) in &mut entities_query.iter().enumerate() {
+            let translate = Mat34f::from_translation(position.0);
+            let rotate = Mat34f::from(rotation.0);
+            let scale = Mat34f::from_scale(scale.0);
+            instances.push(
+                gfx::RaytracingInstanceInfo::<gfx_platform::Device> {
+                    transform: (translate * rotate * scale).m,
+                    instance_id: index as u32,
+                    instance_mask: 0xff,
+                    hit_group_index: 0,
+                    instance_flags: 0,
+                    blas: &blas.blas
+                }
+            );
+        }
         if t.tlas.is_none() {
-            let mut instances = Vec::new();
-            for (index, (position, scale, rotation, blas)) in &mut entities_query.iter().enumerate() {
-                let translate = Mat34f::from_translation(position.0);
-                let rotate = Mat34f::from(rotation.0);
-                let scale = Mat34f::from_scale(scale.0);
-                instances.push(
-                    gfx::RaytracingInstanceInfo::<gfx_platform::Device> {
-                        transform: (translate * rotate * scale).m,
-                        instance_id: index as u32,
-                        instance_mask: 0xff,
-                        hit_group_index: 0,
-                        instance_flags: 0,
-                        blas: &blas.blas
-                    }
-                );
-                let tlas = device.create_raytracing_tlas_with_heap(&gfx::RaytracingTLASInfo {
-                    instances: &instances,
-                    build_flags: gfx::AccelerationStructureBuildFlags::PREFER_FAST_TRACE,
-                    },
-                    &mut pmfx.shader_heap
-                )?;
-                
-                t.tlas = Some(tlas);
-            }
+            let tlas = device.create_raytracing_tlas_with_heap(&gfx::RaytracingTLASInfo {
+                instances: &instances,
+                build_flags: gfx::AccelerationStructureBuildFlags::PREFER_FAST_TRACE |
+                    gfx::AccelerationStructureBuildFlags::ALLOW_UPDATE
+                },
+                &mut pmfx.shader_heap
+            )?;
+            t.tlas = Some(tlas);
         }
     }
 
@@ -246,7 +250,6 @@ pub fn animate_lights (
     let extent = 60.0;
     for (mut position, _) in &mut light_query {
         position.0 = vec3f(sin(time.accumulated), cos(time.accumulated), cos(time.accumulated)) * extent;
-        position.0 += vec3f(100.0, 10.0, 100.0);
     }
 
     Ok(())
@@ -259,16 +262,45 @@ pub fn animate_lights (
 
 #[export_compute_fn]
 pub fn render_meshes_raytraced(
-    device: ResMut<DeviceRes>,
+    mut device: ResMut<DeviceRes>,
     pmfx: &Res<PmfxRes>,
     pass: &pmfx::ComputePass<gfx_platform::Device>,
-    tlas_query: Query<&TLAS>
+    mut entities_query: Query<(&mut Position, &mut Scale, &mut Rotation, &BLAS)>,
+    mut tlas_query: Query<&mut TLAS>,
 ) -> Result<(), hotline_rs::Error> {
     let pmfx = &pmfx.0;
 
+    let mut heap = pmfx.shader_heap.clone();
+
     let output_size = pmfx.get_texture_2d_size("staging_output").expect("expected staging_output");
     let output_tex = pmfx.get_texture("staging_output").expect("expected staging_output");
 
+    // update tlas
+    for mut t in &mut tlas_query {
+        let mut instances = Vec::new();
+        for (index, (position, scale, rotation, blas)) in &mut entities_query.iter().enumerate() {
+            let translate = Mat34f::from_translation(position.0);
+            let rotate = Mat34f::from(rotation.0);
+            let scale = Mat34f::from_scale(scale.0);
+            instances.push(
+                gfx::RaytracingInstanceInfo::<gfx_platform::Device> {
+                    transform: (translate * rotate * scale).m,
+                    instance_id: index as u32,
+                    instance_mask: 0xff,
+                    hit_group_index: 0,
+                    instance_flags: 0,
+                    blas: &blas.blas
+                }
+            );
+        }
+        
+        if let Some(tlas) = t.tlas.as_ref() {
+            let instance_buffer = device.create_raytracing_instance_buffer(&instances, &mut heap)?;
+            pass.cmd_buf.update_raytracing_tlas(tlas, &instance_buffer, instances.len(), gfx::AccelerationStructureRebuildMode::Refit);
+            t.instance_buffer = Some(instance_buffer);
+        }
+    }
+
     let camera = pmfx.get_camera_constants("main_camera");
     if let Ok(camera) = camera {
         for t in &tlas_query {            

src/gfx.rs

@@ -788,6 +788,13 @@ pub enum RaytracingHitGeometry {
     ProceduralPrimitive
 }
 
+pub enum AccelerationStructureRebuildMode {
+    /// Quick refit will update the transforms only
+    Refit,
+    /// Full update will update BLAS topologies
+    Full
+}
+
 #[derive(Serialize, Deserialize, Clone)]
 pub struct RaytracingHitGroup {
     pub name: String,
@@ -1214,6 +1221,12 @@ pub trait Device: 'static + Send + Sync + Sized + Any + Clone {
         data: Option<&[T]>,
         heap: &mut Self::Heap
     ) -> Result<Self::Buffer, Error>;
+    /// Create an upload buffer that can be used specifically for raytracing acceleration structure instances
+    fn create_raytracing_instance_buffer(
+        &mut self,
+        instances: &Vec<RaytracingInstanceInfo<Self>>,
+        heap: &mut Self::Heap // TODO; need to refactor heap handling
+    ) -> Result<Self::Buffer, Error>;  
     /// Create a `Buffer` specifically for reading back data from the GPU mainly for `Query` use
     fn create_read_back_buffer(
         &mut self,
@@ -1398,8 +1411,6 @@ pub trait CmdBuf<D: Device>: Send + Sync + Clone {
     /// Binds the heap with offset (texture srv, uav) on to the `slot` of a pipeline.
     /// this is like a traditional bindful render architecture `cmd.set_binding(pipeline, heap, 0, texture1_id)`
     fn set_binding<T: Pipeline>(&self, pipeline: &T, heap: &D::Heap, slot: u32, offset: usize);
-    // TODO:
-    fn set_tlas(&self, tlas: &D::RaytracingTLAS);
     /// Push a small amount of data into the command buffer for a render pipeline, num values and dest offset are the number of 32bit values
     fn push_render_constants<T: Sized>(&self, slot: u32, num_values: u32, dest_offset: u32, data: &[T]);
     /// Push a small amount of data into the command buffer for a compute pipeline, num values and dest offset are the number of 32bit values
@@ -1435,6 +1446,8 @@ pub trait CmdBuf<D: Device>: Send + Sync + Clone {
     );
     /// Issue dispatch call for ray tracing with the specified `RaytracingShaderBindingTable` which is associated with the bound `RaytracingPipeline`
     fn dispatch_rays(&self, sbt: &D::RaytracingShaderBindingTable, numthreads: Size3);
+    /// Update a raytracing TLAS with instance transform info contained in `instance_buffer` of length `instance_count`. Use `mode` to control quick refit or full rebuild
+    fn update_raytracing_tlas(&self, tlas: &D::RaytracingTLAS, instance_buffer: &D::Buffer, instance_count: usize, mode: AccelerationStructureRebuildMode);
     /// Resolves the `subresource` (mip index, 3d texture slice or array slice)
     fn resolve_texture_subresource(&self, texture: &D::Texture, subresource: u32) -> Result<(), Error>;
     /// Generates a full mip chain for the specified `texture` where `heap` is the shader heap the texture was created on 

src/gfx/d3d12.rs

@@ -812,6 +812,13 @@ fn to_d3d12_raytracing_acceleration_structure_build_flags
     d3d12_flags
 }
 
+fn to_d3d12_raytracing_acceleration_structure_update_flags(mode: AccelerationStructureRebuildMode) -> D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAGS {
+    match mode {
+        AccelerationStructureRebuildMode::Refit => D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_PERFORM_UPDATE,
+        AccelerationStructureRebuildMode::Full => D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BUILD_FLAG_NONE
+    }
+}
+
 fn get_d3d12_error_blob_string(blob: &ID3DBlob) -> String {
     unsafe {
         String::from_raw_parts(
@@ -1720,6 +1727,7 @@ pub struct RaytracingBLAS {
 
 pub struct RaytracingTLAS {
     pub(crate) tlas_buffer: Buffer,
+    pub(crate) scratch_buffer: Option<Buffer>,
     pub(crate) shader_heap_id: u16
 }
 
@@ -3436,19 +3444,19 @@ impl super::Device for Device {
         }
     }
 
-    fn create_raytracing_tlas_with_heap(
+    fn create_raytracing_instance_buffer(
         &mut self,
-        info: &RaytracingTLASInfo<Self>,
-        heap: &mut Heap
-    ) -> result::Result<RaytracingTLAS, super::Error> {
+        instances: &Vec<RaytracingInstanceInfo<Self>>,
+        heap: &mut Heap 
+    ) -> result::Result<Buffer, super::Error> {
         // pack 24: 8 bits
         let pack_24_8 = |a, b| {
             (a & 0x00ffffff) | ((b & 0x000000ff) << 24)
         };
 
-        // create instance descslea
-        let num_instances = info.instances.len();
-        let instance_descs: Vec<D3D12_RAYTRACING_INSTANCE_DESC> = info.instances.iter()
+        // create instance descs
+        let num_instances = instances.len();
+        let instance_descs: Vec<D3D12_RAYTRACING_INSTANCE_DESC> = instances.iter()
             .map(|x| 
                 D3D12_RAYTRACING_INSTANCE_DESC {
                 Transform: x.transform,
@@ -3460,6 +3468,7 @@ impl super::Device for Device {
 
         // create upload buffer for instance descs
         let stride = std::mem::size_of::<D3D12_RAYTRACING_INSTANCE_DESC>();
+
         let instance_buffer = self.create_buffer_with_heap(&BufferInfo {
             usage: super::BufferUsage::UPLOAD,
             cpu_access: super::CpuAccessFlags::NONE,
@@ -3470,13 +3479,25 @@ impl super::Device for Device {
             }, 
             Some(instance_descs.as_slice()),
             heap
-        ).expect(format!("hotline_rs::gfx::d3d12: failed to create a scratch buffer for raytracing blas of size {}", stride * num_instances).as_str());
+        )?;
+
+        Ok(instance_buffer)
+    }
+
+    fn create_raytracing_tlas_with_heap(
+        &mut self,
+        info: &RaytracingTLASInfo<Self>,
+        heap: &mut Heap
+    ) -> result::Result<RaytracingTLAS, super::Error> {
+
+        // create instance buffer
+        let instance_buffer = self.create_raytracing_instance_buffer(info.instances, heap)?;
 
         // create acceleration structure inputs
         let inputs = D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS {
             Type: D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL,
             Flags: to_d3d12_raytracing_acceleration_structure_build_flags(info.build_flags),
-            NumDescs: num_instances as u32,
+            NumDescs: info.instances.len() as u32,
             DescsLayout: D3D12_ELEMENTS_LAYOUT_ARRAY,
             Anonymous: D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS_0 {
                 InstanceDescs: instance_buffer.d3d_virtual_address(),
@@ -3541,6 +3562,12 @@ impl super::Device for Device {
         // return the result
         Ok(RaytracingTLAS {
             tlas_buffer,
+            scratch_buffer: if info.build_flags.contains(AccelerationStructureBuildFlags::ALLOW_UPDATE) {
+                Some(scratch_buffer)
+            }
+            else {
+                None
+            },
             shader_heap_id: self.shader_heap.as_ref().map(|x| x.id).unwrap_or(0)
         })
     }
@@ -4342,13 +4369,6 @@ impl super::CmdBuf<Device> for CmdBuf {
         }
     }
 
-    fn set_tlas(&self, tlas: &RaytracingTLAS) {
-        let cmd = self.cmd().cast::<ID3D12GraphicsCommandList4>().unwrap();
-        unsafe {
-            cmd.SetComputeRootShaderResourceView(0, tlas.tlas_buffer.resource.as_ref().unwrap().GetGPUVirtualAddress());
-        }
-    }
-
     fn dispatch_rays(&self, sbt: &RaytracingShaderBindingTable, numthreads: Size3) {
         unsafe {
             let dispatch_desc = D3D12_DISPATCH_RAYS_DESC {
@@ -4381,6 +4401,39 @@ impl super::CmdBuf<Device> for CmdBuf {
         }
     }
 
+    fn update_raytracing_tlas(&self, tlas: &RaytracingTLAS, instance_buffer: &Buffer, instance_count: usize, mode: AccelerationStructureRebuildMode) {
+
+        // create acceleration structure inputs
+        let inputs = D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS {
+            Type: D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL,
+            Flags: to_d3d12_raytracing_acceleration_structure_update_flags(mode),
+            NumDescs: instance_count as u32,
+            DescsLayout: D3D12_ELEMENTS_LAYOUT_ARRAY,
+            Anonymous: D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS_0 {
+                InstanceDescs: instance_buffer.d3d_virtual_address(),
+            }
+        };
+                
+        // create blas desc
+        let blas_desc = D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_DESC {
+            Inputs: inputs,
+            SourceAccelerationStructureData: tlas.tlas_buffer.d3d_virtual_address(),
+            DestAccelerationStructureData: tlas.tlas_buffer.d3d_virtual_address(),
+            ScratchAccelerationStructureData: tlas.scratch_buffer
+                .as_ref()
+                .expect("hotline_rs::gfx::d3d12: tlas is required to be created with ALLOW_UPDATE in order to update it")
+                .d3d_virtual_address()
+        };
+
+        // build blas
+        unsafe {
+            let bb = self.bb_index;
+            let cmd = self.command_list[bb].cast::<ID3D12GraphicsCommandList4>()
+                .expect("hotline_rs::gfx::d3d12: expected ID3D12GraphicsCommandList4 availability to create raytracing blas");
+            cmd.BuildRaytracingAccelerationStructure(&blas_desc, None);
+        }
+    }
+
     fn read_back_backbuffer(&mut self, swap_chain: &SwapChain) -> result::Result<ReadBackRequest, super::Error> {
         let bb = self.bb_index;
         let bbz = self.bb_index as u32;