Merge pull request #1453 from subspace/parallel-chiapos

Parallel chiapos

crates/subspace-core-primitives/Cargo.toml

@@ -26,7 +26,7 @@ kzg = { git = "https://github.com/subspace/rust-kzg", rev = "1058cc8c8af8461b490
 num-traits = { version = "0.2.15", default-features = false }
 parity-scale-codec = { version = "3.4.0", default-features = false, features = ["derive", "max-encoded-len"] }
 parking_lot = { version = "0.12.1", optional = true }
-rayon = { version = "1.6.1", optional = true }
+rayon = { version = "1.7.0", optional = true }
 scale-info = { version = "2.5.0", default-features = false, features = ["derive"] }
 serde = { version = "1.0.159", optional = true, features = ["alloc", "derive"] }
 serde_arrays = { version = "0.1.0", optional = true }

crates/subspace-farmer-components/Cargo.toml

@@ -26,14 +26,14 @@ lru = "0.10.0"
 parity-scale-codec = "3.4.0"
 parking_lot = "0.12.1"
 rand = "0.8.5"
-rayon = "1.6.1"
+rayon = "1.7.0"
 schnorrkel = "0.9.1"
 serde = { version = "1.0.159", features = ["derive"] }
 static_assertions = "1.1.0"
 subspace-archiving = { version = "0.1.0", path = "../subspace-archiving" }
 subspace-core-primitives = { version = "0.1.0", path = "../subspace-core-primitives" }
 subspace-erasure-coding = { version = "0.1.0", path = "../subspace-erasure-coding" }
-subspace-proof-of-space = { version = "0.1.0", path = "../subspace-proof-of-space" }
+subspace-proof-of-space = { version = "0.1.0", path = "../subspace-proof-of-space", features = ["parallel"] }
 subspace-verification = { version = "0.1.0", path = "../subspace-verification" }
 thiserror = "1.0.38"
 tokio = { version = "1.27.0", features = ["macros", "parking_lot", "rt-multi-thread", "signal", "sync"] }

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

@@ -199,10 +199,8 @@ where
                 }
             }
 
-            // TODO: Can potentially happen in parallel with other work below, saving a bit of
-            //  end-to-end latency
             // Derive PoSpace table
