[wip] Preallocate MSM points

Cargo.toml

@@ -33,7 +33,7 @@ cuda-mobile = []
 [dependencies]
 blst = "~0.3.11"
 semolina = "~0.1.3"
-sppark = "~0.1.2"
+sppark = { git = "https://github.com/lurk-lab/sppark", branch = "preallocate-msm" }
 halo2curves = { version = "0.6.0" }
 pasta_curves = { git = "https://github.com/lurk-lab/pasta_curves", branch = "dev", version = ">=0.3.1, <=0.5", features = ["repr-c"] }
 rand = "^0"

benches/grumpkin_msm.rs

@@ -30,10 +30,21 @@ fn criterion_benchmark(c: &mut Criterion) {
 
     group.bench_function(format!("2**{} points", bench_npow), |b| {
         b.iter(|| {
-            let _ = grumpkin_msm::bn256(&points, &scalars);
+            let _ = grumpkin_msm::bn256::msm(&points, &scalars);
         })
     });
 
+    let context = grumpkin_msm::bn256::init(&points);
+
+    group.bench_function(
+        format!("\"preallocate\" 2**{} points", bench_npow),
+        |b| {
+            b.iter(|| {
+                let _ = grumpkin_msm::bn256::with(&context, &scalars);
+            })
+        },
+    );
+
     group.finish();
 
     #[cfg(feature = "cuda")]
@@ -54,10 +65,22 @@ fn criterion_benchmark(c: &mut Criterion) {
 
         group.bench_function(format!("2**{} points", bench_npow), |b| {
             b.iter(|| {
-                let _ = grumpkin_msm::bn256(&points, &scalars);
+                let _ = grumpkin_msm::bn256::msm(&points, &scalars);
             })
         });
 
+        let context = grumpkin_msm::bn256::init(&points);
+
+        group.bench_function(
+            format!("preallocate 2**{} points", bench_npow),
+            |b| {
+                b.iter(|| {
+                    let _ =
+                        grumpkin_msm::bn256::with(&context, &scalars);
+                })
+            },
+        );
+
         group.finish();
     }
 }

benches/pasta_msm.rs

@@ -30,10 +30,21 @@ fn criterion_benchmark(c: &mut Criterion) {
 
     group.bench_function(format!("2**{} points", bench_npow), |b| {
         b.iter(|| {
-            let _ = grumpkin_msm::pasta::pallas(&points, &scalars);
+            let _ = grumpkin_msm::pasta::pallas::msm(&points, &scalars);
         })
     });
 
+    let context = grumpkin_msm::pasta::pallas::init(&points);
+
+    group.bench_function(
+        format!("\"preallocate\" 2**{} points", bench_npow),
+        |b| {
+            b.iter(|| {
+                let _ = grumpkin_msm::pasta::pallas::with(&context, &scalars);
+            })
+        },
+    );
+
     group.finish();
 
     #[cfg(feature = "cuda")]
@@ -54,10 +65,22 @@ fn criterion_benchmark(c: &mut Criterion) {
 
         group.bench_function(format!("2**{} points", bench_npow), |b| {
             b.iter(|| {
-                let _ = grumpkin_msm::pasta::pallas(&points, &scalars);
+                let _ = grumpkin_msm::pasta::pallas::msm(&points, &scalars);
             })
         });
 
+        let context = grumpkin_msm::pasta::pallas::init(&points);
+
+        group.bench_function(
+            format!("preallocate 2**{} points", bench_npow),
+            |b| {
+                b.iter(|| {
+                    let _ =
+                        grumpkin_msm::pasta::pallas::with(&context, &scalars);
+                })
+            },
+        );
+
         group.finish();
     }
 }

cuda/bn254.cu

@@ -17,8 +17,26 @@ typedef fr_t scalar_t;
 #include <msm/pippenger.cuh>
 
 #ifndef __CUDA_ARCH__
-extern "C"
-RustError cuda_pippenger_bn254(point_t *out, const affine_t points[], size_t npoints,
-                                              const scalar_t scalars[])
-{   return mult_pippenger<bucket_t>(out, points, npoints, scalars);   }
+
+extern "C" void drop_msm_context_bn254(msm_context_t<affine_t::mem_t> &ref) {
+    CUDA_OK(cudaFree(ref.d_points));
+}
+
+extern "C" RustError
+cuda_bn254_init(const affine_t points[], size_t npoints, msm_context_t<affine_t::mem_t> *msm_context)
+{
+    return mult_pippenger_init<bucket_t, point_t, affine_t, scalar_t>(points, npoints, msm_context);
+}
+
+extern "C" RustError cuda_bn254(point_t *out, const affine_t points[], size_t npoints,
+                                          const scalar_t scalars[])
+{
+    return mult_pippenger<bucket_t>(out, points, npoints, scalars);
+}
+
+extern "C" RustError cuda_bn254_with(point_t *out, msm_context_t<affine_t::mem_t> *msm_context, size_t npoints,
+                                                       const scalar_t scalars[])
+{
+    return mult_pippenger_with<bucket_t, point_t, affine_t, scalar_t>(out, msm_context, npoints, scalars);
+}
 #endif

cuda/grumpkin.cu

@@ -17,8 +17,26 @@ typedef fp_t scalar_t;
 #include <msm/pippenger.cuh>
 
 #ifndef __CUDA_ARCH__
