it works

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+    "files.associations": {
+        "__locale": "cpp",
+        "ios": "cpp"
+    }
+}

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "pasta-msm"
 version = "0.1.4"
-edition = "2018"
+edition = "2021"
 license = "Apache-2.0"
 description = "Optimized multiscalar multiplicaton for Pasta moduli for x86_64 and aarch64"
 repository = "https://github.com/supranational/pasta-msm"
@@ -32,6 +32,7 @@ semolina = "~0.1.3"
 sppark = { git = "https://github.com/lurk-lab/sppark.git", branch = "pushing-the-limit" }
 pasta_curves = { git = "https://github.com/lurk-lab/pasta_curves", branch = "dev", version = ">=0.3.1, <=0.5", features = ["repr-c"] }
 paste = "1.0.14"
+rayon = "1.5"
 
 [build-dependencies]
 cc = "^1.0.70"
@@ -44,5 +45,9 @@ rand_chacha = "^0"
 rayon = "1.5"
 
 [[bench]]
-name = "main"
+name = "msm"
+harness = false
+
+[[bench]]
+name = "spmvm"
 harness = false

benches/spmvm.rs

@@ -0,0 +1,110 @@
+// Copyright Supranational LLC
+// Licensed under the Apache License, Version 2.0, see LICENSE for details.
+// SPDX-License-Identifier: Apache-2.0
+
+#![allow(dead_code)]
+#![allow(unused_imports)]
+#![allow(unused_mut)]
+#![allow(non_snake_case)]
+
+use criterion::{criterion_group, criterion_main, Criterion};
+
+use pasta_curves::{group::ff::{PrimeField, Field}, pallas};
+use pasta_msm::{self, utils::SparseMatrix, spmvm::{CudaSparseMatrix, sparse_matrix_witness_pallas, CudaWitness}};
+use rand::Rng;
+
+#[cfg(feature = "cuda")]
+extern "C" {
+    fn cuda_available() -> bool;
+}
+
+pub fn generate_csr<F: PrimeField>(
+    n: usize,
+    m: usize,
+) -> SparseMatrix<F> {
+    let mut rng = rand::thread_rng();
+
+    let mut data = Vec::new();
+    let mut col_idx = Vec::new();
+    let mut row_ptr = Vec::new();
+    row_ptr.push(0);
+
+    for _ in 0..n {
+        let num_elements = rng.gen_range(5..=10); // Random number of elements between 5 to 10
+        for _ in 0..num_elements {
+            data.push(F::random(&mut rng)); // Random data value
+            col_idx.push(rng.gen_range(0..m)); // Random column index
+        }
+        row_ptr.push(data.len()); // Add the index of the next row start
+    }
+
+    data.shrink_to_fit();
+    col_idx.shrink_to_fit();
+    row_ptr.shrink_to_fit();
+
+    let csr = SparseMatrix {
+        data,
+        indices: col_idx,
+        indptr: row_ptr,
+        cols: m,
+    };
+
+    csr
+}
+
+include!("../src/tests.rs");
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let bench_npow: usize = std::env::var("BENCH_NPOW")
+        .unwrap_or("17".to_string())
+        .parse()
+        .unwrap();
+    let n: usize = 1 << bench_npow;
+
+    println!("generating random matrix and scalars, just hang on...");
+    let csr = generate_csr(n, n);
+    let cuda_csr =
+        CudaSparseMatrix::new(&csr.data, &csr.indices, &csr.indptr, n, n);
+    let W = crate::tests::gen_scalars(n - 10);
+    let U = crate::tests::gen_scalars(9);
+    let witness = CudaWitness::new(&W, &pallas::Scalar::ONE, &U);
+    let scalars = [W.clone(), vec![pallas::Scalar::ONE], U.clone()].concat();
+
+    #[cfg(feature = "cuda")]
+    {
+        unsafe { pasta_msm::CUDA_OFF = true };
+    }
+
+    let mut group = c.benchmark_group("CPU");
+    group.sample_size(10);
+
+    group.bench_function(format!("2**{} points", bench_npow), |b| {
+        b.iter(|| {
+            let _ = csr.multiply_vec(&scalars);
+        })
+    });
+
+    group.finish();
+
+    #[cfg(feature = "cuda")]
+    if unsafe { cuda_available() } {
+        unsafe { pasta_msm::CUDA_OFF = false };
+
+        let mut group = c.benchmark_group("GPU");
+        group.sample_size(20);
+
+        let mut cuda_res = vec![pallas::Scalar::ONE; cuda_csr.num_rows];
+        for nthreads in [128, 256, 512, 1024] {
+            group.bench_function(format!("2**{} points, nthreads={}", bench_npow, nthreads), |b| {
+                b.iter(|| {
+                    let _ = sparse_matrix_witness_pallas(&cuda_csr, &witness, &mut cuda_res, nthreads);
+                })
+            });
+        }
+
+        group.finish();
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

cuda/pallas.cu

@@ -16,12 +16,24 @@ typedef vesta_t scalar_t;
 
 #include <msm/pippenger.cuh>
 #include <spmvm/spmvm.cuh>
+#include <spmvm/double.cuh>
 
 #ifndef __CUDA_ARCH__
 
-extern "C" RustError spmvm_pallas(scalar_t *scalars, size_t nscalars)
+extern "C" RustError cuda_double_pallas(double_host_t<scalar_t> *csr, scalar_t *scalars, scalar_t *out)
 {
-    return double_scalars<scalar_t>(scalars, nscalars);
+    return double_scalars<scalar_t>(csr, scalars, out);
+}
+
+extern "C" RustError cuda_sparse_matrix_mul_pallas(spmvm_host_t<scalar_t> *csr, const scalar_t *scalars, scalar_t *out, size_t nthreads)
+{
+    return sparse_matrix_mul<scalar_t>(csr, scalars, out, nthreads);
+}
+
+extern "C" RustError cuda_sparse_matrix_witness_pallas(
+    spmvm_host_t<scalar_t> *csr, const witness_t<scalar_t> *witness, scalar_t *out, size_t nthreads)
+{
+    return sparse_matrix_witness<scalar_t>(csr, witness, out, nthreads);
 }
 
 extern "C" void drop_msm_context_pallas(msm_context_t<affine_t::mem_t> &ref)

examples/spmvm.rs

@@ -1,41 +1,84 @@
 // Copyright Supranational LLC
 // Licensed under the Apache License, Version 2.0, see LICENSE for details.
 // SPDX-License-Identifier: Apache-2.0
+#![allow(non_snake_case)]
 
 use std::time::Instant;
 
-use pasta_curves::group::ff::PrimeField;
-use pasta_msm::spmvm::double_pallas;
+use pasta_curves::{group::ff::{PrimeField, Field}, pallas};
+use pasta_msm::{
+    spmvm::{sparse_matrix_mul_pallas, CudaSparseMatrix, CudaWitness, sparse_matrix_witness_pallas},
+    utils::SparseMatrix,
+};
+use rand::Rng;
 
-pub fn generate_scalars<F: PrimeField>(
-    len: usize,
-) -> Vec<F> {
+pub fn generate_csr<F: PrimeField>(n: usize, m: usize) -> SparseMatrix<F> {
     let mut rng = rand::thread_rng();
-    let scalars = (0..len)
-        .map(|_| F::random(&mut rng))
-        .collect::<Vec<_>>();
+
+    let mut data = Vec::new();
+    let mut col_idx = Vec::new();
+    let mut row_ptr = Vec::new();
+    row_ptr.push(0);
+
+    for _ in 0..n {
+        let num_elements = rng.gen_range(5..=10); // Random number of elements between 5 to 10
+        for _ in 0..num_elements {
+            data.push(F::random(&mut rng)); // Random data value
+            col_idx.push(rng.gen_range(0..m)); // Random column index
+        }
+        row_ptr.push(data.len()); // Add the index of the next row start
+    }
+
+    data.shrink_to_fit();
+    col_idx.shrink_to_fit();
+    row_ptr.shrink_to_fit();
+
+    let csr = SparseMatrix {
+        data,
+        indices: col_idx,
+        indptr: row_ptr,
+        cols: m,
+    };
+
+    csr
+}
+
+pub fn generate_scalars<F: PrimeField>(len: usize) -> Vec<F> {
+    let mut rng = rand::thread_rng();
+    let scalars = (0..len).map(|_| F::random(&mut rng)).collect::<Vec<_>>();
 
     scalars
 }
 
 /// cargo run --release --example spmvm
 fn main() {
     let npow: usize = std::env::var("NPOW")
-        .unwrap_or("23".to_string())
+        .unwrap_or("17".to_string())
         .parse()
         .unwrap();
     let n = 1usize << npow;
+    let nthreads: usize = std::env::var("NTHREADS")
+        .unwrap_or("128".to_string())
+        .parse()
+        .unwrap();
 
-    let mut scalars = generate_scalars(n);
+    let csr = generate_csr(n, n);
+    let cuda_csr =
+        CudaSparseMatrix::new(&csr.data, &csr.indices, &csr.indptr, n, n);
+    let W = generate_scalars(n - 10);
+    let U = generate_scalars(9);
+    let scalars = [W.clone(), vec![pallas::Scalar::ONE], U.clone()].concat();
 
     let start = Instant::now();
-    let double_scalars = scalars.iter().map(|x| x + x).collect::<Vec<_>>();
+    let res = csr.multiply_vec(&scalars);
     println!("cpu took: {:?}", start.elapsed());
 
+    let witness = CudaWitness::new(&W, &pallas::Scalar::ONE, &U);
+    let mut cuda_res = vec![pallas::Scalar::ONE; cuda_csr.num_rows];
     let start = Instant::now();
-    double_pallas(&mut scalars);
+    sparse_matrix_witness_pallas(&cuda_csr, &witness, &mut cuda_res, nthreads);
     println!("gpu took: {:?}", start.elapsed());
 
-    assert_eq!(double_scalars, scalars);
+    assert_eq!(res, cuda_res);
     println!("success!");
 }

src/lib.rs

@@ -3,6 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 pub mod spmvm;
+pub mod utils;
 
 extern crate semolina;