Merge pull request #3004 from autonomys/gpu-plotting-preparation

GPU plotting preparation

Cargo.toml

@@ -141,8 +141,3 @@ sp-weights = { git = "https://github.com/subspace/polkadot-sdk", rev = "5626154d
 staging-xcm = { git = "https://github.com/subspace/polkadot-sdk", rev = "5626154d0781ac9a6ffd5a6207ed237f425ae631" }
 substrate-prometheus-endpoint = { git = "https://github.com/subspace/polkadot-sdk", rev = "5626154d0781ac9a6ffd5a6207ed237f425ae631" }
 xcm-procedural = { git = "https://github.com/subspace/polkadot-sdk", rev = "5626154d0781ac9a6ffd5a6207ed237f425ae631" }
-
-# TODO: Importing https://github.com/supranational/blst/pull/203 to take advantage of optimizations introduced there,
-#  switch to upstream once merged or once similar performance improvements land upstream
-[patch."https://github.com/supranational/blst.git"]
-blst = { git = "https://github.com/subspace/blst", rev = "ab042e18cb3b62e131423380513964e4b2c7b445" }

crates/subspace-core-primitives/Cargo.toml

@@ -17,15 +17,17 @@ bench = false
 
 [dependencies]
 blake3 = { version = "1.5.3", default-features = false }
+# TODO: Remove once we switch to big-endian
+blst = "0.3.13"
 bytes = { version = "1.7.1", default-features = false }
 derive_more = { version = "1.0.0", default-features = false, features = ["full"] }
 hex = { version = "0.4.3", default-features = false, features = ["alloc"] }
-kzg = { git = "https://github.com/sifraitech/rust-kzg", rev = "c34b73916af9b8a699a74bd0186f82f25e72861c", default-features = false }
+kzg = { git = "https://github.com/grandinetech/rust-kzg", rev = "6c8fcc623df3d7e8c0f30951a49bfea764f90bf4", default-features = false }
 num-traits = { version = "0.2.18", default-features = false }
 parity-scale-codec = { version = "3.6.12", default-features = false, features = ["bytes", "derive", "max-encoded-len"] }
 parking_lot = { version = "0.12.2", optional = true }
 rayon = { version = "1.10.0", optional = true }
-rust-kzg-blst = { git = "https://github.com/sifraitech/rust-kzg", rev = "c34b73916af9b8a699a74bd0186f82f25e72861c", default-features = false }
+rust-kzg-blst = { git = "https://github.com/grandinetech/rust-kzg", rev = "6c8fcc623df3d7e8c0f30951a49bfea764f90bf4", default-features = false }
 scale-info = { version = "2.11.2", default-features = false, features = ["derive"] }
 serde = { version = "1.0.206", optional = true, features = ["alloc", "derive"] }
 serde_arrays = { version = "0.1.0", optional = true }

crates/subspace-core-primitives/src/crypto.rs

@@ -21,12 +21,13 @@ extern crate alloc;
 pub mod kzg;
 
 use crate::Blake3Hash;
-use ::kzg::Fr;
 #[cfg(not(feature = "std"))]
 use alloc::format;
 #[cfg(not(feature = "std"))]
 use alloc::string::String;
 #[cfg(not(feature = "std"))]
+use alloc::string::ToString;
+#[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
 use core::cmp::Ordering;
 use core::hash::{Hash, Hasher};
@@ -207,14 +208,36 @@ impl TryFrom<[u8; Self::FULL_BYTES]> for Scalar {
 
     #[inline]
     fn try_from(value: [u8; Self::FULL_BYTES]) -> Result<Self, Self::Error> {
-        FsFr::from_bytes(&value).map(Scalar)
+        // TODO: The whole method should have been just the following line, but upstream `rust-kzg`
+        //  switched to big-endian and we have to maintain little-endian version for now
+        // FsFr::from_bytes(&value).map(Scalar)
+        let mut bls_scalar = blst::blst_scalar::default();
+        let mut fr = blst::blst_fr::default();
+        unsafe {
+            blst::blst_scalar_from_lendian(&mut bls_scalar, value.as_ptr());
+            if !blst::blst_scalar_fr_check(&bls_scalar) {
+                return Err("Invalid scalar".to_string());
+            }
+            blst::blst_fr_from_scalar(&mut fr, &bls_scalar);
+        }
+        Ok(Self(FsFr(fr)))
     }
 }
 
 impl From<&Scalar> for [u8; Scalar::FULL_BYTES] {
     #[inline]
     fn from(value: &Scalar) -> Self {
-        value.0.to_bytes()
+        // TODO: The whole method should have been just the following line, but upstream `rust-kzg`
+        //  switched to big-endian and we have to maintain little-endian version for now
+        // value.0.to_bytes()
+        let mut scalar = blst::blst_scalar::default();
+        let mut bytes = [0u8; 32];
+        unsafe {
+            blst::blst_scalar_from_fr(&mut scalar, &value.0 .0);
+            blst::blst_lendian_from_scalar(bytes.as_mut_ptr(), &scalar);
+        }
+
+        bytes
     }
 }
 

crates/subspace-core-primitives/src/crypto/kzg.rs

@@ -29,7 +29,7 @@ use spin::Mutex;
 use tracing::debug;
 
 /// Embedded KZG settings as bytes, too big for `no_std` in most cases
-/// Generated with with following command (using current Ethereum KZG Summoning Ceremony):
+/// Generated using following command (using current Ethereum KZG Summoning Ceremony):
 /// ```bash
 /// curl -s https://seq.ceremony.ethereum.org/info/current_state | jq '.transcripts[3].powersOfTau' | jq -r '.G1Powers + .G2Powers | map(.[2:]) | join("")' | xxd -r -p - eth-public-parameters.bin
 /// ```
@@ -72,10 +72,13 @@ pub fn bytes_to_kzg_settings(
 
     // Below is the same as `FsKZGSettings::new(&s1, &s2, num_g1_powers, &fft_settings)`, but without
     // extra checks (parameters are static anyway) and without unnecessary allocations
+    // TODO: Switch to `::new()` constructor once
+    //  https://github.com/grandinetech/rust-kzg/issues/264 is resolved
     Ok(FsKZGSettings {
         fs: fft_settings,
         secret_g1,
         secret_g2,
+        precomputation: None,
     })
 }
 

crates/subspace-core-primitives/src/lib.rs

@@ -164,7 +164,7 @@ pub type SlotNumber = u64;
 pub type SolutionRange = u64;
 
 /// Computes the following:
-/// ```
+/// ```text
 /// MAX * slot_probability / (pieces_in_sector * chunks / s_buckets) / sectors
 /// ```
 pub const fn sectors_to_solution_range(
@@ -183,7 +183,7 @@ pub const fn sectors_to_solution_range(
 }
 
 /// Computes the following:
-/// ```
+/// ```text
 /// MAX * slot_probability / (pieces_in_sector * chunks / s_buckets) / solution_range
 /// ```
 pub const fn solution_range_to_sectors(
@@ -274,8 +274,10 @@ impl Default for PosProof {
 }
 
 impl PosProof {
+    /// Constant K used for proof of space
+    pub const K: u8 = 20;
     /// Size of proof of space proof in bytes.
-    pub const SIZE: usize = 20 * 8;
+    pub const SIZE: usize = Self::K as usize * 8;
 
     /// Proof hash.
     pub fn hash(&self) -> Blake3Hash {

crates/subspace-erasure-coding/Cargo.toml

@@ -15,12 +15,12 @@ include = [
 bench = false
 
 [dependencies]
-kzg = { git = "https://github.com/sifraitech/rust-kzg", rev = "c34b73916af9b8a699a74bd0186f82f25e72861c", default-features = false }
-rust-kzg-blst = { git = "https://github.com/sifraitech/rust-kzg", rev = "c34b73916af9b8a699a74bd0186f82f25e72861c", default-features = false }
+kzg = { git = "https://github.com/grandinetech/rust-kzg", rev = "6c8fcc623df3d7e8c0f30951a49bfea764f90bf4", default-features = false }
+rust-kzg-blst = { git = "https://github.com/grandinetech/rust-kzg", rev = "6c8fcc623df3d7e8c0f30951a49bfea764f90bf4", default-features = false }
 subspace-core-primitives = { version = "0.1.0", path = "../subspace-core-primitives", default-features = false }
 
 [dev-dependencies]
-rust-kzg-blst = { git = "https://github.com/sifraitech/rust-kzg", rev = "c34b73916af9b8a699a74bd0186f82f25e72861c" }
+rust-kzg-blst = { git = "https://github.com/grandinetech/rust-kzg", rev = "6c8fcc623df3d7e8c0f30951a49bfea764f90bf4" }
 criterion = "0.5.1"
 rand = "0.8.5"
 

crates/subspace-farmer/src/plotter.rs

@@ -6,6 +6,7 @@
 //! implementations without the rest of the library being aware of implementation details.
 
 pub mod cpu;
+pub mod pool;
 
 use async_trait::async_trait;
 use futures::channel::mpsc;

crates/subspace-farmer/src/plotter/cpu.rs

@@ -354,28 +354,6 @@ where
                         metrics.plotting_capacity_used.inc();
                     }
 
-                    let plotting_fn = || {
-                        let mut sector = Vec::new();
-
-                        let plotted_sector = encode_sector(
-                            downloaded_sector,
-                            EncodeSectorOptions {
-                                sector_index,
-                                sector_output: &mut sector,
-                                records_encoder: &mut CpuRecordsEncoder::<PosTable>::new(
-                                    &mut (0..record_encoding_concurrency.get())
-                                        .map(|_| PosTable::generator())
-                                        .collect::<Vec<_>>(),
-                                    &erasure_coding,
-                                    &global_mutex,
-                                ),
-                                abort_early: &abort_early,
-                            },
-                        )?;
-
-                        Ok((sector, plotted_sector))
-                    };
-
                     let thread_pool = if replotting {
                         &thread_pools.replotting
                     } else {
@@ -400,8 +378,31 @@ where
 
                     let encoding_start = Instant::now();
 
-                    let plotting_result =
-                        tokio::task::block_in_place(|| thread_pool.install(plotting_fn));
+                    let plotting_result = tokio::task::block_in_place(|| {
+                        thread_pool.install(|| {
+                            let mut sector = Vec::new();
+                            let mut generators = (0..record_encoding_concurrency.get())
+                                .map(|_| PosTable::generator())
+                                .collect::<Vec<_>>();
+                            let mut records_encoder = CpuRecordsEncoder::<PosTable>::new(
+                                &mut generators,
+                                &erasure_coding,
+                                &global_mutex,
+                            );
+
+                            let plotted_sector = encode_sector(
+                                downloaded_sector,
+                                EncodeSectorOptions {
+                                    sector_index,
+                                    sector_output: &mut sector,
+                                    records_encoder: &mut records_encoder,
+                                    abort_early: &abort_early,
+                                },
+                            )?;
+
+                            Ok((sector, plotted_sector))
+                        })
+                    });
                     drop(thread_pools);
                     if let Some(metrics) = &metrics {
                         metrics.plotting_capacity_used.dec();