[wgpu-hal] Blas compaction

CHANGELOG.md

@@ -50,6 +50,7 @@ Bottom level categories:
 
 #### General
 
+- Support compaction in wgpu-hal. By @Vecvec in [#7101](https://github.com/gfx-rs/wgpu/pull/7101).
 - Avoid using default features in many dependencies, etc. By Brody in [#7031](https://github.com/gfx-rs/wgpu/pull/7031)
 - Use `hashbrown` to simplify no-std support. By Brody in [#6938](https://github.com/gfx-rs/wgpu/pull/6938) & [#6925](https://github.com/gfx-rs/wgpu/pull/6925).
 - If you use Binding Arrays in a bind group, you may not use Dynamic Offset Buffers or Uniform Buffers in that bind group. By @cwfitzgerald in [#6811](https://github.com/gfx-rs/wgpu/pull/6811)

wgpu-core/src/device/ray_tracing.rs

@@ -85,6 +85,8 @@ impl Device {
                     label: blas_desc.label.as_deref(),
                     size: size_info.acceleration_structure_size,
                     format: hal::AccelerationStructureFormat::BottomLevel,
+                    // change this once compaction is implemented in wgpu-core
+                    allow_compaction: false,
                 })
         }
         .map_err(DeviceError::from_hal)?;
@@ -136,6 +138,7 @@ impl Device {
                     label: desc.label.as_deref(),
                     size: size_info.acceleration_structure_size,
                     format: hal::AccelerationStructureFormat::TopLevel,
+                    allow_compaction: false,
                 })
         }
         .map_err(DeviceError::from_hal)?;

wgpu-hal/examples/ray-traced-triangle/main.rs

@@ -525,6 +525,7 @@ impl<A: hal::Api> Example<A> {
                 label: Some("blas"),
                 size: blas_sizes.acceleration_structure_size,
                 format: hal::AccelerationStructureFormat::BottomLevel,
+                allow_compaction: false,
             })
         }
         .unwrap();
@@ -534,6 +535,7 @@ impl<A: hal::Api> Example<A> {
                 label: Some("tlas"),
                 size: tlas_sizes.acceleration_structure_size,
                 format: hal::AccelerationStructureFormat::TopLevel,
+                allow_compaction: false,
             })
         }
         .unwrap();

wgpu-hal/src/dx12/command.rs

@@ -383,7 +383,9 @@ impl crate::CommandEncoder for super::CommandEncoder {
                     },
                 };
                 self.temp.barriers.push(raw);
-            } else if barrier.usage.from == wgt::BufferUses::STORAGE_READ_WRITE {
+            } else if barrier.usage.from == wgt::BufferUses::STORAGE_READ_WRITE
+                || barrier.usage.from == wgt::BufferUses::ACCELERATION_STRUCTURE_QUERY
+            {
                 let raw = Direct3D12::D3D12_RESOURCE_BARRIER {
                     Type: Direct3D12::D3D12_RESOURCE_BARRIER_TYPE_UAV,
                     Flags: Direct3D12::D3D12_RESOURCE_BARRIER_FLAG_NONE,
@@ -681,6 +683,29 @@ impl crate::CommandEncoder for super::CommandEncoder {
             )
         };
     }
+    unsafe fn read_acceleration_structure_compact_size(
+        &mut self,
+        acceleration_structure: &super::AccelerationStructure,
+        buf: &super::Buffer,
+    ) {
+        let list = self
+            .list
+            .as_ref()
+            .unwrap()
+            .cast::<Direct3D12::ID3D12GraphicsCommandList4>()
+            .unwrap();
+        unsafe {
+            list.EmitRaytracingAccelerationStructurePostbuildInfo(
+                &Direct3D12::D3D12_RAYTRACING_ACCELERATION_STRUCTURE_POSTBUILD_INFO_DESC {
+                    DestBuffer: buf.resource.GetGPUVirtualAddress(),
+                    InfoType: Direct3D12::D3D12_RAYTRACING_ACCELERATION_STRUCTURE_POSTBUILD_INFO_COMPACTED_SIZE,
+                },
+                &[
+                    acceleration_structure.resource.GetGPUVirtualAddress()
+                ],
+            )
+        }
+    }
     unsafe fn reset_queries(&mut self, _set: &super::QuerySet, _range: Range<u32>) {
         // nothing to do here
     }
@@ -1505,4 +1530,25 @@ impl crate::CommandEncoder for super::CommandEncoder {
             }])
         }
     }
+
+    unsafe fn copy_acceleration_structure_to_acceleration_structure(
+        &mut self,
+        src: &super::AccelerationStructure,
+        dst: &super::AccelerationStructure,
+        copy: wgt::AccelerationStructureCopy,
+    ) {
+        let list = self
+            .list
+            .as_ref()
+            .unwrap()
+            .cast::<Direct3D12::ID3D12GraphicsCommandList4>()
+            .unwrap();
+        unsafe {
+            list.CopyRaytracingAccelerationStructure(
+                dst.resource.GetGPUVirtualAddress(),
+                src.resource.GetGPUVirtualAddress(),
+                conv::map_acceleration_structure_copy_mode(copy),
+            )
+        }
+    }
 }

