Support custom instance index

blade-graphics/Cargo.toml

@@ -22,7 +22,7 @@ raw-window-handle = "0.5"
 [target.'cfg(any(target_os = "ios", target_os = "macos"))'.dependencies]
 block = "0.1"
 core-graphics-types = "0.1"
-metal = { git = "https://github.com/gfx-rs/metal-rs", rev = "33c96191390c1f9c69060f9703667e95dd75b475" }
+metal = { git = "https://github.com/gfx-rs/metal-rs", rev = "64d98e700c301aee0a0855ae2f1008750ef33ca9" }
 objc = "0.2.5"
 naga = { workspace = true, features = ["msl-out"] }
 

blade-graphics/src/lib.rs

@@ -376,6 +376,7 @@ pub struct AccelerationStructureInstance {
     pub acceleration_structure_index: u32,
     pub transform: Transform,
     pub mask: u32,
+    pub custom_index: u32,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq)]

blade-graphics/src/metal/command.rs

@@ -337,6 +337,9 @@ impl crate::traits::AccelerationStructureEncoder
             &primitive_acceleration_structures,
         ));
         descriptor.set_instance_count(instance_count as _);
+        descriptor.set_instance_descriptor_type(
+            metal::MTLAccelerationStructureInstanceDescriptorType::UserID,
+        );
         descriptor.set_instance_descriptor_buffer(instance_data.buffer.as_ref());
         descriptor.set_instance_descriptor_buffer_offset(instance_data.offset);
 

blade-graphics/src/metal/resource.rs

