GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git
src/vk/rt.rs (91.2K)
1 //! The tier-1 RT engine (RT-renderer phase 2): an internal triangle+material
2 //! scene, a CPU-built binned-SAH BVH uploaded as storage buffers, and the
3 //! `rt.wgsl` compute path tracer with progressive accumulation.
4 //!
5 //! The stage renders into its own pane-sized storage image and blits it into
6 //! the backdrop image's viewport-pane region — exactly the slot the raster
7 //! `SceneStage` fills — so the swapchain copy, blur plates, and the 2D UI pass
8 //! are untouched. Runs on plain Vulkan compute (no `VK_KHR_ray_*`), which is
9 //! the point: it works on the integrated GPU; a tier-2 ray-query backend can
10 //! later swap out just the traversal.
11 //!
12 //! Scene schema is internal by design — importers (OBJ/glTF) belong in a
13 //! future loader that converts *into* [`RtTriangle`]/[`RtMaterial`].
14
15 use ash::vk;
16 use gpu_allocator::vulkan::{Allocation, AllocationCreateDesc, AllocationScheme, Allocator};
17 use gpu_allocator::MemoryLocation;
18
19 use super::renderer::{
20 compile_wgsl, compile_wgsl_ray_query, create_cpu_buffer, destroy_cpu_buffer, AllocatedBuffer,
21 };
22
23 /// One triangle of an RT scene, in the same space as the camera's `inv_mvp`
24 /// (for the designer: mesh space, the space `Vertex3D` positions live in).
25 #[derive(Debug, Clone, Copy, PartialEq)]
26 pub struct RtTriangle {
27 pub p0: [f32; 3],
28 pub p1: [f32; 3],
29 pub p2: [f32; 3],
30 /// Index into the material slice passed alongside.
31 pub material: u32,
32 }
33
34 /// Lambertian surface + optional emission, linear color (matching the raster
35 /// path, whose vertex colors land in the sRGB attachment as linear values).
36 #[derive(Debug, Clone, Copy, PartialEq)]
37 pub struct RtMaterial {
38 pub albedo: [f32; 3],
39 pub emission: [f32; 3],
40 }
41
42 /// The full camera: the inverse of the raster path's `proj * view * model`.
43 /// Rays are unprojected from NDC through it, so any matrix stack that renders
44 /// the raster viewport drives the tracer unchanged.
45 #[derive(Debug, Clone, Copy, PartialEq)]
46 pub struct RtCamera {
47 pub inv_mvp: [[f32; 4]; 4],
48 }
49
50 // --- GPU layouts (must match rt.wgsl) ---
51
52 #[repr(C)]
53 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
54 struct GpuTriangle {
55 p0: [f32; 4], // w = material index (bitcast)
56 p1: [f32; 4],
57 p2: [f32; 4],
58 }
59
60 #[repr(C)]
61 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
62 struct GpuMaterial {
63 albedo: [f32; 4],
64 emission: [f32; 4],
65 }
66
67 #[repr(C)]
68 #[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
69 pub(crate) struct GpuBvhNode {
70 pub(crate) min: [f32; 3],
71 /// Leaf (`count > 0`): first triangle. Internal: left child; right = +1.
72 pub(crate) left_first: u32,
73 pub(crate) max: [f32; 3],
74 pub(crate) count: u32,
75 }
76
77 #[repr(C)]
78 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
79 struct RtParams {
80 inv_mvp: [[f32; 4]; 4],
81 width: u32,
82 height: u32,
83 sample_index: u32,
84 max_bounces: u32,
85 spp: u32,
86 _pad: [u32; 3],
87 }
88
89 // --- BVH construction (binned SAH) ---
90
91 const BVH_BINS: usize = 8;
92 const BVH_LEAF_MAX: u32 = 4;
93
94 #[derive(Clone, Copy)]
95 struct Aabb {
96 min: [f32; 3],
97 max: [f32; 3],
98 }
99
100 impl Aabb {
101 const EMPTY: Aabb = Aabb { min: [f32::INFINITY; 3], max: [f32::NEG_INFINITY; 3] };
102
103 fn grow(&mut self, p: [f32; 3]) {
104 for a in 0..3 {
105 self.min[a] = self.min[a].min(p[a]);
106 self.max[a] = self.max[a].max(p[a]);
107 }
108 }
109
110 fn grow_aabb(&mut self, other: &Aabb) {
111 self.grow(other.min);
112 self.grow(other.max);
113 }
114
115 fn half_area(&self) -> f32 {
116 let dx = (self.max[0] - self.min[0]).max(0.0);
117 let dy = (self.max[1] - self.min[1]).max(0.0);
118 let dz = (self.max[2] - self.min[2]).max(0.0);
119 dx * dy + dy * dz + dz * dx
120 }
121 }
122
123 fn tri_aabb(t: &RtTriangle) -> Aabb {
124 let mut b = Aabb::EMPTY;
125 b.grow(t.p0);
126 b.grow(t.p1);
127 b.grow(t.p2);
128 b
129 }
130
131 fn tri_centroid(t: &RtTriangle) -> [f32; 3] {
132 let mut c = [0.0f32; 3];
133 for a in 0..3 {
134 c[a] = (t.p0[a] + t.p1[a] + t.p2[a]) / 3.0;
135 }
136 c
137 }
138
139 /// Build a BVH over `triangles`, reordering them so leaves reference
140 /// contiguous ranges. Returns the flat node array (empty input → empty vec).
141 pub(crate) fn build_bvh(triangles: &mut Vec<RtTriangle>) -> Vec<GpuBvhNode> {
142 if triangles.is_empty() {
143 return Vec::new();
144 }
145 let bounds: Vec<Aabb> = triangles.iter().map(tri_aabb).collect();
146 let centroids: Vec<[f32; 3]> = triangles.iter().map(tri_centroid).collect();
147 let mut order: Vec<u32> = (0..triangles.len() as u32).collect();
148
149 fn range_bounds(order: &[u32], bounds: &[Aabb]) -> Aabb {
150 let mut b = Aabb::EMPTY;
151 for &i in order {
152 b.grow_aabb(&bounds[i as usize]);
153 }
154 b
155 }
156
157 let mut nodes: Vec<GpuBvhNode> = Vec::with_capacity(triangles.len() * 2);
158 let root_bounds = range_bounds(&order, &bounds);
159 nodes.push(GpuBvhNode {
160 min: root_bounds.min,
161 left_first: 0,
162 max: root_bounds.max,
163 count: triangles.len() as u32,
164 });
165
166 // (node index, start, count) work list over `order`.
167 let mut work = vec![(0usize, 0usize, triangles.len())];
168 while let Some((node_idx, start, count)) = work.pop() {
169 if (count as u32) <= BVH_LEAF_MAX {
170 continue; // stays a leaf
171 }
172 let slice = &mut order[start..start + count];
173
174 // Centroid bounds pick the split axis.
175 let mut cb = Aabb::EMPTY;
176 for &i in slice.iter() {
177 cb.grow(centroids[i as usize]);
178 }
179 let mut axis = 0;
180 let mut extent = 0.0f32;
181 for a in 0..3 {
182 let e = cb.max[a] - cb.min[a];
183 if e > extent {
184 extent = e;
185 axis = a;
186 }
187 }
188
189 let mut split_at = None;
190 if extent > 1e-12 {
191 // Binned SAH along `axis`.
192 let scale = BVH_BINS as f32 / extent;
193 let bin_of = |i: u32| -> usize {
194 (((centroids[i as usize][axis] - cb.min[axis]) * scale) as usize)
195 .min(BVH_BINS - 1)
196 };
197 let mut bin_bounds = [Aabb::EMPTY; BVH_BINS];
198 let mut bin_counts = [0usize; BVH_BINS];
199 for &i in slice.iter() {
200 let b = bin_of(i);
201 bin_counts[b] += 1;
202 bin_bounds[b].grow_aabb(&bounds[i as usize]);
203 }
204 // Cost of each of the BINS-1 split planes.
205 let mut best_cost = f32::INFINITY;
206 let mut best_plane = 0usize;
207 for plane in 1..BVH_BINS {
208 let (mut lb, mut rb) = (Aabb::EMPTY, Aabb::EMPTY);
209 let (mut lc, mut rc) = (0usize, 0usize);
210 for b in 0..plane {
211 lb.grow_aabb(&bin_bounds[b]);
212 lc += bin_counts[b];
213 }
214 for b in plane..BVH_BINS {
215 rb.grow_aabb(&bin_bounds[b]);
216 rc += bin_counts[b];
217 }
218 if lc == 0 || rc == 0 {
219 continue;
220 }
221 let cost = lb.half_area() * lc as f32 + rb.half_area() * rc as f32;
222 if cost < best_cost {
223 best_cost = cost;
224 best_plane = plane;
225 }
226 }
227 if best_plane > 0 {
228 let mut mid = 0usize;
229 for k in 0..count {
230 if bin_of(slice[k]) < best_plane {
231 slice.swap(k, mid);
232 mid += 1;
233 }
234 }
235 if mid > 0 && mid < count {
236 split_at = Some(mid);
237 }
238 }
239 }
240 // Degenerate centroids or a one-sided SAH result: median split keeps
241 // the tree balanced instead of forcing a giant leaf.
242 let mid = split_at.unwrap_or(count / 2);
243
244 let left_bounds = range_bounds(&slice[..mid], &bounds);
245 let right_bounds = range_bounds(&slice[mid..], &bounds);
246 let left_idx = nodes.len();
247 nodes.push(GpuBvhNode {
248 min: left_bounds.min,
249 left_first: (start) as u32,
250 max: left_bounds.max,
251 count: mid as u32,
252 });
253 nodes.push(GpuBvhNode {
254 min: right_bounds.min,
255 left_first: (start + mid) as u32,
256 max: right_bounds.max,
257 count: (count - mid) as u32,
258 });
259 nodes[node_idx].left_first = left_idx as u32;
260 nodes[node_idx].count = 0;
261 work.push((left_idx, start, mid));
262 work.push((left_idx + 1, start + mid, count - mid));
263 }
264
265 // Apply the final order to the triangle array so leaf ranges are direct.
266 let reordered: Vec<RtTriangle> =
267 order.iter().map(|&i| triangles[i as usize]).collect();
268 *triangles = reordered;
269 nodes
270 }
271
272 // --- The Vulkan stage ---
273
274 const MAX_SAMPLES: u32 = 1024;
275 const MAX_BOUNCES: u32 = 4;
276 const WORKGROUP: u32 = 8;
277
278 /// The two trace backends. They share every shader line except
279 /// `intersect_scene` (rt_bvh.wgsl vs rt_query.wgsl) and binding 1.
280 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
281 enum RtTier {
282 /// CPU-built BVH traversed in compute — runs on any device.
283 Compute,
284 /// Driver acceleration structures + VK_KHR_ray_query — RT cores.
285 RayQuery,
286 }
287
288 /// Tier-2 GPU objects: one BLAS over the triangle buffer, a one-instance
289 /// TLAS over it. Rebuilt wholesale on every scene replacement.
290 struct Accel {
291 blas: vk::AccelerationStructureKHR,
292 blas_buffer: AllocatedBuffer,
293 tlas: vk::AccelerationStructureKHR,
294 tlas_buffer: AllocatedBuffer,
295 instances: AllocatedBuffer,
296 }
297
298 struct RtFrame {
299 uniforms: AllocatedBuffer,
300 descriptor_set: vk::DescriptorSet,
301 }
302
303 #[repr(C)]
304 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
305 struct DenoiseParams {
306 width: u32,
307 height: u32,
308 step: u32,
309 first: u32,
310 last: u32,
311 inv_sqrt_n: f32,
312 _pad: [u32; 2],
313 }
314
315 const DENOISE_ITERATIONS: usize = 3; // à-trous steps 1, 2, 4
316
317 struct DenoiserFrame {
318 /// DENOISE_ITERATIONS dynamic-offset slices of [`DenoiseParams`].
319 uniforms: AllocatedBuffer,
320 /// src = ping, dst = pong.
321 set_a: vk::DescriptorSet,
322 /// src = pong, dst = ping.
323 set_b: vk::DescriptorSet,
324 }
325
326 /// The à-trous denoise pipeline (rt_denoise.wgsl). Owned by [`RtStage`];
327 /// its buffers (features/ping/pong) live on the stage with the other
328 /// pane-sized targets.
329 struct Denoiser {
330 pipeline: vk::Pipeline,
331 pipeline_layout: vk::PipelineLayout,
332 descriptor_set_layout: vk::DescriptorSetLayout,
333 descriptor_pool: vk::DescriptorPool,
334 shader_module: vk::ShaderModule,
335 uniform_stride: vk::DeviceSize,
336 frames: Vec<DenoiserFrame>,
337 }
338
339 impl Denoiser {
340 fn new(
341 device: &ash::Device,
342 allocator: &mut Allocator,
343 frames_in_flight: usize,
344 min_uniform_align: vk::DeviceSize,
345 ) -> Self {
346 unsafe {
347 let bindings = [
348 vk::DescriptorSetLayoutBinding::default()
349 .binding(0)
350 .descriptor_type(vk::DescriptorType::UNIFORM_BUFFER_DYNAMIC)
351 .descriptor_count(1)
352 .stage_flags(vk::ShaderStageFlags::COMPUTE),
353 vk::DescriptorSetLayoutBinding::default()
354 .binding(1)
355 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
356 .descriptor_count(1)
357 .stage_flags(vk::ShaderStageFlags::COMPUTE),
358 vk::DescriptorSetLayoutBinding::default()
359 .binding(2)
360 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
361 .descriptor_count(1)
362 .stage_flags(vk::ShaderStageFlags::COMPUTE),
363 vk::DescriptorSetLayoutBinding::default()
364 .binding(3)
365 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
366 .descriptor_count(1)
367 .stage_flags(vk::ShaderStageFlags::COMPUTE),
368 vk::DescriptorSetLayoutBinding::default()
369 .binding(4)
370 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
371 .descriptor_count(1)
372 .stage_flags(vk::ShaderStageFlags::COMPUTE),
373 vk::DescriptorSetLayoutBinding::default()
374 .binding(5)
375 .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
376 .descriptor_count(1)
377 .stage_flags(vk::ShaderStageFlags::COMPUTE),
378 ];
379 let descriptor_set_layout = device
380 .create_descriptor_set_layout(
381 &vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings),
382 None,
383 )
384 .expect("Failed to create denoise descriptor set layout");
385 let set_layouts_one = [descriptor_set_layout];
386 let pipeline_layout = device
387 .create_pipeline_layout(
388 &vk::PipelineLayoutCreateInfo::default().set_layouts(&set_layouts_one),
389 None,
390 )
391 .expect("Failed to create denoise pipeline layout");
392 let spirv = compile_wgsl(include_str!("rt_denoise.wgsl"));
393 let shader_module = device
394 .create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&spirv), None)
395 .expect("Failed to create denoise shader module");
396 let pipeline = device
397 .create_compute_pipelines(
398 vk::PipelineCache::null(),
399 &[vk::ComputePipelineCreateInfo::default()
400 .stage(
401 vk::PipelineShaderStageCreateInfo::default()
402 .stage(vk::ShaderStageFlags::COMPUTE)
403 .module(shader_module)
404 .name(c"cs_denoise"),
405 )
406 .layout(pipeline_layout)],
407 None,
408 )
409 .expect("Failed to create denoise pipeline")[0];
410
411 let n = frames_in_flight as u32;
412 let pool_sizes = [
413 vk::DescriptorPoolSize::default()
414 .ty(vk::DescriptorType::UNIFORM_BUFFER_DYNAMIC)
415 .descriptor_count(2 * n),
416 vk::DescriptorPoolSize::default()
417 .ty(vk::DescriptorType::STORAGE_BUFFER)
418 .descriptor_count(8 * n),
419 vk::DescriptorPoolSize::default()
420 .ty(vk::DescriptorType::STORAGE_IMAGE)
421 .descriptor_count(2 * n),
422 ];
423 let descriptor_pool = device
424 .create_descriptor_pool(
425 &vk::DescriptorPoolCreateInfo::default()
426 .max_sets(2 * n)
427 .pool_sizes(&pool_sizes),
428 None,
429 )
430 .expect("Failed to create denoise descriptor pool");
431 let set_layouts: Vec<vk::DescriptorSetLayout> =
432 vec![descriptor_set_layout; frames_in_flight * 2];
433 let sets = device
434 .allocate_descriptor_sets(
435 &vk::DescriptorSetAllocateInfo::default()
436 .descriptor_pool(descriptor_pool)
437 .set_layouts(&set_layouts),
438 )
439 .expect("Failed to allocate denoise descriptor sets");
440
441 let uniform_stride = (std::mem::size_of::<DenoiseParams>() as vk::DeviceSize)
442 .next_multiple_of(min_uniform_align.max(1));
443 let frames: Vec<DenoiserFrame> = (0..frames_in_flight)
444 .map(|i| {
445 let uniforms = create_cpu_buffer(
446 device,
447 allocator,
448 uniform_stride * DENOISE_ITERATIONS as vk::DeviceSize,
449 vk::BufferUsageFlags::UNIFORM_BUFFER,
450 "rt-denoise-uniforms",
451 );
452 let (set_a, set_b) = (sets[2 * i], sets[2 * i + 1]);
453 for set in [set_a, set_b] {
454 let infos = [vk::DescriptorBufferInfo::default()
455 .buffer(uniforms.buffer)
456 .range(std::mem::size_of::<DenoiseParams>() as vk::DeviceSize)];
457 device.update_descriptor_sets(
458 &[vk::WriteDescriptorSet::default()
459 .dst_set(set)
460 .dst_binding(0)
461 .descriptor_type(vk::DescriptorType::UNIFORM_BUFFER_DYNAMIC)
462 .buffer_info(&infos)],
463 &[],
464 );
465 }
466 DenoiserFrame { uniforms, set_a, set_b }
467 })
468 .collect();
469
470 Denoiser {
471 pipeline,
472 pipeline_layout,
473 descriptor_set_layout,
474 descriptor_pool,
475 shader_module,
476 uniform_stride,
477 frames,
478 }
479 }
480 }
481
482 /// Re-point the per-target bindings after the pane-sized buffers are
483 /// (re)created. Device is idle (target recreation contract).
484 fn write_target_descriptors(
485 &self,
486 device: &ash::Device,
487 accum: vk::Buffer,
488 features: vk::Buffer,
489 ping: vk::Buffer,
490 pong: vk::Buffer,
491 output_view: vk::ImageView,
492 ) {
493 for frame in &self.frames {
494 for (set, src, dst) in
495 [(frame.set_a, ping, pong), (frame.set_b, pong, ping)]
496 {
497 let buf_infos = [
498 vk::DescriptorBufferInfo::default().buffer(accum).range(vk::WHOLE_SIZE),
499 vk::DescriptorBufferInfo::default().buffer(features).range(vk::WHOLE_SIZE),
500 vk::DescriptorBufferInfo::default().buffer(src).range(vk::WHOLE_SIZE),
501 vk::DescriptorBufferInfo::default().buffer(dst).range(vk::WHOLE_SIZE),
502 ];
503 let image_infos = [vk::DescriptorImageInfo::default()
504 .image_view(output_view)
505 .image_layout(vk::ImageLayout::GENERAL)];
506 let writes: Vec<vk::WriteDescriptorSet> = buf_infos
507 .iter()
508 .enumerate()
509 .map(|(i, info)| {
510 vk::WriteDescriptorSet::default()
511 .dst_set(set)
512 .dst_binding(1 + i as u32)
513 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
514 .buffer_info(std::slice::from_ref(info))
515 })
516 .chain(std::iter::once(
517 vk::WriteDescriptorSet::default()
518 .dst_set(set)
519 .dst_binding(5)
520 .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
521 .image_info(&image_infos),
522 ))
523 .collect();
524 unsafe { device.update_descriptor_sets(&writes, &[]) };
525 }
526 }
527 }
528
529 /// After the frame fence: the per-iteration params. `n_after` is the
530 /// sample count the accumulation will hold once this frame's dispatch
531 /// lands — the color sigma tightens as it grows.
532 fn write_frame_uniforms(&mut self, frame_index: usize, width: u32, height: u32, n_after: u32) {
533 let inv_sqrt_n = 1.0 / (n_after.max(1) as f32).sqrt();
534 let frame = &mut self.frames[frame_index];
535 let mapped = frame.uniforms.allocation.as_mut().unwrap().mapped_slice_mut().unwrap();
536 for i in 0..DENOISE_ITERATIONS {
537 let params = DenoiseParams {
538 width,
539 height,
540 step: 1 << i,
541 first: (i == 0) as u32,
542 last: (i == DENOISE_ITERATIONS - 1) as u32,
543 inv_sqrt_n,
544 _pad: [0; 2],
545 };
546 let offset = self.uniform_stride as usize * i;
547 mapped[offset..offset + std::mem::size_of::<DenoiseParams>()]
548 .copy_from_slice(bytemuck::bytes_of(¶ms));
549 }
550 }
551
552 /// Record the à-trous iterations. The tracer's dispatch has already run
553 /// in this command buffer; the last iteration rewrites `out_img` (still
554 /// in GENERAL). Iteration parity: 0 → set_b (writes ping), 1 → set_a,
555 /// 2 → set_b.
556 fn record(&self, device: &ash::Device, cmd: vk::CommandBuffer, frame_index: usize, w: u32, h: u32) {
557 let frame = &self.frames[frame_index];
558 unsafe {
559 device.cmd_bind_pipeline(cmd, vk::PipelineBindPoint::COMPUTE, self.pipeline);
560 for i in 0..DENOISE_ITERATIONS {
561 // Order this iteration's reads after the previous compute
562 // writes (tracer or prior iteration).
563 device.cmd_pipeline_barrier(
564 cmd,
565 vk::PipelineStageFlags::COMPUTE_SHADER,
566 vk::PipelineStageFlags::COMPUTE_SHADER,
567 vk::DependencyFlags::empty(),
568 &[vk::MemoryBarrier::default()
569 .src_access_mask(vk::AccessFlags::SHADER_WRITE)
570 .dst_access_mask(
571 vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE,
572 )],
573 &[],
574 &[],
575 );
576 let set = if i % 2 == 0 { frame.set_b } else { frame.set_a };
577 device.cmd_bind_descriptor_sets(
578 cmd,
579 vk::PipelineBindPoint::COMPUTE,
580 self.pipeline_layout,
581 0,
582 &[set],
583 &[(self.uniform_stride as u32) * i as u32],
584 );
585 device.cmd_dispatch(cmd, w.div_ceil(WORKGROUP), h.div_ceil(WORKGROUP), 1);
586 }
587 }
588 }
589
590 fn destroy(&mut self, device: &ash::Device, allocator: &mut Allocator) {
591 unsafe {
592 for frame in &mut self.frames {
593 let mut uniforms = std::mem::replace(&mut frame.uniforms, AllocatedBuffer::null());
594 destroy_cpu_buffer(device, allocator, &mut uniforms);
595 }
596 device.destroy_descriptor_pool(self.descriptor_pool, None);
597 device.destroy_descriptor_set_layout(self.descriptor_set_layout, None);
598 device.destroy_pipeline(self.pipeline, None);
599 device.destroy_pipeline_layout(self.pipeline_layout, None);
600 device.destroy_shader_module(self.shader_module, None);
601 }
602 }
603 }
604
605 pub(crate) struct RtStage {
606 tier: RtTier,
607 accel_loader: Option<ash::khr::acceleration_structure::Device>,
608 as_scratch_align: vk::DeviceSize,
609 accel: Option<Accel>,
610
611 pipeline: vk::Pipeline,
612 pipeline_layout: vk::PipelineLayout,
613 descriptor_set_layout: vk::DescriptorSetLayout,
614 descriptor_pool: vk::DescriptorPool,
615 shader_module: vk::ShaderModule,
616 frames: Vec<RtFrame>,
617
618 nodes: AllocatedBuffer,
619 tris: AllocatedBuffer,
620 materials: AllocatedBuffer,
621 tri_count: u32,
622
623 accum: AllocatedBuffer,
624 /// Primary-hit features (2 vec4 per pixel) written by the tracer, read
625 /// by the denoiser.
626 features: AllocatedBuffer,
627 /// À-trous ping-pong color buffers (1 vec4 per pixel each).
628 ping: AllocatedBuffer,
629 pong: AllocatedBuffer,
630 denoiser: Option<Denoiser>,
631 output_image: vk::Image,
632 output_view: vk::ImageView,
633 output_allocation: Option<Allocation>,
634 output_size: (u32, u32),
635 output_initialized: bool,
636
637 pane: (u32, u32, u32, u32),
638 pane_moved: bool,
639 camera: Option<RtCamera>,
640 sample_index: u32,
641 /// Samples per dispatch: 1 interactive, higher for offscreen rendering.
642 spp: u32,
643 staged: bool,
644 }
645
646 impl RtStage {
647 /// `accel_loader` present means the device has the ray-query stack; the
648 /// stage then runs tier 2 unless `CCE_VK_RT=compute` forces the BVH tier.
649 pub(crate) fn new(
650 device: &ash::Device,
651 allocator: &mut Allocator,
652 frames_in_flight: usize,
653 accel_loader: Option<&ash::khr::acceleration_structure::Device>,
654 as_scratch_align: vk::DeviceSize,
655 min_uniform_align: vk::DeviceSize,
656 ) -> Self {
657 let force_compute = std::env::var("CCE_VK_RT").is_ok_and(|v| v == "compute");
658 let denoise_on = !std::env::var("CCE_VK_RT_DENOISE")
659 .is_ok_and(|v| v == "off" || v == "0" || v == "false");
660 let tier = if accel_loader.is_some() && !force_compute {
661 RtTier::RayQuery
662 } else {
663 RtTier::Compute
664 };
665 log::info!(
666 "RT stage: {} tier",
667 match tier {
668 RtTier::Compute => "compute (BVH)",
669 RtTier::RayQuery => "ray-query (hardware)",
670 }
671 );
672 let binding1_type = match tier {
673 RtTier::Compute => vk::DescriptorType::STORAGE_BUFFER,
674 RtTier::RayQuery => vk::DescriptorType::ACCELERATION_STRUCTURE_KHR,
675 };
676 unsafe {
677 let bindings = [
678 vk::DescriptorSetLayoutBinding::default()
679 .binding(0)
680 .descriptor_type(vk::DescriptorType::UNIFORM_BUFFER)
681 .descriptor_count(1)
682 .stage_flags(vk::ShaderStageFlags::COMPUTE),
683 vk::DescriptorSetLayoutBinding::default()
684 .binding(1)
685 .descriptor_type(binding1_type)
686 .descriptor_count(1)
687 .stage_flags(vk::ShaderStageFlags::COMPUTE),
688 vk::DescriptorSetLayoutBinding::default()
689 .binding(2)
690 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
691 .descriptor_count(1)
692 .stage_flags(vk::ShaderStageFlags::COMPUTE),
693 vk::DescriptorSetLayoutBinding::default()
694 .binding(3)
695 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
696 .descriptor_count(1)
697 .stage_flags(vk::ShaderStageFlags::COMPUTE),
698 vk::DescriptorSetLayoutBinding::default()
699 .binding(4)
700 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
701 .descriptor_count(1)
702 .stage_flags(vk::ShaderStageFlags::COMPUTE),
703 vk::DescriptorSetLayoutBinding::default()
704 .binding(5)
705 .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
706 .descriptor_count(1)
707 .stage_flags(vk::ShaderStageFlags::COMPUTE),
708 vk::DescriptorSetLayoutBinding::default()
709 .binding(6)
710 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
711 .descriptor_count(1)
712 .stage_flags(vk::ShaderStageFlags::COMPUTE),
713 ];
714 let descriptor_set_layout = device
715 .create_descriptor_set_layout(
716 &vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings),
717 None,
718 )
719 .expect("Failed to create RT descriptor set layout");
720 let set_layouts_one = [descriptor_set_layout];
721 let pipeline_layout = device
722 .create_pipeline_layout(
723 &vk::PipelineLayoutCreateInfo::default().set_layouts(&set_layouts_one),
724 None,
725 )
726 .expect("Failed to create RT pipeline layout");
727
728 let spirv = match tier {
729 RtTier::Compute => compile_wgsl(&format!(
730 "{}\n{}",
731 include_str!("rt_common.wgsl"),
732 include_str!("rt_bvh.wgsl")
733 )),
734 RtTier::RayQuery => compile_wgsl_ray_query(&format!(
735 "{}\n{}",
736 include_str!("rt_common.wgsl"),
737 include_str!("rt_query.wgsl")
738 )),
739 };
740 let shader_module = device
741 .create_shader_module(&vk::ShaderModuleCreateInfo::default().code(&spirv), None)
742 .expect("Failed to create RT shader module");
743 let pipeline = device
744 .create_compute_pipelines(
745 vk::PipelineCache::null(),
746 &[vk::ComputePipelineCreateInfo::default()
747 .stage(
748 vk::PipelineShaderStageCreateInfo::default()
749 .stage(vk::ShaderStageFlags::COMPUTE)
750 .module(shader_module)
751 .name(c"cs_main"),
752 )
753 .layout(pipeline_layout)],
754 None,
755 )
756 .expect("Failed to create RT compute pipeline")[0];
757
758 let n = frames_in_flight as u32;
759 let mut pool_sizes = vec![
760 vk::DescriptorPoolSize::default()
761 .ty(vk::DescriptorType::UNIFORM_BUFFER)
762 .descriptor_count(n),
763 vk::DescriptorPoolSize::default()
764 .ty(vk::DescriptorType::STORAGE_BUFFER)
765 .descriptor_count(5 * n),
766 vk::DescriptorPoolSize::default()
767 .ty(vk::DescriptorType::STORAGE_IMAGE)
768 .descriptor_count(n),
769 ];
770 if tier == RtTier::RayQuery {
771 pool_sizes.push(
772 vk::DescriptorPoolSize::default()
773 .ty(vk::DescriptorType::ACCELERATION_STRUCTURE_KHR)
774 .descriptor_count(n),
775 );
776 }
777 let descriptor_pool = device
778 .create_descriptor_pool(
779 &vk::DescriptorPoolCreateInfo::default()
780 .max_sets(n)
781 .pool_sizes(&pool_sizes),
782 None,
783 )
784 .expect("Failed to create RT descriptor pool");
785 let set_layouts: Vec<vk::DescriptorSetLayout> =
786 vec![descriptor_set_layout; frames_in_flight];
787 let sets = device
788 .allocate_descriptor_sets(
789 &vk::DescriptorSetAllocateInfo::default()
790 .descriptor_pool(descriptor_pool)
791 .set_layouts(&set_layouts),
792 )
793 .expect("Failed to allocate RT descriptor sets");
794 let frames: Vec<RtFrame> = sets
795 .into_iter()
796 .map(|descriptor_set| {
797 let uniforms = create_cpu_buffer(
798 device,
799 allocator,
800 std::mem::size_of::<RtParams>() as vk::DeviceSize,
801 vk::BufferUsageFlags::UNIFORM_BUFFER,
802 "rt-uniforms",
803 );
804 let buffer_infos = [vk::DescriptorBufferInfo::default()
805 .buffer(uniforms.buffer)
806 .offset(0)
807 .range(std::mem::size_of::<RtParams>() as vk::DeviceSize)];
808 device.update_descriptor_sets(
809 &[vk::WriteDescriptorSet::default()
810 .dst_set(descriptor_set)
811 .dst_binding(0)
812 .descriptor_type(vk::DescriptorType::UNIFORM_BUFFER)
813 .buffer_info(&buffer_infos)],
814 &[],
815 );
816 RtFrame { uniforms, descriptor_set }
817 })
818 .collect();
819
820 let denoiser = denoise_on
821 .then(|| Denoiser::new(device, allocator, frames_in_flight, min_uniform_align));
822
823 RtStage {
824 tier,
825 accel_loader: accel_loader.cloned(),
826 as_scratch_align,
827 accel: None,
828 pipeline,
829 pipeline_layout,
830 descriptor_set_layout,
831 descriptor_pool,
832 shader_module,
833 frames,
834 nodes: AllocatedBuffer::null(),
835 tris: AllocatedBuffer::null(),
836 materials: AllocatedBuffer::null(),
837 tri_count: 0,
838 accum: AllocatedBuffer::null(),
839 features: AllocatedBuffer::null(),
840 ping: AllocatedBuffer::null(),
841 pong: AllocatedBuffer::null(),
842 denoiser,
843 output_image: vk::Image::null(),
844 output_view: vk::ImageView::null(),
845 output_allocation: None,
846 output_size: (0, 0),
847 output_initialized: false,
848 pane: (0, 0, 0, 0),
849 pane_moved: false,
850 camera: None,
851 sample_index: 0,
852 spp: 1,
853 staged: false,
854 }
855 }
856 }
857
858 /// Replace the scene. Tier 1 builds the BVH on the CPU (reordering a copy
859 /// of the triangles); tier 2 builds driver acceleration structures on the
860 /// given queue instead. Caller must have the device idle.
861 pub(crate) fn set_scene(
862 &mut self,
863 device: &ash::Device,
864 allocator: &mut Allocator,
865 queue: vk::Queue,
866 command_pool: vk::CommandPool,
867 triangles: &[RtTriangle],
868 materials: &[RtMaterial],
869 ) {
870 let mut tris: Vec<RtTriangle> = triangles.to_vec();
871 let nodes = match self.tier {
872 RtTier::Compute => build_bvh(&mut tris),
873 RtTier::RayQuery => Vec::new(),
874 };
875 let gpu_tris: Vec<GpuTriangle> = tris
876 .iter()
877 .map(|t| GpuTriangle {
878 p0: [t.p0[0], t.p0[1], t.p0[2], f32::from_bits(t.material)],
879 p1: [t.p1[0], t.p1[1], t.p1[2], 0.0],
880 p2: [t.p2[0], t.p2[1], t.p2[2], 0.0],
881 })
882 .collect();
883 let gpu_mats: Vec<GpuMaterial> = if materials.is_empty() {
884 vec![GpuMaterial { albedo: [0.8, 0.8, 0.8, 0.0], emission: [0.0; 4] }]
885 } else {
886 materials
887 .iter()
888 .map(|m| GpuMaterial {
889 albedo: [m.albedo[0], m.albedo[1], m.albedo[2], 0.0],
890 emission: [m.emission[0], m.emission[1], m.emission[2], 0.0],
891 })
892 .collect()
893 };
894
895 self.destroy_accel(device, allocator);
896 for buf in [&mut self.nodes, &mut self.tris, &mut self.materials] {
897 destroy_cpu_buffer(device, allocator, buf);
898 }
899 let upload = |allocator: &mut Allocator,
900 bytes: &[u8],
901 usage: vk::BufferUsageFlags,
902 name: &str|
903 -> AllocatedBuffer {
904 let mut buf = create_cpu_buffer(
905 device,
906 allocator,
907 (bytes.len() as vk::DeviceSize).max(64),
908 usage,
909 name,
910 );
911 if !bytes.is_empty() {
912 buf.allocation.as_mut().unwrap().mapped_slice_mut().unwrap()[..bytes.len()]
913 .copy_from_slice(bytes);
914 }
915 buf
916 };
917 // Tier 2 reads the same triangle buffer as BLAS build input (the
918 // shading data still comes through the storage binding).
919 let tri_usage = match self.tier {
920 RtTier::Compute => vk::BufferUsageFlags::STORAGE_BUFFER,
921 RtTier::RayQuery => {
922 vk::BufferUsageFlags::STORAGE_BUFFER
923 | vk::BufferUsageFlags::SHADER_DEVICE_ADDRESS
924 | vk::BufferUsageFlags::ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_KHR
925 }
926 };
927 self.tris = upload(allocator, bytemuck::cast_slice(&gpu_tris), tri_usage, "rt-tris");
928 self.materials = upload(
929 allocator,
930 bytemuck::cast_slice(&gpu_mats),
931 vk::BufferUsageFlags::STORAGE_BUFFER,
932 "rt-materials",
933 );
934 self.tri_count = tris.len() as u32;
935 self.sample_index = 0;
936
937 // Binding 1 (per tier), then the shared 2/3.
938 match self.tier {
939 RtTier::Compute => {
940 self.nodes = upload(
941 allocator,
942 bytemuck::cast_slice(&nodes),
943 vk::BufferUsageFlags::STORAGE_BUFFER,
944 "rt-nodes",
945 );
946 for frame in &self.frames {
947 let infos = [vk::DescriptorBufferInfo::default()
948 .buffer(self.nodes.buffer)
949 .range(vk::WHOLE_SIZE)];
950 unsafe {
951 device.update_descriptor_sets(
952 &[vk::WriteDescriptorSet::default()
953 .dst_set(frame.descriptor_set)
954 .dst_binding(1)
955 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
956 .buffer_info(&infos)],
957 &[],
958 );
959 }
960 }
961 }
962 RtTier::RayQuery => {
963 if self.tri_count > 0 {
964 self.build_accel(device, allocator, queue, command_pool);
965 let accel = self.accel.as_ref().unwrap();
966 let handles = [accel.tlas];
967 for frame in &self.frames {
968 let mut as_info =
969 vk::WriteDescriptorSetAccelerationStructureKHR::default()
970 .acceleration_structures(&handles);
971 let mut write = vk::WriteDescriptorSet::default()
972 .dst_set(frame.descriptor_set)
973 .dst_binding(1)
974 .descriptor_type(vk::DescriptorType::ACCELERATION_STRUCTURE_KHR)
975 .push_next(&mut as_info);
976 write.descriptor_count = 1;
977 unsafe { device.update_descriptor_sets(&[write], &[]) };
978 }
979 }
980 }
981 }
982
983 for frame in &self.frames {
984 let infos = [
985 vk::DescriptorBufferInfo::default().buffer(self.tris.buffer).range(vk::WHOLE_SIZE),
986 vk::DescriptorBufferInfo::default()
987 .buffer(self.materials.buffer)
988 .range(vk::WHOLE_SIZE),
989 ];
990 let writes: Vec<vk::WriteDescriptorSet> = infos
991 .iter()
992 .enumerate()
993 .map(|(i, info)| {
994 vk::WriteDescriptorSet::default()
995 .dst_set(frame.descriptor_set)
996 .dst_binding(2 + i as u32)
997 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
998 .buffer_info(std::slice::from_ref(info))
999 })
1000 .collect();
1001 unsafe { device.update_descriptor_sets(&writes, &[]) };
1002 }
1003 }
1004
1005 /// Build the BLAS (over `self.tris`, opaque triangles) and a one-instance
1006 /// TLAS, on the given queue with a blocking one-time submit. Device is
1007 /// idle (set_scene contract), so replacing old structures is safe.
1008 fn build_accel(
1009 &mut self,
1010 device: &ash::Device,
1011 allocator: &mut Allocator,
1012 queue: vk::Queue,
1013 command_pool: vk::CommandPool,
1014 ) {
1015 let loader = self.accel_loader.clone().expect("tier 2 without accel loader");
1016 let create_as_buffer = |allocator: &mut Allocator,
1017 size: vk::DeviceSize,
1018 usage: vk::BufferUsageFlags,
1019 name: &str|
1020 -> AllocatedBuffer {
1021 unsafe {
1022 let buffer = device
1023 .create_buffer(
1024 &vk::BufferCreateInfo::default()
1025 .size(size)
1026 .usage(usage | vk::BufferUsageFlags::SHADER_DEVICE_ADDRESS)
1027 .sharing_mode(vk::SharingMode::EXCLUSIVE),
1028 None,
1029 )
1030 .expect("Failed to create AS buffer");
1031 let requirements = device.get_buffer_memory_requirements(buffer);
1032 let allocation = allocator
1033 .allocate(&AllocationCreateDesc {
1034 name,
1035 requirements,
1036 location: MemoryLocation::GpuOnly,
1037 linear: true,
1038 allocation_scheme: AllocationScheme::GpuAllocatorManaged,
1039 })
1040 .expect("Failed to allocate AS memory");
1041 device
1042 .bind_buffer_memory(buffer, allocation.memory(), allocation.offset())
1043 .expect("Failed to bind AS memory");
1044 AllocatedBuffer { buffer, allocation: Some(allocation), size }
1045 }
1046 };
1047 let addr_of = |buffer: vk::Buffer| unsafe {
1048 device.get_buffer_device_address(&vk::BufferDeviceAddressInfo::default().buffer(buffer))
1049 };
1050
1051 unsafe {
1052 // --- BLAS over the triangle buffer (stride 16: p0/p1/p2 vec4s).
1053 let tri_addr = addr_of(self.tris.buffer);
1054 let blas_geometry = vk::AccelerationStructureGeometryKHR::default()
1055 .geometry_type(vk::GeometryTypeKHR::TRIANGLES)
1056 .flags(vk::GeometryFlagsKHR::OPAQUE)
1057 .geometry(vk::AccelerationStructureGeometryDataKHR {
1058 triangles: vk::AccelerationStructureGeometryTrianglesDataKHR::default()
1059 .vertex_format(vk::Format::R32G32B32_SFLOAT)
1060 .vertex_data(vk::DeviceOrHostAddressConstKHR { device_address: tri_addr })
1061 .vertex_stride(16)
1062 .max_vertex(self.tri_count * 3 - 1)
1063 .index_type(vk::IndexType::NONE_KHR),
1064 });
1065 let blas_geometries = [blas_geometry];
1066 let mut blas_build = vk::AccelerationStructureBuildGeometryInfoKHR::default()
1067 .ty(vk::AccelerationStructureTypeKHR::BOTTOM_LEVEL)
1068 .flags(vk::BuildAccelerationStructureFlagsKHR::PREFER_FAST_TRACE)
1069 .mode(vk::BuildAccelerationStructureModeKHR::BUILD)
1070 .geometries(&blas_geometries);
1071 let blas_sizes = {
1072 let mut sizes = vk::AccelerationStructureBuildSizesInfoKHR::default();
1073 loader.get_acceleration_structure_build_sizes(
1074 vk::AccelerationStructureBuildTypeKHR::DEVICE,
1075 &blas_build,
1076 &[self.tri_count],
1077 &mut sizes,
1078 );
1079 sizes
1080 };
1081 let blas_buffer = create_as_buffer(
1082 allocator,
1083 blas_sizes.acceleration_structure_size,
1084 vk::BufferUsageFlags::ACCELERATION_STRUCTURE_STORAGE_KHR,
1085 "rt-blas",
1086 );
1087 let blas = loader
1088 .create_acceleration_structure(
1089 &vk::AccelerationStructureCreateInfoKHR::default()
1090 .buffer(blas_buffer.buffer)
1091 .size(blas_sizes.acceleration_structure_size)
1092 .ty(vk::AccelerationStructureTypeKHR::BOTTOM_LEVEL),
1093 None,
1094 )
1095 .expect("Failed to create BLAS");
1096
1097 // --- One-instance TLAS.
1098 let blas_addr = loader.get_acceleration_structure_device_address(
1099 &vk::AccelerationStructureDeviceAddressInfoKHR::default()
1100 .acceleration_structure(blas),
1101 );
1102 let instance = vk::AccelerationStructureInstanceKHR {
1103 transform: vk::TransformMatrixKHR {
1104 matrix: [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
1105 },
1106 instance_custom_index_and_mask: vk::Packed24_8::new(0, 0xff),
1107 instance_shader_binding_table_record_offset_and_flags: vk::Packed24_8::new(0, 0),
1108 acceleration_structure_reference: vk::AccelerationStructureReferenceKHR {
1109 device_handle: blas_addr,
1110 },
1111 };
1112 let instance_bytes = std::slice::from_raw_parts(
1113 (&instance as *const vk::AccelerationStructureInstanceKHR).cast::<u8>(),
1114 std::mem::size_of::<vk::AccelerationStructureInstanceKHR>(),
1115 );
1116 let mut instances = create_cpu_buffer(
1117 device,
1118 allocator,
1119 instance_bytes.len() as vk::DeviceSize,
1120 vk::BufferUsageFlags::SHADER_DEVICE_ADDRESS
1121 | vk::BufferUsageFlags::ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_KHR,
1122 "rt-tlas-instances",
1123 );
1124 instances.allocation.as_mut().unwrap().mapped_slice_mut().unwrap()
1125 [..instance_bytes.len()]
1126 .copy_from_slice(instance_bytes);
1127
1128 let tlas_geometry = vk::AccelerationStructureGeometryKHR::default()
1129 .geometry_type(vk::GeometryTypeKHR::INSTANCES)
1130 .geometry(vk::AccelerationStructureGeometryDataKHR {
1131 instances: vk::AccelerationStructureGeometryInstancesDataKHR::default()
1132 .array_of_pointers(false)
1133 .data(vk::DeviceOrHostAddressConstKHR {
1134 device_address: addr_of(instances.buffer),
1135 }),
1136 });
1137 let tlas_geometries = [tlas_geometry];
1138 let mut tlas_build = vk::AccelerationStructureBuildGeometryInfoKHR::default()
1139 .ty(vk::AccelerationStructureTypeKHR::TOP_LEVEL)
1140 .flags(vk::BuildAccelerationStructureFlagsKHR::PREFER_FAST_TRACE)
1141 .mode(vk::BuildAccelerationStructureModeKHR::BUILD)
1142 .geometries(&tlas_geometries);
1143 let tlas_sizes = {
1144 let mut sizes = vk::AccelerationStructureBuildSizesInfoKHR::default();
1145 loader.get_acceleration_structure_build_sizes(
1146 vk::AccelerationStructureBuildTypeKHR::DEVICE,
1147 &tlas_build,
1148 &[1],
1149 &mut sizes,
1150 );
1151 sizes
1152 };
1153 let tlas_buffer = create_as_buffer(
1154 allocator,
1155 tlas_sizes.acceleration_structure_size,
1156 vk::BufferUsageFlags::ACCELERATION_STRUCTURE_STORAGE_KHR,
1157 "rt-tlas",
1158 );
1159 let tlas = loader
1160 .create_acceleration_structure(
1161 &vk::AccelerationStructureCreateInfoKHR::default()
1162 .buffer(tlas_buffer.buffer)
1163 .size(tlas_sizes.acceleration_structure_size)
1164 .ty(vk::AccelerationStructureTypeKHR::TOP_LEVEL),
1165 None,
1166 )
1167 .expect("Failed to create TLAS");
1168
1169 // Shared scratch, aligned to the device's scratch requirement
1170 // (buffer device addresses only guarantee allocation alignment).
1171 let scratch_size =
1172 blas_sizes.build_scratch_size.max(tlas_sizes.build_scratch_size);
1173 let mut scratch = create_as_buffer(
1174 allocator,
1175 scratch_size + self.as_scratch_align,
1176 vk::BufferUsageFlags::STORAGE_BUFFER,
1177 "rt-as-scratch",
1178 );
1179 let scratch_addr =
1180 addr_of(scratch.buffer).next_multiple_of(self.as_scratch_align.max(1));
1181
1182 blas_build = blas_build
1183 .dst_acceleration_structure(blas)
1184 .scratch_data(vk::DeviceOrHostAddressKHR { device_address: scratch_addr });
1185 tlas_build = tlas_build
1186 .dst_acceleration_structure(tlas)
1187 .scratch_data(vk::DeviceOrHostAddressKHR { device_address: scratch_addr });
1188
1189 // One-time submit: BLAS build → barrier → TLAS build.
1190 let cmd = device
1191 .allocate_command_buffers(
1192 &vk::CommandBufferAllocateInfo::default()
1193 .command_pool(command_pool)
1194 .level(vk::CommandBufferLevel::PRIMARY)
1195 .command_buffer_count(1),
1196 )
1197 .expect("Failed to allocate AS build command buffer")[0];
1198 device
1199 .begin_command_buffer(
1200 cmd,
1201 &vk::CommandBufferBeginInfo::default()
1202 .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT),
1203 )
1204 .unwrap();
1205 let blas_range = [vk::AccelerationStructureBuildRangeInfoKHR::default()
1206 .primitive_count(self.tri_count)];
1207 loader.cmd_build_acceleration_structures(cmd, &[blas_build], &[&blas_range]);
1208 device.cmd_pipeline_barrier(
1209 cmd,
1210 vk::PipelineStageFlags::ACCELERATION_STRUCTURE_BUILD_KHR,
1211 vk::PipelineStageFlags::ACCELERATION_STRUCTURE_BUILD_KHR,
1212 vk::DependencyFlags::empty(),
1213 &[vk::MemoryBarrier::default()
1214 .src_access_mask(vk::AccessFlags::ACCELERATION_STRUCTURE_WRITE_KHR)
1215 .dst_access_mask(
1216 vk::AccessFlags::ACCELERATION_STRUCTURE_READ_KHR
1217 | vk::AccessFlags::ACCELERATION_STRUCTURE_WRITE_KHR,
1218 )],
1219 &[],
1220 &[],
1221 );
1222 let tlas_range =
1223 [vk::AccelerationStructureBuildRangeInfoKHR::default().primitive_count(1)];
1224 loader.cmd_build_acceleration_structures(cmd, &[tlas_build], &[&tlas_range]);
1225 device.end_command_buffer(cmd).unwrap();
1226 let cmds = [cmd];
1227 device
1228 .queue_submit(
1229 queue,
1230 &[vk::SubmitInfo::default().command_buffers(&cmds)],
1231 vk::Fence::null(),
1232 )
1233 .expect("AS build submit failed");
1234 let _ = device.queue_wait_idle(queue);
1235 device.free_command_buffers(command_pool, &cmds);
1236 destroy_cpu_buffer(device, allocator, &mut scratch);
1237
1238 self.accel = Some(Accel { blas, blas_buffer, tlas, tlas_buffer, instances });
1239 }
1240 }
1241
1242 fn destroy_accel(&mut self, device: &ash::Device, allocator: &mut Allocator) {
1243 if let Some(mut accel) = self.accel.take() {
1244 let loader = self.accel_loader.as_ref().expect("accel without loader");
1245 unsafe {
1246 loader.destroy_acceleration_structure(accel.tlas, None);
1247 loader.destroy_acceleration_structure(accel.blas, None);
1248 }
1249 destroy_cpu_buffer(device, allocator, &mut accel.tlas_buffer);
1250 destroy_cpu_buffer(device, allocator, &mut accel.blas_buffer);
1251 destroy_cpu_buffer(device, allocator, &mut accel.instances);
1252 }
1253 }
1254
1255 /// Stage an RT frame for the viewport pane (physical pixels). Recreates the
1256 /// pane-sized targets on size change (waits for device idle) and resets the
1257 /// accumulation when the camera, size, or scene changed.
1258 pub(crate) fn stage(
1259 &mut self,
1260 device: &ash::Device,
1261 allocator: &mut Allocator,
1262 pane: (u32, u32, u32, u32),
1263 camera: RtCamera,
1264 ) {
1265 let (_, _, w, h) = pane;
1266 if w == 0 || h == 0 {
1267 self.staged = false;
1268 return;
1269 }
1270 if (w, h) != self.output_size {
1271 unsafe {
1272 let _ = device.device_wait_idle();
1273 }
1274 self.recreate_targets(device, allocator, w, h);
1275 self.sample_index = 0;
1276 }
1277 if pane != self.pane && self.pane != (0, 0, 0, 0) {
1278 // Pane moved or shrank: stale RT pixels sit outside the new
1279 // region; clear the backdrop once before the next blit.
1280 self.pane_moved = true;
1281 }
1282 if self.camera != Some(camera) {
1283 self.camera = Some(camera);
1284 self.sample_index = 0;
1285 }
1286 self.pane = pane;
1287 self.staged = true;
1288 }
1289
1290 fn recreate_targets(&mut self, device: &ash::Device, allocator: &mut Allocator, w: u32, h: u32) {
1291 self.destroy_targets(device, allocator);
1292 self.output_size = (w, h);
1293 self.output_initialized = false;
1294 let gpu_buffer = |allocator: &mut Allocator,
1295 bytes_per_px: vk::DeviceSize,
1296 name: &'static str|
1297 -> AllocatedBuffer {
1298 let size = (w as vk::DeviceSize) * (h as vk::DeviceSize) * bytes_per_px;
1299 unsafe {
1300 let buffer = device
1301 .create_buffer(
1302 &vk::BufferCreateInfo::default()
1303 .size(size)
1304 .usage(vk::BufferUsageFlags::STORAGE_BUFFER)
1305 .sharing_mode(vk::SharingMode::EXCLUSIVE),
1306 None,
1307 )
1308 .expect("Failed to create RT target buffer");
1309 let requirements = device.get_buffer_memory_requirements(buffer);
1310 let allocation = allocator
1311 .allocate(&AllocationCreateDesc {
1312 name,
1313 requirements,
1314 location: MemoryLocation::GpuOnly,
1315 linear: true,
1316 allocation_scheme: AllocationScheme::GpuAllocatorManaged,
1317 })
1318 .expect("Failed to allocate RT target memory");
1319 device
1320 .bind_buffer_memory(buffer, allocation.memory(), allocation.offset())
1321 .expect("Failed to bind RT target memory");
1322 AllocatedBuffer { buffer, allocation: Some(allocation), size }
1323 }
1324 };
1325 self.accum = gpu_buffer(allocator, 16, "rt-accum");
1326 self.features = gpu_buffer(allocator, 32, "rt-features");
1327 if self.denoiser.is_some() {
1328 self.ping = gpu_buffer(allocator, 16, "rt-denoise-ping");
1329 self.pong = gpu_buffer(allocator, 16, "rt-denoise-pong");
1330 }
1331 unsafe {
1332 let image = device
1333 .create_image(
1334 &vk::ImageCreateInfo::default()
1335 .image_type(vk::ImageType::TYPE_2D)
1336 .format(vk::Format::R8G8B8A8_UNORM)
1337 .extent(vk::Extent3D { width: w, height: h, depth: 1 })
1338 .mip_levels(1)
1339 .array_layers(1)
1340 .samples(vk::SampleCountFlags::TYPE_1)
1341 .tiling(vk::ImageTiling::OPTIMAL)
1342 .usage(vk::ImageUsageFlags::STORAGE | vk::ImageUsageFlags::TRANSFER_SRC)
1343 .initial_layout(vk::ImageLayout::UNDEFINED),
1344 None,
1345 )
1346 .expect("Failed to create RT output image");
1347 let requirements = device.get_image_memory_requirements(image);
1348 let allocation = allocator
1349 .allocate(&AllocationCreateDesc {
1350 name: "rt-output",
1351 requirements,
1352 location: MemoryLocation::GpuOnly,
1353 linear: false,
1354 allocation_scheme: AllocationScheme::GpuAllocatorManaged,
1355 })
1356 .expect("Failed to allocate RT output memory");
1357 device
1358 .bind_image_memory(image, allocation.memory(), allocation.offset())
1359 .expect("Failed to bind RT output memory");
1360 let view = device
1361 .create_image_view(
1362 &vk::ImageViewCreateInfo::default()
1363 .image(image)
1364 .view_type(vk::ImageViewType::TYPE_2D)
1365 .format(vk::Format::R8G8B8A8_UNORM)
1366 .subresource_range(
1367 vk::ImageSubresourceRange::default()
1368 .aspect_mask(vk::ImageAspectFlags::COLOR)
1369 .level_count(1)
1370 .layer_count(1),
1371 ),
1372 None,
1373 )
1374 .expect("Failed to create RT output view");
1375 self.output_image = image;
1376 self.output_view = view;
1377 self.output_allocation = Some(allocation);
1378
1379 for frame in &self.frames {
1380 let accum_infos = [vk::DescriptorBufferInfo::default()
1381 .buffer(self.accum.buffer)
1382 .range(vk::WHOLE_SIZE)];
1383 let feature_infos = [vk::DescriptorBufferInfo::default()
1384 .buffer(self.features.buffer)
1385 .range(vk::WHOLE_SIZE)];
1386 let image_infos = [vk::DescriptorImageInfo::default()
1387 .image_view(view)
1388 .image_layout(vk::ImageLayout::GENERAL)];
1389 device.update_descriptor_sets(
1390 &[
1391 vk::WriteDescriptorSet::default()
1392 .dst_set(frame.descriptor_set)
1393 .dst_binding(4)
1394 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1395 .buffer_info(&accum_infos),
1396 vk::WriteDescriptorSet::default()
1397 .dst_set(frame.descriptor_set)
1398 .dst_binding(5)
1399 .descriptor_type(vk::DescriptorType::STORAGE_IMAGE)
1400 .image_info(&image_infos),
1401 vk::WriteDescriptorSet::default()
1402 .dst_set(frame.descriptor_set)
1403 .dst_binding(6)
1404 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1405 .buffer_info(&feature_infos),
1406 ],
1407 &[],
1408 );
1409 }
1410 if let Some(denoiser) = &self.denoiser {
1411 denoiser.write_target_descriptors(
1412 device,
1413 self.accum.buffer,
1414 self.features.buffer,
1415 self.ping.buffer,
1416 self.pong.buffer,
1417 view,
1418 );
1419 }
1420 }
1421 }
1422
1423 fn destroy_targets(&mut self, device: &ash::Device, allocator: &mut Allocator) {
1424 unsafe {
1425 if self.output_view != vk::ImageView::null() {
1426 device.destroy_image_view(self.output_view, None);
1427 device.destroy_image(self.output_image, None);
1428 self.output_view = vk::ImageView::null();
1429 self.output_image = vk::Image::null();
1430 }
1431 }
1432 if let Some(a) = self.output_allocation.take() {
1433 let _ = allocator.free(a);
1434 }
1435 destroy_cpu_buffer(device, allocator, &mut self.accum);
1436 destroy_cpu_buffer(device, allocator, &mut self.features);
1437 destroy_cpu_buffer(device, allocator, &mut self.ping);
1438 destroy_cpu_buffer(device, allocator, &mut self.pong);
1439 self.output_size = (0, 0);
1440 }
1441
1442 /// True while another dispatch would still refine the image.
1443 pub(crate) fn accumulating(&self) -> bool {
1444 self.tri_count > 0 && self.sample_index < MAX_SAMPLES
1445 }
1446
1447 /// After the frame fence: write this frame's params.
1448 pub(crate) fn write_frame_uniforms(&mut self, frame_index: usize) {
1449 if !self.staged || self.tri_count == 0 {
1450 return;
1451 }
1452 let Some(camera) = self.camera else { return };
1453 let params = RtParams {
1454 inv_mvp: camera.inv_mvp,
1455 width: self.output_size.0,
1456 height: self.output_size.1,
1457 sample_index: self.sample_index,
1458 max_bounces: MAX_BOUNCES,
1459 spp: self.spp,
1460 _pad: [0; 3],
1461 };
1462 let frame = &mut self.frames[frame_index];
1463 frame.uniforms.allocation.as_mut().unwrap().mapped_slice_mut().unwrap()
1464 [..std::mem::size_of::<RtParams>()]
1465 .copy_from_slice(bytemuck::bytes_of(¶ms));
1466 if let Some(denoiser) = &mut self.denoiser {
1467 denoiser.write_frame_uniforms(
1468 frame_index,
1469 self.output_size.0,
1470 self.output_size.1,
1471 self.sample_index + self.spp,
1472 );
1473 }
1474 }
1475
1476 /// Record one accumulation dispatch + the blit into the backdrop's pane
1477 /// region. Returns false when there is nothing to do (not staged, empty
1478 /// scene, or converged) — the backdrop then simply keeps its content.
1479 /// On true, the backdrop ends in TRANSFER_SRC (like `SceneStage::record`).
1480 ///
1481 /// `backdrop_in_transfer_src` says the raster scene pass already ran this
1482 /// frame (backdrop in TRANSFER_SRC); otherwise it is in SHADER_READ_ONLY.
1483 pub(crate) fn record(
1484 &mut self,
1485 device: &ash::Device,
1486 cmd: vk::CommandBuffer,
1487 frame_index: usize,
1488 backdrop_image: vk::Image,
1489 backdrop_extent: vk::Extent2D,
1490 backdrop_in_transfer_src: bool,
1491 ) -> bool {
1492 if !self.staged || self.tri_count == 0 || self.output_image == vk::Image::null() {
1493 self.staged = false;
1494 return false;
1495 }
1496 self.staged = false;
1497 if self.sample_index >= MAX_SAMPLES {
1498 return false;
1499 }
1500 let (w, h) = self.output_size;
1501 let color_range = vk::ImageSubresourceRange::default()
1502 .aspect_mask(vk::ImageAspectFlags::COLOR)
1503 .level_count(1)
1504 .layer_count(1);
1505 unsafe {
1506 // Output image to GENERAL for the compute write; order this
1507 // dispatch's accum access after the previous frame's. The src
1508 // stage always includes COMPUTE_SHADER — the accum buffer
1509 // barrier's access flags must be legal for it even on the first
1510 // dispatch, when the image side is still UNDEFINED.
1511 let (old_layout, src_access, src_stage) = if self.output_initialized {
1512 (
1513 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
1514 vk::AccessFlags::TRANSFER_READ,
1515 vk::PipelineStageFlags::TRANSFER | vk::PipelineStageFlags::COMPUTE_SHADER,
1516 )
1517 } else {
1518 (
1519 vk::ImageLayout::UNDEFINED,
1520 vk::AccessFlags::empty(),
1521 vk::PipelineStageFlags::COMPUTE_SHADER,
1522 )
1523 };
1524 device.cmd_pipeline_barrier(
1525 cmd,
1526 src_stage,
1527 vk::PipelineStageFlags::COMPUTE_SHADER,
1528 vk::DependencyFlags::empty(),
1529 &[],
1530 &[vk::BufferMemoryBarrier::default()
1531 .src_access_mask(vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE)
1532 .dst_access_mask(vk::AccessFlags::SHADER_READ | vk::AccessFlags::SHADER_WRITE)
1533 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1534 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1535 .buffer(self.accum.buffer)
1536 .size(vk::WHOLE_SIZE)],
1537 &[vk::ImageMemoryBarrier::default()
1538 .src_access_mask(src_access)
1539 .dst_access_mask(vk::AccessFlags::SHADER_WRITE)
1540 .old_layout(old_layout)
1541 .new_layout(vk::ImageLayout::GENERAL)
1542 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1543 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1544 .image(self.output_image)
1545 .subresource_range(color_range)],
1546 );
1547 self.output_initialized = true;
1548
1549 device.cmd_bind_pipeline(cmd, vk::PipelineBindPoint::COMPUTE, self.pipeline);
1550 device.cmd_bind_descriptor_sets(
1551 cmd,
1552 vk::PipelineBindPoint::COMPUTE,
1553 self.pipeline_layout,
1554 0,
1555 &[self.frames[frame_index].descriptor_set],
1556 &[],
1557 );
1558 device.cmd_dispatch(cmd, w.div_ceil(WORKGROUP), h.div_ceil(WORKGROUP), 1);
1559
1560 // À-trous denoise passes; the last one rewrites out_img (still
1561 // GENERAL), so the transfer barrier below covers either writer.
1562 if let Some(denoiser) = &self.denoiser {
1563 denoiser.record(device, cmd, frame_index, w, h);
1564 }
1565
1566 // Output to TRANSFER_SRC, backdrop to TRANSFER_DST for the blit.
1567 let (bd_old, bd_access, bd_stage) = if backdrop_in_transfer_src {
1568 (
1569 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
1570 vk::AccessFlags::TRANSFER_READ,
1571 vk::PipelineStageFlags::TRANSFER,
1572 )
1573 } else {
1574 (
1575 vk::ImageLayout::SHADER_READ_ONLY_OPTIMAL,
1576 vk::AccessFlags::SHADER_READ,
1577 vk::PipelineStageFlags::FRAGMENT_SHADER,
1578 )
1579 };
1580 device.cmd_pipeline_barrier(
1581 cmd,
1582 vk::PipelineStageFlags::COMPUTE_SHADER | bd_stage,
1583 vk::PipelineStageFlags::TRANSFER,
1584 vk::DependencyFlags::empty(),
1585 &[],
1586 &[],
1587 &[
1588 vk::ImageMemoryBarrier::default()
1589 .src_access_mask(vk::AccessFlags::SHADER_WRITE)
1590 .dst_access_mask(vk::AccessFlags::TRANSFER_READ)
1591 .old_layout(vk::ImageLayout::GENERAL)
1592 .new_layout(vk::ImageLayout::TRANSFER_SRC_OPTIMAL)
1593 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1594 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1595 .image(self.output_image)
1596 .subresource_range(color_range),
1597 vk::ImageMemoryBarrier::default()
1598 .src_access_mask(bd_access)
1599 .dst_access_mask(vk::AccessFlags::TRANSFER_WRITE)
1600 .old_layout(bd_old)
1601 .new_layout(vk::ImageLayout::TRANSFER_DST_OPTIMAL)
1602 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1603 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1604 .image(backdrop_image)
1605 .subresource_range(color_range),
1606 ],
1607 );
1608
1609 if self.pane_moved {
1610 self.pane_moved = false;
1611 device.cmd_clear_color_image(
1612 cmd,
1613 backdrop_image,
1614 vk::ImageLayout::TRANSFER_DST_OPTIMAL,
1615 &vk::ClearColorValue { float32: [0.0; 4] },
1616 &[color_range],
1617 );
1618 }
1619
1620 // Blit (not copy): converts UNORM → the backdrop's sRGB format.
1621 let (px, py, _, _) = self.pane;
1622 let dst_x0 = px.min(backdrop_extent.width);
1623 let dst_y0 = py.min(backdrop_extent.height);
1624 let bw = w.min(backdrop_extent.width - dst_x0);
1625 let bh = h.min(backdrop_extent.height - dst_y0);
1626 if bw > 0 && bh > 0 {
1627 let layers = vk::ImageSubresourceLayers::default()
1628 .aspect_mask(vk::ImageAspectFlags::COLOR)
1629 .layer_count(1);
1630 device.cmd_blit_image(
1631 cmd,
1632 self.output_image,
1633 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
1634 backdrop_image,
1635 vk::ImageLayout::TRANSFER_DST_OPTIMAL,
1636 &[vk::ImageBlit::default()
1637 .src_subresource(layers)
1638 .src_offsets([
1639 vk::Offset3D { x: 0, y: 0, z: 0 },
1640 vk::Offset3D { x: bw as i32, y: bh as i32, z: 1 },
1641 ])
1642 .dst_subresource(layers)
1643 .dst_offsets([
1644 vk::Offset3D { x: dst_x0 as i32, y: dst_y0 as i32, z: 0 },
1645 vk::Offset3D {
1646 x: (dst_x0 + bw) as i32,
1647 y: (dst_y0 + bh) as i32,
1648 z: 1,
1649 },
1650 ])],
1651 vk::Filter::NEAREST,
1652 );
1653 }
1654
1655 // Backdrop to TRANSFER_SRC: the swapchain copy path expects it
1656 // exactly as SceneStage::record leaves it.
1657 device.cmd_pipeline_barrier(
1658 cmd,
1659 vk::PipelineStageFlags::TRANSFER,
1660 vk::PipelineStageFlags::TRANSFER,
1661 vk::DependencyFlags::empty(),
1662 &[],
1663 &[],
1664 &[vk::ImageMemoryBarrier::default()
1665 .src_access_mask(vk::AccessFlags::TRANSFER_WRITE)
1666 .dst_access_mask(vk::AccessFlags::TRANSFER_READ)
1667 .old_layout(vk::ImageLayout::TRANSFER_DST_OPTIMAL)
1668 .new_layout(vk::ImageLayout::TRANSFER_SRC_OPTIMAL)
1669 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1670 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1671 .image(backdrop_image)
1672 .subresource_range(color_range)],
1673 );
1674 }
1675 self.sample_index += self.spp;
1676 true
1677 }
1678
1679 pub(crate) fn destroy(&mut self, device: &ash::Device, allocator: &mut Allocator) {
1680 self.destroy_targets(device, allocator);
1681 self.destroy_accel(device, allocator);
1682 if let Some(mut denoiser) = self.denoiser.take() {
1683 denoiser.destroy(device, allocator);
1684 }
1685 for buf in [&mut self.nodes, &mut self.tris, &mut self.materials] {
1686 destroy_cpu_buffer(device, allocator, buf);
1687 }
1688 unsafe {
1689 for frame in &mut self.frames {
1690 let mut uniforms = std::mem::replace(&mut frame.uniforms, AllocatedBuffer::null());
1691 destroy_cpu_buffer(device, allocator, &mut uniforms);
1692 }
1693 device.destroy_descriptor_pool(self.descriptor_pool, None);
1694 device.destroy_descriptor_set_layout(self.descriptor_set_layout, None);
1695 device.destroy_pipeline(self.pipeline, None);
1696 device.destroy_pipeline_layout(self.pipeline_layout, None);
1697 device.destroy_shader_module(self.shader_module, None);
1698 }
1699 }
1700 }
1701
1702 // --- Headless offscreen rendering (thumbnails, previews) ---
1703
1704 /// One-shot path-traced rendering with no window anywhere: a headless
1705 /// [`super::VkCore`] + an [`RtStage`] whose "backdrop" is a private sRGB
1706 /// target image, read back to CPU pixels. This is the seam consumers like
1707 /// the cce-files thumbnailer sit on.
1708 ///
1709 /// Not `Send`-safe by design intent (owns a device); create it on the worker
1710 /// thread that renders.
1711 pub struct RtOffscreen {
1712 stage: RtStage,
1713 target_image: vk::Image,
1714 target_allocation: Option<Allocation>,
1715 readback: AllocatedBuffer,
1716 size: (u32, u32),
1717 cmd: vk::CommandBuffer,
1718 fence: vk::Fence,
1719 // Declared last: dropped after everything above is destroyed in Drop.
1720 core: super::VkCore,
1721 }
1722
1723 impl RtOffscreen {
1724 /// Samples per submit: keeps each dispatch well under GPU watchdog
1725 /// timeouts even at large sizes; a render loops submits to reach the
1726 /// requested sample count.
1727 const CHUNK_SPP: u32 = 8;
1728
1729 pub fn new() -> Self {
1730 let mut core = super::VkCore::new_headless();
1731 let device = core.device.clone();
1732 let accel_loader = core.accel_loader.clone();
1733 let as_scratch_align = core.as_scratch_align;
1734 let min_uniform_align = core.min_uniform_align;
1735 let allocator = core.allocator.as_mut().unwrap();
1736 let stage = RtStage::new(
1737 &device,
1738 allocator,
1739 1,
1740 accel_loader.as_ref(),
1741 as_scratch_align,
1742 min_uniform_align,
1743 );
1744 unsafe {
1745 let cmd = device
1746 .allocate_command_buffers(
1747 &vk::CommandBufferAllocateInfo::default()
1748 .command_pool(core.command_pool)
1749 .level(vk::CommandBufferLevel::PRIMARY)
1750 .command_buffer_count(1),
1751 )
1752 .expect("Failed to allocate RT offscreen command buffer")[0];
1753 let fence = device
1754 .create_fence(&vk::FenceCreateInfo::default(), None)
1755 .expect("Failed to create RT offscreen fence");
1756 RtOffscreen {
1757 stage,
1758 target_image: vk::Image::null(),
1759 target_allocation: None,
1760 readback: AllocatedBuffer::null(),
1761 size: (0, 0),
1762 cmd,
1763 fence,
1764 core,
1765 }
1766 }
1767 }
1768
1769 /// Replace the scene (same schema as `VkRenderer::set_rt_scene`).
1770 pub fn set_scene(&mut self, triangles: &[RtTriangle], materials: &[RtMaterial]) {
1771 unsafe {
1772 let _ = self.core.device.device_wait_idle();
1773 }
1774 let device = self.core.device.clone();
1775 let queue = self.core.queue;
1776 let command_pool = self.core.command_pool;
1777 self.stage.set_scene(
1778 &device,
1779 self.core.allocator.as_mut().unwrap(),
1780 queue,
1781 command_pool,
1782 triangles,
1783 materials,
1784 );
1785 }
1786
1787 /// Render `samples` paths per pixel and return tightly packed
1788 /// sRGB-encoded RGBA8 pixels (`width * height * 4` bytes). Blocks until
1789 /// the GPU finishes; meant for worker threads, not frame loops.
1790 pub fn render(
1791 &mut self,
1792 camera: RtCamera,
1793 width: u32,
1794 height: u32,
1795 samples: u32,
1796 ) -> Vec<u8> {
1797 let width = width.max(1);
1798 let height = height.max(1);
1799 let samples = samples.clamp(1, MAX_SAMPLES);
1800 let device = self.core.device.clone();
1801 self.ensure_target(width, height);
1802
1803 // Fresh accumulation every render: thumbnails are one-shot.
1804 self.stage.sample_index = 0;
1805 let extent = vk::Extent2D { width, height };
1806 let mut done = 0u32;
1807 while done < samples {
1808 self.stage.spp = Self::CHUNK_SPP.min(samples - done);
1809 self.stage.stage(
1810 &device,
1811 self.core.allocator.as_mut().unwrap(),
1812 (0, 0, width, height),
1813 camera,
1814 );
1815 // stage() resets sample_index when the camera or size changed —
1816 // keep our resume point, not the reset, after the first chunk.
1817 self.stage.sample_index = done;
1818 self.stage.write_frame_uniforms(0);
1819 unsafe {
1820 device
1821 .begin_command_buffer(self.cmd, &vk::CommandBufferBeginInfo::default())
1822 .unwrap();
1823 let recorded =
1824 self.stage.record(&device, self.cmd, 0, self.target_image, extent, true);
1825 device.end_command_buffer(self.cmd).unwrap();
1826 assert!(recorded, "RT offscreen: nothing recorded (empty scene?)");
1827 self.submit_and_wait();
1828 }
1829 done += Self::CHUNK_SPP.min(samples - done);
1830 }
1831
1832 // Copy the sRGB target (left in TRANSFER_SRC by record) to the
1833 // readback buffer and map it.
1834 unsafe {
1835 device
1836 .begin_command_buffer(self.cmd, &vk::CommandBufferBeginInfo::default())
1837 .unwrap();
1838 device.cmd_copy_image_to_buffer(
1839 self.cmd,
1840 self.target_image,
1841 vk::ImageLayout::TRANSFER_SRC_OPTIMAL,
1842 self.readback.buffer,
1843 &[vk::BufferImageCopy::default()
1844 .image_subresource(
1845 vk::ImageSubresourceLayers::default()
1846 .aspect_mask(vk::ImageAspectFlags::COLOR)
1847 .layer_count(1),
1848 )
1849 .image_extent(vk::Extent3D { width, height, depth: 1 })],
1850 );
1851 device.cmd_pipeline_barrier(
1852 self.cmd,
1853 vk::PipelineStageFlags::TRANSFER,
1854 vk::PipelineStageFlags::HOST,
1855 vk::DependencyFlags::empty(),
1856 &[],
1857 &[vk::BufferMemoryBarrier::default()
1858 .src_access_mask(vk::AccessFlags::TRANSFER_WRITE)
1859 .dst_access_mask(vk::AccessFlags::HOST_READ)
1860 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1861 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1862 .buffer(self.readback.buffer)
1863 .size(vk::WHOLE_SIZE)],
1864 &[],
1865 );
1866 device.end_command_buffer(self.cmd).unwrap();
1867 self.submit_and_wait();
1868 }
1869 let len = (width * height * 4) as usize;
1870 self.readback.allocation.as_ref().unwrap().mapped_slice().unwrap()[..len].to_vec()
1871 }
1872
1873 unsafe fn submit_and_wait(&mut self) {
1874 let device = &self.core.device;
1875 let cmds = [self.cmd];
1876 device
1877 .queue_submit(
1878 self.core.queue,
1879 &[vk::SubmitInfo::default().command_buffers(&cmds)],
1880 self.fence,
1881 )
1882 .expect("RT offscreen submit failed");
1883 device
1884 .wait_for_fences(&[self.fence], true, u64::MAX)
1885 .expect("RT offscreen fence wait failed");
1886 device.reset_fences(&[self.fence]).unwrap();
1887 }
1888
1889 fn ensure_target(&mut self, width: u32, height: u32) {
1890 if (width, height) == self.size {
1891 return;
1892 }
1893 let device = self.core.device.clone();
1894 unsafe {
1895 let _ = device.device_wait_idle();
1896 }
1897 self.destroy_target();
1898 let allocator = self.core.allocator.as_mut().unwrap();
1899 unsafe {
1900 let image = device
1901 .create_image(
1902 &vk::ImageCreateInfo::default()
1903 .image_type(vk::ImageType::TYPE_2D)
1904 .format(vk::Format::R8G8B8A8_SRGB)
1905 .extent(vk::Extent3D { width, height, depth: 1 })
1906 .mip_levels(1)
1907 .array_layers(1)
1908 .samples(vk::SampleCountFlags::TYPE_1)
1909 .tiling(vk::ImageTiling::OPTIMAL)
1910 .usage(
1911 vk::ImageUsageFlags::TRANSFER_DST | vk::ImageUsageFlags::TRANSFER_SRC,
1912 )
1913 .initial_layout(vk::ImageLayout::UNDEFINED),
1914 None,
1915 )
1916 .expect("Failed to create RT offscreen target");
1917 let requirements = device.get_image_memory_requirements(image);
1918 let allocation = allocator
1919 .allocate(&AllocationCreateDesc {
1920 name: "rt-offscreen-target",
1921 requirements,
1922 location: MemoryLocation::GpuOnly,
1923 linear: false,
1924 allocation_scheme: AllocationScheme::GpuAllocatorManaged,
1925 })
1926 .expect("Failed to allocate RT offscreen target memory");
1927 device
1928 .bind_image_memory(image, allocation.memory(), allocation.offset())
1929 .expect("Failed to bind RT offscreen target memory");
1930 self.target_image = image;
1931 self.target_allocation = Some(allocation);
1932
1933 self.readback = create_cpu_buffer(
1934 &device,
1935 allocator,
1936 (width as vk::DeviceSize) * (height as vk::DeviceSize) * 4,
1937 vk::BufferUsageFlags::TRANSFER_DST,
1938 "rt-readback",
1939 );
1940
1941 // RtStage::record expects the blit destination in TRANSFER_SRC
1942 // (the steady state SceneStage leaves the backdrop in).
1943 device
1944 .begin_command_buffer(self.cmd, &vk::CommandBufferBeginInfo::default())
1945 .unwrap();
1946 device.cmd_pipeline_barrier(
1947 self.cmd,
1948 vk::PipelineStageFlags::TOP_OF_PIPE,
1949 vk::PipelineStageFlags::TRANSFER,
1950 vk::DependencyFlags::empty(),
1951 &[],
1952 &[],
1953 &[vk::ImageMemoryBarrier::default()
1954 .src_access_mask(vk::AccessFlags::empty())
1955 .dst_access_mask(vk::AccessFlags::TRANSFER_READ)
1956 .old_layout(vk::ImageLayout::UNDEFINED)
1957 .new_layout(vk::ImageLayout::TRANSFER_SRC_OPTIMAL)
1958 .src_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1959 .dst_queue_family_index(vk::QUEUE_FAMILY_IGNORED)
1960 .image(image)
1961 .subresource_range(
1962 vk::ImageSubresourceRange::default()
1963 .aspect_mask(vk::ImageAspectFlags::COLOR)
1964 .level_count(1)
1965 .layer_count(1),
1966 )],
1967 );
1968 device.end_command_buffer(self.cmd).unwrap();
1969 self.submit_and_wait();
1970 }
1971 self.size = (width, height);
1972 }
1973
1974 fn destroy_target(&mut self) {
1975 unsafe {
1976 if self.target_image != vk::Image::null() {
1977 self.core.device.destroy_image(self.target_image, None);
1978 self.target_image = vk::Image::null();
1979 }
1980 }
1981 if let Some(a) = self.target_allocation.take() {
1982 let _ = self.core.allocator.as_mut().unwrap().free(a);
1983 }
1984 let device = self.core.device.clone();
1985 destroy_cpu_buffer(&device, self.core.allocator.as_mut().unwrap(), &mut self.readback);
1986 self.size = (0, 0);
1987 }
1988 }
1989
1990 impl Default for RtOffscreen {
1991 fn default() -> Self {
1992 Self::new()
1993 }
1994 }
1995
1996 impl Drop for RtOffscreen {
1997 fn drop(&mut self) {
1998 unsafe {
1999 let _ = self.core.device.device_wait_idle();
2000 }
2001 self.destroy_target();
2002 let device = self.core.device.clone();
2003 self.stage.destroy(&device, self.core.allocator.as_mut().unwrap());
2004 unsafe {
2005 self.core.device.destroy_fence(self.fence, None);
2006 // The command buffer dies with the pool in VkCore's Drop.
2007 }
2008 }
2009 }
2010
2011 #[cfg(test)]
2012 mod tests {
2013 use super::*;
2014
2015 // CPU mirror of the shader's traversal, for parity testing.
2016 fn intersect_tri_cpu(ro: [f32; 3], rd: [f32; 3], t: &RtTriangle, t_limit: f32) -> f32 {
2017 let sub = |a: [f32; 3], b: [f32; 3]| [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
2018 let cross = |a: [f32; 3], b: [f32; 3]| {
2019 [
2020 a[1] * b[2] - a[2] * b[1],
2021 a[2] * b[0] - a[0] * b[2],
2022 a[0] * b[1] - a[1] * b[0],
2023 ]
2024 };
2025 let dot = |a: [f32; 3], b: [f32; 3]| a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
2026 let e1 = sub(t.p1, t.p0);
2027 let e2 = sub(t.p2, t.p0);
2028 let h = cross(rd, e2);
2029 let a = dot(e1, h);
2030 if a.abs() < 1e-8 {
2031 return 1e30;
2032 }
2033 let f = 1.0 / a;
2034 let s = sub(ro, t.p0);
2035 let u = f * dot(s, h);
2036 if !(0.0..=1.0).contains(&u) {
2037 return 1e30;
2038 }
2039 let q = cross(s, e1);
2040 let v = f * dot(rd, q);
2041 if v < 0.0 || u + v > 1.0 {
2042 return 1e30;
2043 }
2044 let tt = f * dot(e2, q);
2045 if tt > 1e-4 && tt < t_limit {
2046 return tt;
2047 }
2048 1e30
2049 }
2050
2051 fn traverse_bvh_cpu(
2052 nodes: &[GpuBvhNode],
2053 tris: &[RtTriangle],
2054 ro: [f32; 3],
2055 rd: [f32; 3],
2056 ) -> (f32, Option<usize>) {
2057 if nodes.is_empty() {
2058 return (1e30, None);
2059 }
2060 let inv = [1.0 / rd[0], 1.0 / rd[1], 1.0 / rd[2]];
2061 let hit_aabb = |min: [f32; 3], max: [f32; 3], t_limit: f32| -> bool {
2062 let mut tn = f32::NEG_INFINITY;
2063 let mut tf = f32::INFINITY;
2064 for a in 0..3 {
2065 let t1 = (min[a] - ro[a]) * inv[a];
2066 let t2 = (max[a] - ro[a]) * inv[a];
2067 tn = tn.max(t1.min(t2));
2068 tf = tf.min(t1.max(t2));
2069 }
2070 tf >= tn.max(0.0) && tn < t_limit
2071 };
2072 let mut best = 1e30f32;
2073 let mut best_tri = None;
2074 let mut stack = vec![0u32];
2075 while let Some(idx) = stack.pop() {
2076 let node = &nodes[idx as usize];
2077 if !hit_aabb(node.min, node.max, best) {
2078 continue;
2079 }
2080 if node.count > 0 {
2081 for i in node.left_first..node.left_first + node.count {
2082 let t = intersect_tri_cpu(ro, rd, &tris[i as usize], best);
2083 if t < best {
2084 best = t;
2085 best_tri = Some(i as usize);
2086 }
2087 }
2088 } else {
2089 stack.push(node.left_first);
2090 stack.push(node.left_first + 1);
2091 }
2092 }
2093 (best, best_tri)
2094 }
2095
2096 fn brute_force(tris: &[RtTriangle], ro: [f32; 3], rd: [f32; 3]) -> (f32, Option<usize>) {
2097 let mut best = 1e30f32;
2098 let mut best_tri = None;
2099 for (i, t) in tris.iter().enumerate() {
2100 let tt = intersect_tri_cpu(ro, rd, t, best);
2101 if tt < best {
2102 best = tt;
2103 best_tri = Some(i);
2104 }
2105 }
2106 (best, best_tri)
2107 }
2108
2109 // Deterministic LCG so the test needs no rand dependency.
2110 struct Lcg(u64);
2111 impl Lcg {
2112 fn next_f32(&mut self) -> f32 {
2113 self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
2114 ((self.0 >> 33) as f32) / (u32::MAX >> 1) as f32
2115 }
2116 fn point(&mut self, scale: f32) -> [f32; 3] {
2117 [
2118 (self.next_f32() - 0.5) * scale,
2119 (self.next_f32() - 0.5) * scale,
2120 (self.next_f32() - 0.5) * scale,
2121 ]
2122 }
2123 }
2124
2125 fn random_scene(n: usize, seed: u64) -> Vec<RtTriangle> {
2126 let mut rng = Lcg(seed);
2127 (0..n)
2128 .map(|i| {
2129 let c = rng.point(20.0);
2130 let jitter = |rng: &mut Lcg, c: [f32; 3]| {
2131 let d = rng.point(2.0);
2132 [c[0] + d[0], c[1] + d[1], c[2] + d[2]]
2133 };
2134 RtTriangle {
2135 p0: jitter(&mut rng, c),
2136 p1: jitter(&mut rng, c),
2137 p2: jitter(&mut rng, c),
2138 material: (i % 5) as u32,
2139 }
2140 })
2141 .collect()
2142 }
2143
2144 #[test]
2145 fn test_bvh_matches_brute_force() {
2146 let mut tris = random_scene(500, 42);
2147 let nodes = build_bvh(&mut tris);
2148 assert!(!nodes.is_empty());
2149 let mut rng = Lcg(7);
2150 let mut hits = 0;
2151 for _ in 0..200 {
2152 let ro = rng.point(40.0);
2153 let target = rng.point(10.0);
2154 let d = [target[0] - ro[0], target[1] - ro[1], target[2] - ro[2]];
2155 let len = (d[0] * d[0] + d[1] * d[1] + d[2] * d[2]).sqrt().max(1e-6);
2156 let rd = [d[0] / len, d[1] / len, d[2] / len];
2157 let (t_bvh, tri_bvh) = traverse_bvh_cpu(&nodes, &tris, ro, rd);
2158 let (t_ref, tri_ref) = brute_force(&tris, ro, rd);
2159 assert_eq!(tri_bvh, tri_ref, "different triangle hit");
2160 assert!((t_bvh - t_ref).abs() < 1e-4, "t mismatch: {t_bvh} vs {t_ref}");
2161 if tri_bvh.is_some() {
2162 hits += 1;
2163 }
2164 }
2165 assert!(hits > 20, "test rays barely hit the scene ({hits}/200)");
2166 }
2167
2168 #[test]
2169 fn test_bvh_leaf_ranges_cover_all_triangles() {
2170 let mut tris = random_scene(300, 9);
2171 let nodes = build_bvh(&mut tris);
2172 let mut seen = vec![false; tris.len()];
2173 for node in &nodes {
2174 if node.count > 0 {
2175 for i in node.left_first..node.left_first + node.count {
2176 assert!(!seen[i as usize], "triangle {i} in two leaves");
2177 seen[i as usize] = true;
2178 }
2179 }
2180 }
2181 assert!(seen.iter().all(|&s| s), "not every triangle is in a leaf");
2182 }
2183
2184 #[test]
2185 fn test_bvh_degenerate_identical_centroids() {
2186 // All triangles share one centroid: SAH can't split, the median
2187 // fallback must still terminate and cover everything.
2188 let tri = RtTriangle {
2189 p0: [0.0, 0.0, 0.0],
2190 p1: [1.0, 0.0, 0.0],
2191 p2: [0.0, 1.0, 0.0],
2192 material: 0,
2193 };
2194 let mut tris = vec![tri; 100];
2195 let nodes = build_bvh(&mut tris);
2196 let covered: u32 = nodes.iter().filter(|n| n.count > 0).map(|n| n.count).sum();
2197 assert_eq!(covered, 100);
2198 let (t, hit) = traverse_bvh_cpu(&nodes, &tris, [0.2, 0.2, -5.0], [0.0, 0.0, 1.0]);
2199 assert!(hit.is_some());
2200 assert!((t - 5.0).abs() < 1e-3);
2201 }
2202
2203 #[test]
2204 fn test_bvh_empty_and_single() {
2205 let mut empty: Vec<RtTriangle> = Vec::new();
2206 assert!(build_bvh(&mut empty).is_empty());
2207
2208 let mut single = vec![RtTriangle {
2209 p0: [-1.0, -1.0, 0.0],
2210 p1: [1.0, -1.0, 0.0],
2211 p2: [0.0, 1.0, 0.0],
2212 material: 3,
2213 }];
2214 let nodes = build_bvh(&mut single);
2215 assert_eq!(nodes.len(), 1);
2216 assert_eq!(nodes[0].count, 1);
2217 let (t, hit) = traverse_bvh_cpu(&nodes, &single, [0.0, 0.0, -3.0], [0.0, 0.0, 1.0]);
2218 assert_eq!(hit, Some(0));
2219 assert!((t - 3.0).abs() < 1e-4);
2220 }
2221
2222 #[test]
2223 fn test_rt_shaders_compile() {
2224 // naga parse + validate + SPIR-V write for both tiers; panics on failure.
2225 let tier1 = compile_wgsl(&format!(
2226 "{}\n{}",
2227 include_str!("rt_common.wgsl"),
2228 include_str!("rt_bvh.wgsl")
2229 ));
2230 assert!(!tier1.is_empty());
2231 let tier2 = compile_wgsl_ray_query(&format!(
2232 "{}\n{}",
2233 include_str!("rt_common.wgsl"),
2234 include_str!("rt_query.wgsl")
2235 ));
2236 assert!(!tier2.is_empty());
2237 let denoise = compile_wgsl(include_str!("rt_denoise.wgsl"));
2238 assert!(!denoise.is_empty());
2239 }
2240
2241 /// End-to-end GPU test — needs a Vulkan device, so ignored by default.
2242 /// Run with: cargo test --lib vk::rt -- --ignored
2243 #[test]
2244 #[ignore = "requires a Vulkan device"]
2245 fn test_offscreen_render_smoke() {
2246 let mut off = RtOffscreen::new();
2247 // A red triangle filling the view center, camera looking down -Z.
2248 off.set_scene(
2249 &[RtTriangle {
2250 p0: [-1.0, -1.0, 0.0],
2251 p1: [1.0, -1.0, 0.0],
2252 p2: [0.0, 1.5, 0.0],
2253 material: 0,
2254 }],
2255 &[RtMaterial { albedo: [0.9, 0.1, 0.1], emission: [0.0; 3] }],
2256 );
2257 let proj = glam::Mat4::perspective_rh(0.9, 1.0, 0.1, 100.0);
2258 let view = glam::Mat4::look_at_rh(
2259 glam::Vec3::new(0.0, 0.0, 3.0),
2260 glam::Vec3::ZERO,
2261 glam::Vec3::Y,
2262 );
2263 let camera = RtCamera { inv_mvp: (proj * view).inverse().to_cols_array_2d() };
2264 let (w, h) = (64u32, 64u32);
2265 let px = off.render(camera, w, h, 16);
2266 assert_eq!(px.len(), (w * h * 4) as usize);
2267 // Center pixel hits the triangle: red-dominant. Corner pixel is sky:
2268 // blue >= red. Alpha opaque everywhere.
2269 let at = |x: u32, y: u32| {
2270 let i = ((y * w + x) * 4) as usize;
2271 (px[i], px[i + 1], px[i + 2], px[i + 3])
2272 };
2273 let (cr, _cg, cb, ca) = at(w / 2, h / 2);
2274 assert!(ca == 255, "alpha not opaque: {ca}");
2275 assert!(cr > cb, "center not red-dominant: r={cr} b={cb}");
2276 let (sr, _sg, sb, _sa) = at(1, 1);
2277 assert!(sb >= sr, "corner sky not blue-ish: r={sr} b={sb}");
2278 }
2279 }