Ray Tracing support - stage 1

Cargo.toml

@@ -8,7 +8,7 @@ members = [
 exclude = []
 
 [workspace.dependencies]
-naga = { version = "0.11", features = ["wgsl-in", "span", "validate"] }
+naga = { git = "https://github.com/kvark/naga", branch = "ray-query", features = ["wgsl-in", "span", "validate"] }
 web-sys = "0.3.60"
 
 [lib]
@@ -32,6 +32,7 @@ blade-graphics = { version = "0.1.0", path = "blade-graphics" }
 
 [dev-dependencies]
 bytemuck = { version = "1", features = ["derive"] }
+del-msh = "0.1"
 egui = { version = "0.20", features = ["bytemuck"] }
 naga = { workspace = true }
 nanorand = { version = "0.7", default-features = false, features = ["wyrand"] }

blade-graphics/Cargo.toml

@@ -15,14 +15,14 @@ bytemuck = "1"
 codespan-reporting = "0.11"
 hidden-trait = "0.1"
 log = "0.4"
+mint = "0.5"
 naga = { workspace = true, features = ["clone"] }
 raw-window-handle = "0.5"
 
 [target.'cfg(any(target_os = "ios", target_os = "macos"))'.dependencies]
 block = "0.1"
 core-graphics-types = "0.1"
-foreign-types = "0.3"
-metal = "0.24"
+metal = { git = "https://github.com/kvark/metal-rs", branch = "rt" }
 objc = "0.2.5"
 naga = { workspace = true, features = ["msl-out"] }
 

blade-graphics/src/gles/command.rs

@@ -59,6 +59,13 @@ impl crate::ShaderBindable for crate::BufferPiece {
         }
     }
 }