@@ -110,18 +110,19 @@ impl super::Context {
         let mut instance_descriptors = Vec::with_capacity(instances.len());
         for instance in instances {
             let transposed = mint::ColumnMatrix3x4::from(instance.transform);
-            instance_descriptors.push(metal::MTLAccelerationStructureInstanceDescriptor {
+            instance_descriptors.push(metal::MTLAccelerationStructureUserIDInstanceDescriptor {
                 acceleration_structure_index: instance.acceleration_structure_index,
                 mask: instance.mask,
                 transformation_matrix: transposed.into(),
                 options: metal::MTLAccelerationStructureInstanceOptions::None,
                 intersection_function_table_offset: 0,
+                user_id: instance.custom_index,
             });
         }
         let buffer = self.device.lock().unwrap().new_buffer_with_data(
             instance_descriptors.as_ptr() as *const _,
-            (mem::size_of::<metal::MTLAccelerationStructureInstanceDescriptor>() * instances.len())
-                as _,
+            (mem::size_of::<metal::MTLAccelerationStructureUserIDInstanceDescriptor>()
+                * instances.len()) as _,
             metal::MTLResourceOptions::StorageModeShared,
         );
         super::Buffer {

blade-graphics/src/vulkan/resource.rs

@@ -155,7 +155,7 @@ impl super::Context {
             let vk_instance = vk::AccelerationStructureInstanceKHR {
                 transform: unsafe { mem::transmute(instance.transform) },
                 instance_custom_index_and_mask: vk::Packed24_8::new(
-                    0, //TODO? Metal doesn't support it
+                    instance.custom_index,
                     instance.mask as u8,
                 ),
                 instance_shader_binding_table_record_offset_and_flags: vk::Packed24_8::new(0, 0),

blade-render/src/gltf_loader.rs

@@ -112,18 +112,21 @@ impl super::Scene {
         let (doc, buffers, _images) = gltf::import(path).unwrap();
         let mut scene = super::Scene::default();
         let g_scene = doc.scenes().next().unwrap();
-        let mut loader = LoadContext {
-            gltf_buffers: &buffers,
-            gpu,
-            encoder: encoder.transfer(),
-            temp_buffers: Vec::new(),
-            scene: &mut scene,
+
+        let mut temp_buffers = {
+            let mut loader = LoadContext {
+                gltf_buffers: &buffers,
+                gpu,
+                encoder: encoder.transfer(),
+                temp_buffers: Vec::new(),
+                scene: &mut scene,
+            };
+            for g_node in g_scene.nodes() {
+                loader.populate(g_node, glam::Mat4::IDENTITY);
+            }
+            loader.temp_buffers
         };
-        for g_node in g_scene.nodes() {
-            loader.populate(g_node, glam::Mat4::IDENTITY);
-        }
 
-        let mut temp_buffers = loader.temp_buffers;
         scene.populate_bottom_level_acceleration_structures(gpu, encoder, &mut temp_buffers);
         (scene, temp_buffers)
     }

blade-render/src/lib.rs

@@ -32,6 +32,7 @@ pub struct Renderer {
     draw_pipeline: blade::RenderPipeline,
     scene: Scene,
     acceleration_structure: blade::AccelerationStructure,
+    data_buffers: Vec<blade::Buffer>,
     is_tlas_dirty: bool,
     screen_size: blade::Extent,
 }
@@ -40,4 +41,5 @@ pub struct Camera {
     pub pos: mint::Vector3<f32>,
     pub rot: mint::Quaternion<f32>,
     pub fov_y: f32,
+    pub depth: f32,
 }

blade-render/src/renderer.rs

@@ -84,14 +84,17 @@ impl super::Scene {
     ) -> blade::AccelerationStructure {
         let mut instances = Vec::with_capacity(self.objects.len());
         let mut blases = Vec::with_capacity(self.objects.len());
+        let mut custom_index = 0;
 
         for object in self.objects.iter() {
             instances.push(blade::AccelerationStructureInstance {
                 acceleration_structure_index: blases.len() as u32,
                 transform: object.transform.into(),
                 mask: 0xFF,
+                custom_index,
             });
             blases.push(object.acceleration_structure);
+            custom_index += object.geometries.len() as u32;
         }
 
         // Needs to be a separate encoder in order to force synchronization
@@ -184,6 +187,7 @@ impl super::Renderer {
             rt_pipeline,
             draw_pipeline,
             acceleration_structure: blade::AccelerationStructure::default(),
+            data_buffers: Vec::new(),
             is_tlas_dirty: true,
             screen_size,
         }
@@ -215,16 +219,21 @@ impl super::Renderer {
         temp_buffers: &mut Vec<blade::Buffer>,
     ) {
         if self.is_tlas_dirty {
+            self.is_tlas_dirty = false;
+            //TODO: properly remove the old TLAS and buffers
             self.acceleration_structure = self.scene.build_top_level_acceleration_structure(
                 command_encoder,
                 gpu,
                 temp_buffers,
             );
-            self.is_tlas_dirty = false;
+            temp_buffers.extend(self.data_buffers.drain(..));
+            //TODO
         }
     }
 
     pub fn ray_trace(&self, command_encoder: &mut blade::CommandEncoder, camera: &super::Camera) {
+        assert!(!self.is_tlas_dirty);
+
         let mut pass = command_encoder.compute();
         let mut pc = pass.with(&self.rt_pipeline);
         let wg_size = self.rt_pipeline.get_workgroup_size();
@@ -240,7 +249,7 @@ impl super::Renderer {
             &ShaderData {
                 parameters: Parameters {
                     cam_position: camera.pos.into(),
-                    depth: 100.0,
+                    depth: camera.depth,
                     cam_orientation: camera.rot.into(),
                     fov: [fov_x, camera.fov_y],
                     pad: [0.0; 2],

examples/ray-query/main.rs

@@ -172,6 +172,7 @@ impl Example {
                 ]
                 .into(),
                 mask: 0xFF,
+                custom_index: 0,
             },
             blade::AccelerationStructureInstance {
                 acceleration_structure_index: 0,
@@ -182,6 +183,7 @@ impl Example {
                 ]
                 .into(),
                 mask: 0xFF,
+                custom_index: 0,
             },
         ];
         let tlas_sizes = context.get_top_level_acceleration_structure_sizes(instances.len() as u32);

examples/scene/main.rs

@@ -119,6 +119,7 @@ fn main() {
         pos: [0.0, 1.0, 5.0].into(),
         rot: [0.0, 1.0, 0.0, 0.0].into(),
         fov_y: 0.3,
+        depth: 100.0,
     };
     let mut example = Example::new(&window, "examples/scene/data/cornellBox.gltf", camera);