Skip to content

Commit

Permalink
make icicle purely optional
Browse files Browse the repository at this point in the history
  • Loading branch information
sagar-a16z committed Nov 15, 2024
1 parent a32f4b3 commit b6201b7
Show file tree
Hide file tree
Showing 5 changed files with 198 additions and 180 deletions.
10 changes: 5 additions & 5 deletions jolt-core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"
Expand All @@ -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"] }
Expand Down Expand Up @@ -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]
Expand Down
87 changes: 11 additions & 76 deletions jolt-core/src/msm/icicle.rs → jolt-core/src/msm/icicle/adapter.rs
Original file line number Diff line number Diff line change
@@ -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,
};
Expand All @@ -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;
Expand Down Expand Up @@ -165,7 +99,7 @@ pub fn icicle_msm<V: VariableBaseMSM + Icicle>(
&cfg,
&mut msm_result[..],
)
.unwrap();
.unwrap();

drop(_guard);
drop(span);
Expand Down Expand Up @@ -212,7 +146,7 @@ pub fn icicle_batch_msm<V: VariableBaseMSM + Icicle>(
DeviceVec::<<<V as Icicle>::C as Curve>::ScalarField>::device_malloc_async(
len, &stream,
)
.unwrap();
.unwrap();
let scalars_mont = unsafe {
&*(&scalars[..] as *const _ as *const [<<V as Icicle>::C as Curve>::ScalarField])
};
Expand All @@ -233,7 +167,7 @@ pub fn icicle_batch_msm<V: VariableBaseMSM + Icicle>(
&cfg,
&mut msm_result[..],
)
.unwrap();
.unwrap();

msm_result
.copy_to_host_async(
Expand Down Expand Up @@ -278,7 +212,7 @@ pub fn icicle_variable_batch_msm<V: VariableBaseMSM + Icicle>(
scalars.len(),
&stream,
)
.unwrap();
.unwrap();
let scalars_mont = unsafe {
&*(&scalars[..] as *const _ as *const [<<V as Icicle>::C as Curve>::ScalarField])
};
Expand All @@ -299,7 +233,7 @@ pub fn icicle_variable_batch_msm<V: VariableBaseMSM + Icicle>(
&cfg,
&mut msm_result[..],
)
.unwrap();
.unwrap();

msm_result
.copy_to_host_async(
Expand Down Expand Up @@ -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`
Expand Down Expand Up @@ -395,4 +330,4 @@ mod tests {
let total = total_memory_bits();
assert!(total > 0);
}
}
}
94 changes: 94 additions & 0 deletions jolt-core/src/msm/icicle/mod.rs
Original file line number Diff line number Diff line change
@@ -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)
}
Loading

0 comments on commit b6201b7

Please sign in to comment.