feat: ICICLE msm integration

.github/workflows/build.yml

@@ -64,7 +64,7 @@ jobs:
       - name: Check all tests and binaries compilation
         run: |
           cargo check --workspace --tests --lib --bins
-          cargo check --workspace --all-features
+          cargo check --workspace --features 'std parallel test-srs test-apis'
 
       - name: Check no_std support and WASM compilation
         env:

flake.nix

@@ -21,12 +21,15 @@
   outputs = { self, nixpkgs, flake-utils, rust-overlay, pre-commit-hooks, ... }:
     flake-utils.lib.eachDefaultSystem (system:
       let
-        overlays = [
-          (import rust-overlay)
-        ];
+        overlays = [ (import rust-overlay) ];
         pkgs = import nixpkgs { inherit system overlays; };
-        nightlyToolchain = pkgs.rust-bin.selectLatestNightlyWith
-          (toolchain: toolchain.minimal.override { extensions = [ "rustfmt" ]; });
+        pkgsAllowUnfree = import nixpkgs {
+          inherit system;
+          config.allowUnfree = true;
+        };
+        gcc11 = pkgs.overrideCC pkgs.stdenv pkgs.gcc11;
+        nightlyToolchain = pkgs.rust-bin.selectLatestNightlyWith (toolchain:
+          toolchain.minimal.override { extensions = [ "rustfmt" ]; });
 
         stableToolchain = pkgs.rust-bin.stable.latest.minimal.override {
           extensions = [ "clippy" "llvm-tools-preview" "rust-src" ];
@@ -42,8 +45,52 @@
           fi
           exec ${stableToolchain}/bin/cargo "$@"
         '';
-      in with pkgs;
-      {
+        baseShell = with pkgs;
+          clang15Stdenv.mkDerivation {
+            name = "clang15-nix-shell";
+            buildInputs = [
+              argbash
+              openssl
+              pkg-config
+              git
+              nixpkgs-fmt
+
+              cargo-with-nightly
+              stableToolchain
+              nightlyToolchain
+              cargo-sort
+              clang-tools_15
+              clangStdenv
+              llvm_15
+            ] ++ lib.optionals stdenv.isDarwin
+              [ darwin.apple_sdk.frameworks.Security ];
+
+            CARGO_TARGET_DIR = "target/nix_rustc";
+
+            shellHook = ''
+              export RUST_BACKTRACE=full
+              export PATH="$PATH:$(pwd)/target/debug:$(pwd)/target/release"
+              # Prevent cargo aliases from using programs in `~/.cargo` to avoid conflicts with local rustup installations.
+              export CARGO_HOME=$HOME/.cargo-nix
+
+              # Ensure `cargo fmt` uses `rustfmt` from nightly.
+              export RUSTFMT="${nightlyToolchain}/bin/rustfmt"
+
+              export C_INCLUDE_PATH="${llvmPackages_15.libclang.lib}/lib/clang/${llvmPackages_15.libclang.version}/include"
+              export LIBCLANG_PATH=
+              export CC="${clang-tools_15.clang}/bin/clang"
+              export CXX="${clang-tools_15.clang}/bin/clang++"
+              export AR="${llvm_15}/bin/llvm-ar"
+              export CFLAGS="-mcpu=generic"
+
+              # by default choose u64_backend
+              export RUSTFLAGS='--cfg curve25519_dalek_backend="u64"'
+            ''
+            # install pre-commit hooks
+            + self.check.${system}.pre-commit-check.shellHook;
+          };
+      in
+      with pkgs; {
         check = {
           pre-commit-check = pre-commit-hooks.lib.${system}.run {
             src = ./.;
@@ -72,48 +119,34 @@
                 entry = "cargo sort -w";
                 pass_filenames = false;
               };
+              nixpkgs-fmt.enable = true;
             };
           };
         };
-        devShell = clang15Stdenv.mkDerivation {
-          name = "clang15-nix-shell";
-          buildInputs = [
-            argbash
-            openssl
-            pkg-config
-            git
-
-            cargo-with-nightly
-            stableToolchain
-            nightlyToolchain
-            cargo-sort
-            clang-tools_15
-            clangStdenv
-            llvm_15
-          ] ++ lib.optionals stdenv.isDarwin [ darwin.apple_sdk.frameworks.Security ];
-
-          CARGO_TARGET_DIR = "target/nix_rustc";
-
-          shellHook = ''
-            export RUST_BACKTRACE=full
-            export PATH="$PATH:$(pwd)/target/debug:$(pwd)/target/release"
-            # Prevent cargo aliases from using programs in `~/.cargo` to avoid conflicts with local rustup installations.
-            export CARGO_HOME=$HOME/.cargo-nix
-
-            # Ensure `cargo fmt` uses `rustfmt` from nightly.
-            export RUSTFMT="${nightlyToolchain}/bin/rustfmt"
-
-            export C_INCLUDE_PATH="${llvmPackages_15.libclang.lib}/lib/clang/${llvmPackages_15.libclang.version}/include"
-            export CC="${clang-tools_15.clang}/bin/clang"
-            export AR="${llvm_15}/bin/llvm-ar"
-            export CFLAGS="-mcpu=generic"
+        devShell = baseShell;
+        # extra dev shells
+        devShells = {
+          # run with `nix develop .#cudaShell`
+          cudaShell =
+            let cudatoolkit = pkgsAllowUnfree.cudaPackages_12_3.cudatoolkit;
+            in baseShell.overrideAttrs (oldAttrs: {
+              # for GPU/CUDA env (e.g. to run ICICLE code)
+              name = "cuda-env-shell";
+              buildInputs = oldAttrs.buildInputs
+                ++ [ cmake cudatoolkit util-linux gcc11 ];
+              # CXX is overridden to use gcc as icicle-curves's build scripts need them
+              shellHook = oldAttrs.shellHook + ''
 
-            # by default choose u64_backend
-            export RUSTFLAGS='--cfg curve25519_dalek_backend="u64"'
-          ''
-          # install pre-commit hooks
-          + self.check.${system}.pre-commit-check.shellHook;
+                export PATH="${pkgs.gcc11}/bin:${cudatoolkit}/bin:${cudatoolkit}/nvvm/bin:$PATH"
+                export LD_LIBRARY_PATH=${cudatoolkit}/lib
+                export CUDA_PATH=${cudatoolkit}
+                export CPATH="${cudatoolkit}/include"
+                export LIBRARY_PATH="$LIBRARY_PATH:/lib"
+                export CMAKE_CUDA_COMPILER=$CUDA_PATH/bin/nvcc
+                export LIBCLANG_PATH=${llvmPackages_15.libclang.lib}/lib
+                export CFLAGS=""
+              '';
+            });
         };
-      }
-    );
+      });
 }

