ab_proof_of_space_gpu/
host.rs

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

ab_proof_of_space_gpu/
host.rs