wgpu-hal/src/dx12/conv.rs

@@ -4,7 +4,9 @@ pub fn map_buffer_usage_to_resource_flags(
     usage: wgt::BufferUses,
 ) -> Direct3D12::D3D12_RESOURCE_FLAGS {
     let mut flags = Direct3D12::D3D12_RESOURCE_FLAG_NONE;
-    if usage.contains(wgt::BufferUses::STORAGE_READ_WRITE) {
+    if usage.contains(wgt::BufferUses::STORAGE_READ_WRITE)
+        || usage.contains(wgt::BufferUses::ACCELERATION_STRUCTURE_QUERY)
+    {
         flags |= Direct3D12::D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS;
     }
     flags
@@ -141,6 +143,9 @@ pub fn map_buffer_usage_to_state(usage: wgt::BufferUses) -> Direct3D12::D3D12_RE
     if usage.intersects(Bu::INDIRECT) {
         state |= Direct3D12::D3D12_RESOURCE_STATE_INDIRECT_ARGUMENT;
     }
+    if usage.intersects(Bu::ACCELERATION_STRUCTURE_QUERY) {
+        state |= Direct3D12::D3D12_RESOURCE_STATE_UNORDERED_ACCESS;
+    }
     state
 }
 
@@ -398,3 +403,16 @@ pub(crate) fn map_acceleration_structure_geometry_flags(
     }
     d3d_flags
 }
+
+pub(crate) fn map_acceleration_structure_copy_mode(
+    mode: wgt::AccelerationStructureCopy,
+) -> Direct3D12::D3D12_RAYTRACING_ACCELERATION_STRUCTURE_COPY_MODE {
+    match mode {
+        wgt::AccelerationStructureCopy::Clone => {
+            Direct3D12::D3D12_RAYTRACING_ACCELERATION_STRUCTURE_COPY_MODE_CLONE
+        }
+        wgt::AccelerationStructureCopy::Compact => {
+            Direct3D12::D3D12_RAYTRACING_ACCELERATION_STRUCTURE_COPY_MODE_COMPACT
+        }
+    }
+}

wgpu-hal/src/dynamic/command.rs

@@ -179,6 +179,18 @@ pub trait DynCommandEncoder: DynResource + std::fmt::Debug {
         &mut self,
         barrier: AccelerationStructureBarrier,
     );