primitives/Cargo.toml

@@ -39,6 +39,11 @@ generic-array = { version = "0", features = [
         "serde",
 ] } # not a direct dependency, but we need serde
 hashbrown = "0.14.3"
+icicle-bls12-377 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true, features = ["arkworks"] }
+icicle-bls12-381 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true, features = ["arkworks"] }
+icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true, features = ["arkworks"] }
+icicle-core = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true }
+icicle-cuda-runtime = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true }
 itertools = { workspace = true, features = ["use_alloc"] }
 jf-relation = { path = "../relation", default-features = false }
 jf-utils = { path = "../utilities" }
@@ -84,6 +89,12 @@ path = "benches/pcs_size.rs"
 harness = false
 required-features = ["test-srs"]
 
+[[bench]]
+name = "kzg-gpu"
+path = "benches/kzg_gpu.rs"
+harness = false
+required-features = ["test-srs"]
+
 [[bench]]
 name = "reed-solomon"
 path = "benches/reed_solomon.rs"
@@ -135,3 +146,12 @@ parallel = [
 ]
 test-srs = []
 seq-fk-23 = [] # FK23 without parallelism
+icicle = [
+       "icicle-cuda-runtime",
+       "icicle-core",
+       "icicle-core/arkworks",
+       "icicle-bn254",
+       "icicle-bls12-381",
+       "icicle-bls12-377",
+       "parallel",
+]