+impl crate::ShaderBindable for super::AccelerationStructure {
+    fn bind_to(&self, ctx: &mut super::PipelineContext, index: u32) {
+        for _ in ctx.targets[index as usize].iter() {
+            unimplemented!()
+        }
+    }
+}
 
 impl super::CommandEncoder {
     pub fn start(&mut self) {
@@ -84,6 +91,10 @@ impl super::CommandEncoder {
         }
     }
 
+    pub fn acceleration_structure(&mut self) -> super::PassEncoder<()> {
+        unimplemented!()
+    }
+
     pub fn compute(&mut self) -> super::PassEncoder<super::ComputePipeline> {
         super::PassEncoder {
             commands: &mut self.commands,
@@ -215,7 +226,7 @@ impl<T> Drop for super::PassEncoder<'_, T> {
                 .push(super::Command::InvalidateAttachment(attachment));
         }
         match self.kind {
-            super::PassKind::Transfer => {}
+            super::PassKind::Transfer | super::PassKind::AccelerationStructure => {}
             super::PassKind::Compute => {
                 self.commands.push(super::Command::ResetAllSamplers);
             }
@@ -293,6 +304,27 @@ impl crate::traits::TransferEncoder for super::PassEncoder<'_, ()> {
     }
 }
 
+impl super::PassEncoder<'_, ()> {
+    pub fn build_bottom_level(
+        &mut self,
+        _acceleration_structure: super::AccelerationStructure,
+        _meshes: &[crate::AccelerationStructureMesh],
+        _scratch_data: crate::BufferPiece,
+    ) {
+        unimplemented!()
+    }
+
+    pub fn build_top_level(
+        &mut self,
+        _acceleration_structure: super::AccelerationStructure,
+        _instance_count: u32,
+        _instance_data: crate::BufferPiece,
+        _scratch_data: crate::BufferPiece,
+    ) {
+        unimplemented!()
+    }
+}
+
 #[hidden_trait::expose]
 impl crate::traits::PipelineEncoder for super::PipelineEncoder<'_> {
     fn bind<D: crate::ShaderData>(&mut self, group: u32, data: &D) {

blade-graphics/src/gles/mod.rs

@@ -77,6 +77,9 @@ pub struct Sampler {
     raw: glow::Sampler,
 }
 
+#[derive(Clone, Copy, Debug, Hash, PartialEq)]
+pub struct AccelerationStructure {}
+
 type SlotList = Vec<u32>;
 
 struct BindGroupInfo {
@@ -312,6 +315,7 @@ pub struct CommandEncoder {
 
 enum PassKind {
     Transfer,
+    AccelerationStructure,
     Compute,
     Render,
 }
@@ -351,6 +355,12 @@ struct ExecutionContext {
     plain_buffer: glow::Buffer,
 }
 
+impl Context {
+    pub fn capabilities(&self) -> crate::Capabilities {
+        crate::Capabilities { ray_query: false }
+    }
+}
+
 #[hidden_trait::expose]
 impl crate::traits::CommandDevice for Context {
     type CommandEncoder = CommandEncoder;
@@ -439,6 +449,7 @@ fn describe_texture_format(format: crate::TextureFormat) -> FormatInfo {
         Tf::Rgba8Unorm => (glow::RGBA8, glow::RGBA, glow::UNSIGNED_BYTE),
         Tf::Rgba8UnormSrgb => (glow::SRGB8_ALPHA8, glow::RGBA, glow::UNSIGNED_BYTE),
         Tf::Bgra8UnormSrgb => (glow::SRGB8_ALPHA8, glow::BGRA, glow::UNSIGNED_BYTE),
+        Tf::Rgba16Float => (glow::RGBA16F, glow::RGBA, glow::FLOAT),
         Tf::Depth32Float => (glow::DEPTH_COMPONENT32F, glow::DEPTH_COMPONENT, glow::FLOAT),
     };
     FormatInfo {

blade-graphics/src/gles/pipeline.rs

@@ -137,6 +137,9 @@ impl super::Context {
                             targets.push(params[0] as u32);
                         }
                     }
+                    crate::ShaderBinding::AccelerationStructure => {
+                        unimplemented!()
+                    }
                     crate::ShaderBinding::Plain { size } => {
                         if let Some(index) = gl.get_uniform_block_index(program, glsl_name) {
                             let expected_size = gl.get_active_uniform_block_parameter_i32(

blade-graphics/src/gles/resource.rs

@@ -1,6 +1,43 @@
 use glow::HasContext as _;
 use std::{ptr, slice};
 
+impl super::Context {
+    pub fn get_bottom_level_acceleration_structure_sizes(
+        &self,
+        _meshes: &[crate::AccelerationStructureMesh],
+    ) -> crate::AccelerationStructureSizes {
+        unimplemented!()
+    }
+
+    pub fn get_top_level_acceleration_structure_sizes(
+        &self,
+        _instance_count: u32,
+    ) -> crate::AccelerationStructureSizes {
+        unimplemented!()
+    }
+
+    pub fn create_acceleration_structure_instance_buffer(
+        &self,
+        _instances: &[crate::AccelerationStructureInstance],
+    ) -> super::Buffer {
+        unimplemented!()
+    }
+
+    pub fn create_acceleration_structure(
+        &self,
+        _desc: crate::AccelerationStructureDesc,
+    ) -> super::AccelerationStructure {
+        unimplemented!()
+    }
+
+    pub fn destroy_acceleration_structure(
+        &self,
+        _acceleration_structure: super::AccelerationStructure,
+    ) {
+        unimplemented!()
+    }
+}
+
 #[hidden_trait::expose]
 impl crate::traits::ResourceDevice for super::Context {
     type Buffer = super::Buffer;

blade-graphics/src/lib.rs

@@ -24,6 +24,7 @@
 )]
 
 pub use naga::{StorageAccess, VectorSize};
+pub type Transform = mint::RowMatrix3x4<f32>;
 
 #[cfg_attr(
     all(not(vulkan), not(gles), any(target_os = "ios", target_os = "macos")),
@@ -44,6 +45,9 @@ pub mod util;
 pub mod limits {
     pub const PLAIN_DATA_SIZE: u32 = 256;
     pub const RESOURCES_IN_GROUP: u32 = 8;
+    pub const STORAGE_BUFFER_ALIGNMENT: u64 = 256;
+    pub const ACCELERATION_STRUCTURE_BUFFER_ALIGNMENT: u64 = 256;
+    pub const ACCELERATION_STRUCTURE_SCRATCH_ALIGNMENT: u64 = 256;
 }
 
 pub use hal::*;
@@ -59,6 +63,11 @@ pub struct ContextDesc {
 #[derive(Debug)]
 pub struct NotSupportedError;
 
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct Capabilities {
+    pub ray_query: bool,
+}
+
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub enum Memory {
     /// Device-local memory. Fast for GPU operations.
@@ -139,6 +148,7 @@ pub enum TextureFormat {
     Rgba8Unorm,
     Rgba8UnormSrgb,
     Bgra8UnormSrgb,
+    Rgba16Float,
     Depth32Float,
 }
 
@@ -329,6 +339,58 @@ pub struct SamplerDesc<'a> {
     pub border_color: Option<TextureColor>,
 }
 
+#[derive(Debug)]
+pub enum AccelerationStructureType {
+    TopLevel,
+    BottomLevel,
+}
+
+#[derive(Debug)]
+pub struct AccelerationStructureDesc<'a> {
+    pub name: &'a str,
+    pub ty: AccelerationStructureType,
+    pub buffer: Buffer,
+    /// Offset into the buffer where AS is placed.
+    /// Must be a multiple of `limits::ACCELERATION_STRUCTURE_BUFFER_ALIGNMENT`.
+    pub offset: u64,
+    pub size: u64,
+}
+
+#[non_exhaustive]
+#[derive(Clone, Copy, Debug, Hash, Eq, Ord, PartialEq, PartialOrd)]
+pub enum VertexFormat {
+    Rgb32Float,
+}
+
+#[derive(Clone, Debug)]
+pub struct AccelerationStructureMesh {
+    pub vertex_data: BufferPiece,
+    pub vertex_format: VertexFormat,
+    pub vertex_stride: u32,
+    pub vertex_count: u32,
+    pub index_data: BufferPiece,
+    pub index_type: Option<IndexType>,
+    pub triangle_count: u32,
+    pub transform_data: BufferPiece,
+    pub is_opaque: bool,
+}
+
+#[derive(Clone, Debug)]
+pub struct AccelerationStructureInstance {
+    pub acceleration_structure: AccelerationStructure,
+    pub transform: Transform,
+    pub mask: u32,
+    pub custom_index: u32,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub struct AccelerationStructureSizes {
+    /// Size of the permanent GPU data
+    pub data: u64,
+    /// Size of the scratch space
+    pub scratch: u64,
+}
+
 pub struct Shader {
     module: naga::Module,
     info: naga::valid::ModuleInfo,
@@ -356,6 +418,7 @@ pub enum ShaderBinding {
     Texture,
     Sampler,
     Buffer,
+    AccelerationStructure,
     Plain { size: u32 },
 }
 
@@ -380,6 +443,9 @@ impl HasShaderBinding for Sampler {
 impl HasShaderBinding for BufferPiece {
     const TYPE: ShaderBinding = ShaderBinding::Buffer;
 }
+impl HasShaderBinding for AccelerationStructure {
+    const TYPE: ShaderBinding = ShaderBinding::AccelerationStructure;
+}
 
 #[derive(Clone, Debug, Default, PartialEq)]
 pub struct ShaderDataLayout {

blade-graphics/src/metal/command.rs

@@ -62,6 +62,21 @@ impl crate::ShaderBindable for crate::BufferPiece {
         }
     }
 }
+impl crate::ShaderBindable for crate::AccelerationStructure {
+    fn bind_to(&self, ctx: &mut super::PipelineContext, index: u32) {
+        let slot = ctx.targets[index as usize] as _;
+        let value = Some(self.as_ref());
+        if let Some(encoder) = ctx.vs_encoder {
+            encoder.set_vertex_acceleration_structure(slot, value);
+        }
+        if let Some(encoder) = ctx.fs_encoder {
+            encoder.set_fragment_acceleration_structure(slot, value);
+        }
+        if let Some(encoder) = ctx.cs_encoder {
+            encoder.set_acceleration_structure(slot, value);
+        }
+    }
+}
 
 impl super::CommandEncoder {
     pub fn start(&mut self) {
@@ -95,6 +110,20 @@ impl super::CommandEncoder {
         }
     }
 
+    pub fn acceleration_structure(&mut self) -> super::AccelerationStructureCommandEncoder {
+        let raw = objc::rc::autoreleasepool(|| {
+            self.raw
+                .as_mut()
+                .unwrap()
+                .new_acceleration_structure_command_encoder()
+                .to_owned()
+        });
+        super::AccelerationStructureCommandEncoder {
+            raw,
+            phantom: PhantomData,
+        }
+    }
+
     pub fn compute(&mut self) -> super::ComputeCommandEncoder {
         let raw = objc::rc::autoreleasepool(|| {
             self.raw
@@ -272,6 +301,34 @@ impl Drop for super::TransferCommandEncoder<'_> {
     }
 }
 
+impl<'a> super::AccelerationStructureCommandEncoder<'a> {
+    //TODO: move into the trait
+    pub fn build_bottom_level(
+        &mut self,
+        _acceleration_structure: super::AccelerationStructure,
+        _meshes: &[crate::AccelerationStructureMesh],
+        _scratch_data: crate::BufferPiece,
+    ) {
+        unimplemented!()
+    }
+
+    pub fn build_top_level(
+        &mut self,
+        _acceleration_structure: super::AccelerationStructure,
+        _instance_count: u32,
+        _instance_data: crate::BufferPiece,
+        _scratch_data: crate::BufferPiece,
+    ) {
+        unimplemented!()
+    }
+}
+
+impl Drop for super::AccelerationStructureCommandEncoder<'_> {
+    fn drop(&mut self) {
+        self.raw.end_encoding();
+    }
+}
+
 impl super::ComputeCommandEncoder<'_> {
     pub fn with<'p>(
         &'p mut self,