1use alloc::{sync::Arc, vec::Vec};
2use core::{
3 cmp::max,
4 num::NonZeroU64,
5 ops::{Deref, Range},
6};
7
8use wgt::{math::align_to, BufferUsages, BufferUses, Features};
9
10use crate::{
11 command::encoder::EncodingState,
12 ray_tracing::{AsAction, AsBuild, TlasBuild, ValidateAsActionsError},
13 resource::InvalidResourceError,
14};
15use crate::{command::EncoderStateError, device::resource::CommandIndices};
16use crate::{
17 command::{ArcCommand, ArcReferences, CommandBufferMutable},
18 device::queue::TempResource,
19 global::Global,
20 id::CommandEncoderId,
21 init_tracker::MemoryInitKind,
22 ray_tracing::{
23 ArcBlasAabbGeometry, ArcBlasBuildEntry, ArcBlasGeometries, ArcBlasTriangleGeometry,
24 ArcTlasInstance, ArcTlasPackage, BlasBuildEntry, BlasGeometries,
25 BuildAccelerationStructureError, OwnedBlasBuildEntry, OwnedTlasPackage, TlasPackage,
26 },
27 resource::{Blas, BlasCompactState, Labeled, StagingBuffer, Tlas},
28 scratch::ScratchBuffer,
29 snatch::SnatchGuard,
30};
31use crate::{lock::RwLockWriteGuard, resource::RawResourceAccess};
32
33use crate::id::{BlasId, TlasId};
34
35struct BlasStore<'a> {
36 blas: Arc<Blas>,
37 entries: hal::AccelerationStructureEntries<'a, dyn hal::DynBuffer>,
38 scratch_buffer_offset: u64,
39}
40
41struct UnsafeTlasStore<'a> {
42 tlas: Arc<Tlas>,
43 entries: hal::AccelerationStructureEntries<'a, dyn hal::DynBuffer>,
44 scratch_buffer_offset: u64,
45}
46
47struct TlasStore<'a> {
48 internal: UnsafeTlasStore<'a>,
49 range: Range<usize>,
50}
51
52impl Global {
53 fn resolve_blas_id(&self, blas_id: BlasId) -> Result<Arc<Blas>, InvalidResourceError> {
54 self.hub.blas_s.get(blas_id).get()
55 }
56
57 fn resolve_tlas_id(&self, tlas_id: TlasId) -> Result<Arc<Tlas>, InvalidResourceError> {
58 self.hub.tlas_s.get(tlas_id).get()
59 }
60
61 pub fn command_encoder_mark_acceleration_structures_built(
62 &self,
63 command_encoder_id: CommandEncoderId,
64 blas_ids: &[BlasId],
65 tlas_ids: &[TlasId],
66 ) -> Result<(), EncoderStateError> {
67 profiling::scope!("CommandEncoder::mark_acceleration_structures_built");
68
69 let hub = &self.hub;
70
71 let cmd_enc = hub.command_encoders.get(command_encoder_id);
72
73 let mut cmd_buf_data = cmd_enc.data.lock();
74 cmd_buf_data.with_buffer(
75 crate::command::EncodingApi::Raw,
76 |cmd_buf_data| -> Result<(), BuildAccelerationStructureError> {
77 let device = &cmd_enc.device;
78 device.check_is_valid()?;
79 device.require_features(Features::EXPERIMENTAL_RAY_QUERY)?;
80
81 let mut build_command = AsBuild::with_capacity(blas_ids.len(), tlas_ids.len());
82
83 for blas in blas_ids {
84 let blas = hub.blas_s.get(*blas).get()?;
85 build_command.blas_s_built.push(blas);
86 }
87
88 for tlas in tlas_ids {
89 let tlas = hub.tlas_s.get(*tlas).get()?;
90 build_command.tlas_s_built.push(TlasBuild {
91 tlas,
92 dependencies: Vec::new(),
93 });
94 }
95
96 cmd_buf_data.as_actions.push(AsAction::Build(build_command));
97 Ok(())
98 },
99 )
100 }
101
102 pub fn command_encoder_build_acceleration_structures<'a>(
103 &self,
104 command_encoder_id: CommandEncoderId,
105 blas_iter: impl Iterator<Item = BlasBuildEntry<'a>>,
106 tlas_iter: impl Iterator<Item = TlasPackage<'a>>,
107 ) -> Result<(), EncoderStateError> {
108 profiling::scope!("CommandEncoder::build_acceleration_structures");
109
110 let hub = &self.hub;
111
112 let cmd_enc = hub.command_encoders.get(command_encoder_id);
113 let mut cmd_buf_data = cmd_enc.data.lock();
114
115 cmd_buf_data.push_with(|| -> Result<_, BuildAccelerationStructureError> {
116 let blas = blas_iter
117 .map(|blas_entry| {
118 let geometries = match blas_entry.geometries {
119 BlasGeometries::TriangleGeometries(triangle_geometries) => {
120 let tri_geo = triangle_geometries
121 .map(|tg| {
122 Ok(ArcBlasTriangleGeometry {
123 size: tg.size.clone(),
124 vertex_buffer: self.resolve_buffer_id(tg.vertex_buffer)?,
125 index_buffer: tg
126 .index_buffer
127 .map(|id| self.resolve_buffer_id(id))
128 .transpose()?,
129 transform_buffer: tg
130 .transform_buffer
131 .map(|id| self.resolve_buffer_id(id))
132 .transpose()?,
133 first_vertex: tg.first_vertex,
134 vertex_stride: tg.vertex_stride,
135 first_index: tg.first_index,
136 transform_buffer_offset: tg.transform_buffer_offset,
137 })
138 })
139 .collect::<Result<_, BuildAccelerationStructureError>>()?;
140 ArcBlasGeometries::TriangleGeometries(tri_geo)
141 }
142 BlasGeometries::AabbGeometries(aabb_geometries) => {
143 let aabb_geo = aabb_geometries
144 .map(|ag| {
145 Ok(ArcBlasAabbGeometry {
146 size: ag.size.clone(),
147 stride: ag.stride,
148 aabb_buffer: self.resolve_buffer_id(ag.aabb_buffer)?,
149 primitive_offset: ag.primitive_offset,
150 })
151 })
152 .collect::<Result<_, BuildAccelerationStructureError>>()?;
153 ArcBlasGeometries::AabbGeometries(aabb_geo)
154 }
155 };
156 Ok(ArcBlasBuildEntry {
157 blas: self.resolve_blas_id(blas_entry.blas_id)?,
158 geometries,
159 })
160 })
161 .collect::<Result<_, BuildAccelerationStructureError>>()?;
162
163 let tlas = tlas_iter
164 .map(|tlas_package| {
165 let instances = tlas_package
166 .instances
167 .map(|instance| {
168 instance
169 .as_ref()
170 .map(|instance| {
171 Ok(ArcTlasInstance {
172 blas: self.resolve_blas_id(instance.blas_id)?,
173 transform: *instance.transform,
174 custom_data: instance.custom_data,
175 mask: instance.mask,
176 })
177 })
178 .transpose()
179 })
180 .collect::<Result<_, BuildAccelerationStructureError>>()?;
181 Ok(ArcTlasPackage {
182 tlas: self.resolve_tlas_id(tlas_package.tlas_id)?,
183 instances,
184 lowest_unmodified: tlas_package.lowest_unmodified,
185 })
186 })
187 .collect::<Result<_, BuildAccelerationStructureError>>()?;
188
189 Ok(ArcCommand::BuildAccelerationStructures { blas, tlas })
190 })
191 }
192}
193
194#[inline(never)]
195pub(crate) fn build_acceleration_structures(
196 state: &mut EncodingState,
197 blas: Vec<OwnedBlasBuildEntry<ArcReferences>>,
198 tlas: Vec<OwnedTlasPackage<ArcReferences>>,
199) -> Result<(), BuildAccelerationStructureError> {
200 state
201 .device
202 .require_features(Features::EXPERIMENTAL_RAY_QUERY)?;
203
204 let mut build_command = AsBuild::with_capacity(blas.len(), tlas.len());
205 let mut input_barriers = Vec::<hal::BufferBarrier<dyn hal::DynBuffer>>::new();
206 let mut scratch_buffer_blas_size = 0;
207 let mut blas_storage = Vec::with_capacity(blas.len());
208 iter_blas(
209 blas.iter(),
210 &mut build_command,
211 &mut input_barriers,
212 &mut scratch_buffer_blas_size,
213 &mut blas_storage,
214 state,
215 )?;
216
217 let mut scratch_buffer_tlas_size: u64 = 0;
218 let mut tlas_storage = Vec::<TlasStore>::with_capacity(tlas.len());
219 let mut instance_buffer_staging_source = Vec::<u8>::new();
220
221 for package in tlas.iter() {
222 let tlas = &package.tlas;
223 state.tracker.tlas_s.insert_single(tlas.clone());
224
225 let scratch_buffer_offset = scratch_buffer_tlas_size;
226 scratch_buffer_tlas_size = scratch_buffer_tlas_size.saturating_add(align_to(
227 tlas.size_info.build_scratch_size,
228 u64::from(state.device.alignments.ray_tracing_scratch_buffer_alignment),
229 ));
230
231 let first_byte_index = instance_buffer_staging_source.len();
232
233 let mut dependencies = Vec::new();
234
235 let mut instance_count = 0;
236 for instance in package.instances.iter().flatten() {
237 if instance.custom_data >= (1u32 << 24u32) {
238 return Err(BuildAccelerationStructureError::TlasInvalidCustomIndex(
239 tlas.error_ident(),
240 ));
241 }
242 let blas = &instance.blas;
243 state.tracker.blas_s.insert_single(blas.clone());
244
245 instance_buffer_staging_source.extend(state.device.raw().tlas_instance_to_bytes(
246 hal::TlasInstance {
247 transform: instance.transform,
248 custom_data: instance.custom_data,
249 mask: instance.mask,
250 blas_address: blas.handle,
251 pipeline_intersection_data_offset: 0,
252 },
253 ));
254
255 if tlas
256 .flags
257 .contains(wgpu_types::AccelerationStructureFlags::ALLOW_RAY_HIT_VERTEX_RETURN)
258 && !blas
259 .flags
260 .contains(wgpu_types::AccelerationStructureFlags::ALLOW_RAY_HIT_VERTEX_RETURN)
261 {
262 return Err(
263 BuildAccelerationStructureError::TlasDependentMissingVertexReturn(
264 tlas.error_ident(),
265 blas.error_ident(),
266 ),
267 );
268 }
269
270 instance_count += 1;
271
272 dependencies.push(blas.clone());
273 }
274
275 build_command.tlas_s_built.push(TlasBuild {
276 tlas: tlas.clone(),
277 dependencies,
278 });
279
280 if instance_count > tlas.max_instance_count {
281 return Err(BuildAccelerationStructureError::TlasInstanceCountExceeded(
282 tlas.error_ident(),
283 instance_count,
284 tlas.max_instance_count,
285 ));
286 }
287
288 tlas_storage.push(TlasStore {
289 internal: UnsafeTlasStore {
290 tlas: tlas.clone(),
291 entries: hal::AccelerationStructureEntries::Instances(
292 hal::AccelerationStructureInstances {
293 buffer: Some(tlas.instance_buffer.as_ref()),
294 offset: 0,
295 count: instance_count,
296 },
297 ),
298 scratch_buffer_offset,
299 },
300 range: first_byte_index..instance_buffer_staging_source.len(),
301 });
302 }
303
304 let Some(scratch_size) =
305 wgt::BufferSize::new(max(scratch_buffer_blas_size, scratch_buffer_tlas_size))
306 else {
307 return Ok(());
309 };
310
311 if scratch_size.get() == u64::MAX {
312 return Err(crate::device::DeviceError::OutOfMemory.into());
313 }
314
315 let scratch_buffer = ScratchBuffer::new(state.device, scratch_size)?;
316
317 let scratch_buffer_barrier = hal::BufferBarrier::<dyn hal::DynBuffer> {
318 buffer: scratch_buffer.raw(),
319 usage: hal::StateTransition {
320 from: BufferUses::ACCELERATION_STRUCTURE_SCRATCH,
321 to: BufferUses::ACCELERATION_STRUCTURE_SCRATCH,
322 },
323 };
324
325 let mut tlas_descriptors = Vec::with_capacity(tlas_storage.len());
326
327 for &TlasStore {
328 internal:
329 UnsafeTlasStore {
330 ref tlas,
331 ref entries,
332 ref scratch_buffer_offset,
333 },
334 ..
335 } in &tlas_storage
336 {
337 if tlas.update_mode == wgt::AccelerationStructureUpdateMode::PreferUpdate {
338 log::warn!("build_acceleration_structures called with PreferUpdate, but only rebuild is implemented");
339 }
340 tlas_descriptors.push(hal::BuildAccelerationStructureDescriptor {
341 entries,
342 mode: hal::AccelerationStructureBuildMode::Build,
343 flags: tlas.flags,
344 source_acceleration_structure: None,
345 destination_acceleration_structure: tlas.try_raw(state.snatch_guard)?,
346 scratch_buffer: scratch_buffer.raw(),
347 scratch_buffer_offset: *scratch_buffer_offset,
348 })
349 }
350
351 let blas_present = !blas_storage.is_empty();
352 let tlas_present = !tlas_storage.is_empty();
353
354 let raw_encoder = &mut state.raw_encoder;
355
356 let mut blas_s_compactable = Vec::new();
357 let mut descriptors = Vec::with_capacity(blas.len());
358
359 for storage in &blas_storage {
360 descriptors.push(map_blas(
361 storage,
362 scratch_buffer.raw(),
363 state.snatch_guard,
364 &mut blas_s_compactable,
365 )?);
366 }
367
368 build_blas(
369 *raw_encoder,
370 blas_present,
371 tlas_present,
372 input_barriers,
373 &descriptors,
374 scratch_buffer_barrier,
375 blas_s_compactable,
376 );
377
378 if tlas_present {
379 let staging_buffer = if !instance_buffer_staging_source.is_empty() {
380 let mut staging_buffer = StagingBuffer::new(
381 state.device,
382 wgt::BufferSize::new(instance_buffer_staging_source.len() as u64).unwrap(),
383 )?;
384 staging_buffer.write(&instance_buffer_staging_source);
385 let flushed = staging_buffer.flush();
386 Some(flushed)
387 } else {
388 None
389 };
390
391 unsafe {
392 if let Some(ref staging_buffer) = staging_buffer {
393 raw_encoder.transition_buffers(&[hal::BufferBarrier::<dyn hal::DynBuffer> {
394 buffer: staging_buffer.raw(),
395 usage: hal::StateTransition {
396 from: BufferUses::MAP_WRITE,
397 to: BufferUses::COPY_SRC,
398 },
399 }]);
400 }
401 }
402
403 let mut instance_buffer_barriers = Vec::new();
404 for &TlasStore {
405 internal: UnsafeTlasStore { ref tlas, .. },
406 ref range,
407 } in &tlas_storage
408 {
409 let size = match wgt::BufferSize::new((range.end - range.start) as u64) {
410 None => continue,
411 Some(size) => size,
412 };
413 instance_buffer_barriers.push(hal::BufferBarrier::<dyn hal::DynBuffer> {
414 buffer: tlas.instance_buffer.as_ref(),
415 usage: hal::StateTransition {
416 from: BufferUses::COPY_DST,
417 to: BufferUses::TOP_LEVEL_ACCELERATION_STRUCTURE_INPUT,
418 },
419 });
420 unsafe {
421 raw_encoder.transition_buffers(&[hal::BufferBarrier::<dyn hal::DynBuffer> {
422 buffer: tlas.instance_buffer.as_ref(),
423 usage: hal::StateTransition {
424 from: BufferUses::TOP_LEVEL_ACCELERATION_STRUCTURE_INPUT,
425 to: BufferUses::COPY_DST,
426 },
427 }]);
428 let temp = hal::BufferCopy {
429 src_offset: range.start as u64,
430 dst_offset: 0,
431 size,
432 };
433 raw_encoder.copy_buffer_to_buffer(
434 staging_buffer.as_ref().unwrap().raw(),
435 tlas.instance_buffer.as_ref(),
436 &[temp],
437 );
438 }
439 }
440
441 unsafe {
442 raw_encoder.transition_buffers(&instance_buffer_barriers);
443
444 raw_encoder.build_acceleration_structures(&tlas_descriptors);
445
446 raw_encoder.place_acceleration_structure_barrier(hal::AccelerationStructureBarrier {
447 usage: hal::StateTransition {
448 from: hal::AccelerationStructureUses::BUILD_OUTPUT,
449 to: hal::AccelerationStructureUses::SHADER_INPUT,
450 },
451 });
452 }
453
454 if let Some(staging_buffer) = staging_buffer {
455 state
456 .temp_resources
457 .push(TempResource::StagingBuffer(staging_buffer));
458 }
459 }
460
461 state
462 .temp_resources
463 .push(TempResource::ScratchBuffer(scratch_buffer));
464
465 state.as_actions.push(AsAction::Build(build_command));
466
467 Ok(())
468}
469
470impl CommandBufferMutable {
471 pub(crate) fn validate_acceleration_structure_actions(
472 &self,
473 snatch_guard: &SnatchGuard,
474 command_index_guard: &mut RwLockWriteGuard<CommandIndices>,
475 ) -> Result<(), ValidateAsActionsError> {
476 profiling::scope!("CommandEncoder::[submission]::validate_as_actions");
477 for action in &self.as_actions {
478 match action {
479 AsAction::Build(build) => {
480 let build_command_index = NonZeroU64::new(
481 command_index_guard.next_acceleration_structure_build_command_index,
482 )
483 .unwrap();
484
485 command_index_guard.next_acceleration_structure_build_command_index += 1;
486 for blas in build.blas_s_built.iter() {
487 let mut state_lock = blas.compacted_state.lock();
488 *state_lock = match *state_lock {
489 BlasCompactState::Compacted => {
490 unreachable!("Should be validated out in build.")
491 }
492 _ => BlasCompactState::Idle,
495 };
496 *blas.built_index.write() = Some(build_command_index);
497 }
498
499 for tlas_build in build.tlas_s_built.iter() {
500 for blas in &tlas_build.dependencies {
501 if blas.built_index.read().is_none() {
502 return Err(ValidateAsActionsError::UsedUnbuiltBlas(
503 blas.error_ident(),
504 tlas_build.tlas.error_ident(),
505 ));
506 }
507 }
508 *tlas_build.tlas.built_index.write() = Some(build_command_index);
509 tlas_build
510 .tlas
511 .dependencies
512 .write()
513 .clone_from(&tlas_build.dependencies)
514 }
515 }
516 AsAction::UseTlas(tlas) => {
517 let tlas_build_index = tlas.built_index.read();
518 let dependencies = tlas.dependencies.read();
519
520 if (*tlas_build_index).is_none() {
521 return Err(ValidateAsActionsError::UsedUnbuiltTlas(tlas.error_ident()));
522 }
523 for blas in dependencies.deref() {
524 let blas_build_index = *blas.built_index.read();
525 if blas_build_index.is_none() {
526 return Err(ValidateAsActionsError::UsedUnbuiltBlas(
527 tlas.error_ident(),
528 blas.error_ident(),
529 ));
530 }
531 if blas_build_index.unwrap() > tlas_build_index.unwrap() {
532 return Err(ValidateAsActionsError::BlasNewerThenTlas(
533 blas.error_ident(),
534 tlas.error_ident(),
535 ));
536 }
537 blas.try_raw(snatch_guard)?;
538 }
539 }
540 }
541 }
542 Ok(())
543 }
544
545 pub(crate) fn set_acceleration_structure_dependencies(&self, snatch_guard: &SnatchGuard) {
546 profiling::scope!("CommandEncoder::[submission]::set_acceleration_structure_dependencies");
547 let tlas_dependencies_locks: Vec<_> = self
548 .as_actions
549 .iter()
550 .filter_map(|action| {
551 if let AsAction::UseTlas(tlas) = action {
552 Some(tlas.dependencies.read())
553 } else {
554 None
555 }
556 })
557 .collect();
558 let mut tlas_dependencies_lock_iter = tlas_dependencies_locks.iter();
559 let mut dependencies = Vec::new();
560 for action in &self.as_actions {
561 match action {
562 AsAction::Build(build) => {
563 for tlas_build in build.tlas_s_built.iter() {
564 for dependency in &tlas_build.dependencies {
565 if let Some(dependency) = dependency.raw(snatch_guard) {
566 dependencies.push(dependency);
567 }
568 }
569 }
570 }
571 AsAction::UseTlas(_tlas) => {
572 let tlas_dependencies = tlas_dependencies_lock_iter.next().unwrap(); for dependency in tlas_dependencies.iter() {
574 if let Some(dependency) = dependency.raw(snatch_guard) {
575 dependencies.push(dependency);
576 }
577 }
578 }
579 }
580 }
581 if !dependencies.is_empty() {
582 unsafe {
583 self.encoder
584 .raw
585 .set_acceleration_structure_dependencies(&self.encoder.list, &dependencies);
586 }
587 }
588 }
589}
590
591fn iter_blas<'snatch_guard: 'buffers, 'buffers>(
593 blas_iter: impl Iterator<Item = &'buffers OwnedBlasBuildEntry<ArcReferences>>,
594 build_command: &mut AsBuild,
595 input_barriers: &mut Vec<hal::BufferBarrier<'buffers, dyn hal::DynBuffer>>,
596 scratch_buffer_blas_size: &mut u64,
597 blas_storage: &mut Vec<BlasStore<'buffers>>,
598 state: &mut EncodingState<'snatch_guard, '_>,
599) -> Result<(), BuildAccelerationStructureError> {
600 for entry in blas_iter {
601 let blas = &entry.blas;
602 state.tracker.blas_s.insert_single(blas.clone());
603
604 build_command.blas_s_built.push(blas.clone());
605
606 match &entry.geometries {
607 ArcBlasGeometries::TriangleGeometries(triangle_geometries) => {
608 let mut triangle_entries =
609 Vec::<hal::AccelerationStructureTriangles<dyn hal::DynBuffer>>::new();
610
611 for (i, mesh) in triangle_geometries.iter().enumerate() {
612 let size_desc = match &blas.sizes {
613 wgt::BlasGeometrySizeDescriptors::Triangles { descriptors } => descriptors,
614 _ => {
615 return Err(BuildAccelerationStructureError::BlasGeometryKindMismatch(
616 blas.error_ident(),
617 ));
618 }
619 };
620 if i >= size_desc.len() {
621 return Err(BuildAccelerationStructureError::IncompatibleBlasBuildSizes(
622 blas.error_ident(),
623 ));
624 }
625 let size_desc = &size_desc[i];
626
627 if size_desc.flags != mesh.size.flags {
628 return Err(BuildAccelerationStructureError::IncompatibleBlasFlags(
629 blas.error_ident(),
630 size_desc.flags,
631 mesh.size.flags,
632 ));
633 }
634
635 if size_desc.vertex_count < mesh.size.vertex_count {
636 return Err(
637 BuildAccelerationStructureError::IncompatibleBlasVertexCount(
638 blas.error_ident(),
639 size_desc.vertex_count,
640 mesh.size.vertex_count,
641 ),
642 );
643 }
644
645 if size_desc.vertex_format != mesh.size.vertex_format {
646 return Err(BuildAccelerationStructureError::DifferentBlasVertexFormats(
647 blas.error_ident(),
648 size_desc.vertex_format,
649 mesh.size.vertex_format,
650 ));
651 }
652
653 if size_desc
654 .vertex_format
655 .min_acceleration_structure_vertex_stride()
656 > mesh.vertex_stride
657 {
658 return Err(BuildAccelerationStructureError::VertexStrideTooSmall(
659 blas.error_ident(),
660 size_desc
661 .vertex_format
662 .min_acceleration_structure_vertex_stride(),
663 mesh.vertex_stride,
664 ));
665 }
666
667 if mesh.vertex_stride
668 % size_desc
669 .vertex_format
670 .acceleration_structure_stride_alignment()
671 != 0
672 {
673 return Err(BuildAccelerationStructureError::VertexStrideUnaligned(
674 blas.error_ident(),
675 size_desc
676 .vertex_format
677 .acceleration_structure_stride_alignment(),
678 mesh.vertex_stride,
679 ));
680 }
681
682 match (size_desc.index_count, mesh.size.index_count) {
683 (Some(_), None) | (None, Some(_)) => {
684 return Err(
685 BuildAccelerationStructureError::BlasIndexCountProvidedMismatch(
686 blas.error_ident(),
687 ),
688 )
689 }
690 (Some(create), Some(build)) if create < build => {
691 return Err(
692 BuildAccelerationStructureError::IncompatibleBlasIndexCount(
693 blas.error_ident(),
694 create,
695 build,
696 ),
697 )
698 }
699 _ => {}
700 }
701
702 if size_desc.index_format != mesh.size.index_format {
703 return Err(BuildAccelerationStructureError::DifferentBlasIndexFormats(
704 blas.error_ident(),
705 size_desc.index_format,
706 mesh.size.index_format,
707 ));
708 }
709
710 if size_desc.index_count.is_some() && mesh.index_buffer.is_none() {
711 return Err(BuildAccelerationStructureError::MissingIndexBuffer(
712 blas.error_ident(),
713 ));
714 }
715 let vertex_buffer = mesh.vertex_buffer.clone();
716 let vertex_pending = state.tracker.buffers.set_single(
717 &vertex_buffer,
718 BufferUses::BOTTOM_LEVEL_ACCELERATION_STRUCTURE_INPUT,
719 );
720 let vertex_buffer = {
721 let vertex_raw = mesh.vertex_buffer.as_ref().try_raw(state.snatch_guard)?;
722 let vertex_buffer = &mesh.vertex_buffer;
723 vertex_buffer.check_usage(BufferUsages::BLAS_INPUT)?;
724
725 if let Some(barrier) = vertex_pending.map(|pending| {
726 pending.into_hal(vertex_buffer.as_ref(), state.snatch_guard)
727 }) {
728 input_barriers.push(barrier);
729 }
730 if u64::from(mesh.size.vertex_count)
731 .checked_add(u64::from(mesh.first_vertex))
732 .and_then(|end_vertex| end_vertex.checked_mul(mesh.vertex_stride))
733 .is_none_or(|end| vertex_buffer.size < end)
734 {
735 return Err(BuildAccelerationStructureError::InsufficientBufferSize {
736 buffer_ident: vertex_buffer.error_ident(),
737 offset: u64::from(mesh.first_vertex)
738 .saturating_mul(mesh.vertex_stride),
739 region_size: u64::from(mesh.size.vertex_count)
740 .saturating_mul(mesh.vertex_stride),
741 buffer_size: vertex_buffer.size,
742 });
743 }
744 let vertex_buffer_offset = mesh.first_vertex as u64 * mesh.vertex_stride;
745 state.buffer_memory_init_actions.extend(
746 vertex_buffer.initialization_status.read().create_action(
747 vertex_buffer,
748 vertex_buffer_offset
749 ..(vertex_buffer_offset
750 + mesh.size.vertex_count as u64 * mesh.vertex_stride),
751 MemoryInitKind::NeedsInitializedMemory,
752 ),
753 );
754 vertex_raw
755 };
756 let index_buffer = if let Some(ref index_buffer) = mesh.index_buffer {
757 if mesh.first_index.is_none()
758 || mesh.size.index_count.is_none()
759 || mesh.size.index_count.is_none()
760 {
761 return Err(BuildAccelerationStructureError::MissingAssociatedData(
762 index_buffer.error_ident(),
763 ));
764 }
765 let index_pending = state.tracker.buffers.set_single(
766 index_buffer,
767 BufferUses::BOTTOM_LEVEL_ACCELERATION_STRUCTURE_INPUT,
768 );
769 let index_raw = index_buffer.try_raw(state.snatch_guard)?;
770 index_buffer.check_usage(BufferUsages::BLAS_INPUT)?;
771
772 if let Some(barrier) = index_pending.map(|pending| {
773 pending.into_hal(index_buffer.as_ref(), state.snatch_guard)
774 }) {
775 input_barriers.push(barrier);
776 }
777 let index_stride = mesh.size.index_format.unwrap().byte_size();
778 let Some(vertex_offset) =
780 mesh.first_index.unwrap().checked_mul(index_stride)
781 else {
782 return Err(BuildAccelerationStructureError::OffsetLimitedTo4GB {
783 buffer_ident: index_buffer.error_ident(),
784 offset: u64::from(mesh.first_index.unwrap())
785 .saturating_mul(u64::from(index_stride)),
786 count: u64::from(mesh.first_index.unwrap()),
787 stride: u64::from(index_stride),
788 });
789 };
790 let Some(indexes_size) =
791 mesh.size.index_count.unwrap().checked_mul(index_stride)
792 else {
793 return Err(BuildAccelerationStructureError::OffsetLimitedTo4GB {
794 buffer_ident: index_buffer.error_ident(),
795 offset: u64::from(mesh.size.index_count.unwrap())
796 .saturating_mul(u64::from(index_stride)),
797 count: u64::from(mesh.size.index_count.unwrap()),
798 stride: u64::from(index_stride),
799 });
800 };
801
802 if mesh.size.index_count.unwrap() % 3 != 0 {
803 return Err(BuildAccelerationStructureError::InvalidIndexCount(
804 index_buffer.error_ident(),
805 mesh.size.index_count.unwrap(),
806 ));
807 }
808 if index_buffer.size < u64::from(vertex_offset)
809 || index_buffer.size - u64::from(vertex_offset)
810 < u64::from(indexes_size)
811 {
812 return Err(BuildAccelerationStructureError::InsufficientBufferSize {
813 buffer_ident: index_buffer.error_ident(),
814 offset: u64::from(vertex_offset),
815 region_size: u64::from(indexes_size),
816 buffer_size: index_buffer.size,
817 });
818 }
819
820 state.buffer_memory_init_actions.extend(
821 index_buffer.initialization_status.read().create_action(
822 index_buffer,
823 u64::from(vertex_offset)
824 ..(u64::from(vertex_offset) + u64::from(indexes_size)),
825 MemoryInitKind::NeedsInitializedMemory,
826 ),
827 );
828 Some((index_raw, vertex_offset))
829 } else {
830 None
831 };
832 let transform_buffer = if let Some(ref transform_buffer) = mesh.transform_buffer
833 {
834 if !blas
835 .flags
836 .contains(wgt::AccelerationStructureFlags::USE_TRANSFORM)
837 {
838 return Err(BuildAccelerationStructureError::UseTransformMissing(
839 blas.error_ident(),
840 ));
841 }
842 if mesh.transform_buffer_offset.is_none() {
843 return Err(BuildAccelerationStructureError::MissingAssociatedData(
844 transform_buffer.error_ident(),
845 ));
846 }
847 let transform_pending = state.tracker.buffers.set_single(
848 transform_buffer,
849 BufferUses::BOTTOM_LEVEL_ACCELERATION_STRUCTURE_INPUT,
850 );
851 if mesh.transform_buffer_offset.is_none() {
852 return Err(BuildAccelerationStructureError::MissingAssociatedData(
853 transform_buffer.error_ident(),
854 ));
855 }
856 let transform_raw = transform_buffer.try_raw(state.snatch_guard)?;
857 transform_buffer.check_usage(BufferUsages::BLAS_INPUT)?;
858
859 if let Some(barrier) = transform_pending.map(|pending| {
860 pending.into_hal(transform_buffer.as_ref(), state.snatch_guard)
861 }) {
862 input_barriers.push(barrier);
863 }
864
865 let Ok(offset) = u32::try_from(mesh.transform_buffer_offset.unwrap())
867 else {
868 return Err(BuildAccelerationStructureError::OffsetLimitedTo4GB {
869 buffer_ident: transform_buffer.error_ident(),
870 offset: mesh.transform_buffer_offset.unwrap(),
871 count: mesh.transform_buffer_offset.unwrap(),
872 stride: 1,
873 });
874 };
875
876 if offset % wgt::TRANSFORM_BUFFER_ALIGNMENT as u32 != 0 {
877 return Err(
878 BuildAccelerationStructureError::UnalignedTransformBufferOffset(
879 transform_buffer.error_ident(),
880 ),
881 );
882 }
883 if transform_buffer.size < 48
884 || transform_buffer.size - 48 < u64::from(offset)
885 {
886 return Err(BuildAccelerationStructureError::InsufficientBufferSize {
887 buffer_ident: transform_buffer.error_ident(),
888 region_size: 48,
889 offset: u64::from(offset),
890 buffer_size: transform_buffer.size,
891 });
892 }
893 state.buffer_memory_init_actions.extend(
894 transform_buffer.initialization_status.read().create_action(
895 transform_buffer,
896 u64::from(offset)..(u64::from(offset) + 48),
897 MemoryInitKind::NeedsInitializedMemory,
898 ),
899 );
900 Some((transform_raw, offset))
901 } else {
902 if blas
903 .flags
904 .contains(wgt::AccelerationStructureFlags::USE_TRANSFORM)
905 {
906 return Err(BuildAccelerationStructureError::TransformMissing(
907 blas.error_ident(),
908 ));
909 }
910 None
911 };
912
913 let triangles = hal::AccelerationStructureTriangles {
914 vertex_buffer: Some(vertex_buffer),
915 vertex_format: mesh.size.vertex_format,
916 first_vertex: mesh.first_vertex,
917 vertex_count: mesh.size.vertex_count,
918 vertex_stride: mesh.vertex_stride,
919 indices: index_buffer.map(|(index_raw, vertex_offset)| {
920 hal::AccelerationStructureTriangleIndices::<dyn hal::DynBuffer> {
921 format: mesh.size.index_format.unwrap(),
922 buffer: Some(index_raw),
923 offset: vertex_offset,
924 count: mesh.size.index_count.unwrap(),
925 }
926 }),
927 transform: transform_buffer.map(|(transform_raw, offset)| {
928 hal::AccelerationStructureTriangleTransform {
929 buffer: transform_raw,
930 offset,
931 }
932 }),
933 flags: mesh.size.flags,
934 };
935 triangle_entries.push(triangles);
936 }
937
938 {
939 let scratch_buffer_offset = *scratch_buffer_blas_size;
940 *scratch_buffer_blas_size = scratch_buffer_blas_size.saturating_add(align_to(
941 blas.size_info.build_scratch_size,
942 u64::from(state.device.alignments.ray_tracing_scratch_buffer_alignment),
943 ));
944
945 blas_storage.push(BlasStore {
946 blas: blas.clone(),
947 entries: hal::AccelerationStructureEntries::Triangles(triangle_entries),
948 scratch_buffer_offset,
949 });
950 }
951 }
952 ArcBlasGeometries::AabbGeometries(aabb_geometries) => {
953 let mut aabb_entries =
954 Vec::<hal::AccelerationStructureAABBs<dyn hal::DynBuffer>>::new();
955
956 for (i, aabb) in aabb_geometries.iter().enumerate() {
957 let size_desc = match &blas.sizes {
958 wgt::BlasGeometrySizeDescriptors::AABBs { descriptors } => descriptors,
959 _ => {
960 return Err(BuildAccelerationStructureError::BlasGeometryKindMismatch(
961 blas.error_ident(),
962 ));
963 }
964 };
965 if i >= size_desc.len() {
966 return Err(BuildAccelerationStructureError::IncompatibleBlasBuildSizes(
967 blas.error_ident(),
968 ));
969 }
970 let size_desc = &size_desc[i];
971
972 if size_desc.flags != aabb.size.flags {
973 return Err(BuildAccelerationStructureError::IncompatibleBlasFlags(
974 blas.error_ident(),
975 size_desc.flags,
976 aabb.size.flags,
977 ));
978 }
979
980 if size_desc.primitive_count < aabb.size.primitive_count {
981 return Err(
982 BuildAccelerationStructureError::IncompatibleBlasAabbPrimitiveCount(
983 blas.error_ident(),
984 size_desc.primitive_count,
985 aabb.size.primitive_count,
986 ),
987 );
988 }
989
990 if aabb.primitive_offset % 8 != 0 {
991 return Err(
992 BuildAccelerationStructureError::UnalignedAabbPrimitiveOffset(
993 blas.error_ident(),
994 ),
995 );
996 }
997
998 if aabb.stride < wgt::AABB_GEOMETRY_MIN_STRIDE || aabb.stride % 8 != 0 {
999 return Err(BuildAccelerationStructureError::InvalidAabbStride(
1000 blas.error_ident(),
1001 aabb.stride,
1002 ));
1003 }
1004
1005 let aabb_buffer = aabb.aabb_buffer.clone();
1006 let aabb_pending = state.tracker.buffers.set_single(
1007 &aabb_buffer,
1008 BufferUses::BOTTOM_LEVEL_ACCELERATION_STRUCTURE_INPUT,
1009 );
1010 let aabb_raw = {
1011 let aabb_raw = aabb.aabb_buffer.as_ref().try_raw(state.snatch_guard)?;
1012 let aabb_buffer = &aabb.aabb_buffer;
1013 aabb_buffer.check_usage(BufferUsages::BLAS_INPUT)?;
1014
1015 if let Some(barrier) = aabb_pending.map(|pending| {
1016 pending.into_hal(aabb_buffer.as_ref(), state.snatch_guard)
1017 }) {
1018 input_barriers.push(barrier);
1019 }
1020
1021 let Some(aabb_size) = u32::try_from(aabb.stride)
1023 .ok()
1024 .and_then(|stride| aabb.size.primitive_count.checked_mul(stride))
1025 else {
1026 return Err(BuildAccelerationStructureError::OffsetLimitedTo4GB {
1027 buffer_ident: aabb_buffer.error_ident(),
1028 offset: u64::from(aabb.size.primitive_count)
1029 .saturating_mul(aabb.stride),
1030 count: u64::from(aabb.size.primitive_count),
1031 stride: aabb.stride,
1032 });
1033 };
1034
1035 if aabb_buffer.size < u64::from(aabb.primitive_offset)
1036 || aabb_buffer.size - u64::from(aabb.primitive_offset)
1037 < u64::from(aabb_size)
1038 {
1039 return Err(BuildAccelerationStructureError::InsufficientBufferSize {
1040 buffer_ident: aabb_buffer.error_ident(),
1041 region_size: u64::from(aabb_size),
1042 offset: u64::from(aabb.primitive_offset),
1043 buffer_size: aabb_buffer.size,
1044 });
1045 }
1046
1047 state.buffer_memory_init_actions.extend(
1048 aabb_buffer.initialization_status.read().create_action(
1049 aabb_buffer,
1050 u64::from(aabb.primitive_offset)
1051 ..u64::from(aabb.primitive_offset) + u64::from(aabb_size),
1052 MemoryInitKind::NeedsInitializedMemory,
1053 ),
1054 );
1055 aabb_raw
1056 };
1057
1058 aabb_entries.push(hal::AccelerationStructureAABBs {
1059 buffer: Some(aabb_raw),
1060 offset: aabb.primitive_offset,
1061 count: aabb.size.primitive_count,
1062 stride: aabb.stride,
1063 flags: aabb.size.flags,
1064 });
1065 }
1066
1067 {
1068 let scratch_buffer_offset = *scratch_buffer_blas_size;
1069 *scratch_buffer_blas_size = scratch_buffer_blas_size.saturating_add(align_to(
1070 blas.size_info.build_scratch_size,
1071 u64::from(state.device.alignments.ray_tracing_scratch_buffer_alignment),
1072 ));
1073
1074 blas_storage.push(BlasStore {
1075 blas: blas.clone(),
1076 entries: hal::AccelerationStructureEntries::AABBs(aabb_entries),
1077 scratch_buffer_offset,
1078 });
1079 }
1080 }
1081 }
1082 }
1083 Ok(())
1084}
1085
1086fn map_blas<'a>(
1087 storage: &'a BlasStore<'_>,
1088 scratch_buffer: &'a dyn hal::DynBuffer,
1089 snatch_guard: &'a SnatchGuard,
1090 blases_compactable: &mut Vec<(
1091 &'a dyn hal::DynBuffer,
1092 &'a dyn hal::DynAccelerationStructure,
1093 )>,
1094) -> Result<
1095 hal::BuildAccelerationStructureDescriptor<
1096 'a,
1097 dyn hal::DynBuffer,
1098 dyn hal::DynAccelerationStructure,
1099 >,
1100 BuildAccelerationStructureError,
1101> {
1102 let BlasStore {
1103 blas,
1104 entries,
1105 scratch_buffer_offset,
1106 } = storage;
1107 if blas.update_mode == wgt::AccelerationStructureUpdateMode::PreferUpdate {
1108 log::debug!("only rebuild implemented")
1109 }
1110 let raw = blas.try_raw(snatch_guard)?;
1111
1112 let state_lock = blas.compacted_state.lock();
1113 if let BlasCompactState::Compacted = *state_lock {
1114 return Err(BuildAccelerationStructureError::CompactedBlas(
1115 blas.error_ident(),
1116 ));
1117 }
1118
1119 if blas
1120 .flags
1121 .contains(wgpu_types::AccelerationStructureFlags::ALLOW_COMPACTION)
1122 {
1123 blases_compactable.push((blas.compaction_buffer.as_ref().unwrap().as_ref(), raw));
1124 }
1125 Ok(hal::BuildAccelerationStructureDescriptor {
1126 entries,
1127 mode: hal::AccelerationStructureBuildMode::Build,
1128 flags: blas.flags,
1129 source_acceleration_structure: None,
1130 destination_acceleration_structure: raw,
1131 scratch_buffer,
1132 scratch_buffer_offset: *scratch_buffer_offset,
1133 })
1134}
1135
1136fn build_blas<'a>(
1137 cmd_buf_raw: &mut dyn hal::DynCommandEncoder,
1138 blas_present: bool,
1139 tlas_present: bool,
1140 input_barriers: Vec<hal::BufferBarrier<dyn hal::DynBuffer>>,
1141 blas_descriptors: &[hal::BuildAccelerationStructureDescriptor<
1142 'a,
1143 dyn hal::DynBuffer,
1144 dyn hal::DynAccelerationStructure,
1145 >],
1146 scratch_buffer_barrier: hal::BufferBarrier<dyn hal::DynBuffer>,
1147 blas_s_for_compaction: Vec<(
1148 &'a dyn hal::DynBuffer,
1149 &'a dyn hal::DynAccelerationStructure,
1150 )>,
1151) {
1152 unsafe {
1153 cmd_buf_raw.transition_buffers(&input_barriers);
1154 }
1155
1156 if blas_present {
1157 unsafe {
1158 cmd_buf_raw.place_acceleration_structure_barrier(hal::AccelerationStructureBarrier {
1159 usage: hal::StateTransition {
1160 from: hal::AccelerationStructureUses::BUILD_INPUT,
1161 to: hal::AccelerationStructureUses::BUILD_OUTPUT,
1162 },
1163 });
1164
1165 cmd_buf_raw.build_acceleration_structures(blas_descriptors);
1166 }
1167 }
1168
1169 if blas_present && tlas_present {
1170 unsafe {
1171 cmd_buf_raw.transition_buffers(&[scratch_buffer_barrier]);
1172 }
1173 }
1174
1175 let mut source_usage = hal::AccelerationStructureUses::empty();
1176 let mut destination_usage = hal::AccelerationStructureUses::empty();
1177 for &(buf, blas) in blas_s_for_compaction.iter() {
1178 unsafe {
1179 cmd_buf_raw.transition_buffers(&[hal::BufferBarrier {
1180 buffer: buf,
1181 usage: hal::StateTransition {
1182 from: BufferUses::ACCELERATION_STRUCTURE_QUERY,
1183 to: BufferUses::ACCELERATION_STRUCTURE_QUERY,
1184 },
1185 }])
1186 }
1187 unsafe { cmd_buf_raw.read_acceleration_structure_compact_size(blas, buf) }
1188 destination_usage |= hal::AccelerationStructureUses::COPY_SRC;
1189 }
1190
1191 if blas_present {
1192 source_usage |= hal::AccelerationStructureUses::BUILD_OUTPUT;
1193 destination_usage |= hal::AccelerationStructureUses::BUILD_INPUT
1194 }
1195 if tlas_present {
1196 source_usage |= hal::AccelerationStructureUses::SHADER_INPUT;
1197 destination_usage |= hal::AccelerationStructureUses::BUILD_OUTPUT;
1198 }
1199 unsafe {
1200 cmd_buf_raw.place_acceleration_structure_barrier(hal::AccelerationStructureBarrier {
1201 usage: hal::StateTransition {
1202 from: source_usage,
1203 to: destination_usage,
1204 },
1205 });
1206 }
1207}