primitives/benches/kzg_gpu.rs

@@ -0,0 +1,99 @@
+//! This benchmark meant for MSM speed comparison between arkworks and
+//! GPU-accelerated code We use `UnivariateKzgPCS::commit()` as a proxy for MSM
+//!
+//! Run `cargo bench --bench kzg-gpu --features "test-srs icicle"`
+use ark_bn254::Bn254;
+#[cfg(feature = "icicle")]
+use ark_ec::models::{short_weierstrass::Affine, CurveConfig};
+use ark_ec::pairing::Pairing;
+use ark_poly::{univariate::DensePolynomial, DenseUVPolynomial};
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+#[cfg(feature = "icicle")]
+use jf_primitives::icicle_deps::*;
+use jf_primitives::pcs::{
+    prelude::{PolynomialCommitmentScheme, UnivariateKzgPCS},
+    StructuredReferenceString,
+};
+use jf_utils::test_rng;
+
+const MIN_LOG_DEGREE: usize = 19;
+const MAX_LOG_DEGREE: usize = 23;
+
+/// running MSM using arkworks backend
+pub fn kzg_ark<E: Pairing>(c: &mut Criterion) {
+    let mut group = c.benchmark_group("MSM with arkworks");
+    let mut rng = test_rng();
+
+    let supported_degree = 2usize.pow(MAX_LOG_DEGREE as u32);
+    let pp = UnivariateKzgPCS::<E>::gen_srs_for_testing(&mut rng, supported_degree).unwrap();
+
+    // setup for commit first
+    for log_degree in MIN_LOG_DEGREE..MAX_LOG_DEGREE {
+        let degree = 2usize.pow(log_degree as u32);
+        let (ck, _vk) = pp.trim(degree).unwrap();
+        let p = <DensePolynomial<E::ScalarField> as DenseUVPolynomial<E::ScalarField>>::rand(
+            degree, &mut rng,
+        );
+
+        group.bench_with_input(
+            BenchmarkId::from_parameter(log_degree),
+            &log_degree,
+            |b, _log_degree| b.iter(|| UnivariateKzgPCS::<E>::commit(&ck, &p).unwrap()),
+        );
+    }
+    group.finish();
+}
+
+/// running MSM using ICICLE backends
+#[cfg(feature = "icicle")]
+pub fn kzg_icicle<E, C>(c: &mut Criterion)
+where
+    C: IcicleCurve + MSM<C>,
+    C::ScalarField: ArkConvertible<ArkEquivalent = E::ScalarField>,
+    C::BaseField: ArkConvertible<ArkEquivalent = <C::ArkSWConfig as CurveConfig>::BaseField>,
+    E: Pairing<G1Affine = Affine<<C as IcicleCurve>::ArkSWConfig>>,
+{
+    let mut group = c.benchmark_group("MSM with ICICLE");
+    let mut rng = test_rng();
+
+    let supported_degree = 2usize.pow(MAX_LOG_DEGREE as u32);
+    let pp = UnivariateKzgPCS::<E>::gen_srs_for_testing(&mut rng, supported_degree).unwrap();
+    // TODO: (alex) figure out load longer SRS first, and only use part of it later
+    // currently it will error if the `scalars.len() % points.len() != 0`
+    // while we can tap into a slice behind the reference via
+    // `HostOrDeviceSlice[..]` which gives `&[T]` however msm() doesn't accept
+    // &[T], only `&HostOrDeviceSlice`
+
+    // setup for commit first
+    for log_degree in MIN_LOG_DEGREE..MAX_LOG_DEGREE {
+        let degree = 2usize.pow(log_degree as u32);
+        let (ck, _vk) = pp.trim(degree).unwrap();
+        let p = <DensePolynomial<E::ScalarField> as DenseUVPolynomial<E::ScalarField>>::rand(
+            degree, &mut rng,
+        );
+
+        group.bench_with_input(
+            BenchmarkId::from_parameter(log_degree),
+            &log_degree,
+            |b, _log_degree| {
+                b.iter(|| UnivariateKzgPCS::<E>::commit_with_gpu::<C>(&ck, &p).unwrap())
+            },
+        );
+    }
+    group.finish();
+}
+
+fn kzg_gpu_bn254(c: &mut Criterion) {
+    kzg_ark::<Bn254>(c);
+    #[cfg(feature = "icicle")]
+    kzg_icicle::<Bn254, icicle_bn254::curve::CurveCfg>(c);
+}
+
+criterion_group! {
+    name = kzg_gpu_benches;
+    config = Criterion::default().sample_size(10);
+    targets =
+        kzg_gpu_bn254,
+}
+
+criterion_main!(kzg_gpu_benches);

