Merge pull request #1686 from subspace/faster-chiapos-part-4

Introduce and use stateful proof of space generator API

crates/pallet-subspace/src/mock.rs

@@ -53,6 +53,7 @@ use subspace_farmer_components::plotting::{plot_sector, PieceGetterRetryPolicy};
 use subspace_farmer_components::sector::{sector_size, SectorMetadata};
 use subspace_farmer_components::FarmerProtocolInfo;
 use subspace_proof_of_space::shim::ShimTable;
+use subspace_proof_of_space::Table;
 use subspace_solving::REWARD_SIGNING_CONTEXT;
 
 type PosTable = ShimTable;
@@ -412,6 +413,8 @@ pub fn create_signed_vote(
     let pieces_in_sector = farmer_protocol_info.max_pieces_in_sector;
     let sector_size = sector_size(pieces_in_sector);
 
+    let mut table_generator = PosTable::generator();
+
     for sector_index in iter::from_fn(|| Some(rand::random())) {
         let mut plotted_sector_bytes = vec![0; sector_size];
         let mut plotted_sector_metadata_bytes = vec![0; SectorMetadata::encoded_size()];
@@ -427,6 +430,7 @@ pub fn create_signed_vote(
             pieces_in_sector,
             &mut plotted_sector_bytes,
             &mut plotted_sector_metadata_bytes,
+            &mut table_generator,
         ))
         .unwrap();
 
@@ -445,7 +449,7 @@ pub fn create_signed_vote(
         };
 
         let solution = solution_candidates
-            .into_iter::<_, PosTable>(&reward_address, kzg, erasure_coding)
+            .into_iter::<_, PosTable>(&reward_address, kzg, erasure_coding, &mut table_generator)
             .unwrap()
             .next()
             .unwrap()

crates/sp-lightclient/src/tests.rs

@@ -33,6 +33,7 @@ use subspace_farmer_components::auditing::audit_sector;
 use subspace_farmer_components::plotting::{plot_sector, PieceGetterRetryPolicy};
 use subspace_farmer_components::sector::{sector_size, SectorMetadata};
 use subspace_farmer_components::FarmerProtocolInfo;
+use subspace_proof_of_space::Table;
 use subspace_solving::REWARD_SIGNING_CONTEXT;
 use subspace_verification::{
     calculate_block_weight, derive_randomness, verify_solution, VerifySolutionParams,
@@ -154,6 +155,8 @@ fn valid_header(
     let pieces_in_sector = farmer_parameters.farmer_protocol_info.max_pieces_in_sector;
     let sector_size = sector_size(pieces_in_sector);
 
+    let mut table_generator = PosTable::generator();
+
     for sector_index in iter::from_fn(|| Some(rand::random())) {
         let mut plotted_sector_bytes = vec![0; sector_size];
         let mut plotted_sector_metadata_bytes = vec![0; SectorMetadata::encoded_size()];
@@ -169,6 +172,7 @@ fn valid_header(
             pieces_in_sector,
             &mut plotted_sector_bytes,
             &mut plotted_sector_metadata_bytes,
+            &mut table_generator,
         ))
         .unwrap();
 
@@ -193,6 +197,7 @@ fn valid_header(
                 &public_key,
                 &farmer_parameters.kzg,
                 &farmer_parameters.erasure_coding,
+                &mut table_generator,
             )
             .unwrap()
             .next()

crates/subspace-farmer-components/benches/auditing.rs

@@ -21,6 +21,7 @@ use subspace_farmer_components::plotting::{plot_sector, PieceGetterRetryPolicy,
 use subspace_farmer_components::sector::{sector_size, SectorContentsMap, SectorMetadata};
 use subspace_farmer_components::FarmerProtocolInfo;
 use subspace_proof_of_space::chia::ChiaTable;
+use subspace_proof_of_space::Table;
 
 type PosTable = ChiaTable;
 
@@ -51,6 +52,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         NonZeroUsize::new(Record::NUM_S_BUCKETS.next_power_of_two().ilog2() as usize).unwrap(),
     )
     .unwrap();
+    let mut table_generator = PosTable::generator();
     let archived_history_segment = archiver
         .add_block(
             AsRef::<[u8]>::as_ref(input.as_ref()).to_vec(),
@@ -123,6 +125,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             pieces_in_sector,
             &mut plotted_sector_bytes,
             &mut plotted_sector_metadata_bytes,
+            &mut table_generator,
         ))
         .unwrap();
 

crates/subspace-farmer-components/benches/plotting.rs

@@ -12,6 +12,7 @@ use subspace_farmer_components::plotting::{plot_sector, PieceGetterRetryPolicy};
 use subspace_farmer_components::sector::{sector_size, SectorMetadata};
 use subspace_farmer_components::FarmerProtocolInfo;
 use subspace_proof_of_space::chia::ChiaTable;
+use subspace_proof_of_space::Table;
 
 type PosTable = ChiaTable;
 
@@ -33,6 +34,7 @@ fn criterion_benchmark(c: &mut Criterion) {
         NonZeroUsize::new(Record::NUM_S_BUCKETS.next_power_of_two().ilog2() as usize).unwrap(),
     )
     .unwrap();
+    let mut table_generator = PosTable::generator();
     let archived_history_segment = archiver
         .add_block(
             AsRef::<[u8]>::as_ref(input.as_ref()).to_vec(),
@@ -73,6 +75,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                 black_box(pieces_in_sector),
                 black_box(&mut sector_bytes),
                 black_box(&mut sector_metadata_bytes),
+                black_box(&mut table_generator),
             ))
             .unwrap();
         })

