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

ab_proof_of_space_gpu/
host.rs