primitives/src/lib.rs

@@ -39,4 +39,34 @@ pub mod vdf;
 pub mod vid;
 pub mod vrf;
 
+/// dependecies required for ICICLE-related code, group import for convenience
+#[cfg(feature = "icicle")]
+pub mod icicle_deps {
+    pub use icicle_core::{
+        curve::{Affine as IcicleAffine, Curve as IcicleCurve, Projective as IcicleProjective},
+        msm::{MSMConfig, MSM},
+        traits::{ArkConvertible, FieldImpl},
+    };
+    pub use icicle_cuda_runtime::{memory::HostOrDeviceSlice, stream::CudaStream};
+
+    /// curve-specific types both from arkworks and from ICICLE
+    /// including Pairing, CurveCfg, Fr, Fq etc.
+    pub mod curves {
+        pub use ark_bls12_381::Bls12_381;
+        pub use ark_bn254::Bn254;
+        pub use icicle_bls12_381::curve::CurveCfg as IcicleBls12_381;
+        pub use icicle_bn254::curve::CurveCfg as IcicleBn254;
+    }
+
+    // TODO: remove this after `warmup()` is added upstream
+    // https://github.com/ingonyama-zk/icicle/pull/422#issuecomment-1980881638
+    /// Create a new stream and warmup
+    pub fn warmup_new_stream() -> Result<CudaStream, ()> {
+        let stream = CudaStream::create().unwrap();
+        // TODO: consider using an error type?
+        let _warmup_bytes = HostOrDeviceSlice::<'_, u8>::cuda_malloc_async(1024, &stream).unwrap();
+        Ok(stream)
+    }
+}
+
 pub(crate) mod utils;

primitives/src/pcs/errors.rs

@@ -11,6 +11,10 @@ use crate::errors::PrimitivesError;
 use ark_serialize::SerializationError;
 use ark_std::string::{String, ToString};
 use displaydoc::Display;
+#[cfg(feature = "icicle")]
+use icicle_core::error::IcicleError;
+#[cfg(feature = "icicle")]
+use icicle_cuda_runtime::error::CudaError;
 
 /// A `enum` specifying the possible failure modes of the PCS.
 #[derive(Display, Debug)]
@@ -29,6 +33,9 @@ pub enum PCSError {
     TranscriptError(TranscriptError),
     /// Error from upstream dependencies: {0}
     UpstreamError(String),
+    #[cfg(feature = "icicle")]
+    /// Error from ICICLE: {0}
+    IcicleError(String),
 }
 
 impl ark_std::error::Error for PCSError {}
@@ -50,3 +57,17 @@ impl From<PrimitivesError> for PCSError {
         Self::UpstreamError(e.to_string())
     }
 }
+
+#[cfg(feature = "icicle")]
+impl From<IcicleError> for PCSError {
+    fn from(e: IcicleError) -> Self {
+        Self::IcicleError(ark_std::format!("{:?}", e))
+    }
+}
+#[cfg(feature = "icicle")]
+impl From<CudaError> for PCSError {
+    fn from(e: CudaError) -> Self {
+        let icicle_err = IcicleError::from_cuda_error(e);
+        icicle_err.into()
+    }
+}