Add experimental use of icicle GPU Acceleration on pse feature

README.md

@@ -47,6 +47,30 @@ cargo bench --bench inner_product
 
 These benchmarks use the `criterion` crate to run `create_proof` 10 times for statistical analysis. Note the benchmark circuits perform more than a one multiplication / inner product per circuit.
 
+### GPU Acceleration
+
+If you have access to NVIDIA GPUs, you can enable acceleration by building with the feature `halo2-icicle` and setting the following environment variable:
+
+```sh
+export ENABLE_ICICLE_GPU=true
+```
+
+GPU acceleration is provided by [Icicle](https://github.com/ingonyama-zk/icicle)
+
+To go back to running with CPU, the previous environment variable must be **unset** instead of being switched to a value of false:
+
+```sh
+unset ENABLE_ICICLE_GPU
+```
+
+> [!NOTE] 
+> Even with the above environment variable set, for circuits where k <= 8, icicle is only enabled in certain areas where batching MSMs will help; all other places will fallback to using CPU MSM. To change the value of `k` where icicle is enabled, you can set the environment variable `ICICLE_SMALL_CIRCUIT`.
+> 
+> Example: The following will cause icicle single MSM to be used throughout when k > 10 and CPU single MSM with certain locations using icicle batched MSM when k <= 10
+> ```sh
+> export ICICLE_SMALL_CIRCUIT=10
+> ```
+
 ## halo2-ecc
 
 This crate uses `halo2-base` to provide a library of elliptic curve cryptographic primitives. In particular, we support elliptic curves over base fields that are larger than the scalar field used in the proving system (e.g., `F_r` for bn254 when using Halo 2 with a KZG backend).

halo2-base/Cargo.toml

@@ -21,6 +21,8 @@ ark-std = { version = "0.3.0", features = ["print-trace"], optional = true }
 halo2_proofs_axiom = { version = "0.4", package = "halo2-axiom", optional = true }
 # Use PSE halo2 and halo2curves for compatibility when feature = "halo2-pse" is on
 halo2_proofs = { git = "https://github.com/privacy-scaling-explorations/halo2.git", rev = "7a21656", optional = true }
+# Use PSE halo2 and halo2curves with Icicle GPU Acceleration when feature = "halo2-icicle" is on
+halo2_proofs_icicle = { git = "https://github.com/ingonyama-zk/halo2.git", branch = "axiom-icicle", package = "halo2_proofs", optional = true }
 
 # This is Scroll's audited poseidon circuit. We only use it for the Native Poseidon spec. We do not use the halo2 circuit at all (and it wouldn't even work because the halo2_proofs tag is not compatbile).
 # We forked it to upgrade to ff v0.13 and removed the circuit module
@@ -55,7 +57,10 @@ default = ["halo2-axiom", "display", "test-utils"]
 asm = ["halo2_proofs_axiom?/asm"]
 dev-graph = ["halo2_proofs/dev-graph", "plotters"] # only works with halo2-pse for now
 halo2-pse = ["halo2_proofs/circuit-params"]
+halo2-icicle = ["halo2_proofs_icicle/icicle_gpu", "halo2_proofs_icicle/circuit-params"]
+profile-icicle = ["halo2_proofs_icicle/profile"]
 halo2-axiom = ["halo2_proofs_axiom"]
+halo2-axiom-icicle = ["halo2_proofs_axiom"]
 display = []
 profile = ["halo2_proofs_axiom?/profile"]
 test-utils = ["dep:rand", "ark-std"]

halo2-base/src/gates/flex_gate/threads/single_phase.rs

@@ -212,9 +212,9 @@ pub fn assign_with_constraints<F: ScalarField, const ROTATIONS: usize>(
         for (i, (advice, &q)) in ctx.advice.iter().zip(ctx.selector.iter()).enumerate() {
             let column = basic_gate.value;
             let value = if use_unknown { Value::unknown() } else { Value::known(advice) };
-            #[cfg(feature = "halo2-axiom")]
+            #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
             let cell = region.assign_advice(column, row_offset, value).cell();
-            #[cfg(not(feature = "halo2-axiom"))]
+            #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
             let cell = region
                 .assign_advice(|| "", column, row_offset, || value.map(|v| *v))
                 .unwrap()
@@ -244,9 +244,9 @@ pub fn assign_with_constraints<F: ScalarField, const ROTATIONS: usize>(
                         .get(gate_index)
                         .unwrap_or_else(|| panic!("NOT ENOUGH ADVICE COLUMNS. Perhaps blinding factors were not taken into account. The max non-poisoned rows is {max_rows}"));
                 let column = basic_gate.value;
-                #[cfg(feature = "halo2-axiom")]
+                #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
                 let ncell = region.assign_advice(column, row_offset, value);
-                #[cfg(not(feature = "halo2-axiom"))]
+                #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
                 let ncell =
                     region.assign_advice(|| "", column, row_offset, || value.map(|v| *v)).unwrap();
                 raw_constrain_equal(region, ncell.cell(), cell);

halo2-base/src/lib.rs

@@ -25,17 +25,26 @@ use mimalloc::MiMalloc;
 #[global_allocator]
 static GLOBAL: MiMalloc = MiMalloc;
 
-#[cfg(all(feature = "halo2-pse", feature = "halo2-axiom"))]
+#[cfg(any(
+    all(feature = "halo2-pse", feature = "halo2-axiom"),
+    all(feature = "halo2-pse", feature = "halo2-icicle"),
+    all(feature = "halo2-pse", feature = "halo2-axiom-icicle"),
+    all(feature = "halo2-axiom", feature = "halo2-icicle"),
+    all(feature = "halo2-axiom", feature = "halo2-axiom-icicle"),
+    all(feature = "halo2-icicle", feature = "halo2-axiom-icicle")
+))]
 compile_error!(
-    "Cannot have both \"halo2-pse\" and \"halo2-axiom\" features enabled at the same time!"
+    "Cannot have multiple of \"halo2-pse\", \"halo2-axiom\", \"halo2-axiom-icicle\", or \"halo2-icicle\" features enabled at the same time!"
 );
-#[cfg(not(any(feature = "halo2-pse", feature = "halo2-axiom")))]
-compile_error!("Must enable exactly one of \"halo2-pse\" or \"halo2-axiom\" features to choose which halo2_proofs crate to use.");
+#[cfg(not(any(feature = "halo2-pse", feature = "halo2-axiom", feature = "halo2-icicle", feature = "halo2-axiom-icicle")))]
+compile_error!("Must enable exactly one of \"halo2-pse\", \"halo2-axiom\", \"halo2-axiom-icicle\", or \"halo2-icicle\" features to choose which halo2_proofs crate to use.");
 
 // use gates::flex_gate::MAX_PHASE;
 #[cfg(feature = "halo2-pse")]
 pub use halo2_proofs;
-#[cfg(feature = "halo2-axiom")]
+#[cfg(feature = "halo2-icicle")]
+pub use halo2_proofs_icicle as halo2_proofs;
+#[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
 pub use halo2_proofs_axiom as halo2_proofs;
 
 use halo2_proofs::halo2curves::ff;
@@ -55,10 +64,10 @@ pub mod utils;
 pub mod virtual_region;
 
 /// Constant representing whether the Layouter calls `synthesize` once just to get region shape.
-#[cfg(feature = "halo2-axiom")]
+#[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
 pub const SKIP_FIRST_PASS: bool = false;
 /// Constant representing whether the Layouter calls `synthesize` once just to get region shape.
-#[cfg(feature = "halo2-pse")]
+#[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
 pub const SKIP_FIRST_PASS: bool = true;
 
 /// Convenience Enum which abstracts the scenarios under a value is added to an advice column.

halo2-base/src/utils/halo2.rs

@@ -1,16 +1,16 @@
 use crate::ff::Field;
 use crate::halo2_proofs::{
     circuit::{AssignedCell, Cell, Region, Value},
-    plonk::{Advice, Assigned, Column, Fixed},
+    plonk::{Advice, Assigned, Circuit, Column, Fixed},
 };
 use crate::virtual_region::copy_constraints::{CopyConstraintManager, SharedCopyConstraintManager};
 use crate::AssignedValue;
 
 /// Raw (physical) assigned cell in Plonkish arithmetization.
-#[cfg(feature = "halo2-axiom")]
+#[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
 pub type Halo2AssignedCell<'v, F> = AssignedCell<&'v Assigned<F>, F>;
 /// Raw (physical) assigned cell in Plonkish arithmetization.
-#[cfg(not(feature = "halo2-axiom"))]
+#[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
 pub type Halo2AssignedCell<'v, F> = AssignedCell<Assigned<F>, F>;
 
 /// Assign advice to physical region.
@@ -21,11 +21,11 @@ pub fn raw_assign_advice<'v, F: Field>(
     offset: usize,
     value: Value<impl Into<Assigned<F>>>,
 ) -> Halo2AssignedCell<'v, F> {
