-
Notifications
You must be signed in to change notification settings - Fork 175
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
a32f4b3
commit b6201b7
Showing
5 changed files
with
198 additions
and
180 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | ||
} |
Oops, something went wrong.