Skip to main content

ab_proof_of_space_gpu/
host.rs

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