crates/subspace-farmer-components/benches/proving.rs

@@ -21,6 +21,7 @@ use subspace_farmer_components::plotting::{plot_sector, PieceGetterRetryPolicy,
 use subspace_farmer_components::sector::{sector_size, SectorContentsMap, SectorMetadata};
 use subspace_farmer_components::FarmerProtocolInfo;
 use subspace_proof_of_space::chia::ChiaTable;
+use subspace_proof_of_space::Table;
 
 type PosTable = ChiaTable;
 
@@ -53,6 +54,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         NonZeroUsize::new(Record::NUM_S_BUCKETS.next_power_of_two().ilog2() as usize).unwrap(),
     )
     .unwrap();
+    let mut table_generator = PosTable::generator();
     let archived_history_segment = archiver
         .add_block(
             AsRef::<[u8]>::as_ref(input.as_ref()).to_vec(),
@@ -125,6 +127,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             pieces_in_sector,
             &mut plotted_sector_bytes,
             &mut plotted_sector_metadata_bytes,
+            &mut table_generator,
         ))
         .unwrap();
 
@@ -164,7 +167,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
 
         let num_actual_solutions = solution_candidates
             .clone()
-            .into_iter::<_, PosTable>(&reward_address, &kzg, &erasure_coding)
+            .into_iter::<_, PosTable>(&reward_address, &kzg, &erasure_coding, &mut table_generator)
             .unwrap()
             .len();
 
@@ -194,6 +197,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
                         black_box(&reward_address),
                         black_box(&kzg),
                         black_box(&erasure_coding),
+                        black_box(&mut table_generator),
                     )
                     .unwrap()
                     // Process just one solution
@@ -260,6 +264,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
                                 black_box(&reward_address),
                                 black_box(&kzg),
                                 black_box(&erasure_coding),
+                                black_box(&mut table_generator),
                             )
                             .unwrap()
                             // Process just one solution

crates/subspace-farmer-components/benches/reading.rs

@@ -20,6 +20,7 @@ use subspace_farmer_components::reading::read_piece;
 use subspace_farmer_components::sector::{sector_size, SectorContentsMap, SectorMetadata};
 use subspace_farmer_components::FarmerProtocolInfo;
 use subspace_proof_of_space::chia::ChiaTable;
+use subspace_proof_of_space::Table;
 
 type PosTable = ChiaTable;
 
@@ -50,6 +51,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         NonZeroUsize::new(Record::NUM_S_BUCKETS.next_power_of_two().ilog2() as usize).unwrap(),
     )
     .unwrap();
