subspace_proof_of_space_wgpu/
host.rs

1//! Host side: runs GPU proof generation and encodes records with subspace's KZG scheme, so a
2//! GPU-plotted sector reads back byte-for-byte the same as the CPU encoder.
3
4use ab_proof_of_space_gpu::GpuRecordsEncoderInstance;
5use std::ops::DerefMut;
6use std::simd::Simd;
7use subspace_core_primitives::ScalarBytes;
8use subspace_core_primitives::pieces::Record;
9use subspace_core_primitives::pos::PosSeed;
10use subspace_erasure_coding::ErasureCoding;
11use subspace_kzg::Scalar;
12
13/// A single wgpu proof-of-space encoder (one GPU queue) plus the erasure coding used to encode
14/// records host-side.
15pub struct WgpuDevice {
16    instance: GpuRecordsEncoderInstance,
17    erasure_coding: ErasureCoding,
18}
19
20impl WgpuDevice {
21    /// Create a new device from an abundance proof encoder instance.
22    pub fn new(instance: GpuRecordsEncoderInstance, erasure_coding: ErasureCoding) -> Self {
23        Self {
24            instance,
25            erasure_coding,
26        }
27    }
28
29    /// Generate proofs on the GPU and encode a record with subspace's KZG scheme.
30    ///
31    /// Mirrors `record_encoding` in `subspace-farmer-components`, so output is byte-identical to
32    /// the CPU encoder.
33    pub fn generate_and_encode_pospace(
34        &mut self,
35        seed: &PosSeed,
36        record: &mut Record,
37        encoded_chunks_used_output: impl ExactSizeIterator<Item = impl DerefMut<Target = bool>>,
38    ) -> Result<(), String> {
39        let proofs = self
40            .instance
41            .create_proofs(seed)
42            .map_err(|error| error.to_string())?;
43        let proofs = proofs.proofs();
44
45        let source_record_chunks = record.to_vec();
46        let parity_record_chunks = self
47            .erasure_coding
48            .extend(
49                &source_record_chunks
50                    .iter()
51                    .map(|scalar_bytes| {
52                        Scalar::try_from(scalar_bytes)
53                            .expect("Record chunks are valid scalar bytes; qed")
54                    })
55                    .collect::<Vec<_>>(),
56            )
57            .expect("Erasure coding instance supports this many shards; qed")
58            .into_iter()
59            .map(<[u8; ScalarBytes::FULL_BYTES]>::from)
60            .collect::<Vec<_>>();
61
62        let mut encoded_chunks_used = vec![false; Record::NUM_S_BUCKETS];
63        let mut chunks_scratch =
64            Vec::<[u8; ScalarBytes::FULL_BYTES]>::with_capacity(Record::NUM_S_BUCKETS);
65        for s_bucket in 0..Record::NUM_S_BUCKETS {
66            let record_chunk = if s_bucket % 2 == 0 {
67                &source_record_chunks[s_bucket / 2]
68            } else {
69                &parity_record_chunks[s_bucket / 2]
70            };
71
72            let proof_found = (proofs.found_proofs[s_bucket / u8::BITS as usize]
73                >> (s_bucket % u8::BITS as usize))
74                & 1
75                == 1;
76            let encoded_chunk = if proof_found {
77                (Simd::from(*record_chunk) ^ Simd::from(*proofs.proofs[s_bucket].hash())).to_array()
78            } else {
79                // Dummy value indicating no proof
80                [0; ScalarBytes::FULL_BYTES]
81            };
82            chunks_scratch.push(encoded_chunk);
83        }
84
85        let num_successfully_encoded_chunks = chunks_scratch
86            .drain(..)
87            .zip(encoded_chunks_used.iter_mut())
88            .filter_map(|(maybe_encoded_chunk, encoded_chunk_used)| {
89                if maybe_encoded_chunk == [0; ScalarBytes::FULL_BYTES] {
90                    None
91                } else {
92                    *encoded_chunk_used = true;
93                    Some(maybe_encoded_chunk)
94                }
95            })
96            .take(record.len())
97            .zip(record.iter_mut())
98            .map(|(input_chunk, output_chunk)| {
99                *output_chunk = input_chunk;
100            })
101            .count();
102
103        source_record_chunks
104            .iter()
105            .zip(&parity_record_chunks)
106            .flat_map(|(a, b)| [a, b])
107            .zip(encoded_chunks_used.iter())
108            .filter_map(|(record_chunk, encoded_chunk_used)| {
109                if *encoded_chunk_used {
110                    None
111                } else {
112                    Some(record_chunk)
113                }
114            })
115            .zip(record.iter_mut().skip(num_successfully_encoded_chunks))
116            .for_each(|(input_chunk, output_chunk)| {
117                *output_chunk = *input_chunk;
118            });
119
120        encoded_chunks_used_output
121            .zip(&encoded_chunks_used)
122            .for_each(|(mut output, input)| *output = *input);
123
124        Ok(())
125    }
126}
subspace_proof_of_space_wgpu/host.rs

subspace_proof_of_space_wgpu/
host.rs