-extern "C"
-RustError cuda_pippenger_grumpkin(point_t *out, const affine_t points[], size_t npoints,
-                                              const scalar_t scalars[])
-{   return mult_pippenger<bucket_t>(out, points, npoints, scalars);   }
+
+extern "C" void drop_msm_context_grumpkin(msm_context_t<affine_t::mem_t> &ref) {
+    CUDA_OK(cudaFree(ref.d_points));
+}
+
+extern "C" RustError
+cuda_grumpkin_init(const affine_t points[], size_t npoints, msm_context_t<affine_t::mem_t> *msm_context)
+{
+    return mult_pippenger_init<bucket_t, point_t, affine_t, scalar_t>(points, npoints, msm_context);
+}
+
+extern "C" RustError cuda_grumpkin(point_t *out, const affine_t points[], size_t npoints,
+                                          const scalar_t scalars[])
+{
+    return mult_pippenger<bucket_t>(out, points, npoints, scalars);
+}
+
+extern "C" RustError cuda_grumpkin_with(point_t *out, msm_context_t<affine_t::mem_t> *msm_context, size_t npoints,
+                                                       const scalar_t scalars[])
+{
+    return mult_pippenger_with<bucket_t, point_t, affine_t, scalar_t>(out, msm_context, npoints, scalars);
+}
 #endif

cuda/pallas.cu

@@ -17,8 +17,27 @@ typedef vesta_t scalar_t;
 #include <msm/pippenger.cuh>
 
 #ifndef __CUDA_ARCH__
-extern "C"
-RustError cuda_pippenger_pallas(point_t *out, const affine_t points[], size_t npoints,
-                                              const scalar_t scalars[])
-{   return mult_pippenger<bucket_t>(out, points, npoints, scalars);   }
+
+extern "C" void drop_msm_context_pallas(msm_context_t<affine_t::mem_t> &ref) {
+    CUDA_OK(cudaFree(ref.d_points));
+}
+
+extern "C" RustError
+cuda_pallas_init(const affine_t points[], size_t npoints, msm_context_t<affine_t::mem_t> *msm_context)
+{
+    return mult_pippenger_init<bucket_t, point_t, affine_t, scalar_t>(points, npoints, msm_context);
+}
+
+extern "C" RustError cuda_pallas(point_t *out, const affine_t points[], size_t npoints,
+                                          const scalar_t scalars[])
+{
+    return mult_pippenger<bucket_t>(out, points, npoints, scalars);
+}
+
+extern "C" RustError cuda_pallas_with(point_t *out, msm_context_t<affine_t::mem_t> *msm_context, size_t npoints,
+                                                       const scalar_t scalars[])
+{
+    return mult_pippenger_with<bucket_t, point_t, affine_t, scalar_t>(out, msm_context, npoints, scalars);
+}
+
 #endif

cuda/vesta.cu

@@ -17,8 +17,27 @@ typedef pallas_t scalar_t;
 #include <msm/pippenger.cuh>
 
 #ifndef __CUDA_ARCH__
-extern "C"
-RustError cuda_pippenger_vesta(point_t *out, const affine_t points[], size_t npoints,
-                                              const scalar_t scalars[])
-{   return mult_pippenger<bucket_t>(out, points, npoints, scalars);   }
+
+extern "C" void drop_msm_context_vesta(msm_context_t<affine_t::mem_t> &ref) {
+    CUDA_OK(cudaFree(ref.d_points));
+}
+
+extern "C" RustError
+cuda_vesta_init(const affine_t points[], size_t npoints, msm_context_t<affine_t::mem_t> *msm_context)
+{
+    return mult_pippenger_init<bucket_t, point_t, affine_t, scalar_t>(points, npoints, msm_context);
+}
+
+extern "C" RustError cuda_vesta(point_t *out, const affine_t points[], size_t npoints,
+                                          const scalar_t scalars[])
+{
+    return mult_pippenger<bucket_t>(out, points, npoints, scalars);
+}
+
+extern "C" RustError cuda_vesta_with(point_t *out, msm_context_t<affine_t::mem_t> *msm_context, size_t npoints,
+                                                       const scalar_t scalars[])
+{
+    return mult_pippenger_with<bucket_t, point_t, affine_t, scalar_t>(out, msm_context, npoints, scalars);
+}
+
 #endif

examples/grumpkin_msm.rs

@@ -20,7 +20,7 @@ fn main() {
         unsafe { grumpkin_msm::CUDA_OFF = false };
     }
 
-    let res = grumpkin_msm::bn256(&points, &scalars).to_affine();
+    let res = grumpkin_msm::bn256::msm(&points, &scalars).to_affine();
     let native = naive_multiscalar_mul(&points, &scalars);
     assert_eq!(res, native);
     println!("success!")

examples/pasta_msm.rs

@@ -8,7 +8,7 @@ use pasta_curves::group::Curve;
 
 fn main() {
     let bench_npow: usize = std::env::var("BENCH_NPOW")
-        .unwrap_or("17".to_string())
+        .unwrap_or("22".to_string())
         .parse()
         .unwrap();
     let npoints: usize = 1 << bench_npow;
@@ -22,8 +22,10 @@ fn main() {
         unsafe { grumpkin_msm::CUDA_OFF = false };
     }
 
-    let res = grumpkin_msm::pasta::pallas(&points, &scalars).to_affine();
     let native = naive_multiscalar_mul(&points, &scalars);
+    let context = grumpkin_msm::pasta::pallas::init(&points);
+    let res = grumpkin_msm::pasta::pallas::with(&context, &scalars).to_affine();
+
     assert_eq!(res, native);
     println!("success!")
 }