+    let mut table_generator = PosTable::generator();
     let archived_history_segment = archiver
         .add_block(
             AsRef::<[u8]>::as_ref(input.as_ref()).to_vec(),
@@ -120,6 +122,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             pieces_in_sector,
             &mut plotted_sector_bytes,
             &mut plotted_sector_metadata_bytes,
+            &mut table_generator,
         ))
         .unwrap();
 
@@ -148,6 +151,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
                 black_box(&plotted_sector.sector_metadata),
                 black_box(&plotted_sector_bytes),
                 black_box(&erasure_coding),
+                black_box(&mut table_generator),
             )
             .unwrap();
         })
@@ -195,6 +199,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
                             black_box(&plotted_sector.sector_metadata),
                             black_box(sector),
                             black_box(&erasure_coding),
+                            black_box(&mut table_generator),
                         )
                         .unwrap();
                     }

crates/subspace-farmer-components/src/plotting.rs

@@ -11,7 +11,6 @@ use futures::stream::FuturesOrdered;
 use futures::StreamExt;
 use parity_scale_codec::Encode;
 use parking_lot::Mutex;
-use rayon::prelude::*;
 use std::error::Error;
 use std::simd::Simd;
 use std::sync::Arc;
@@ -23,7 +22,7 @@ use subspace_core_primitives::{
     RecordWitness, SBucket, SectorId, SectorIndex,
 };
 use subspace_erasure_coding::ErasureCoding;
-use subspace_proof_of_space::{Quality, Table};
+use subspace_proof_of_space::{Quality, Table, TableGenerator};
 use thiserror::Error;
 use tokio::sync::Semaphore;
 use tracing::{debug, warn};
@@ -169,6 +168,7 @@ pub async fn plot_sector<PG, PosTable>(
     pieces_in_sector: u16,
     sector_output: &mut [u8],
     sector_metadata_output: &mut [u8],
+    table_generator: &mut PosTable::Generator,
 ) -> Result<PlottedSector, PlottingError>
 where
     PG: PieceGetter,
@@ -248,11 +248,13 @@ where
     (PieceOffset::ZERO..)
         .zip(raw_sector.records.iter_mut())
         .zip(sector_contents_map.iter_record_bitfields_mut())
-        // TODO: Doesn't work without a bridge: https://github.com/ferrilab/bitvec/issues/143
-        .par_bridge()
+        // TODO: Ideally, we'd use parallelism here, but using `.par_bridge()` causes Chia table
+        //  derivation to only use a single thread, which slows everything to essentially
+        //  single-threaded
         .for_each(|((piece_offset, record), mut encoded_chunks_used)| {
-            // Derive PoSpace table
-            let pos_table = PosTable::generate(
+            // Derive PoSpace table (use parallel mode because multiple tables concurrently will use
+            // too much RAM)
+            let pos_table = table_generator.generate_parallel(
                 &sector_id.derive_evaluation_seed(piece_offset, farmer_protocol_info.history_size),
             );
 

crates/subspace-farmer-components/src/proving.rs

@@ -2,14 +2,13 @@ use crate::auditing::ChunkCandidate;
 use crate::reading::{read_record_metadata, read_sector_record_chunks, ReadingError};
 use crate::sector::{SectorContentsMap, SectorContentsMapFromBytesError, SectorMetadata};
 use std::collections::VecDeque;
-use std::marker::PhantomData;
 use subspace_core_primitives::crypto::kzg::{Commitment, Kzg, Witness};
 use subspace_core_primitives::crypto::Scalar;
 use subspace_core_primitives::{
     PieceOffset, PosProof, PublicKey, Record, SBucket, SectorId, SectorIndex, Solution,
 };
 use subspace_erasure_coding::ErasureCoding;
-use subspace_proof_of_space::{Quality, Table};
+use subspace_proof_of_space::{Quality, Table, TableGenerator};
 use thiserror::Error;
 
 /// Errors that happen during proving
@@ -116,6 +115,7 @@ impl<'a> SolutionCandidates<'a> {
         reward_address: &'a RewardAddress,
         kzg: &'a Kzg,
         erasure_coding: &'a ErasureCoding,
+        table_generator: &'a mut PosTable::Generator,
     ) -> Result<
         impl ExactSizeIterator<Item = Result<Solution<PublicKey, RewardAddress>, ProvingError>> + 'a,
         ProvingError,
@@ -135,6 +135,7 @@ impl<'a> SolutionCandidates<'a> {
             kzg,
             erasure_coding,
             self.chunk_candidates,
+            table_generator,
         )
     }
 }
@@ -148,7 +149,10 @@ struct ChunkCache {
     proof_of_space: PosProof,
 }
 
-struct SolutionCandidatesIterator<'a, RewardAddress, PosTable> {
+struct SolutionCandidatesIterator<'a, RewardAddress, PosTable>
+where
+    PosTable: Table,
+{
     public_key: &'a PublicKey,
     reward_address: &'a RewardAddress,
     sector_index: SectorIndex,
@@ -163,7 +167,7 @@ struct SolutionCandidatesIterator<'a, RewardAddress, PosTable> {
     winning_chunks: VecDeque<WinningChunk>,
     count: usize,
     chunk_cache: Option<ChunkCache>,
-    _pos_table: PhantomData<PosTable>,
+    table_generator: &'a mut PosTable::Generator,
 }
 
 // TODO: This can be potentially parallelized with rayon
@@ -200,7 +204,7 @@ where
             }
 
             // Derive PoSpace table
