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

ab_proof_of_space_gpu/
host.rs
