1use pfx_gpu::Gpu;
2use wgpu::util::DeviceExt;
3
4pub const RT_WGSL: &str = include_str!("rt.wgsl");
5
6#[derive(Clone, Debug, Default, PartialEq)]
7pub struct RtGeometry {
8 pub opaque: Vec<[f32; 3]>,
9 pub filtered: Vec<[f32; 3]>,
10 pub map: Vec<u32>,
11 pub boxes: Vec<[f32; 6]>,
12}
13
14impl RtGeometry {
15 pub fn new(triangles: &[[[f32; 3]; 3]], opaque: &[bool], boxes: Vec<(u32, [f32; 6])>) -> Self {
16 let mut first = Vec::new();
17 let mut rest = Vec::new();
18 let mut opaque_slots = Vec::new();
19 let mut filtered_slots = Vec::new();
20 for (slot, (vertices, &solid)) in triangles.iter().zip(opaque).enumerate() {
21 if solid {
22 first.extend_from_slice(vertices);
23 opaque_slots.push(slot as u32);
24 } else {
25 rest.extend_from_slice(vertices);
26 filtered_slots.push(slot as u32);
27 }
28 }
29 let leading = if opaque_slots.is_empty() {
30 filtered_slots.len()
31 } else {
32 opaque_slots.len()
33 };
34 let mut map = vec![leading as u32, triangles.len() as u32];
35 map.extend(opaque_slots);
36 map.extend(filtered_slots);
37 map.extend(boxes.iter().map(|(index, _)| *index));
38 Self {
39 opaque: first,
40 filtered: rest,
41 map,
42 boxes: boxes.into_iter().map(|(_, bounds)| bounds).collect(),
43 }
44 }
45
46 pub fn is_empty(&self) -> bool {
47 self.opaque.is_empty() && self.filtered.is_empty() && self.boxes.is_empty()
48 }
49}
50
51pub fn boxes() -> bool {
52 std::env::var("PFX_TRACE_RT_BOXES").is_ok_and(|value| value == "1")
53}
54
55pub fn offered(gpu: &Gpu) -> bool {
56 gpu.device
57 .features()
58 .contains(wgpu::Features::EXPERIMENTAL_RAY_QUERY)
59}
60
61pub fn wanted(gpu: &Gpu) -> bool {
62 offered(gpu) && std::env::var("PFX_TRACE_RT").map_or(true, |value| value != "0")
63}
64
65pub struct Rt {
66 pub tlas: wgpu::Tlas,
67 pub map: wgpu::Buffer,
68 _blases: Vec<wgpu::Blas>,
69 _inputs: Vec<wgpu::Buffer>,
70}
71
72const IDENTITY: [f32; 12] = [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0];
73
74fn input(gpu: &Gpu, label: &'static str, bytes: &[u8]) -> wgpu::Buffer {
75 gpu.device
76 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
77 label: Some(label),
78 contents: bytes,
79 usage: wgpu::BufferUsages::BLAS_INPUT,
80 })
81}
82
83impl Rt {
84 pub fn build(gpu: &Gpu, geometry: &RtGeometry) -> Self {
85 let device = &gpu.device;
86 let lists: Vec<(&Vec<[f32; 3]>, wgpu::AccelerationStructureGeometryFlags)> = [
87 (
88 &geometry.opaque,
89 wgpu::AccelerationStructureGeometryFlags::OPAQUE,
90 ),
91 (
92 &geometry.filtered,
93 wgpu::AccelerationStructureGeometryFlags::NO_DUPLICATE_ANY_HIT_INVOCATION,
94 ),
95 ]
96 .into_iter()
97 .filter(|(vertices, _)| !vertices.is_empty())
98 .collect();
99 let sizes: Vec<wgpu::BlasTriangleGeometrySizeDescriptor> = lists
100 .iter()
101 .map(
102 |(vertices, flags)| wgpu::BlasTriangleGeometrySizeDescriptor {
103 vertex_format: wgpu::VertexFormat::Float32x3,
104 vertex_count: vertices.len() as u32,
105 index_format: None,
106 index_count: None,
107 flags: *flags,
108 },
109 )
110 .collect();
111 let vertex_buffers: Vec<wgpu::Buffer> = lists
112 .iter()
113 .map(|(vertices, _)| input(gpu, "trace rt vertices", bytemuck::cast_slice(vertices)))
114 .collect();
115 let flags = wgpu::AccelerationStructureFlags::PREFER_FAST_TRACE;
116 let tri_blas = (!sizes.is_empty()).then(|| {
117 device.create_blas(
118 &wgpu::CreateBlasDescriptor {
119 label: Some("trace rt triangles"),
120 flags,
121 update_mode: wgpu::AccelerationStructureUpdateMode::Build,
122 },
123 wgpu::BlasGeometrySizeDescriptors::Triangles {
124 descriptors: sizes.clone(),
125 },
126 )
127 });
128 let box_size = wgpu::BlasAABBGeometrySizeDescriptor {
129 primitive_count: geometry.boxes.len() as u32,
130 flags: wgpu::AccelerationStructureGeometryFlags::NO_DUPLICATE_ANY_HIT_INVOCATION,
131 };
132 let box_buffer = (!geometry.boxes.is_empty())
133 .then(|| input(gpu, "trace rt boxes", bytemuck::cast_slice(&geometry.boxes)));
134 let box_blas = box_buffer.as_ref().map(|_| {
135 device.create_blas(
136 &wgpu::CreateBlasDescriptor {
137 label: Some("trace rt shapes"),
138 flags,
139 update_mode: wgpu::AccelerationStructureUpdateMode::Build,
140 },
141 wgpu::BlasGeometrySizeDescriptors::AABBs {
142 descriptors: vec![box_size.clone()],
143 },
144 )
145 });
146 let mut tlas = device.create_tlas(&wgpu::CreateTlasDescriptor {
147 label: Some("trace rt world"),
148 max_instances: 2,
149 flags,
150 update_mode: wgpu::AccelerationStructureUpdateMode::Build,
151 });
152 for (slot, blas) in [&tri_blas, &box_blas].into_iter().flatten().enumerate() {
153 *tlas.get_mut_single(slot).unwrap() =
154 Some(wgpu::TlasInstance::new(blas, IDENTITY, slot as u32, 0xff));
155 }
156 let mut entries = Vec::new();
157 if let Some(blas) = &tri_blas {
158 entries.push(wgpu::BlasBuildEntry {
159 blas,
160 geometry: wgpu::BlasGeometries::TriangleGeometries(
161 sizes
162 .iter()
163 .zip(&vertex_buffers)
164 .map(|(size, buffer)| wgpu::BlasTriangleGeometry {
165 size,
166 vertex_buffer: buffer,
167 first_vertex: 0,
168 vertex_stride: 12,
169 index_buffer: None,
170 first_index: None,
171 transform_buffer: None,
172 transform_buffer_offset: None,
173 })
174 .collect(),
175 ),
176 });
177 }
178 if let (Some(blas), Some(buffer)) = (&box_blas, &box_buffer) {
179 entries.push(wgpu::BlasBuildEntry {
180 blas,
181 geometry: wgpu::BlasGeometries::AabbGeometries(vec![wgpu::BlasAabbGeometry {
182 size: &box_size,
183 stride: 24,
184 aabb_buffer: buffer,
185 primitive_offset: 0,
186 }]),
187 });
188 }
189 let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
190 label: Some("trace rt build"),
191 });
192 encoder.build_acceleration_structures(entries.iter(), [&tlas]);
193 gpu.queue.submit([encoder.finish()]);
194 let map = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
195 label: Some("trace rt map"),
196 contents: bytemuck::cast_slice(&geometry.map),
197 usage: wgpu::BufferUsages::STORAGE,
198 });
199 let mut inputs = vertex_buffers;
200 inputs.extend(box_buffer);
201 Self {
202 tlas,
203 map,
204 _blases: [tri_blas, box_blas].into_iter().flatten().collect(),
205 _inputs: inputs,
206 }
207 }
208}
209
210#[cfg(test)]
211mod tests {
212 use super::*;
213
214 #[test]
215 fn the_map_leads_with_the_first_geometry_then_slots_then_shapes() {
216 let tri = |x: f32| [[x, 0.0, 0.0], [x, 1.0, 0.0], [x, 0.0, 1.0]];
217 let geometry = RtGeometry::new(
218 &[tri(0.0), tri(1.0), tri(2.0)],
219 &[false, true, false],
220 vec![(4, [0.0; 6])],
221 );
222 assert_eq!(geometry.map, vec![1, 3, 1, 0, 2, 4]);
223 assert_eq!(geometry.opaque.len(), 3);
224 assert_eq!(geometry.filtered.len(), 6);
225 let none_solid = RtGeometry::new(&[tri(0.0), tri(1.0)], &[false, false], Vec::new());
226 assert_eq!(none_solid.map, vec![2, 2, 0, 1]);
227 }
228
229 #[test]
230 fn the_ray_query_shader_validates() {
231 for lit in [false, true] {
232 let source = crate::gpu::shader_rt(lit);
233 let module = naga::front::wgsl::parse_str(&source)
234 .unwrap_or_else(|error| panic!("{}", error.emit_to_string(&source)));
235 naga::valid::Validator::new(
236 naga::valid::ValidationFlags::all(),
237 naga::valid::Capabilities::all(),
238 )
239 .validate(&module)
240 .unwrap();
241 }
242 }
243}