-            let pos_table = PosTable::generate_parallel(
+            let pos_table = self.table_generator.generate_parallel(
                 &self
                     .sector_id
                     .derive_evaluation_seed(piece_offset, self.sector_metadata.history_size),
@@ -351,6 +355,7 @@ where
         kzg: &'a Kzg,
         erasure_coding: &'a ErasureCoding,
         chunk_candidates: VecDeque<ChunkCandidate>,
+        table_generator: &'a mut PosTable::Generator,
     ) -> Result<Self, ProvingError> {
         if erasure_coding.max_shards() < Record::NUM_S_BUCKETS {
             return Err(ProvingError::InvalidErasureCodingInstance);
@@ -411,7 +416,7 @@ where
             winning_chunks,
             count,
             chunk_cache: None,
-            _pos_table: PhantomData,
+            table_generator,
         })
     }
 }

crates/subspace-farmer-components/src/reading.rs

@@ -10,7 +10,7 @@ use subspace_core_primitives::{
     Piece, PieceOffset, Record, RecordCommitment, RecordWitness, SBucket, SectorId,
 };
 use subspace_erasure_coding::ErasureCoding;
-use subspace_proof_of_space::{Quality, Table};
+use subspace_proof_of_space::{Quality, Table, TableGenerator};
 use thiserror::Error;
 
 /// Errors that happen during reading
@@ -257,6 +257,7 @@ pub fn read_piece<PosTable>(
     sector_metadata: &SectorMetadata,
     sector: &[u8],
     erasure_coding: &ErasureCoding,
+    table_generator: &mut PosTable::Generator,
 ) -> Result<Piece, ReadingError>
 where
     PosTable: Table,
@@ -284,7 +285,7 @@ where
             pieces_in_sector,
             &sector_metadata.s_bucket_offsets(),
             &sector_contents_map,
-            &PosTable::generate(
+            &table_generator.generate(
                 &sector_id.derive_evaluation_seed(piece_offset, sector_metadata.history_size),
             ),
             sector,

crates/subspace-farmer/src/single_disk_plot/farming.rs

@@ -83,6 +83,8 @@ where
     NC: NodeClient,
     PosTable: Table,
 {
+    let mut table_generator = PosTable::generator();
+
     while let Some(slot_info) = slot_info_notifications.next().await {
         let slot = slot_info.slot_number;
         let sectors_metadata = sectors_metadata.read();
@@ -120,6 +122,7 @@ where
                 &reward_address,
                 &kzg,
                 &erasure_coding,
+                &mut table_generator,
             )? {
                 let solution = match maybe_solution {
                     Ok(solution) => solution,