feat(icicle): Add icicle MSM support (#503)

* feat(icicle): Add icicle MSM support * use icicle-jolt to compile cuda deps * remove icicle from CI * address review comments * fmt
a16z · Jan 7, 2025 · 1145ec4 · 1145ec4
1 parent f1e5ab3
commit 1145ec4
Show file tree

Hide file tree

Showing 35 changed files with 1,790 additions and 216 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -54,6 +54,7 @@ members = [
 
 [features]
 host = ["jolt-sdk/host"]
+icicle = ["jolt-core/icicle"]
 
 [lib]
 path = "./src/lib.rs"

diff --git a/README.md b/README.md
@@ -69,6 +69,29 @@ Examples in the [`examples`](./examples/) directory can be run using e.g.
 
 ```cargo run --release -p sha2-chain```
 
+## CUDA Support
+
+JOLT supports CUDA acceleration via [icicle](https://github.com/ingonyama-zk/icicle-jolt).
+
+Dependencies:
+1. Install [CUDA Toolkit](https://developer.nvidia.com/cuda-toolkit)
+2. Install [CMake](https://cmake.org/)
+
+Now you may build Jolt with CUDA acceleration using the `--features icicle` flag.
+
+### Build
+
+```cargo build -p jolt-core --features icicle```
+
+### Bench
+
+```
+# Set the icicle backend path - this won't be needed in the future
+export ICICLE_BACKEND_INSTALL_DIR=$(pwd)/target/debug/deps/icicle/lib/backend
+cargo bench --bench msm_batch --no-fail-fast -p jolt-core --features icicle
+```
+
+Note - NVIDIA doesn't support cross compilation on MacOS. Only Windows or Linux.
 
 ## Performance profiling
 

diff --git a/examples/alloc/Cargo.toml b/examples/alloc/Cargo.toml
@@ -6,3 +6,6 @@ edition = "2021"
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "alloc-guest", path = "./guest" }
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/collatz/Cargo.toml b/examples/collatz/Cargo.toml
@@ -6,3 +6,6 @@ edition = "2021"
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "collatz-guest", path = "./guest" }
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/fibonacci/Cargo.toml b/examples/fibonacci/Cargo.toml
@@ -6,3 +6,6 @@ edition = "2021"
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "fibonacci-guest", path = "./guest" }
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/muldiv/Cargo.toml b/examples/muldiv/Cargo.toml
@@ -7,3 +7,5 @@ edition = "2021"
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "muldiv-guest", path = "./guest" }
 
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/multi-function/Cargo.toml b/examples/multi-function/Cargo.toml
@@ -6,3 +6,6 @@ edition = "2021"
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "multi-function-guest", path = "./guest" }
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/overflow/Cargo.toml b/examples/overflow/Cargo.toml
@@ -5,4 +5,7 @@ edition = "2021"
 
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
-guest = { package = "overflow-guest", path = "./guest" }
+guest = { package = "overflow-guest", path = "./guest" }
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/sha2-chain/Cargo.toml b/examples/sha2-chain/Cargo.toml
@@ -8,3 +8,6 @@ jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "sha2-chain-guest", path = "./guest" }
 
 hex = "0.4.3"
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/sha2-ex/Cargo.toml b/examples/sha2-ex/Cargo.toml
@@ -8,3 +8,6 @@ jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "sha2-guest", path = "./guest" }
 
 hex = "0.4.3"
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/sha3-chain/Cargo.toml b/examples/sha3-chain/Cargo.toml
@@ -6,5 +6,7 @@ edition = "2021"
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "sha3-chain-guest", path = "./guest" }
+hex = "0.4.3"
 
-hex = "0.4.3"
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/sha3-ex/Cargo.toml b/examples/sha3-ex/Cargo.toml
@@ -6,5 +6,7 @@ edition = "2021"
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "sha3-guest", path = "./guest" }
-
 hex = "0.4.3"
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/examples/stdlib/Cargo.toml b/examples/stdlib/Cargo.toml
@@ -6,3 +6,6 @@ edition = "2021"
 [dependencies]
 jolt-sdk = { path = "../../jolt-sdk", features = ["host"] }
 guest = { package = "stdlib-guest", path = "./guest" }
+
+[features]
+icicle = ["jolt-sdk/icicle"]
diff --git a/jolt-core/Cargo.toml b/jolt-core/Cargo.toml
@@ -19,6 +19,18 @@ repository = "https://github.com/a16z/jolt"
 license-file = "LICENSE"
 keywords = ["SNARK", "cryptography", "proofs"]
 
+[features]
+default = [
+    "ark-ec/parallel",
+    "ark-ff/parallel",
+    "ark-std/parallel",
+    "ark-ff/asm",
+    "host",
+    "rayon",
+]
+host = ["dep:reqwest", "dep:tokio"]
+icicle = ["default", "dep:icicle-runtime", "dep:icicle-core", "dep:icicle-bn254"]
+
 [dependencies]
 ark-bn254 = "0.4.0"
 ark-ec = { version = "0.4.2", default-features = false }
@@ -65,6 +77,7 @@ bytemuck = "1.19.0"
 tokio = { version = "1.38.0", optional = true }
 alloy-primitives = "0.7.6"
 alloy-sol-types = "0.7.6"
+once_cell = "1.19.0"
 
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["html_reports"] }
@@ -93,23 +106,24 @@ harness = false
 name = "compute_cubic"
 harness = false
 
+[[bench]]
+name = "msm"
+harness = false
+
+[[bench]]
+name = "msm_batch"
+harness = false
+
 [lib]
 name = "jolt_core"
 path = "src/lib.rs"
 
-[features]
-default = [
-    "ark-ec/parallel",
-    "ark-ff/parallel",
-    "ark-std/parallel",
-    "ark-ff/asm",
-    "host",
-    "rayon",
-]
-host = ["dep:reqwest", "dep:tokio"]
-
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+icicle-runtime = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", features = ["cuda_backend"], rev = "ed93e21", optional = true }
+icicle-core = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", rev = "ed93e21", optional = true }
+icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", features = ["cuda_backend"], rev = "ed93e21", optional = true }
 memory-stats = "1.0.0"
+sys-info = "0.9.1"
 tokio = { version = "1.38.0", optional = true, features = ["rt-multi-thread"] }
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]

diff --git a/jolt-core/benches/iai.rs b/jolt-core/benches/iai.rs
@@ -35,7 +35,7 @@ fn eval_poly_setup<F: JoltField>(size: usize) -> (DensePolynomial<F>, Vec<F>) {
 #[library_benchmark]
 #[bench::long(msm_setup::<G1Projective>(4096))]
 fn bench_msm<G: CurveGroup>(input: (Vec<G>, Vec<G::ScalarField>)) -> G {
-    black_box(VariableBaseMSM::msm(&G::normalize_batch(&input.0), &input.1).unwrap())
+    black_box(VariableBaseMSM::msm(&G::normalize_batch(&input.0), None, &input.1).unwrap())
 }
 
 #[library_benchmark]