-            let pos_table = PosTable::generate(
+            let pos_table = PosTable::generate_parallel(
                 &self
                     .sector_id
                     .evaluation_seed(piece_offset, self.sector_metadata.history_size),

crates/subspace-farmer/Cargo.toml

@@ -61,4 +61,4 @@ zeroize = "1.6.0"
 jemallocator = "0.5.0"
 
 [dev-dependencies]
-rayon = "1.6.1"
+rayon = "1.7.0"

crates/subspace-proof-of-space/Cargo.toml

@@ -19,6 +19,7 @@ bench = false
 bitvec = { version = "1.0.1", default-features = false, features = ["alloc", "atomic"], optional = true }
 blake3 = { version = "1.3.3", default-features = false, optional = true }
 chacha20 = { version = "0.9.1", default-features = false, optional = true }
+rayon = { version = "1.7.0", optional = true }
 subspace-chiapos = { git = "https://github.com/subspace/chiapos", rev = "3b1ab3ca24764d25da30e0c8243e0bf304b776a5", optional = true }
 subspace-core-primitives = { version = "0.1.0", path = "../subspace-core-primitives", default-features = false }
 
@@ -52,6 +53,9 @@ std = [
     "chacha20?/std",
     "subspace-core-primitives/std",
 ]
+parallel = [
+    "dep:rayon",
+]
 # Enable Chia proof of space support (legacy implementation uses C++ chiapos), only works in `std` environment for now
 chia-legacy = [
     "std",

crates/subspace-proof-of-space/benches/pos.rs

@@ -3,6 +3,8 @@
 #[cfg(any(feature = "chia-legacy", feature = "chia", feature = "shim"))]
 use criterion::black_box;
 use criterion::{criterion_group, criterion_main, Criterion};
+#[cfg(feature = "parallel")]
+use rayon::ThreadPoolBuilder;
 #[cfg(any(feature = "chia-legacy", feature = "chia", feature = "shim"))]
 use subspace_core_primitives::PosSeed;
 #[cfg(any(feature = "chia-legacy", feature = "chia", feature = "shim"))]
@@ -23,14 +25,32 @@ fn pos_bench<PosTable>(
 ) where
     PosTable: Table,
 {
+    #[cfg(feature = "parallel")]
+    {
+        // Repeated initialization is not supported, we just ignore errors here because of it
+        let _ = ThreadPoolBuilder::new()
+            // Change number of threads if necessary
+            .num_threads(4)
+            .build_global();
+    }
+
     let mut group = c.benchmark_group(name);
 
-    group.bench_function("table", |b| {
+    group.bench_function("table/single", |b| {
         b.iter(|| {
             PosTable::generate(black_box(&SEED));
         });
     });
 
+    #[cfg(feature = "parallel")]
+    {
+        group.bench_function("table/parallel", |b| {
+            b.iter(|| {
+                PosTable::generate_parallel(black_box(&SEED));
+            });
+        });
+    }
+
     let table = PosTable::generate(&SEED);
 
     group.bench_function("quality/no-solution", |b| {

crates/subspace-proof-of-space/src/chia.rs

@@ -1,5 +1,7 @@
 //! Chia proof of space implementation
 use crate::chiapos::Tables;
+#[cfg(any(feature = "parallel", test))]
+use crate::chiapos::TablesCache;
 use crate::{PosTableType, Quality, Table};
 use core::mem;
 use subspace_core_primitives::{PosProof, PosQualityBytes, PosSeed};
@@ -50,6 +52,13 @@ impl Table for ChiaTable {
         }
     }
 
+    #[cfg(any(feature = "parallel", test))]
+    fn generate_parallel(seed: &PosSeed) -> ChiaTable {
+        Self {
+            tables: Tables::<K>::create_parallel((*seed).into(), &mut TablesCache::default()),
+        }
+    }
+
     fn find_quality(&self, challenge_index: u32) -> Option<Self::Quality<'_>> {
         let mut challenge = [0; 32];
         challenge[..mem::size_of::<u32>()].copy_from_slice(&challenge_index.to_le_bytes());
@@ -84,12 +93,21 @@ mod tests {
     #[test]
     fn basic() {
         let table = ChiaTable::generate(&SEED);
+        let table_parallel = ChiaTable::generate_parallel(&SEED);
 
         assert!(table.find_quality(1232460437).is_none());
+        assert!(table_parallel.find_quality(1232460437).is_none());
 
         {
             let challenge_index = 124537303;
             let quality = table.find_quality(challenge_index).unwrap();
+            assert_eq!(
+                quality.to_bytes(),
+                table_parallel
+                    .find_quality(challenge_index)
+                    .unwrap()
+                    .to_bytes()
+            );
             let proof = quality.create_proof();
             let maybe_quality = ChiaTable::is_proof_valid(&SEED, challenge_index, &proof);
             assert_eq!(maybe_quality, Some(quality.to_bytes()));

crates/subspace-proof-of-space/src/chiapos.rs

@@ -37,7 +37,8 @@ macro_rules! impl_any {
     ($($k: expr$(,)? )*) => {
         $(
 impl Tables<$k> {
-    /// Create Chia proof of space tables.
+    /// Create Chia proof of space tables. There also exists [`Self::create_parallel()`] that trades
+    /// CPU efficiency and memory usage for lower latency.
     ///
     /// Advanced version of [`Self::create_simple`] that allows to reuse cache.
     pub fn create(seed: Seed, cache: &mut TablesCache<$k>) -> Self {
@@ -46,6 +47,16 @@ impl Tables<$k> {
         ))
     }
 
+    /// Almost the same as [`Self::create()`], but uses parallelism internally for better
+    /// performance (though not efficiency of CPU and memory usage), if you create multiple tables
+    /// in parallel, prefer [`Self::create()`] for better overall performance.
+    #[cfg(any(feature = "parallel", test))]
+    pub fn create_parallel(seed: Seed, cache: &mut TablesCache<$k>) -> Self {
+        Self(TablesGeneric::<$k>::create_parallel(
+            seed, cache,
+        ))
+    }
+
     /// Create Chia proof of space tables.
     ///
     /// Simpler version of [`Self::create`].