From 5ea422700e821519ec90ef444d4566bb213ce7a9 Mon Sep 17 00:00:00 2001 From: Sagar Dhawan Date: Thu, 14 Nov 2024 18:05:08 -0800 Subject: [PATCH] make icicle purely optional --- jolt-core/Cargo.toml | 10 +- .../src/msm/{icicle.rs => icicle/adapter.rs} | 87 ++---------- jolt-core/src/msm/icicle/mod.rs | 94 +++++++++++++ jolt-core/src/msm/mod.rs | 133 ++++++++++++------ jolt-evm-verifier/script/Cargo.lock | 54 ------- 5 files changed, 198 insertions(+), 180 deletions(-) rename jolt-core/src/msm/{icicle.rs => icicle/adapter.rs} (81%) create mode 100644 jolt-core/src/msm/icicle/mod.rs diff --git a/jolt-core/Cargo.toml b/jolt-core/Cargo.toml index 02f4076d7..38a1f12be 100644 --- a/jolt-core/Cargo.toml +++ b/jolt-core/Cargo.toml @@ -29,7 +29,7 @@ default = [ "rayon", ] host = ["dep:reqwest", "dep:tokio"] -icicle = ["default"] +icicle = ["default", "dep:icicle-runtime", "dep:icicle-core", "dep:icicle-bn254"] [dependencies] ark-bn254 = "0.4.0" @@ -53,9 +53,6 @@ rand = "0.7.3" rand_chacha = { version = "0.3.0", default-features = false } rand_core = { version = "0.6.4", default-features = false } rayon = { version = "^1.8.0", optional = true } -icicle-runtime = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v3.1.0" } -icicle-core = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v3.1.0" } -icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v3.1.0" } serde = { version = "1.0.*", default-features = false } sha3 = "0.10.8" strum = "0.26.3" @@ -82,7 +79,6 @@ tokio = { version = "1.38.0", optional = true } alloy-primitives = "0.7.6" alloy-sol-types = "0.7.6" once_cell = "1.19.0" -sys-info = "0.9.1" [dev-dependencies] criterion = { version = "0.5.1", features = ["html_reports"] } @@ -120,7 +116,11 @@ name = "jolt_core" path = "src/lib.rs" [target.'cfg(not(target_arch = "wasm32"))'.dependencies] +icicle-runtime = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v3.1.0", optional = true } +icicle-core = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v3.1.0", optional = true } +icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v3.1.0", 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/src/msm/icicle.rs b/jolt-core/src/msm/icicle/adapter.rs similarity index 81% rename from jolt-core/src/msm/icicle.rs rename to jolt-core/src/msm/icicle/adapter.rs index cf92ce340..0feb4bcf0 100644 --- a/jolt-core/src/msm/icicle.rs +++ b/jolt-core/src/msm/icicle/adapter.rs @@ -1,10 +1,8 @@ use crate::msm::VariableBaseMSM; -use ark_bn254::G1Projective; -use ark_ec::{CurveGroup, ScalarMul}; +use ark_ec::{ScalarMul, CurveGroup}; use ark_ff::{BigInteger, Field, PrimeField}; -use icicle_bn254::curve::CurveCfg as IcicleBn254; +use icicle_bn254::curve::{CurveCfg as IcicleBn254}; use icicle_core::{ - curve::{Affine, Curve, Projective}, msm::{msm, MSMConfig, MSM}, traits::FieldImpl, }; @@ -13,72 +11,8 @@ use icicle_runtime::{ memory::{DeviceVec, HostSlice}, stream::IcicleStream, }; -use std::sync::Once; - -pub static ICICLE_INIT: Once = Once::new(); -static mut ICICLE_READY: bool = false; - -/// Initializes the icicle backend and sets the CUDA device as active and returns true if successful. -/// -/// Safe to call multiple times; will only initialize the backend once. -pub fn icicle_init() -> bool { - let mut initialized = false; - - ICICLE_INIT.call_once(|| { - #[cfg(feature = "icicle")] - if icicle_runtime::load_backend_from_env_or_default().is_ok() { - if let Ok(devices) = icicle_runtime::get_registered_devices() { - println!("Initializing icicle: available devices {:?}", devices); - - // Attempt to set the CUDA device as active - let device = icicle_runtime::Device::new("CUDA", 0); - if icicle_runtime::set_device(&device).is_ok() { - println!("icicle using device: {:?}", device); - initialized = true; - } else { - println!("Failed to set CUDA device; falling back to CPU."); - } - } - } - - #[cfg(not(feature = "icicle"))] - { - initialized = false; - } - - #[cfg(feature = "icicle")] - if !initialized { - println!("Failed to initialize icicle backend; using JOLT CPU implementations."); - } - - unsafe { ICICLE_READY = initialized }; - }); - - unsafe { ICICLE_READY } -} - -/// Returns the total memory available on the system in bits. -/// -/// If icicle is enabled, it will return the total memory available on the GPU in bits. -pub fn total_memory_bits() -> usize { - const DEFAULT_MEM_GB: usize = 30; - const BITS_PER_BYTE: usize = 8; - const BYTES_PER_KB: usize = 1024; - const BYTES_PER_GB: usize = 1024 * 1024 * 1024; - - if let Ok((total_bytes, _)) = icicle_runtime::get_available_memory() { - // If icicle is enabled and memory is available, return the total memory in bits. - return total_bytes * BITS_PER_BYTE; - } - - // Fallback to system memory if icicle is unavailable or not enabled. - if let Ok(mem_info) = sys_info::mem_info() { - return (mem_info.total as usize * BYTES_PER_KB) * BITS_PER_BYTE; - } - - // Fallback to "default" memory if system memory retrieval fails. - DEFAULT_MEM_GB * BYTES_PER_GB * BITS_PER_BYTE -} +use ark_bn254::G1Projective; +use icicle_core::curve::{Affine, Curve, Projective}; impl Icicle for G1Projective { type C = IcicleBn254; @@ -165,7 +99,7 @@ pub fn icicle_msm( &cfg, &mut msm_result[..], ) - .unwrap(); + .unwrap(); drop(_guard); drop(span); @@ -212,7 +146,7 @@ pub fn icicle_batch_msm( DeviceVec::<<::C as Curve>::ScalarField>::device_malloc_async( len, &stream, ) - .unwrap(); + .unwrap(); let scalars_mont = unsafe { &*(&scalars[..] as *const _ as *const [<::C as Curve>::ScalarField]) }; @@ -233,7 +167,7 @@ pub fn icicle_batch_msm( &cfg, &mut msm_result[..], ) - .unwrap(); + .unwrap(); msm_result .copy_to_host_async( @@ -278,7 +212,7 @@ pub fn icicle_variable_batch_msm( scalars.len(), &stream, ) - .unwrap(); + .unwrap(); let scalars_mont = unsafe { &*(&scalars[..] as *const _ as *const [<::C as Curve>::ScalarField]) }; @@ -299,7 +233,7 @@ pub fn icicle_variable_batch_msm( &cfg, &mut msm_result[..], ) - .unwrap(); + .unwrap(); msm_result .copy_to_host_async( @@ -348,6 +282,7 @@ mod tests { use icicle_bn254::curve::ScalarField as GPUScalar; use rand_core::SeedableRng; use rayon::prelude::*; + use crate::msm::total_memory_bits; // Note due to contention for the gpu device testing using multiple threads leads to unit tests that intermittently fail. // To avoid this run `cargo t --features icicle -- --test-threads=1` @@ -395,4 +330,4 @@ mod tests { let total = total_memory_bits(); assert!(total > 0); } -} +} \ No newline at end of file diff --git a/jolt-core/src/msm/icicle/mod.rs b/jolt-core/src/msm/icicle/mod.rs new file mode 100644 index 000000000..d02ee8506 --- /dev/null +++ b/jolt-core/src/msm/icicle/mod.rs @@ -0,0 +1,94 @@ +#[allow(dead_code)] + +use ark_ec::{CurveGroup, ScalarMul}; +use std::sync::Once; +#[cfg(not(feature = "icicle"))] +use ark_bn254::G1Projective; + +#[cfg(feature = "icicle")] +pub(crate) mod adapter; +#[cfg(feature = "icicle")] +pub use adapter::*; + +static ICICLE_INIT: Once = Once::new(); +static mut ICICLE_READY: bool = false; + +#[cfg(feature = "icicle")] +pub trait CurveGroupConfig: CurveGroup + Icicle {} +#[cfg(not(feature = "icicle"))] +pub trait CurveGroupConfig: CurveGroup {} + +#[cfg(feature = "icicle")] +pub trait ScalarMulConfig: ScalarMul + Icicle {} +#[cfg(not(feature = "icicle"))] +pub trait ScalarMulConfig: ScalarMul {} + +#[cfg(not(feature = "icicle"))] +pub trait Icicle {} +#[cfg(not(feature = "icicle"))] +impl Icicle for G1Projective {} + +/// Initializes the icicle backend and sets the CUDA device as active and returns true if successful. +/// +/// Safe to call multiple times; will only initialize the backend once. +pub fn icicle_init() -> bool { + let mut initialized = false; + + ICICLE_INIT.call_once(|| { + #[cfg(feature = "icicle")] + if icicle_runtime::load_backend_from_env_or_default().is_ok() { + if let Ok(devices) = icicle_runtime::get_registered_devices() { + println!("Initializing icicle: available devices {:?}", devices); + + // Attempt to set the CUDA device as active + let device = icicle_runtime::Device::new("CUDA", 0); + if icicle_runtime::set_device(&device).is_ok() { + println!("icicle using device: {:?}", device); + initialized = true; + } else { + println!("Failed to set CUDA device; falling back to CPU."); + } + } + } + + #[cfg(not(feature = "icicle"))] + { + initialized = false; + } + + #[cfg(feature = "icicle")] + if !initialized { + println!("Failed to initialize icicle backend; using JOLT CPU implementations."); + } + + unsafe { ICICLE_READY = initialized }; + }); + + unsafe { ICICLE_READY } +} + +/// Returns the total memory available on the system in bits. +/// +/// If icicle is enabled, it will return the total memory available on the GPU in bits. +#[allow(dead_code)] +pub fn total_memory_bits() -> usize { + const DEFAULT_MEM_GB: usize = 30; + const BITS_PER_BYTE: usize = 8; + const BYTES_PER_KB: usize = 1024; + const BYTES_PER_GB: usize = 1024 * 1024 * 1024; + + #[cfg(feature = "icicle")] + if let Ok((total_bytes, _)) = icicle_runtime::get_available_memory() { + // If icicle is enabled and memory is available, return the total memory in bits. + return total_bytes.checked_mul(BITS_PER_BYTE).unwrap_or(usize::MAX); + } + + // Fallback to system memory if icicle is unavailable or not enabled. + #[cfg(not(target_arch = "wasm32"))] + if let Ok(mem_info) = sys_info::mem_info() { + return (mem_info.total as usize * BYTES_PER_KB).checked_mul(BITS_PER_BYTE).unwrap_or(usize::MAX); + } + + // Fallback to "default" memory if system memory retrieval fails. + DEFAULT_MEM_GB.checked_mul(BYTES_PER_GB.checked_mul(BITS_PER_BYTE).unwrap_or(usize::MAX)).unwrap_or(usize::MAX) +} diff --git a/jolt-core/src/msm/mod.rs b/jolt-core/src/msm/mod.rs index 9d9ba9ca1..f139c195b 100644 --- a/jolt-core/src/msm/mod.rs +++ b/jolt-core/src/msm/mod.rs @@ -34,11 +34,18 @@ pub trait VariableBaseMSM: ScalarMul + Icicle { } 11..=64 => { if icicle_init() { - let gpu_bases = bases - .par_iter() - .map(|base| ::from_ark_affine(base)) - .collect::>(); - return icicle_msm::(&gpu_bases, scalars, 64); + #[cfg(feature = "icicle")] + { + let gpu_bases = bases + .par_iter() + .map(|base| ::from_ark_affine(base)) + .collect::>(); + return icicle_msm::(&gpu_bases, scalars, 64); + } + #[cfg(not(feature = "icicle"))] + { + unreachable!("icicle_init must not return true without the icicle feature"); + } } let scalars_u64 = &map_field_elements_to_u64::(scalars); @@ -50,11 +57,18 @@ pub trait VariableBaseMSM: ScalarMul + Icicle { } _ => { if icicle_init() { - let gpu_bases = bases - .par_iter() - .map(|base| ::from_ark_affine(base)) - .collect::>(); - return icicle_msm::(&gpu_bases, scalars, 256); + #[cfg(feature = "icicle")] + { + let gpu_bases = bases + .par_iter() + .map(|base| ::from_ark_affine(base)) + .collect::>(); + return icicle_msm::(&gpu_bases, scalars, 256); + } + #[cfg(not(feature = "icicle"))] + { + unreachable!("icicle_init must not return true without the icicle feature"); + } } let scalars = scalars @@ -77,6 +91,7 @@ pub trait VariableBaseMSM: ScalarMul + Icicle { assert!(scalars.iter().all(|s| s.len() == scalars[0].len())); assert_eq!(bases.len(), scalars[0].len()); + #[cfg(feature = "icicle")] let gpu_bases = if icicle_init() { bases .par_iter() @@ -137,16 +152,23 @@ pub trait VariableBaseMSM: ScalarMul + Icicle { MsmType::Small => { if use_icicle { - let scalar_batches: Vec<&[Self::ScalarField]> = - indices.iter().map(|i| scalars[*i]).collect(); - let batch_results = - icicle_batch_msm::(&gpu_bases, &scalar_batches, 10); - assert_eq!(batch_results.len(), scalar_batches.len()); - batch_results - .into_iter() - .enumerate() - .map(|(batch_index, result)| (indices[batch_index], result)) - .collect() + #[cfg(feature = "icicle")] + { + let scalar_batches: Vec<&[Self::ScalarField]> = + indices.iter().map(|i| scalars[*i]).collect(); + let batch_results = + icicle_batch_msm::(&gpu_bases, &scalar_batches, 10); + assert_eq!(batch_results.len(), scalar_batches.len()); + batch_results + .into_iter() + .enumerate() + .map(|(batch_index, result)| (indices[batch_index], result)) + .collect() + } + #[cfg(not(feature = "icicle"))] + { + unreachable!("icicle_init must not return true without the icicle feature"); + } } else { indices .into_par_iter() @@ -161,16 +183,23 @@ pub trait VariableBaseMSM: ScalarMul + Icicle { MsmType::Medium => { if use_icicle { - let scalar_batches: Vec<&[Self::ScalarField]> = - indices.iter().map(|i| scalars[*i]).collect(); - let batch_results = - icicle_batch_msm::(&gpu_bases, &scalar_batches, 64); - assert_eq!(batch_results.len(), scalar_batches.len()); - batch_results - .into_iter() - .enumerate() - .map(|(batch_index, result)| (indices[batch_index], result)) - .collect() + #[cfg(feature = "icicle")] + { + let scalar_batches: Vec<&[Self::ScalarField]> = + indices.iter().map(|i| scalars[*i]).collect(); + let batch_results = + icicle_batch_msm::(&gpu_bases, &scalar_batches, 64); + assert_eq!(batch_results.len(), scalar_batches.len()); + batch_results + .into_iter() + .enumerate() + .map(|(batch_index, result)| (indices[batch_index], result)) + .collect() + } + #[cfg(not(feature = "icicle"))] + { + unreachable!("icicle_init must not return true without the icicle feature"); + } } else { indices .into_par_iter() @@ -190,16 +219,23 @@ pub trait VariableBaseMSM: ScalarMul + Icicle { MsmType::Large => { if use_icicle { - let scalar_batches: Vec<&[Self::ScalarField]> = - indices.iter().map(|i| scalars[*i]).collect(); - let batch_results = - icicle_batch_msm::(&gpu_bases, &scalar_batches, 256); - assert_eq!(batch_results.len(), scalar_batches.len()); - batch_results - .into_iter() - .enumerate() - .map(|(batch_index, result)| (indices[batch_index], result)) - .collect() + #[cfg(feature = "icicle")] + { + let scalar_batches: Vec<&[Self::ScalarField]> = + indices.iter().map(|i| scalars[*i]).collect(); + let batch_results = + icicle_batch_msm::(&gpu_bases, &scalar_batches, 256); + assert_eq!(batch_results.len(), scalar_batches.len()); + batch_results + .into_iter() + .enumerate() + .map(|(batch_index, result)| (indices[batch_index], result)) + .collect() + } + #[cfg(not(feature = "icicle"))] + { + unreachable!("icicle_init must not return true without the icicle feature"); + } } else { indices .into_par_iter() @@ -277,11 +313,18 @@ pub trait VariableBaseMSM: ScalarMul + Icicle { let use_icicle = icicle_init(); if use_icicle { - let gpu_bases = bases - .par_iter() - .map(|base| ::from_ark_affine(base)) - .collect::>(); - icicle_variable_batch_msm::(&gpu_bases, scalar_batches, 256) + #[cfg(feature = "icicle")] + { + let gpu_bases = bases + .par_iter() + .map(|base| ::from_ark_affine(base)) + .collect::>(); + icicle_variable_batch_msm::(&gpu_bases, scalar_batches, 256) + } + #[cfg(not(feature = "icicle"))] + { + unreachable!("icicle_init must not return true without the icicle feature"); + } } else { scalar_batches .par_iter() diff --git a/jolt-evm-verifier/script/Cargo.lock b/jolt-evm-verifier/script/Cargo.lock index f94514c28..1e6a95473 100644 --- a/jolt-evm-verifier/script/Cargo.lock +++ b/jolt-evm-verifier/script/Cargo.lock @@ -599,15 +599,6 @@ version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1462739cb27611015575c0c11df5df7601141071f07518d56fcc1be504cbec97" -[[package]] -name = "cmake" -version = "0.1.51" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fb1e43aa7fd152b1f968787f7dbcdeb306d1867ff373c69955211876c053f91a" -dependencies = [ - "cc", -] - [[package]] name = "cobs" version = "0.2.3" @@ -1326,48 +1317,6 @@ dependencies = [ "tracing", ] -[[package]] -name = "icicle-bn254" -version = "3.1.0" -source = "git+https://github.com/ingonyama-zk/icicle.git?tag=v3.1.0#38712a95af4a118a124321c81383daa93f59f1e4" -dependencies = [ - "cmake", - "icicle-core", - "icicle-hash", - "icicle-runtime", -] - -[[package]] -name = "icicle-core" -version = "3.1.0" -source = "git+https://github.com/ingonyama-zk/icicle.git?tag=v3.1.0#38712a95af4a118a124321c81383daa93f59f1e4" -dependencies = [ - "hex", - "icicle-runtime", - "once_cell", - "rand 0.8.5", - "rayon", -] - -[[package]] -name = "icicle-hash" -version = "3.1.0" -source = "git+https://github.com/ingonyama-zk/icicle.git?tag=v3.1.0#38712a95af4a118a124321c81383daa93f59f1e4" -dependencies = [ - "cmake", - "icicle-core", - "icicle-runtime", - "rand 0.8.5", -] - -[[package]] -name = "icicle-runtime" -version = "3.1.0" -source = "git+https://github.com/ingonyama-zk/icicle.git?tag=v3.1.0#38712a95af4a118a124321c81383daa93f59f1e4" -dependencies = [ - "cmake", -] - [[package]] name = "icu_collections" version = "1.5.0" @@ -1622,9 +1571,6 @@ dependencies = [ "eyre", "fixedbitset", "getrandom 0.2.15", - "icicle-bn254", - "icicle-core", - "icicle-runtime", "indicatif", "itertools 0.10.5", "lazy_static",