speed up weierstrass

succinctlabs · Feb 29, 2024 · 6403426 · 6403426
1 parent 5026c25
commit 6403426
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 50 deletions.
diff --git a/core/src/syscall/precompiles/weierstrass/weierstrass_double.rs b/core/src/syscall/precompiles/weierstrass/weierstrass_double.rs
@@ -27,9 +27,12 @@ use p3_field::AbstractField;
 use p3_field::PrimeField32;
 use p3_matrix::dense::RowMajorMatrix;
 use p3_matrix::MatrixRowSlices;
+use p3_maybe_rayon::prelude::ParallelIterator;
+use p3_maybe_rayon::prelude::ParallelSlice;
 use sp1_derive::AlignedBorrow;
 use std::fmt::Debug;
 use std::marker::PhantomData;
+use tracing::instrument;
 
 pub const NUM_WEIERSTRASS_DOUBLE_COLS: usize = size_of::<WeierstrassDoubleAssignCols<u8>>();
 
@@ -162,41 +165,64 @@ impl<F: PrimeField32, E: EllipticCurve + WeierstrassParameters> MachineAir<F>
         "WeierstrassDoubleAssign".to_string()
     }
 
+    #[instrument(name = "generate WeierstrassDoubleAssign trace", skip_all)]
     fn generate_trace(
         &self,
         input: &ExecutionRecord,
         output: &mut ExecutionRecord,
     ) -> RowMajorMatrix<F> {
-        let mut rows = Vec::new();
-
-        let mut new_field_events = Vec::new();
-
-        for i in 0..input.weierstrass_double_events.len() {
-            let event = input.weierstrass_double_events[i];
-            let mut row = [F::zero(); NUM_WEIERSTRASS_DOUBLE_COLS];
-            let cols: &mut WeierstrassDoubleAssignCols<F> = row.as_mut_slice().borrow_mut();
-
-            // Decode affine points.
-            let p = &event.p;
-            let p = AffinePoint::<E>::from_words_le(p);
-            let (p_x, p_y) = (p.x, p.y);
-
-            // Populate basic columns.
-            cols.is_real = F::one();
-            cols.shard = F::from_canonical_u32(event.shard);
-            cols.clk = F::from_canonical_u32(event.clk);
-            cols.p_ptr = F::from_canonical_u32(event.p_ptr);
-
-            Self::populate_field_ops(cols, p_x, p_y);
-
-            // Populate the memory access columns.
-            for i in 0..NUM_WORDS_EC_POINT {
-                cols.p_access[i].populate(event.p_memory_records[i], &mut new_field_events);
-            }
+        let chunk_size = std::cmp::max(
+            input.weierstrass_double_events.len() / (num_cpus::get() * 2),
+            1,
+        );
 
-            rows.push(row);
+        // Generate the trace rows & corresponding records for each chunk of events in parallel.
+        let rows_and_records = input
+            .weierstrass_double_events
+            .par_chunks(chunk_size)
+            .map(|events| {
+                let mut record = ExecutionRecord::default();
+                let mut new_field_events = Vec::new();
+
+                let rows = events
+                    .iter()
+                    .map(|event| {
+                        let mut row = [F::zero(); NUM_WEIERSTRASS_DOUBLE_COLS];
+                        let cols: &mut WeierstrassDoubleAssignCols<F> =
+                            row.as_mut_slice().borrow_mut();
+
+                        // Decode affine points.
+                        let p = &event.p;
+                        let p = AffinePoint::<E>::from_words_le(p);
+                        let (p_x, p_y) = (p.x, p.y);
+
+                        // Populate basic columns.
+                        cols.is_real = F::one();
+                        cols.shard = F::from_canonical_u32(event.shard);
+                        cols.clk = F::from_canonical_u32(event.clk);
+                        cols.p_ptr = F::from_canonical_u32(event.p_ptr);
+
+                        Self::populate_field_ops(cols, p_x, p_y);
+
+                        // Populate the memory access columns.
+                        for i in 0..NUM_WORDS_EC_POINT {
+                            cols.p_access[i]
+                                .populate(event.p_memory_records[i], &mut new_field_events);
+                        }
+                        row
+                    })
+                    .collect::<Vec<_>>();
+                record.add_field_events(&new_field_events);
+                (rows, record)
+            })
+            .collect::<Vec<_>>();
+
+        // Generate the trace rows for each event.
+        let mut rows = Vec::new();
+        for mut row_and_record in rows_and_records {
+            rows.extend(row_and_record.0);
+            output.append(&mut row_and_record.1);
         }
-        output.add_field_events(&new_field_events);
 
         pad_rows(&mut rows, || {
             let mut row = [F::zero(); NUM_WEIERSTRASS_DOUBLE_COLS];

diff --git a/tests/secp256k1-double/elf/riscv32im-succinct-zkvm-elf b/tests/secp256k1-double/elf/riscv32im-succinct-zkvm-elf
diff --git a/tests/secp256k1-double/src/main.rs b/tests/secp256k1-double/src/main.rs
@@ -6,31 +6,33 @@ extern "C" {
 }
 
 pub fn main() {
-    // generator.
-    // 55066263022277343669578718895168534326250603453777594175500187360389116729240
-    // 32670510020758816978083085130507043184471273380659243275938904335757337482424
-    let mut a: [u8; 64] = [
-        152, 23, 248, 22, 91, 129, 242, 89, 217, 40, 206, 45, 219, 252, 155, 2, 7, 11, 135, 206,
-        149, 98, 160, 85, 172, 187, 220, 249, 126, 102, 190, 121, 184, 212, 16, 251, 143, 208, 71,
-        156, 25, 84, 133, 166, 72, 180, 23, 253, 168, 8, 17, 14, 252, 251, 164, 93, 101, 196, 163,
-        38, 119, 218, 58, 72,
-    ];
+    for _ in 0..10i64.pow(3) {
+        // generator.
+        // 55066263022277343669578718895168534326250603453777594175500187360389116729240
+        // 32670510020758816978083085130507043184471273380659243275938904335757337482424
+        let mut a: [u8; 64] = [
+            152, 23, 248, 22, 91, 129, 242, 89, 217, 40, 206, 45, 219, 252, 155, 2, 7, 11, 135,
+            206, 149, 98, 160, 85, 172, 187, 220, 249, 126, 102, 190, 121, 184, 212, 16, 251, 143,
+            208, 71, 156, 25, 84, 133, 166, 72, 180, 23, 253, 168, 8, 17, 14, 252, 251, 164, 93,
+            101, 196, 163, 38, 119, 218, 58, 72,
+        ];
 
-    unsafe {
-        syscall_secp256k1_double(a.as_mut_ptr() as *mut u32);
-    }
+        unsafe {
+            syscall_secp256k1_double(a.as_mut_ptr() as *mut u32);
+        }
 
-    // 2 * generator.
-    // 89565891926547004231252920425935692360644145829622209833684329913297188986597
-    // 12158399299693830322967808612713398636155367887041628176798871954788371653930
-    let b: [u8; 64] = [
-        229, 158, 112, 92, 185, 9, 172, 171, 167, 60, 239, 140, 75, 142, 119, 92, 216, 124, 192,
-        149, 110, 64, 69, 48, 109, 125, 237, 65, 148, 127, 4, 198, 42, 229, 207, 80, 169, 49, 100,
-        35, 225, 208, 102, 50, 101, 50, 246, 247, 238, 234, 108, 70, 25, 132, 197, 163, 57, 195,
-        61, 166, 254, 104, 225, 26,
-    ];
+        // 2 * generator.
+        // 89565891926547004231252920425935692360644145829622209833684329913297188986597
+        // 12158399299693830322967808612713398636155367887041628176798871954788371653930
+        let b: [u8; 64] = [
+            229, 158, 112, 92, 185, 9, 172, 171, 167, 60, 239, 140, 75, 142, 119, 92, 216, 124,
+            192, 149, 110, 64, 69, 48, 109, 125, 237, 65, 148, 127, 4, 198, 42, 229, 207, 80, 169,
+            49, 100, 35, 225, 208, 102, 50, 101, 50, 246, 247, 238, 234, 108, 70, 25, 132, 197,
+            163, 57, 195, 61, 166, 254, 104, 225, 26,
+        ];
 
-    assert_eq!(a, b);
+        assert_eq!(a, b);
+    }
 
     println!("done");
 }