From cccd6c1679042d5ed0b7149871276f6025bf41a7 Mon Sep 17 00:00:00 2001 From: Yuval Shekel Date: Wed, 31 Jul 2024 19:35:40 +0300 Subject: [PATCH] add ntt benchmark for rust --- wrappers/rust_v3/icicle-core/src/ntt/mod.rs | 136 +++++++++++++++++- .../icicle-curves/icicle-bls12-377/Cargo.toml | 7 + .../icicle-bls12-377/benches/ntt.rs | 5 + .../icicle-curves/icicle-bls12-381/Cargo.toml | 4 + .../icicle-bls12-381/benches/ntt.rs | 5 + .../icicle-curves/icicle-bn254/Cargo.toml | 4 + .../icicle-curves/icicle-bn254/benches/ntt.rs | 5 + .../icicle-curves/icicle-bw6-761/Cargo.toml | 3 + .../icicle-bw6-761/benches/ntt.rs | 5 + 9 files changed, 173 insertions(+), 1 deletion(-) create mode 100644 wrappers/rust_v3/icicle-curves/icicle-bls12-377/benches/ntt.rs create mode 100644 wrappers/rust_v3/icicle-curves/icicle-bls12-381/benches/ntt.rs create mode 100644 wrappers/rust_v3/icicle-curves/icicle-bn254/benches/ntt.rs create mode 100644 wrappers/rust_v3/icicle-curves/icicle-bw6-761/benches/ntt.rs diff --git a/wrappers/rust_v3/icicle-core/src/ntt/mod.rs b/wrappers/rust_v3/icicle-core/src/ntt/mod.rs index b1b24ed3..5ef5d790 100644 --- a/wrappers/rust_v3/icicle-core/src/ntt/mod.rs +++ b/wrappers/rust_v3/icicle-core/src/ntt/mod.rs @@ -427,4 +427,138 @@ macro_rules! impl_ntt_tests { }; } -// TODO Yuval : becnhmarks +#[macro_export] +macro_rules! impl_ntt_bench { + ( + $field_prefix:literal, + $field:ident + ) => { + use std::{env, sync::OnceLock}; + use criterion::{black_box, criterion_group, criterion_main, Criterion}; + use icicle_runtime::{memory::{HostSlice,HostOrDeviceSlice},device::Device,is_device_available, get_active_device, set_device, runtime::load_backend_from_env_or_default}; + use icicle_core::{ + ntt::{NTTConfig, NTTInitDomainConfig, NTTDir, NttAlgorithm, Ordering, NTTDomain, ntt, NTT}, + traits::{GenerateRandom,FieldImpl}, + vec_ops::VecOps, + }; + + fn ntt_for_bench( + input: &(impl HostOrDeviceSlice + ?Sized), + mut batch_ntt_result: &mut (impl HostOrDeviceSlice + ?Sized), + dir: NTTDir, + config: &mut NTTConfig, + _seed: u32, + ) where + ::Config: NTT + GenerateRandom, + ::Config: VecOps, + { + ntt(input, dir, config, batch_ntt_result).unwrap(); + } + + static INIT: OnceLock<()> = OnceLock::new(); + + fn load_and_init_backend_device() { + // Attempt to load the backends + load_backend_from_env_or_default(); // try loading from /opt/icicle/backend or env ${ICICLE_BACKEND_INSTALL_DIR} + + // Check if BENCH_TARGET is defined + let target = env::var("BENCH_TARGET").unwrap_or_else(|_| { + // If not defined, try CUDA first, fallback to CPU + if is_device_available(&Device::new("CUDA", 0)) { + "CUDA".to_string() + } else { + "CPU".to_string() + } + }); + + // Initialize the device with the determined target + let device = Device::new(&target, 0); + set_device(&device).unwrap(); + + println!("ICICLE benchmark with {:?}", device); + } + + fn benchmark_ntt(c: &mut Criterion) + where + ::Config: NTT + GenerateRandom, + ::Config: VecOps, + { + use criterion::SamplingMode; + use icicle_core::ntt::tests::init_domain; + use std::env; + + load_and_init_backend_device(); + + let group_id = format!("{} NTT", $field_prefix); + let mut group = c.benchmark_group(&group_id); + group.sampling_mode(SamplingMode::Flat); + group.sample_size(10); + + const MAX_LOG2: u32 = 25; // max length = 2 ^ MAX_LOG2 + + let max_log2 = env::var("MAX_LOG2") + .unwrap_or_else(|_| MAX_LOG2.to_string()) + .parse::() + .unwrap_or(MAX_LOG2); + + const FAST_TWIDDLES_MODE: bool = false; + + INIT.get_or_init(move || init_domain::<$field>(1 << max_log2, FAST_TWIDDLES_MODE)); + + let coset_generators = [F::one(), F::Config::generate_random(1)[0]]; + let mut config = NTTConfig::::default(); + + for test_size_log2 in (13u32..=max_log2) { + for batch_size_log2 in (7u32..17u32) { + let test_size = 1 << test_size_log2; + let batch_size = 1 << batch_size_log2; + let full_size = batch_size * test_size; + + if full_size > 1 << max_log2 { + continue; + } + + let scalars = F::Config::generate_random(full_size); + let input = HostSlice::from_slice(&scalars); + + let mut batch_ntt_result = vec![F::zero(); batch_size * test_size]; + let batch_ntt_result = HostSlice::from_mut_slice(&mut batch_ntt_result); + let mut config = NTTConfig::::default(); + for dir in [NTTDir::kForward, NTTDir::kInverse ] { + for ordering in [ + Ordering::kNN, + Ordering::kNR, + Ordering::kRN, + Ordering::kRR, + Ordering::kNM, + Ordering::kMN, + ] { + config.ordering = ordering; + config.batch_size = batch_size as i32; + let bench_descr = format!( + "{:?} {:?} {} x {}", + ordering, dir, test_size, batch_size + ); + group.bench_function(&bench_descr, |b| { + b.iter(|| { + ntt_for_bench::( + input, + batch_ntt_result, + dir, + &mut config, + black_box(1), + ) + }) + }); + } + } + } + } + + group.finish(); + } + + criterion_group!(benches, benchmark_ntt<$field, $field>); + criterion_main!(benches); + }; +} diff --git a/wrappers/rust_v3/icicle-curves/icicle-bls12-377/Cargo.toml b/wrappers/rust_v3/icicle-curves/icicle-bls12-377/Cargo.toml index 5d3b8c8b..9f56aa59 100644 --- a/wrappers/rust_v3/icicle-curves/icicle-bls12-377/Cargo.toml +++ b/wrappers/rust_v3/icicle-curves/icicle-bls12-377/Cargo.toml @@ -27,3 +27,10 @@ g2 = ["icicle-core/g2"] ec_ntt = ["icicle-core/ec_ntt"] cuda_backend = ["icicle-runtime/cuda_backend"] pull_cuda_backend = ["icicle-runtime/pull_cuda_backend"] + + +[[bench]] +name = "ntt" +harness = false + + diff --git a/wrappers/rust_v3/icicle-curves/icicle-bls12-377/benches/ntt.rs b/wrappers/rust_v3/icicle-curves/icicle-bls12-377/benches/ntt.rs new file mode 100644 index 00000000..bef2f78d --- /dev/null +++ b/wrappers/rust_v3/icicle-curves/icicle-bls12-377/benches/ntt.rs @@ -0,0 +1,5 @@ +use icicle_bls12_377::curve::ScalarField; + +use icicle_core::impl_ntt_bench; + +impl_ntt_bench!("bls12_377", ScalarField); diff --git a/wrappers/rust_v3/icicle-curves/icicle-bls12-381/Cargo.toml b/wrappers/rust_v3/icicle-curves/icicle-bls12-381/Cargo.toml index 0f345e11..d9a0cb14 100644 --- a/wrappers/rust_v3/icicle-curves/icicle-bls12-381/Cargo.toml +++ b/wrappers/rust_v3/icicle-curves/icicle-bls12-381/Cargo.toml @@ -25,3 +25,7 @@ g2 = ["icicle-core/g2"] ec_ntt = ["icicle-core/ec_ntt"] cuda_backend = ["icicle-runtime/cuda_backend"] pull_cuda_backend = ["icicle-runtime/pull_cuda_backend"] + +[[bench]] +name = "ntt" +harness = false diff --git a/wrappers/rust_v3/icicle-curves/icicle-bls12-381/benches/ntt.rs b/wrappers/rust_v3/icicle-curves/icicle-bls12-381/benches/ntt.rs new file mode 100644 index 00000000..2cf2a351 --- /dev/null +++ b/wrappers/rust_v3/icicle-curves/icicle-bls12-381/benches/ntt.rs @@ -0,0 +1,5 @@ +use icicle_bls12_381::curve::ScalarField; + +use icicle_core::impl_ntt_bench; + +impl_ntt_bench!("bls12_381", ScalarField); diff --git a/wrappers/rust_v3/icicle-curves/icicle-bn254/Cargo.toml b/wrappers/rust_v3/icicle-curves/icicle-bn254/Cargo.toml index 7dad8ffc..8d17852d 100644 --- a/wrappers/rust_v3/icicle-curves/icicle-bn254/Cargo.toml +++ b/wrappers/rust_v3/icicle-curves/icicle-bn254/Cargo.toml @@ -25,3 +25,7 @@ g2 = ["icicle-core/g2"] ec_ntt = ["icicle-core/ec_ntt"] cuda_backend = ["icicle-runtime/cuda_backend"] pull_cuda_backend = ["icicle-runtime/pull_cuda_backend"] + +[[bench]] +name = "ntt" +harness = false diff --git a/wrappers/rust_v3/icicle-curves/icicle-bn254/benches/ntt.rs b/wrappers/rust_v3/icicle-curves/icicle-bn254/benches/ntt.rs new file mode 100644 index 00000000..b0eac3e3 --- /dev/null +++ b/wrappers/rust_v3/icicle-curves/icicle-bn254/benches/ntt.rs @@ -0,0 +1,5 @@ +use icicle_bn254::curve::ScalarField; + +use icicle_core::impl_ntt_bench; + +impl_ntt_bench!("bn254", ScalarField); diff --git a/wrappers/rust_v3/icicle-curves/icicle-bw6-761/Cargo.toml b/wrappers/rust_v3/icicle-curves/icicle-bw6-761/Cargo.toml index 65167c28..2b7c7030 100644 --- a/wrappers/rust_v3/icicle-curves/icicle-bw6-761/Cargo.toml +++ b/wrappers/rust_v3/icicle-curves/icicle-bw6-761/Cargo.toml @@ -24,3 +24,6 @@ cmake = "0.1.50" default = [] g2 = ["icicle-bls12-377/bw6-761-g2"] +[[bench]] +name = "ntt" +harness = false \ No newline at end of file diff --git a/wrappers/rust_v3/icicle-curves/icicle-bw6-761/benches/ntt.rs b/wrappers/rust_v3/icicle-curves/icicle-bw6-761/benches/ntt.rs new file mode 100644 index 00000000..dd02bbdf --- /dev/null +++ b/wrappers/rust_v3/icicle-curves/icicle-bw6-761/benches/ntt.rs @@ -0,0 +1,5 @@ +use icicle_bw6_761::curve::ScalarField; + +use icicle_core::impl_ntt_bench; + +impl_ntt_bench!("bw6_761", ScalarField);