-    #[cfg(feature = "halo2-axiom")]
+    #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
     {
         region.assign_advice(column, offset, value)
     }
-    #[cfg(feature = "halo2-pse")]
+    #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
     {
         let value = value.map(|a| Into::<Assigned<F>>::into(a));
         region
@@ -47,11 +47,11 @@ pub fn raw_assign_fixed<F: Field>(
     offset: usize,
     value: F,
 ) -> Cell {
-    #[cfg(feature = "halo2-axiom")]
+    #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
     {
         region.assign_fixed(column, offset, value)
     }
-    #[cfg(feature = "halo2-pse")]
+    #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
     {
         region
             .assign_fixed(
@@ -68,9 +68,9 @@ pub fn raw_assign_fixed<F: Field>(
 /// Constrain two physical cells to be equal.
 #[inline(always)]
 pub fn raw_constrain_equal<F: Field>(region: &mut Region<F>, left: Cell, right: Cell) {
-    #[cfg(feature = "halo2-axiom")]
+    #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
     region.constrain_equal(left, right);
-    #[cfg(not(feature = "halo2-axiom"))]
+    #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
     region.constrain_equal(left, right).unwrap();
 }
 

halo2-base/src/utils/mod.rs

@@ -1,10 +1,10 @@
 use core::hash::Hash;
 
 use crate::ff::{FromUniformBytes, PrimeField};
-#[cfg(not(feature = "halo2-axiom"))]
+#[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
 use crate::halo2_proofs::arithmetic::CurveAffine;
 use crate::halo2_proofs::circuit::Value;
-#[cfg(feature = "halo2-axiom")]
+#[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
 pub use crate::halo2_proofs::halo2curves::CurveAffineExt;
 
 use num_bigint::BigInt;
@@ -19,7 +19,7 @@ pub mod halo2;
 pub mod testing;
 
 /// Helper trait to convert to and from a [BigPrimeField] by converting a list of [u64] digits
-#[cfg(feature = "halo2-axiom")]
+#[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
 pub trait BigPrimeField: ScalarField {
     /// Converts a slice of [u64] to [BigPrimeField]
     /// * `val`: the slice of u64
@@ -29,7 +29,7 @@ pub trait BigPrimeField: ScalarField {
     /// * The integer value of `val` is already less than the modulus of `Self`
     fn from_u64_digits(val: &[u64]) -> Self;
 }
-#[cfg(feature = "halo2-axiom")]
+#[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
 impl<F> BigPrimeField for F
 where
     F: ScalarField + From<[u64; 4]>, // Assume [u64; 4] is little-endian. We only implement ScalarField when this is true.
@@ -94,7 +94,7 @@ pub trait ScalarField: PrimeField + FromUniformBytes<64> + From<bool> + Hash + O
 // Later: will need to separate BigPrimeField from ScalarField when Goldilocks is introduced
 
 /// [ScalarField] that is ~256 bits long
-#[cfg(feature = "halo2-pse")]
+#[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
 pub trait BigPrimeField = PrimeField<Repr = [u8; 32]> + ScalarField;
 
 /// Converts an [Iterator] of u64 digits into `number_of_limbs` limbs of `bit_len` bits returned as a [Vec].
@@ -245,12 +245,12 @@ pub fn decompose_fe_to_u64_limbs<F: ScalarField>(
     number_of_limbs: usize,
     bit_len: usize,
 ) -> Vec<u64> {
-    #[cfg(feature = "halo2-axiom")]
+    #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
     {
         e.to_u64_limbs(number_of_limbs, bit_len)
     }
 
-    #[cfg(feature = "halo2-pse")]
+    #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
     {
         decompose_u64_digits_to_limbs(fe_to_biguint(e).iter_u64_digits(), number_of_limbs, bit_len)
     }
@@ -351,15 +351,15 @@ pub fn compose(input: Vec<BigUint>, bit_len: usize) -> BigUint {
 }
 
 /// Helper trait
-#[cfg(feature = "halo2-pse")]
+#[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
 pub trait CurveAffineExt: CurveAffine {
     /// Returns the raw affine (X, Y) coordinantes
     fn into_coordinates(self) -> (Self::Base, Self::Base) {
         let coordinates = self.coordinates().unwrap();
         (*coordinates.x(), *coordinates.y())
     }
 }
-#[cfg(feature = "halo2-pse")]
+#[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
 impl<C: CurveAffine> CurveAffineExt for C {}
 
 mod scalar_field_impls {

halo2-base/src/virtual_region/copy_constraints.rs

@@ -61,11 +61,11 @@ impl<F: Field + Ord> CopyConstraintManager<F> {
         let context_cell = self.load_external_cell(assigned_cell.cell());
         let mut value = Assigned::Trivial(F::ZERO);
         assigned_cell.value().map(|v| {
-            #[cfg(feature = "halo2-axiom")]
+            #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
             {
                 value = **v;
             }
-            #[cfg(not(feature = "halo2-axiom"))]
+            #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
             {
                 value = *v;
             }

halo2-base/src/virtual_region/lookups/basic.rs

@@ -93,17 +93,17 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         keys: impl IntoIterator<Item = [AssignedValue<F>; KEY_COL]>,
         copy_manager: Option<&SharedCopyConstraintManager<F>>,
     ) {
-        #[cfg(not(feature = "halo2-axiom"))]
+        #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
         let keys = keys.into_iter().collect::<Vec<_>>();
         layouter
             .assign_region(
                 || "[BasicDynLookupConfig] Advice cells to lookup",
                 |mut region| {
                     self.assign_virtual_to_lookup_to_raw_from_offset(
                         &mut region,
-                        #[cfg(feature = "halo2-axiom")]
+                        #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
                         keys,
-                        #[cfg(not(feature = "halo2-axiom"))]
+                        #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
                         keys.clone(),
                         0,
                         copy_manager,
@@ -157,17 +157,17 @@ impl<const KEY_COL: usize> BasicDynLookupConfig<KEY_COL> {
         rows: impl IntoIterator<Item = [AssignedValue<F>; KEY_COL]>,
         copy_manager: Option<&SharedCopyConstraintManager<F>>,
     ) {
-        #[cfg(not(feature = "halo2-axiom"))]
+        #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
         let rows = rows.into_iter().collect::<Vec<_>>();
         layouter
             .assign_region(
                 || "[BasicDynLookupConfig] Dynamic Lookup Table",
                 |mut region| {
                     self.assign_virtual_table_to_raw_from_offset(
                         &mut region,
-                        #[cfg(feature = "halo2-axiom")]
+                        #[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
                         rows,
-                        #[cfg(not(feature = "halo2-axiom"))]
+                        #[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
                         rows.clone(),
                         0,
                         copy_manager,

halo2-base/src/virtual_region/tests/lookups/memory.rs

@@ -4,12 +4,12 @@ use crate::{
         circuit::{Layouter, SimpleFloorPlanner},
         dev::MockProver,
         halo2curves::bn256::Fr,
-        plonk::{keygen_pk, keygen_vk, Assigned, Circuit, ConstraintSystem, Error},
+        plonk::{keygen_pk, keygen_vk, Assigned, Circuit, ConstraintSystem, Error, FirstPhase},
     },
     virtual_region::lookups::basic::BasicDynLookupConfig,
     AssignedValue, ContextCell,
 };
-use halo2_proofs_axiom::plonk::FirstPhase;
+
 use rand::{rngs::StdRng, Rng, SeedableRng};
 use test_log::test;
 

halo2-ecc/Cargo.toml

@@ -37,6 +37,10 @@ display = ["halo2-base/display"]
 asm = ["halo2-base/asm"]
 halo2-pse = ["halo2-base/halo2-pse"]
 halo2-axiom = ["halo2-base/halo2-axiom"]
+halo2-icicle = ["halo2-base/halo2-icicle"]
+profile = ["halo2-base/profile"]
+profile-icicle = ["halo2-base/profile-icicle"]
+halo2-axiom-icicle = ["halo2-base/halo2-axiom-icicle"]
 jemallocator = ["halo2-base/jemallocator"]
 mimalloc = ["halo2-base/mimalloc"]
 

halo2-ecc/src/bn254/tests/bls_signature.rs

@@ -4,11 +4,15 @@ use std::{
 };
 
 use super::*;
-use crate::halo2curves::pairing::{group::ff::Field, MillerLoopResult};
+use crate::halo2curves::pairing::MillerLoopResult;
+#[cfg(any(feature = "halo2-axiom", feature = "halo2-axiom-icicle"))]
+use crate::halo2curves::pairing::group::ff::Field;
 use crate::{
     bn254::bls_signature::BlsSignatureChip, fields::FpStrategy,
     halo2_proofs::halo2curves::bn256::G2Affine,
 };
+#[cfg(any(feature = "halo2-pse", feature = "halo2-icicle"))]
+use halo2_base::halo2_proofs::arithmetic::Field;
 use halo2_base::{
     gates::RangeChip,
     halo2_proofs::halo2curves::bn256::{multi_miller_loop, G2Prepared, Gt},

halo2-ecc/src/bn254/tests/msm.rs

@@ -72,8 +72,23 @@ fn bench_msm() -> Result<(), Box<dyn std::error::Error>> {
     fs::create_dir_all("data").unwrap();
 
     let results_path = "results/bn254/msm_bench.csv";
-    let mut fs_results = File::create(results_path).unwrap();
-    writeln!(fs_results, "degree,num_advice,num_lookup,num_fixed,lookup_bits,limb_bits,num_limbs,batch_size,window_bits,proof_time,proof_size,verify_time")?;
+    let mut fs_results = match File::options().append(true).open(results_path) {
+        Ok(file) => file,
+        Err(_) => {
+            let mut file = File::create(results_path).unwrap();
+            writeln!(file, "halo2_feature,degree,num_advice,num_lookup,num_fixed,lookup_bits,limb_bits,num_limbs,batch_size,window_bits,proof_time,proof_size,verify_time")?;
+            file
+        }
+    };
+
+    #[cfg(feature = "halo2-icicle")]
+    let halo2_feature = "pse-icicle";
+    #[cfg(feature = "halo2-axiom-icicle")]
+    let halo2_feature = "axiom-icicle";
+    #[cfg(feature = "halo2-axiom")]
+    let halo2_feature = "axiom";
+    #[cfg(feature = "halo2-pse")]
+    let halo2_feature = "pse";
 
     let bench_params_reader = BufReader::new(bench_params_file);
     for line in bench_params_reader.lines() {
@@ -93,7 +108,8 @@ fn bench_msm() -> Result<(), Box<dyn std::error::Error>> {
 
         writeln!(
             fs_results,
-            "{},{},{},{},{},{},{},{},{},{:?},{},{:?}",
+            "{},{},{},{},{},{},{},{},{},{},{:?},{},{:?}",
+            halo2_feature,
             bench_params.degree,
             bench_params.num_advice,
             bench_params.num_lookup_advice,

halo2-ecc/src/secp256k1/tests/ecdsa_tests.rs

@@ -2,7 +2,7 @@ use crate::halo2_proofs::{
     arithmetic::CurveAffine,
     halo2curves::secp256k1::{Fq, Secp256k1Affine},
 };
-
+use halo2_base::halo2_proofs::arithmetic::Field;
 use halo2_base::utils::{biguint_to_fe, fe_to_biguint, modulus};
 use rand::random;
 use test_case::test_case;