+
+    unsafe fn copy_acceleration_structure_to_acceleration_structure(
+        &mut self,
+        src: &dyn DynAccelerationStructure,
+        dst: &dyn DynAccelerationStructure,
+        copy: wgt::AccelerationStructureCopy,
+    );
+    unsafe fn read_acceleration_structure_compact_size(
+        &mut self,
+        acceleration_structure: &dyn DynAccelerationStructure,
+        buf: &dyn DynBuffer,
+    );
 }
 
 impl<C: CommandEncoder + DynResource> DynCommandEncoder for C {
@@ -611,6 +623,26 @@ impl<C: CommandEncoder + DynResource> DynCommandEncoder for C {
     ) {
         unsafe { C::place_acceleration_structure_barrier(self, barrier) };
     }
+
+    unsafe fn copy_acceleration_structure_to_acceleration_structure(
+        &mut self,
+        src: &dyn DynAccelerationStructure,
+        dst: &dyn DynAccelerationStructure,
+        copy: wgt::AccelerationStructureCopy,
+    ) {
+        let src = src.expect_downcast_ref();
+        let dst = dst.expect_downcast_ref();
+        unsafe { C::copy_acceleration_structure_to_acceleration_structure(self, src, dst, copy) };
+    }
+    unsafe fn read_acceleration_structure_compact_size(
+        &mut self,
+        acceleration_structure: &dyn DynAccelerationStructure,
+        buf: &dyn DynBuffer,
+    ) {
+        let acceleration_structure = acceleration_structure.expect_downcast_ref();
+        let buf = buf.expect_downcast_ref();
+        unsafe { C::read_acceleration_structure_compact_size(self, acceleration_structure, buf) }
+    }
 }
 
 impl<'a> PassTimestampWrites<'a, dyn DynQuerySet> {

wgpu-hal/src/empty.rs

@@ -378,6 +378,12 @@ impl crate::CommandEncoder for Encoder {
     unsafe fn begin_query(&mut self, set: &Resource, index: u32) {}
     unsafe fn end_query(&mut self, set: &Resource, index: u32) {}
     unsafe fn write_timestamp(&mut self, set: &Resource, index: u32) {}
+    unsafe fn read_acceleration_structure_compact_size(
+        &mut self,
+        acceleration_structure: &Resource,
+        buf: &Resource,
+    ) {
+    }
     unsafe fn reset_queries(&mut self, set: &Resource, range: Range<u32>) {}
     unsafe fn copy_query_results(
         &mut self,
@@ -510,4 +516,12 @@ impl crate::CommandEncoder for Encoder {
         _barriers: crate::AccelerationStructureBarrier,
     ) {
     }
+
+    unsafe fn copy_acceleration_structure_to_acceleration_structure(
+        &mut self,
+        src: &Resource,
+        dst: &Resource,
+        copy: wgt::AccelerationStructureCopy,
+    ) {
+    }
 }

wgpu-hal/src/gles/command.rs

@@ -1209,4 +1209,21 @@ impl crate::CommandEncoder for super::CommandEncoder {
     ) {
         unimplemented!()
     }
+
+    unsafe fn copy_acceleration_structure_to_acceleration_structure(
+        &mut self,
+        _src: &super::AccelerationStructure,
+        _dst: &super::AccelerationStructure,
+        _copy: wgt::AccelerationStructureCopy,
+    ) {
+        unimplemented!()
+    }
+
+    unsafe fn read_acceleration_structure_compact_size(
+        &mut self,
+        _acceleration_structure: &super::AccelerationStructure,
+        _buf: &super::Buffer,
+    ) {
+        unimplemented!()
+    }
 }

wgpu-hal/src/lib.rs

@@ -1249,6 +1249,12 @@ pub trait CommandEncoder: WasmNotSendSync + fmt::Debug {
     ) where
         T: Iterator<Item = BufferTextureCopy>;
 
+    unsafe fn copy_acceleration_structure_to_acceleration_structure(
+        &mut self,
+        src: &<Self::A as Api>::AccelerationStructure,
+        dst: &<Self::A as Api>::AccelerationStructure,
+        copy: wgt::AccelerationStructureCopy,
+    );
     // pass common
 
     /// Sets the bind group at `index` to `group`.
@@ -1509,6 +1515,12 @@ pub trait CommandEncoder: WasmNotSendSync + fmt::Debug {
         &mut self,
         barrier: AccelerationStructureBarrier,
     );
+    // modeled off dx12, because this is able to be polyfilled in vulkan as opposed to the other way round
+    unsafe fn read_acceleration_structure_compact_size(
+        &mut self,
+        acceleration_structure: &<Self::A as Api>::AccelerationStructure,
+        buf: &<Self::A as Api>::Buffer,
+    );
 }
 
 bitflags!(
@@ -2311,6 +2323,7 @@ pub struct AccelerationStructureDescriptor<'a> {
     pub label: Label<'a>,
     pub size: wgt::BufferAddress,
     pub format: AccelerationStructureFormat,
+    pub allow_compaction: bool,
 }
 
 #[derive(Debug, Clone, Copy, Eq, PartialEq)]
@@ -2397,6 +2410,11 @@ pub struct AccelerationStructureAABBs<'a, B: DynBuffer + ?Sized> {
     pub flags: AccelerationStructureGeometryFlags,
 }
 
+pub struct AccelerationStructureCopy {
+    pub copy_flags: wgt::AccelerationStructureCopy,
+    pub type_flags: wgt::AccelerationStructureType,
+}
+
 /// * `offset` - offset in bytes
 #[derive(Clone, Debug)]
 pub struct AccelerationStructureInstances<'a, B: DynBuffer + ?Sized> {
@@ -2433,6 +2451,12 @@ bitflags::bitflags! {
         const BUILD_OUTPUT = 1 << 1;
         // Tlas used in a shader
         const SHADER_INPUT = 1 << 2;
+        // Blas used to query compacted size
+        const QUERY_INPUT = 1 << 3;
+        // BLAS used as a src for a copy operation
+        const COPY_SRC = 1 << 4;
+        // BLAS used as a dst for a copy operation
+        const COPY_DST = 1 << 5;
     }
 }
 

wgpu-hal/src/metal/command.rs

@@ -392,6 +392,15 @@ impl crate::CommandEncoder for super::CommandEncoder {
         }
     }
 
+    unsafe fn copy_acceleration_structure_to_acceleration_structure(
+        &mut self,
+        _src: &super::AccelerationStructure,
+        _dst: &super::AccelerationStructure,
+        _copy: wgt::AccelerationStructureCopy,
+    ) {
+        unimplemented!()
+    }
+
     unsafe fn begin_query(&mut self, set: &super::QuerySet, index: u32) {
         match set.ty {
             wgt::QueryType::Occlusion => {
@@ -1292,6 +1301,14 @@ impl crate::CommandEncoder for super::CommandEncoder {
     ) {
         unimplemented!()
     }
+
+    unsafe fn read_acceleration_structure_compact_size(
+        &mut self,
+        _acceleration_structure: &super::AccelerationStructure,
+        _buf: &super::Buffer,
+    ) {
+        unimplemented!()
+    }
 }
 
 impl Drop for super::CommandEncoder {