|
| 1 | +use std::{borrow::Cow, num::NonZeroU32}; |
| 2 | + |
| 3 | +use wgpu::util::DeviceExt; |
| 4 | +use wgpu::*; |
| 5 | +use wgpu_test::{ |
| 6 | + gpu_test, FailureCase, GpuTestConfiguration, GpuTestInitializer, TestParameters, TestingContext, |
| 7 | +}; |
| 8 | + |
| 9 | +pub fn all_tests(tests: &mut Vec<GpuTestInitializer>) { |
| 10 | + tests.push(BINDING_ARRAY_TLAS); |
| 11 | +} |
| 12 | + |
| 13 | +#[gpu_test] |
| 14 | +static BINDING_ARRAY_TLAS: GpuTestConfiguration = GpuTestConfiguration::new() |
| 15 | + .parameters( |
| 16 | + TestParameters::default() |
| 17 | + .instance_flags(wgpu::InstanceFlags::GPU_BASED_VALIDATION) |
| 18 | + // Ray queries + acceleration structure bindings are gated behind this experimental feature. |
| 19 | + .features( |
| 20 | + Features::EXPERIMENTAL_RAY_QUERY | Features::ACCELERATION_STRUCTURE_BINDING_ARRAY, |
| 21 | + ) |
| 22 | + .limits({ |
| 23 | + let mut limits = |
| 24 | + Limits::default().using_minimum_supported_acceleration_structure_values(); |
| 25 | + // Keep this small; we only need a couple of array elements. |
| 26 | + limits.max_binding_array_elements_per_shader_stage = 8; |
| 27 | + limits.max_acceleration_structures_per_shader_stage = 8; |
| 28 | + limits |
| 29 | + }) |
| 30 | + // As of writing, Metal's HAL does not implement binding acceleration structures. |
| 31 | + .skip(FailureCase::backend(Backends::METAL)), |
| 32 | + ) |
| 33 | + .run_async(|ctx| async move { binding_array_tlas(ctx).await }); |
| 34 | + |
| 35 | +async fn binding_array_tlas(ctx: TestingContext) { |
| 36 | + // Minimal shader that consumes a TLAS binding array. |
| 37 | + // |
| 38 | + // We don't need to actually "trace" anything for this test. We only need: |
| 39 | + // - Pipeline compilation to accept `binding_array<acceleration_structure>` |
| 40 | + // - Bind group creation to accept `BindingResource::AccelerationStructureArray` |
| 41 | + // - Encoder to successfully set the bind group and submit. |
| 42 | + // |
| 43 | + // Creating a `ray_query` and initializing it against element 0 forces the binding to be used. |
| 44 | + let shader = r#" |
| 45 | + enable wgpu_ray_query; |
| 46 | +
|
| 47 | + @group(0) @binding(0) |
| 48 | + var tlas_array: binding_array<acceleration_structure>; |
| 49 | +
|
| 50 | + @compute |
| 51 | + @workgroup_size(1, 1, 1) |
| 52 | + fn main() { |
| 53 | + var rq: ray_query; |
| 54 | + rayQueryInitialize( |
| 55 | + &rq, |
| 56 | + tlas_array[0], |
| 57 | + RayDesc( |
| 58 | + 0u, |
| 59 | + 0xffu, |
| 60 | + 0.001, |
| 61 | + 1000.0, |
| 62 | + vec3f(0.0, 0.0, 0.0), |
| 63 | + vec3f(0.0, 0.0, 1.0) |
| 64 | + ) |
| 65 | + ); |
| 66 | + } |
| 67 | + "#; |
| 68 | + |
| 69 | + let module = ctx.device.create_shader_module(ShaderModuleDescriptor { |
| 70 | + label: Some("Binding Array TLAS"), |
| 71 | + source: ShaderSource::Wgsl(Cow::Borrowed(shader)), |
| 72 | + }); |
| 73 | + |
| 74 | + // Build a minimal BLAS + two TLAS so we can bind an array of TLAS. |
| 75 | + // |
| 76 | + // This follows the shapes used in the ray tracing examples. |
| 77 | + let vertex_data: [[f32; 3]; 3] = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]; |
| 78 | + let index_data: [u16; 3] = [0, 1, 2]; |
| 79 | + |
| 80 | + let vertex_buf = ctx |
| 81 | + .device |
| 82 | + .create_buffer_init(&wgpu::util::BufferInitDescriptor { |
| 83 | + label: Some("RT Vertex Buffer"), |
| 84 | + contents: bytemuck::cast_slice(&vertex_data), |
| 85 | + usage: BufferUsages::VERTEX | BufferUsages::BLAS_INPUT, |
| 86 | + }); |
| 87 | + |
| 88 | + let index_buf = ctx |
| 89 | + .device |
| 90 | + .create_buffer_init(&wgpu::util::BufferInitDescriptor { |
| 91 | + label: Some("RT Index Buffer"), |
| 92 | + contents: bytemuck::cast_slice(&index_data), |
| 93 | + usage: BufferUsages::INDEX | BufferUsages::BLAS_INPUT, |
| 94 | + }); |
| 95 | + |
| 96 | + let blas_geo_size_desc = wgpu::BlasTriangleGeometrySizeDescriptor { |
| 97 | + vertex_format: wgpu::VertexFormat::Float32x3, |
| 98 | + vertex_count: vertex_data.len() as u32, |
| 99 | + index_format: Some(wgpu::IndexFormat::Uint16), |
| 100 | + index_count: Some(index_data.len() as u32), |
| 101 | + flags: wgpu::AccelerationStructureGeometryFlags::OPAQUE, |
| 102 | + }; |
| 103 | + |
| 104 | + let blas = ctx.device.create_blas( |
| 105 | + &wgpu::CreateBlasDescriptor { |
| 106 | + label: Some("BLAS"), |
| 107 | + flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE, |
| 108 | + update_mode: wgpu::AccelerationStructureUpdateMode::Build, |
| 109 | + }, |
| 110 | + wgpu::BlasGeometrySizeDescriptors::Triangles { |
| 111 | + descriptors: vec![blas_geo_size_desc.clone()], |
| 112 | + }, |
| 113 | + ); |
| 114 | + |
| 115 | + let mut tlas_a = ctx.device.create_tlas(&wgpu::CreateTlasDescriptor { |
| 116 | + label: Some("TLAS A"), |
| 117 | + flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE, |
| 118 | + update_mode: wgpu::AccelerationStructureUpdateMode::Build, |
| 119 | + max_instances: 1, |
| 120 | + }); |
| 121 | + |
| 122 | + let mut tlas_b = ctx.device.create_tlas(&wgpu::CreateTlasDescriptor { |
| 123 | + label: Some("TLAS B"), |
| 124 | + flags: wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE, |
| 125 | + update_mode: wgpu::AccelerationStructureUpdateMode::Build, |
| 126 | + max_instances: 1, |
| 127 | + }); |
| 128 | + |
| 129 | + // Put a single instance into each TLAS. Both reference the same BLAS. |
| 130 | + // |
| 131 | + // NOTE: This indexing API is how TLAS instances are populated in the examples. |
| 132 | + tlas_a[0] = Some(wgpu::TlasInstance::new( |
| 133 | + &blas, |
| 134 | + [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], |
| 135 | + 0, |
| 136 | + 0xff, |
| 137 | + )); |
| 138 | + tlas_b[0] = Some(wgpu::TlasInstance::new( |
| 139 | + &blas, |
| 140 | + [1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], |
| 141 | + 0, |
| 142 | + 0xff, |
| 143 | + )); |
| 144 | + |
| 145 | + // Build BLAS and TLASes. |
| 146 | + let mut encoder = ctx |
| 147 | + .device |
| 148 | + .create_command_encoder(&wgpu::CommandEncoderDescriptor { |
| 149 | + label: Some("RT Build"), |
| 150 | + }); |
| 151 | + |
| 152 | + encoder.build_acceleration_structures( |
| 153 | + std::iter::once(&wgpu::BlasBuildEntry { |
| 154 | + blas: &blas, |
| 155 | + geometry: wgpu::BlasGeometries::TriangleGeometries(vec![wgpu::BlasTriangleGeometry { |
| 156 | + size: &blas_geo_size_desc, |
| 157 | + vertex_buffer: &vertex_buf, |
| 158 | + first_vertex: 0, |
| 159 | + vertex_stride: std::mem::size_of::<[f32; 3]>() as u64, |
| 160 | + index_buffer: Some(&index_buf), |
| 161 | + first_index: Some(0), |
| 162 | + transform_buffer: None, |
| 163 | + transform_buffer_offset: None, |
| 164 | + }]), |
| 165 | + }), |
| 166 | + std::iter::empty::<&wgpu::Tlas>() |
| 167 | + .chain(std::iter::once(&tlas_a)) |
| 168 | + .chain(std::iter::once(&tlas_b)), |
| 169 | + ); |
| 170 | + |
| 171 | + ctx.queue.submit(Some(encoder.finish())); |
| 172 | + |
| 173 | + // Bind group layout with a TLAS array binding. |
| 174 | + let bgl = ctx |
| 175 | + .device |
| 176 | + .create_bind_group_layout(&BindGroupLayoutDescriptor { |
| 177 | + label: Some("TLAS array BGL"), |
| 178 | + entries: &[BindGroupLayoutEntry { |
| 179 | + binding: 0, |
| 180 | + visibility: ShaderStages::COMPUTE, |
| 181 | + ty: BindingType::AccelerationStructure { |
| 182 | + vertex_return: false, |
| 183 | + }, |
| 184 | + count: Some(NonZeroU32::new(2).unwrap()), |
| 185 | + }], |
| 186 | + }); |
| 187 | + |
| 188 | + let tlas_refs: [&Tlas; 2] = [&tlas_a, &tlas_b]; |
| 189 | + |
| 190 | + let bg = ctx.device.create_bind_group(&BindGroupDescriptor { |
| 191 | + label: Some("TLAS array BG"), |
| 192 | + layout: &bgl, |
| 193 | + entries: &[BindGroupEntry { |
| 194 | + binding: 0, |
| 195 | + resource: BindingResource::AccelerationStructureArray(&tlas_refs), |
| 196 | + }], |
| 197 | + }); |
| 198 | + |
| 199 | + let pipeline_layout = ctx |
| 200 | + .device |
| 201 | + .create_pipeline_layout(&PipelineLayoutDescriptor { |
| 202 | + label: Some("TLAS array pipeline layout"), |
| 203 | + bind_group_layouts: &[Some(&bgl)], |
| 204 | + immediate_size: 0, |
| 205 | + }); |
| 206 | + |
| 207 | + let pipeline = ctx |
| 208 | + .device |
| 209 | + .create_compute_pipeline(&ComputePipelineDescriptor { |
| 210 | + label: Some("TLAS array pipeline"), |
| 211 | + layout: Some(&pipeline_layout), |
| 212 | + module: &module, |
| 213 | + entry_point: Some("main"), |
| 214 | + compilation_options: Default::default(), |
| 215 | + cache: None, |
| 216 | + }); |
| 217 | + |
| 218 | + let mut encoder = ctx |
| 219 | + .device |
| 220 | + .create_command_encoder(&CommandEncoderDescriptor { |
| 221 | + label: Some("Dispatch"), |
| 222 | + }); |
| 223 | + |
| 224 | + { |
| 225 | + let mut pass = encoder.begin_compute_pass(&ComputePassDescriptor { |
| 226 | + label: Some("Compute pass"), |
| 227 | + timestamp_writes: None, |
| 228 | + }); |
| 229 | + pass.set_pipeline(&pipeline); |
| 230 | + pass.set_bind_group(0, &bg, &[]); |
| 231 | + pass.dispatch_workgroups(1, 1, 1); |
| 232 | + } |
| 233 | + |
| 234 | + ctx.queue.submit(Some(encoder.finish())); |
| 235 | +} |
0 commit comments