1use crate::shader::constants::{
2 MAX_BUCKET_SIZE, MAX_TABLE_SIZE, NUM_BUCKETS, NUM_MATCH_BUCKETS, NUM_S_BUCKETS,
3 REDUCED_MATCHES_COUNT,
4};
5use crate::shader::find_matches_and_compute_f7::{NUM_ELEMENTS_PER_S_BUCKET, ProofTargets};
6use crate::shader::find_proofs::ProofsHost;
7use crate::shader::types::{Metadata, Position, PositionR};
8use crate::shader::{compute_f1, find_proofs, select_shader_features_limits};
9use ab_chacha8::{ChaCha8Block, ChaCha8State, block_to_bytes};
10use futures::stream::FuturesOrdered;
11use futures::{StreamExt, TryStreamExt};
12use parking_lot::Mutex;
13use rclite::Arc;
14use std::num::NonZeroU8;
15use std::{fmt, iter};
16use subspace_core_primitives::pos::PosSeed;
17use tracing::{debug, warn};
18use wgpu::{
19 AdapterInfo, Backend, BackendOptions, Backends, BindGroup, BindGroupDescriptor, BindGroupEntry,
20 BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingType, Buffer, BufferAddress,
21 BufferAsyncError, BufferBindingType, BufferDescriptor, BufferUsages, CommandEncoderDescriptor,
22 ComputePassDescriptor, ComputePipeline, ComputePipelineDescriptor, DeviceDescriptor,
23 DeviceType, Instance, InstanceDescriptor, InstanceFlags, MapMode, MemoryBudgetThresholds,
24 PipelineCompilationOptions, PipelineLayoutDescriptor, PollError, PollType, Queue,
25 RequestDeviceError, ShaderModule, ShaderRuntimeChecks, ShaderStages,
26};
27
28#[derive(Debug, thiserror::Error)]
30pub enum RecordEncodingError {
31 #[error("Proof creation failed previously and the device is now considered broken")]
33 DeviceBroken,
34 #[error("Failed to map buffer: {0}")]
36 BufferMapping(#[from] BufferAsyncError),
37 #[error("Poll error: {0}")]
39 DevicePoll(#[from] PollError),
40}
41
42pub struct ProofsHostWrapper<'a> {
44 proofs: &'a ProofsHost,
45 proofs_host: &'a Buffer,
46}
47
48impl ProofsHostWrapper<'_> {
49 pub fn proofs(&self) -> &ProofsHost {
51 self.proofs
52 }
53}
54
55impl fmt::Debug for ProofsHostWrapper<'_> {
56 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57 f.debug_struct("ProofsHostWrapper").finish_non_exhaustive()
58 }
59}
60
61impl Drop for ProofsHostWrapper<'_> {
62 fn drop(&mut self) {
63 self.proofs_host.unmap();
64 }
65}
66
67#[derive(Clone)]
69pub struct Device {
70 id: u32,
71 devices: Vec<(wgpu::Device, Queue, ShaderModule)>,
72 adapter_info: AdapterInfo,
73}
74
75impl fmt::Debug for Device {
76 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77 f.debug_struct("Device")
78 .field("id", &self.id)
79 .field("name", &self.adapter_info.name)
80 .field("device_type", &self.adapter_info.device_type)
81 .field("driver", &self.adapter_info.driver)
82 .field("driver_info", &self.adapter_info.driver_info)
83 .field("backend", &self.adapter_info.backend)
84 .finish_non_exhaustive()
85 }
86}
87
88impl Device {
89 pub async fn enumerate<NOQ>(number_of_queues: NOQ) -> Vec<Self>
91 where
92 NOQ: Fn(DeviceType) -> NonZeroU8,
93 {
94 let backends = Backends::from_env().unwrap_or(Backends::METAL | Backends::VULKAN);
95 let instance = Instance::new(InstanceDescriptor {
96 backends,
97 flags: if cfg!(debug_assertions) {
98 InstanceFlags::debugging().with_env()
99 } else {
100 InstanceFlags::from_env_or_default()
101 },
102 memory_budget_thresholds: MemoryBudgetThresholds::default(),
103 backend_options: BackendOptions::from_env_or_default(),
104 display: None,
105 });
106
107 let adapters = instance.enumerate_adapters(backends).await;
108 let number_of_queues = &number_of_queues;
109
110 adapters
111 .into_iter()
112 .zip(0..)
113 .map(|(adapter, id)| async move {
114 let adapter_info = adapter.get_info();
115
116 let (shader, required_features, required_limits) =
117 if let Some((shader, required_features, required_limits)) =
118 select_shader_features_limits(&adapter)
119 {
120 debug!(
121 %id,
122 adapter_info = ?adapter_info,
123 "Compatible adapter found"
124 );
125
126 (shader, required_features, required_limits)
127 } else {
128 debug!(
129 %id,
130 adapter_info = ?adapter_info,
131 "Incompatible adapter found"
132 );
133
134 return None;
135 };
136
137 let devices = iter::repeat_with(|| async {
140 let (device, queue) = adapter
141 .request_device(&DeviceDescriptor {
142 label: None,
143 required_features,
144 required_limits: required_limits.clone(),
145 ..DeviceDescriptor::default()
146 })
147 .await
148 .inspect_err(|error| {
149 warn!(%id, ?adapter_info, %error, "Failed to request the device");
150 })?;
151 let module = if cfg!(debug_assertions) {
152 device.create_shader_module(shader.clone())
153 } else {
154 unsafe {
158 device.create_shader_module_trusted(
159 shader.clone(),
160 ShaderRuntimeChecks::unchecked(),
161 )
162 }
163 };
164
165 Ok::<_, RequestDeviceError>((device, queue, module))
166 })
167 .take(usize::from(
168 number_of_queues(adapter_info.device_type).get(),
169 ))
170 .collect::<FuturesOrdered<_>>()
171 .try_collect::<Vec<_>>()
172 .await
173 .ok()?;
174
175 Some(Self {
176 id,
177 devices,
178 adapter_info,
179 })
180 })
181 .collect::<FuturesOrdered<_>>()
182 .filter_map(|device| async move { device })
183 .collect()
184 .await
185 }
186
187 pub fn id(&self) -> u32 {
189 self.id
190 }
191
192 pub fn name(&self) -> &str {
194 &self.adapter_info.name
195 }
196
197 pub fn device_type(&self) -> DeviceType {
199 self.adapter_info.device_type
200 }
201
202 pub fn driver(&self) -> &str {
204 &self.adapter_info.driver
205 }
206
207 pub fn driver_info(&self) -> &str {
209 &self.adapter_info.driver_info
210 }
211
212 pub fn backend(&self) -> Backend {
214 self.adapter_info.backend
215 }
216
217 pub fn create_proofs_encoder_instances(&self) -> Vec<GpuRecordsEncoderInstance> {
220 self.devices
221 .clone()
222 .into_iter()
223 .map(|(device, queue, module)| GpuRecordsEncoderInstance::new(device, queue, module))
224 .collect()
225 }
226}
227
228pub struct GpuRecordsEncoderInstance {
229 device: wgpu::Device,
230 queue: Queue,
231 mapping_error: Arc<Mutex<Option<BufferAsyncError>>>,
232 tainted: bool,
233 initial_state_host: Buffer,
234 initial_state_gpu: Buffer,
235 proofs_host: Buffer,
236 proofs_gpu: Buffer,
237 bind_group_compute_f1: BindGroup,
238 compute_pipeline_compute_f1: ComputePipeline,
239 bind_group_sort_buckets_a: BindGroup,
240 compute_pipeline_sort_buckets_a: ComputePipeline,
241 bind_group_sort_buckets_b: BindGroup,
242 compute_pipeline_sort_buckets_b: ComputePipeline,
243 bind_group_find_matches_and_compute_f2: BindGroup,
244 compute_pipeline_find_matches_and_compute_f2: ComputePipeline,
245 bind_group_find_matches_and_compute_f3: BindGroup,
246 compute_pipeline_find_matches_and_compute_f3: ComputePipeline,
247 bind_group_find_matches_and_compute_f4: BindGroup,
248 compute_pipeline_find_matches_and_compute_f4: ComputePipeline,
249 bind_group_find_matches_and_compute_f5: BindGroup,
250 compute_pipeline_find_matches_and_compute_f5: ComputePipeline,
251 bind_group_find_matches_and_compute_f6: BindGroup,
252 compute_pipeline_find_matches_and_compute_f6: ComputePipeline,
253 bind_group_find_matches_and_compute_f7: BindGroup,
254 compute_pipeline_find_matches_and_compute_f7: ComputePipeline,
255 bind_group_find_proofs: BindGroup,
256 compute_pipeline_find_proofs: ComputePipeline,
257}
258
259impl fmt::Debug for GpuRecordsEncoderInstance {
260 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
261 f.debug_struct("GpuRecordsEncoderInstance")
262 .finish_non_exhaustive()
263 }
264}
265
266impl GpuRecordsEncoderInstance {
267 fn new(device: wgpu::Device, queue: Queue, module: ShaderModule) -> Self {
268 let initial_state_host = device.create_buffer(&BufferDescriptor {
269 label: Some("initial_state_host"),
270 size: size_of::<ChaCha8Block>() as BufferAddress,
271 usage: BufferUsages::MAP_WRITE | BufferUsages::COPY_SRC,
272 mapped_at_creation: true,
273 });
274
275 let initial_state_gpu = device.create_buffer(&BufferDescriptor {
276 label: Some("initial_state_gpu"),
277 size: initial_state_host.size(),
278 usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
279 mapped_at_creation: false,
280 });
281
282 let bucket_sizes_gpu_buffer_size = size_of::<[u32; NUM_BUCKETS]>() as BufferAddress;
283 let table_6_proof_targets_sizes_gpu_buffer_size =
284 size_of::<[u32; NUM_S_BUCKETS]>() as BufferAddress;
285 let bucket_sizes_gpu = device.create_buffer(&BufferDescriptor {
289 label: Some("bucket_sizes_gpu"),
290 size: bucket_sizes_gpu_buffer_size.max(table_6_proof_targets_sizes_gpu_buffer_size),
291 usage: BufferUsages::STORAGE,
292 mapped_at_creation: false,
293 });
294 let table_6_proof_targets_sizes_gpu = bucket_sizes_gpu.clone();
296
297 let buckets_a_gpu = device.create_buffer(&BufferDescriptor {
298 label: Some("buckets_a_gpu"),
299 size: size_of::<[[PositionR; MAX_BUCKET_SIZE]; NUM_BUCKETS]>() as BufferAddress,
300 usage: BufferUsages::STORAGE,
301 mapped_at_creation: false,
302 });
303
304 let buckets_b_gpu = device.create_buffer(&BufferDescriptor {
305 label: Some("buckets_b_gpu"),
306 size: buckets_a_gpu.size(),
307 usage: BufferUsages::STORAGE,
308 mapped_at_creation: false,
309 });
310
311 let positions_f2_gpu = device.create_buffer(&BufferDescriptor {
312 label: Some("positions_f2_gpu"),
313 size: size_of::<[[[Position; 2]; REDUCED_MATCHES_COUNT]; NUM_MATCH_BUCKETS]>()
314 as BufferAddress,
315 usage: BufferUsages::STORAGE,
316 mapped_at_creation: false,
317 });
318
319 let positions_f3_gpu = device.create_buffer(&BufferDescriptor {
320 label: Some("positions_f3_gpu"),
321 size: positions_f2_gpu.size(),
322 usage: BufferUsages::STORAGE,
323 mapped_at_creation: false,
324 });
325
326 let positions_f4_gpu = device.create_buffer(&BufferDescriptor {
327 label: Some("positions_f4_gpu"),
328 size: positions_f2_gpu.size(),
329 usage: BufferUsages::STORAGE,
330 mapped_at_creation: false,
331 });
332
333 let positions_f5_gpu = device.create_buffer(&BufferDescriptor {
334 label: Some("positions_f5_gpu"),
335 size: positions_f2_gpu.size(),
336 usage: BufferUsages::STORAGE,
337 mapped_at_creation: false,
338 });
339
340 let positions_f6_gpu = device.create_buffer(&BufferDescriptor {
341 label: Some("positions_f6_gpu"),
342 size: positions_f2_gpu.size(),
343 usage: BufferUsages::STORAGE,
344 mapped_at_creation: false,
345 });
346
347 let metadatas_gpu_buffer_size =
348 size_of::<[[Metadata; REDUCED_MATCHES_COUNT]; NUM_MATCH_BUCKETS]>() as BufferAddress;
349 let table_6_proof_targets_gpu_buffer_size = size_of::<
350 [[ProofTargets; NUM_ELEMENTS_PER_S_BUCKET]; NUM_S_BUCKETS],
351 >() as BufferAddress;
352 let metadatas_a_gpu = device.create_buffer(&BufferDescriptor {
353 label: Some("metadatas_a_gpu"),
354 size: metadatas_gpu_buffer_size.max(table_6_proof_targets_gpu_buffer_size),
355 usage: BufferUsages::STORAGE,
356 mapped_at_creation: false,
357 });
358 let table_6_proof_targets_gpu = metadatas_a_gpu.clone();
360
361 let metadatas_b_gpu = device.create_buffer(&BufferDescriptor {
362 label: Some("metadatas_b_gpu"),
363 size: metadatas_gpu_buffer_size,
364 usage: BufferUsages::STORAGE,
365 mapped_at_creation: false,
366 });
367
368 let proofs_host = device.create_buffer(&BufferDescriptor {
369 label: Some("proofs_host"),
370 size: size_of::<ProofsHost>() as BufferAddress,
371 usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
372 mapped_at_creation: false,
373 });
374
375 let proofs_gpu = device.create_buffer(&BufferDescriptor {
376 label: Some("proofs_gpu"),
377 size: proofs_host.size(),
378 usage: BufferUsages::STORAGE | BufferUsages::COPY_SRC,
379 mapped_at_creation: false,
380 });
381
382 let (bind_group_compute_f1, compute_pipeline_compute_f1) =
383 bind_group_and_pipeline_compute_f1(
384 &device,
385 &module,
386 &initial_state_gpu,
387 &bucket_sizes_gpu,
388 &buckets_a_gpu,
389 );
390 let (bind_group_sort_buckets_a, compute_pipeline_sort_buckets_a) =
391 bind_group_and_pipeline_sort_buckets(
392 &device,
393 &module,
394 &bucket_sizes_gpu,
395 &buckets_a_gpu,
396 );
397
398 let (bind_group_sort_buckets_b, compute_pipeline_sort_buckets_b) =
399 bind_group_and_pipeline_sort_buckets(
400 &device,
401 &module,
402 &bucket_sizes_gpu,
403 &buckets_b_gpu,
404 );
405
406 let (bind_group_find_matches_and_compute_f2, compute_pipeline_find_matches_and_compute_f2) =
407 bind_group_and_pipeline_find_matches_and_compute_f2(
408 &device,
409 &module,
410 &buckets_a_gpu,
411 &bucket_sizes_gpu,
412 &buckets_b_gpu,
413 &positions_f2_gpu,
414 &metadatas_b_gpu,
415 );
416
417 let (bind_group_find_matches_and_compute_f3, compute_pipeline_find_matches_and_compute_f3) =
418 bind_group_and_pipeline_find_matches_and_compute_fn::<3>(
419 &device,
420 &module,
421 &buckets_b_gpu,
422 &metadatas_b_gpu,
423 &bucket_sizes_gpu,
424 &buckets_a_gpu,
425 &positions_f3_gpu,
426 &metadatas_a_gpu,
427 );
428
429 let (bind_group_find_matches_and_compute_f4, compute_pipeline_find_matches_and_compute_f4) =
430 bind_group_and_pipeline_find_matches_and_compute_fn::<4>(
431 &device,
432 &module,
433 &buckets_a_gpu,
434 &metadatas_a_gpu,
435 &bucket_sizes_gpu,
436 &buckets_b_gpu,
437 &positions_f4_gpu,
438 &metadatas_b_gpu,
439 );
440
441 let (bind_group_find_matches_and_compute_f5, compute_pipeline_find_matches_and_compute_f5) =
442 bind_group_and_pipeline_find_matches_and_compute_fn::<5>(
443 &device,
444 &module,
445 &buckets_b_gpu,
446 &metadatas_b_gpu,
447 &bucket_sizes_gpu,
448 &buckets_a_gpu,
449 &positions_f5_gpu,
450 &metadatas_a_gpu,
451 );
452
453 let (bind_group_find_matches_and_compute_f6, compute_pipeline_find_matches_and_compute_f6) =
454 bind_group_and_pipeline_find_matches_and_compute_fn::<6>(
455 &device,
456 &module,
457 &buckets_a_gpu,
458 &metadatas_a_gpu,
459 &bucket_sizes_gpu,
460 &buckets_b_gpu,
461 &positions_f6_gpu,
462 &metadatas_b_gpu,
463 );
464
465 let (bind_group_find_matches_and_compute_f7, compute_pipeline_find_matches_and_compute_f7) =
466 bind_group_and_pipeline_find_matches_and_compute_f7(
467 &device,
468 &module,
469 &buckets_b_gpu,
470 &metadatas_b_gpu,
471 &table_6_proof_targets_sizes_gpu,
472 &table_6_proof_targets_gpu,
473 );
474
475 let (bind_group_find_proofs, compute_pipeline_find_proofs) =
476 bind_group_and_pipeline_find_proofs(
477 &device,
478 &module,
479 &positions_f2_gpu,
480 &positions_f3_gpu,
481 &positions_f4_gpu,
482 &positions_f5_gpu,
483 &positions_f6_gpu,
484 &table_6_proof_targets_sizes_gpu,
485 &table_6_proof_targets_gpu,
486 &proofs_gpu,
487 );
488
489 Self {
490 device,
491 queue,
492 mapping_error: Arc::new(Mutex::new(None)),
493 tainted: false,
494 initial_state_host,
495 initial_state_gpu,
496 proofs_host,
497 proofs_gpu,
498 bind_group_compute_f1,
499 compute_pipeline_compute_f1,
500 bind_group_sort_buckets_a,
501 compute_pipeline_sort_buckets_a,
502 bind_group_sort_buckets_b,
503 compute_pipeline_sort_buckets_b,
504 bind_group_find_matches_and_compute_f2,
505 compute_pipeline_find_matches_and_compute_f2,
506 bind_group_find_matches_and_compute_f3,
507 compute_pipeline_find_matches_and_compute_f3,
508 bind_group_find_matches_and_compute_f4,
509 compute_pipeline_find_matches_and_compute_f4,
510 bind_group_find_matches_and_compute_f5,
511 compute_pipeline_find_matches_and_compute_f5,
512 bind_group_find_matches_and_compute_f6,
513 compute_pipeline_find_matches_and_compute_f6,
514 bind_group_find_matches_and_compute_f7,
515 compute_pipeline_find_matches_and_compute_f7,
516 bind_group_find_proofs,
517 compute_pipeline_find_proofs,
518 }
519 }
520
521 pub fn create_proofs(
522 &mut self,
523 seed: &PosSeed,
524 ) -> Result<ProofsHostWrapper<'_>, RecordEncodingError> {
525 if self.tainted {
526 return Err(RecordEncodingError::DeviceBroken);
527 }
528 self.tainted = true;
529
530 let mut encoder = self
531 .device
532 .create_command_encoder(&CommandEncoderDescriptor {
533 label: Some("create_proofs"),
534 });
535
536 self.initial_state_host
538 .get_mapped_range_mut(..)
539 .copy_from_slice(&block_to_bytes(
540 &ChaCha8State::init(seed, &[0; _]).to_repr(),
541 ));
542 self.initial_state_host.unmap();
543
544 encoder.copy_buffer_to_buffer(
545 &self.initial_state_host,
546 0,
547 &self.initial_state_gpu,
548 0,
549 self.initial_state_host.size(),
550 );
551
552 {
553 let mut cpass = encoder.begin_compute_pass(&ComputePassDescriptor {
554 label: Some("create_proofs"),
555 timestamp_writes: None,
556 });
557
558 cpass.set_bind_group(0, &self.bind_group_compute_f1, &[]);
559 cpass.set_pipeline(&self.compute_pipeline_compute_f1);
560 cpass.dispatch_workgroups(
561 MAX_TABLE_SIZE
562 .div_ceil(compute_f1::WORKGROUP_SIZE * compute_f1::ELEMENTS_PER_INVOCATION),
563 1,
564 1,
565 );
566
567 cpass.set_bind_group(0, &self.bind_group_sort_buckets_a, &[]);
568 cpass.set_pipeline(&self.compute_pipeline_sort_buckets_a);
569 cpass.dispatch_workgroups(NUM_BUCKETS as u32, 1, 1);
570
571 cpass.set_bind_group(0, &self.bind_group_find_matches_and_compute_f2, &[]);
572 cpass.set_pipeline(&self.compute_pipeline_find_matches_and_compute_f2);
573 cpass.dispatch_workgroups(NUM_MATCH_BUCKETS as u32, 1, 1);
574
575 cpass.set_bind_group(0, &self.bind_group_sort_buckets_b, &[]);
576 cpass.set_pipeline(&self.compute_pipeline_sort_buckets_b);
577 cpass.dispatch_workgroups(NUM_BUCKETS as u32, 1, 1);
578
579 cpass.set_bind_group(0, &self.bind_group_find_matches_and_compute_f3, &[]);
580 cpass.set_pipeline(&self.compute_pipeline_find_matches_and_compute_f3);
581 cpass.dispatch_workgroups(NUM_MATCH_BUCKETS as u32, 1, 1);
582
583 cpass.set_bind_group(0, &self.bind_group_sort_buckets_a, &[]);
584 cpass.set_pipeline(&self.compute_pipeline_sort_buckets_a);
585 cpass.dispatch_workgroups(NUM_BUCKETS as u32, 1, 1);
586
587 cpass.set_bind_group(0, &self.bind_group_find_matches_and_compute_f4, &[]);
588 cpass.set_pipeline(&self.compute_pipeline_find_matches_and_compute_f4);
589 cpass.dispatch_workgroups(NUM_MATCH_BUCKETS as u32, 1, 1);
590
591 cpass.set_bind_group(0, &self.bind_group_sort_buckets_b, &[]);
592 cpass.set_pipeline(&self.compute_pipeline_sort_buckets_b);
593 cpass.dispatch_workgroups(NUM_BUCKETS as u32, 1, 1);
594
595 cpass.set_bind_group(0, &self.bind_group_find_matches_and_compute_f5, &[]);
596 cpass.set_pipeline(&self.compute_pipeline_find_matches_and_compute_f5);
597 cpass.dispatch_workgroups(NUM_MATCH_BUCKETS as u32, 1, 1);
598
599 cpass.set_bind_group(0, &self.bind_group_sort_buckets_a, &[]);
600 cpass.set_pipeline(&self.compute_pipeline_sort_buckets_a);
601 cpass.dispatch_workgroups(NUM_BUCKETS as u32, 1, 1);
602
603 cpass.set_bind_group(0, &self.bind_group_find_matches_and_compute_f6, &[]);
604 cpass.set_pipeline(&self.compute_pipeline_find_matches_and_compute_f6);
605 cpass.dispatch_workgroups(NUM_MATCH_BUCKETS as u32, 1, 1);
606
607 cpass.set_bind_group(0, &self.bind_group_sort_buckets_b, &[]);
608 cpass.set_pipeline(&self.compute_pipeline_sort_buckets_b);
609 cpass.dispatch_workgroups(NUM_BUCKETS as u32, 1, 1);
610
611 cpass.set_bind_group(0, &self.bind_group_find_matches_and_compute_f7, &[]);
612 cpass.set_pipeline(&self.compute_pipeline_find_matches_and_compute_f7);
613 cpass.dispatch_workgroups(NUM_MATCH_BUCKETS as u32, 1, 1);
614
615 cpass.set_bind_group(0, &self.bind_group_find_proofs, &[]);
616 cpass.set_pipeline(&self.compute_pipeline_find_proofs);
617 cpass.dispatch_workgroups(NUM_S_BUCKETS as u32 / find_proofs::WORKGROUP_SIZE, 1, 1);
618 }
619
620 encoder.copy_buffer_to_buffer(
621 &self.proofs_gpu,
622 0,
623 &self.proofs_host,
624 0,
625 self.proofs_host.size(),
626 );
627
628 encoder.map_buffer_on_submit(&self.initial_state_host, MapMode::Write, .., {
630 let mapping_error = Arc::clone(&self.mapping_error);
631
632 move |r| {
633 if let Err(error) = r {
634 mapping_error.lock().replace(error);
635 }
636 }
637 });
638 encoder.map_buffer_on_submit(&self.proofs_host, MapMode::Read, .., {
639 let mapping_error = Arc::clone(&self.mapping_error);
640
641 move |r| {
642 if let Err(error) = r {
643 mapping_error.lock().replace(error);
644 }
645 }
646 });
647
648 let submission_index = self.queue.submit([encoder.finish()]);
649
650 self.device.poll(PollType::Wait {
651 submission_index: Some(submission_index),
652 timeout: None,
653 })?;
654
655 if let Some(error) = self.mapping_error.lock().take() {
656 return Err(RecordEncodingError::BufferMapping(error));
657 }
658
659 let proofs = {
660 let proofs_host_ptr = self
661 .proofs_host
662 .get_mapped_range(..)
663 .as_ptr()
664 .cast::<ProofsHost>();
665 unsafe { &*proofs_host_ptr }
667 };
668
669 self.tainted = false;
670
671 Ok(ProofsHostWrapper {
672 proofs,
673 proofs_host: &self.proofs_host,
674 })
675 }
676}
677
678fn bind_group_and_pipeline_compute_f1(
679 device: &wgpu::Device,
680 module: &ShaderModule,
681 initial_state_gpu: &Buffer,
682 bucket_sizes_gpu: &Buffer,
683 buckets_gpu: &Buffer,
684) -> (BindGroup, ComputePipeline) {
685 let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
686 label: Some("compute_f1"),
687 entries: &[
688 BindGroupLayoutEntry {
689 binding: 0,
690 count: None,
691 visibility: ShaderStages::COMPUTE,
692 ty: BindingType::Buffer {
693 has_dynamic_offset: false,
694 min_binding_size: None,
695 ty: BufferBindingType::Uniform,
696 },
697 },
698 BindGroupLayoutEntry {
699 binding: 1,
700 count: None,
701 visibility: ShaderStages::COMPUTE,
702 ty: BindingType::Buffer {
703 has_dynamic_offset: false,
704 min_binding_size: None,
705 ty: BufferBindingType::Storage { read_only: false },
706 },
707 },
708 BindGroupLayoutEntry {
709 binding: 2,
710 count: None,
711 visibility: ShaderStages::COMPUTE,
712 ty: BindingType::Buffer {
713 has_dynamic_offset: false,
714 min_binding_size: None,
715 ty: BufferBindingType::Storage { read_only: false },
716 },
717 },
718 ],
719 });
720
721 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
722 label: Some("compute_f1"),
723 bind_group_layouts: &[Some(&bind_group_layout)],
724 immediate_size: 0,
725 });
726
727 let compute_pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
728 compilation_options: PipelineCompilationOptions {
729 constants: &[],
730 zero_initialize_workgroup_memory: false,
731 },
732 cache: None,
733 label: Some("compute_f1"),
734 layout: Some(&pipeline_layout),
735 module,
736 entry_point: Some("compute_f1"),
737 });
738
739 let bind_group = device.create_bind_group(&BindGroupDescriptor {
740 label: Some("compute_f1"),
741 layout: &bind_group_layout,
742 entries: &[
743 BindGroupEntry {
744 binding: 0,
745 resource: initial_state_gpu.as_entire_binding(),
746 },
747 BindGroupEntry {
748 binding: 1,
749 resource: bucket_sizes_gpu.as_entire_binding(),
750 },
751 BindGroupEntry {
752 binding: 2,
753 resource: buckets_gpu.as_entire_binding(),
754 },
755 ],
756 });
757
758 (bind_group, compute_pipeline)
759}
760
761fn bind_group_and_pipeline_sort_buckets(
762 device: &wgpu::Device,
763 module: &ShaderModule,
764 bucket_sizes_gpu: &Buffer,
765 buckets_gpu: &Buffer,
766) -> (BindGroup, ComputePipeline) {
767 let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
768 label: Some("sort_buckets"),
769 entries: &[
770 BindGroupLayoutEntry {
771 binding: 0,
772 count: None,
773 visibility: ShaderStages::COMPUTE,
774 ty: BindingType::Buffer {
775 has_dynamic_offset: false,
776 min_binding_size: None,
777 ty: BufferBindingType::Storage { read_only: false },
778 },
779 },
780 BindGroupLayoutEntry {
781 binding: 1,
782 count: None,
783 visibility: ShaderStages::COMPUTE,
784 ty: BindingType::Buffer {
785 has_dynamic_offset: false,
786 min_binding_size: None,
787 ty: BufferBindingType::Storage { read_only: false },
788 },
789 },
790 ],
791 });
792
793 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
794 label: Some("sort_buckets"),
795 bind_group_layouts: &[Some(&bind_group_layout)],
796 immediate_size: 0,
797 });
798
799 let compute_pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
800 compilation_options: PipelineCompilationOptions {
801 constants: &[],
802 zero_initialize_workgroup_memory: false,
803 },
804 cache: None,
805 label: Some("sort_buckets"),
806 layout: Some(&pipeline_layout),
807 module,
808 entry_point: Some("sort_buckets"),
809 });
810
811 let bind_group = device.create_bind_group(&BindGroupDescriptor {
812 label: Some("sort_buckets"),
813 layout: &bind_group_layout,
814 entries: &[
815 BindGroupEntry {
816 binding: 0,
817 resource: bucket_sizes_gpu.as_entire_binding(),
818 },
819 BindGroupEntry {
820 binding: 1,
821 resource: buckets_gpu.as_entire_binding(),
822 },
823 ],
824 });
825
826 (bind_group, compute_pipeline)
827}
828
829fn bind_group_and_pipeline_find_matches_and_compute_f2(
830 device: &wgpu::Device,
831 module: &ShaderModule,
832 parent_buckets_gpu: &Buffer,
833 bucket_sizes_gpu: &Buffer,
834 buckets_gpu: &Buffer,
835 positions_gpu: &Buffer,
836 metadatas_gpu: &Buffer,
837) -> (BindGroup, ComputePipeline) {
838 let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
839 label: Some("find_matches_and_compute_f2"),
840 entries: &[
841 BindGroupLayoutEntry {
842 binding: 0,
843 count: None,
844 visibility: ShaderStages::COMPUTE,
845 ty: BindingType::Buffer {
846 has_dynamic_offset: false,
847 min_binding_size: None,
848 ty: BufferBindingType::Storage { read_only: true },
849 },
850 },
851 BindGroupLayoutEntry {
852 binding: 1,
853 count: None,
854 visibility: ShaderStages::COMPUTE,
855 ty: BindingType::Buffer {
856 has_dynamic_offset: false,
857 min_binding_size: None,
858 ty: BufferBindingType::Storage { read_only: false },
859 },
860 },
861 BindGroupLayoutEntry {
862 binding: 2,
863 count: None,
864 visibility: ShaderStages::COMPUTE,
865 ty: BindingType::Buffer {
866 has_dynamic_offset: false,
867 min_binding_size: None,
868 ty: BufferBindingType::Storage { read_only: false },
869 },
870 },
871 BindGroupLayoutEntry {
872 binding: 3,
873 count: None,
874 visibility: ShaderStages::COMPUTE,
875 ty: BindingType::Buffer {
876 has_dynamic_offset: false,
877 min_binding_size: None,
878 ty: BufferBindingType::Storage { read_only: false },
879 },
880 },
881 BindGroupLayoutEntry {
882 binding: 4,
883 count: None,
884 visibility: ShaderStages::COMPUTE,
885 ty: BindingType::Buffer {
886 has_dynamic_offset: false,
887 min_binding_size: None,
888 ty: BufferBindingType::Storage { read_only: false },
889 },
890 },
891 ],
892 });
893
894 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
895 label: Some("find_matches_and_compute_f2"),
896 bind_group_layouts: &[Some(&bind_group_layout)],
897 immediate_size: 0,
898 });
899
900 let compute_pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
901 compilation_options: PipelineCompilationOptions {
902 constants: &[],
903 zero_initialize_workgroup_memory: true,
904 },
905 cache: None,
906 label: Some("find_matches_and_compute_f2"),
907 layout: Some(&pipeline_layout),
908 module,
909 entry_point: Some("find_matches_and_compute_f2"),
910 });
911
912 let bind_group = device.create_bind_group(&BindGroupDescriptor {
913 label: Some("find_matches_and_compute_f2"),
914 layout: &bind_group_layout,
915 entries: &[
916 BindGroupEntry {
917 binding: 0,
918 resource: parent_buckets_gpu.as_entire_binding(),
919 },
920 BindGroupEntry {
921 binding: 1,
922 resource: bucket_sizes_gpu.as_entire_binding(),
923 },
924 BindGroupEntry {
925 binding: 2,
926 resource: buckets_gpu.as_entire_binding(),
927 },
928 BindGroupEntry {
929 binding: 3,
930 resource: positions_gpu.as_entire_binding(),
931 },
932 BindGroupEntry {
933 binding: 4,
934 resource: metadatas_gpu.as_entire_binding(),
935 },
936 ],
937 });
938
939 (bind_group, compute_pipeline)
940}
941
942#[expect(
943 clippy::too_many_arguments,
944 reason = "Both I/O and Vulkan stuff together take a lot of arguments"
945)]
946fn bind_group_and_pipeline_find_matches_and_compute_fn<const TABLE_NUMBER: u8>(
947 device: &wgpu::Device,
948 module: &ShaderModule,
949 parent_buckets_gpu: &Buffer,
950 parent_metadatas_gpu: &Buffer,
951 bucket_sizes_gpu: &Buffer,
952 buckets_gpu: &Buffer,
953 positions_gpu: &Buffer,
954 metadatas_gpu: &Buffer,
955) -> (BindGroup, ComputePipeline) {
956 let label = format!("find_matches_and_compute_f{TABLE_NUMBER}");
957 let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
958 label: Some(&label),
959 entries: &[
960 BindGroupLayoutEntry {
961 binding: 0,
962 count: None,
963 visibility: ShaderStages::COMPUTE,
964 ty: BindingType::Buffer {
965 has_dynamic_offset: false,
966 min_binding_size: None,
967 ty: BufferBindingType::Storage { read_only: true },
968 },
969 },
970 BindGroupLayoutEntry {
971 binding: 1,
972 count: None,
973 visibility: ShaderStages::COMPUTE,
974 ty: BindingType::Buffer {
975 has_dynamic_offset: false,
976 min_binding_size: None,
977 ty: BufferBindingType::Storage { read_only: true },
978 },
979 },
980 BindGroupLayoutEntry {
981 binding: 2,
982 count: None,
983 visibility: ShaderStages::COMPUTE,
984 ty: BindingType::Buffer {
985 has_dynamic_offset: false,
986 min_binding_size: None,
987 ty: BufferBindingType::Storage { read_only: false },
988 },
989 },
990 BindGroupLayoutEntry {
991 binding: 3,
992 count: None,
993 visibility: ShaderStages::COMPUTE,
994 ty: BindingType::Buffer {
995 has_dynamic_offset: false,
996 min_binding_size: None,
997 ty: BufferBindingType::Storage { read_only: false },
998 },
999 },
1000 BindGroupLayoutEntry {
1001 binding: 4,
1002 count: None,
1003 visibility: ShaderStages::COMPUTE,
1004 ty: BindingType::Buffer {
1005 has_dynamic_offset: false,
1006 min_binding_size: None,
1007 ty: BufferBindingType::Storage { read_only: false },
1008 },
1009 },
1010 BindGroupLayoutEntry {
1011 binding: 5,
1012 count: None,
1013 visibility: ShaderStages::COMPUTE,
1014 ty: BindingType::Buffer {
1015 has_dynamic_offset: false,
1016 min_binding_size: None,
1017 ty: BufferBindingType::Storage { read_only: false },
1018 },
1019 },
1020 ],
1021 });
1022
1023 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
1024 label: Some(&label),
1025 bind_group_layouts: &[Some(&bind_group_layout)],
1026 immediate_size: 0,
1027 });
1028
1029 let compute_pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
1030 compilation_options: PipelineCompilationOptions {
1031 constants: &[],
1032 zero_initialize_workgroup_memory: true,
1033 },
1034 cache: None,
1035 label: Some(&label),
1036 layout: Some(&pipeline_layout),
1037 module,
1038 entry_point: Some(&format!("find_matches_and_compute_f{TABLE_NUMBER}")),
1039 });
1040
1041 let bind_group = device.create_bind_group(&BindGroupDescriptor {
1042 label: Some(&label),
1043 layout: &bind_group_layout,
1044 entries: &[
1045 BindGroupEntry {
1046 binding: 0,
1047 resource: parent_buckets_gpu.as_entire_binding(),
1048 },
1049 BindGroupEntry {
1050 binding: 1,
1051 resource: parent_metadatas_gpu.as_entire_binding(),
1052 },
1053 BindGroupEntry {
1054 binding: 2,
1055 resource: bucket_sizes_gpu.as_entire_binding(),
1056 },
1057 BindGroupEntry {
1058 binding: 3,
1059 resource: buckets_gpu.as_entire_binding(),
1060 },
1061 BindGroupEntry {
1062 binding: 4,
1063 resource: positions_gpu.as_entire_binding(),
1064 },
1065 BindGroupEntry {
1066 binding: 5,
1067 resource: metadatas_gpu.as_entire_binding(),
1068 },
1069 ],
1070 });
1071
1072 (bind_group, compute_pipeline)
1073}
1074
1075fn bind_group_and_pipeline_find_matches_and_compute_f7(
1076 device: &wgpu::Device,
1077 module: &ShaderModule,
1078 parent_buckets_gpu: &Buffer,
1079 parent_metadatas_gpu: &Buffer,
1080 table_6_proof_targets_sizes_gpu: &Buffer,
1081 table_6_proof_targets_gpu: &Buffer,
1082) -> (BindGroup, ComputePipeline) {
1083 let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
1084 label: Some("find_matches_and_compute_f7"),
1085 entries: &[
1086 BindGroupLayoutEntry {
1087 binding: 0,
1088 count: None,
1089 visibility: ShaderStages::COMPUTE,
1090 ty: BindingType::Buffer {
1091 has_dynamic_offset: false,
1092 min_binding_size: None,
1093 ty: BufferBindingType::Storage { read_only: true },
1094 },
1095 },
1096 BindGroupLayoutEntry {
1097 binding: 1,
1098 count: None,
1099 visibility: ShaderStages::COMPUTE,
1100 ty: BindingType::Buffer {
1101 has_dynamic_offset: false,
1102 min_binding_size: None,
1103 ty: BufferBindingType::Storage { read_only: true },
1104 },
1105 },
1106 BindGroupLayoutEntry {
1107 binding: 2,
1108 count: None,
1109 visibility: ShaderStages::COMPUTE,
1110 ty: BindingType::Buffer {
1111 has_dynamic_offset: false,
1112 min_binding_size: None,
1113 ty: BufferBindingType::Storage { read_only: false },
1114 },
1115 },
1116 BindGroupLayoutEntry {
1117 binding: 3,
1118 count: None,
1119 visibility: ShaderStages::COMPUTE,
1120 ty: BindingType::Buffer {
1121 has_dynamic_offset: false,
1122 min_binding_size: None,
1123 ty: BufferBindingType::Storage { read_only: false },
1124 },
1125 },
1126 ],
1127 });
1128
1129 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
1130 label: Some("find_matches_and_compute_f7"),
1131 bind_group_layouts: &[Some(&bind_group_layout)],
1132 immediate_size: 0,
1133 });
1134
1135 let compute_pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
1136 compilation_options: PipelineCompilationOptions {
1137 constants: &[],
1138 zero_initialize_workgroup_memory: true,
1139 },
1140 cache: None,
1141 label: Some("find_matches_and_compute_f7"),
1142 layout: Some(&pipeline_layout),
1143 module,
1144 entry_point: Some("find_matches_and_compute_f7"),
1145 });
1146
1147 let bind_group = device.create_bind_group(&BindGroupDescriptor {
1148 label: Some("find_matches_and_compute_f7"),
1149 layout: &bind_group_layout,
1150 entries: &[
1151 BindGroupEntry {
1152 binding: 0,
1153 resource: parent_buckets_gpu.as_entire_binding(),
1154 },
1155 BindGroupEntry {
1156 binding: 1,
1157 resource: parent_metadatas_gpu.as_entire_binding(),
1158 },
1159 BindGroupEntry {
1160 binding: 2,
1161 resource: table_6_proof_targets_sizes_gpu.as_entire_binding(),
1162 },
1163 BindGroupEntry {
1164 binding: 3,
1165 resource: table_6_proof_targets_gpu.as_entire_binding(),
1166 },
1167 ],
1168 });
1169
1170 (bind_group, compute_pipeline)
1171}
1172
1173#[expect(
1174 clippy::too_many_arguments,
1175 reason = "Both I/O and Vulkan stuff together take a lot of arguments"
1176)]
1177fn bind_group_and_pipeline_find_proofs(
1178 device: &wgpu::Device,
1179 module: &ShaderModule,
1180 table_2_positions_gpu: &Buffer,
1181 table_3_positions_gpu: &Buffer,
1182 table_4_positions_gpu: &Buffer,
1183 table_5_positions_gpu: &Buffer,
1184 table_6_positions_gpu: &Buffer,
1185 bucket_sizes_gpu: &Buffer,
1186 buckets_gpu: &Buffer,
1187 proofs_gpu: &Buffer,
1188) -> (BindGroup, ComputePipeline) {
1189 let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
1190 label: Some("find_proofs"),
1191 entries: &[
1192 BindGroupLayoutEntry {
1193 binding: 0,
1194 count: None,
1195 visibility: ShaderStages::COMPUTE,
1196 ty: BindingType::Buffer {
1197 has_dynamic_offset: false,
1198 min_binding_size: None,
1199 ty: BufferBindingType::Storage { read_only: true },
1200 },
1201 },
1202 BindGroupLayoutEntry {
1203 binding: 1,
1204 count: None,
1205 visibility: ShaderStages::COMPUTE,
1206 ty: BindingType::Buffer {
1207 has_dynamic_offset: false,
1208 min_binding_size: None,
1209 ty: BufferBindingType::Storage { read_only: true },
1210 },
1211 },
1212 BindGroupLayoutEntry {
1213 binding: 2,
1214 count: None,
1215 visibility: ShaderStages::COMPUTE,
1216 ty: BindingType::Buffer {
1217 has_dynamic_offset: false,
1218 min_binding_size: None,
1219 ty: BufferBindingType::Storage { read_only: true },
1220 },
1221 },
1222 BindGroupLayoutEntry {
1223 binding: 3,
1224 count: None,
1225 visibility: ShaderStages::COMPUTE,
1226 ty: BindingType::Buffer {
1227 has_dynamic_offset: false,
1228 min_binding_size: None,
1229 ty: BufferBindingType::Storage { read_only: true },
1230 },
1231 },
1232 BindGroupLayoutEntry {
1233 binding: 4,
1234 count: None,
1235 visibility: ShaderStages::COMPUTE,
1236 ty: BindingType::Buffer {
1237 has_dynamic_offset: false,
1238 min_binding_size: None,
1239 ty: BufferBindingType::Storage { read_only: true },
1240 },
1241 },
1242 BindGroupLayoutEntry {
1243 binding: 5,
1244 count: None,
1245 visibility: ShaderStages::COMPUTE,
1246 ty: BindingType::Buffer {
1247 has_dynamic_offset: false,
1248 min_binding_size: None,
1249 ty: BufferBindingType::Storage { read_only: false },
1250 },
1251 },
1252 BindGroupLayoutEntry {
1253 binding: 6,
1254 count: None,
1255 visibility: ShaderStages::COMPUTE,
1256 ty: BindingType::Buffer {
1257 has_dynamic_offset: false,
1258 min_binding_size: None,
1259 ty: BufferBindingType::Storage { read_only: true },
1260 },
1261 },
1262 BindGroupLayoutEntry {
1263 binding: 7,
1264 count: None,
1265 visibility: ShaderStages::COMPUTE,
1266 ty: BindingType::Buffer {
1267 has_dynamic_offset: false,
1268 min_binding_size: None,
1269 ty: BufferBindingType::Storage { read_only: false },
1270 },
1271 },
1272 ],
1273 });
1274
1275 let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
1276 label: Some("find_proofs"),
1277 bind_group_layouts: &[Some(&bind_group_layout)],
1278 immediate_size: 0,
1279 });
1280
1281 let compute_pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
1282 compilation_options: PipelineCompilationOptions {
1283 constants: &[],
1284 zero_initialize_workgroup_memory: false,
1285 },
1286 cache: None,
1287 label: Some("find_proofs"),
1288 layout: Some(&pipeline_layout),
1289 module,
1290 entry_point: Some("find_proofs"),
1291 });
1292
1293 let bind_group = device.create_bind_group(&BindGroupDescriptor {
1294 label: Some("find_proofs"),
1295 layout: &bind_group_layout,
1296 entries: &[
1297 BindGroupEntry {
1298 binding: 0,
1299 resource: table_2_positions_gpu.as_entire_binding(),
1300 },
1301 BindGroupEntry {
1302 binding: 1,
1303 resource: table_3_positions_gpu.as_entire_binding(),
1304 },
1305 BindGroupEntry {
1306 binding: 2,
1307 resource: table_4_positions_gpu.as_entire_binding(),
1308 },
1309 BindGroupEntry {
1310 binding: 3,
1311 resource: table_5_positions_gpu.as_entire_binding(),
1312 },
1313 BindGroupEntry {
1314 binding: 4,
1315 resource: table_6_positions_gpu.as_entire_binding(),
1316 },
1317 BindGroupEntry {
1318 binding: 5,
1319 resource: bucket_sizes_gpu.as_entire_binding(),
1320 },
1321 BindGroupEntry {
1322 binding: 6,
1323 resource: buckets_gpu.as_entire_binding(),
1324 },
1325 BindGroupEntry {
1326 binding: 7,
1327 resource: proofs_gpu.as_entire_binding(),
1328 },
1329 ],
1330 });
1331
1332 (bind_group, compute_pipeline)
1333}