Skip to content

Fix benchmarks for ring-gsw over BGG+ encodings #135

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SoraSuegami merged 5 commits into from
Jun 30, 2026
Merged

Fix benchmarks for ring-gsw over BGG+ encodings #135

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
70b5844
Add lattice estimator checks to GPU benchmark tests
SoraSuegami Jun 28, 2026
105df43
Add Ring-GSW bench lattice estimator skip
SoraSuegami Jun 28, 2026
82db042
Add pubkey preprocess bench estimate
SoraSuegami Jun 28, 2026
82a486d
Log PRF final-FE breakdown in IO estimators
SoraSuegami Jun 28, 2026
076dc37
init storage in we
SoraSuegami Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions src/io/aky24_io/bench_estimator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,20 @@ impl<'a, PKBE, EncBE, NBE> Aky24IOBenchEstimator<'a, PKBE, EncBE, NBE> {
repeat_sequential_summary(refresh_decoder_work_per_round, shape.prf_round_count)
}
};
let homomorphic_prf = sequential_summaries(&[round_summary.clone(), final_prg.clone()]);
tracing::info!(
?mode,
homomorphic_prf_total_time_nanos = %homomorphic_prf.total_time,
homomorphic_prf_latency = homomorphic_prf.latency,
homomorphic_prf_max_parallelism = %homomorphic_prf.max_parallelism,
final_mask_total_time_nanos = %final_mask_decrypt.total_time,
final_mask_latency = final_mask_decrypt.latency,
final_mask_max_parallelism = %final_mask_decrypt.max_parallelism,
function_outputs_total_time_nanos = %final_function_decrypt.total_time,
function_outputs_latency = final_function_decrypt.latency,
function_outputs_max_parallelism = %final_function_decrypt.max_parallelism,
"estimated AKY24IO PRF final-FE 3-way breakdown"
);
let final_summary = sequential_summaries(&[final_prg.clone(), final_output_decrypt]);
let compute_without_refresh_decoder =
sequential_summaries(&[round_compute, final_summary.clone()]);
Expand Down
14 changes: 14 additions & 0 deletions src/io/diamond_io/bench_estimator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1080,6 +1080,20 @@ where
shape.prf_round_count,
),
};
let homomorphic_prf = sequential_summaries(&[round_summary.clone(), final_prg.clone()]);
info!(
?mode,
homomorphic_prf_total_time_nanos = %homomorphic_prf.total_time,
homomorphic_prf_latency = homomorphic_prf.latency,
homomorphic_prf_max_parallelism = %homomorphic_prf.max_parallelism,
final_mask_total_time_nanos = %final_mask_decrypt.total_time,
final_mask_latency = final_mask_decrypt.latency,
final_mask_max_parallelism = %final_mask_decrypt.max_parallelism,
function_outputs_total_time_nanos = %final_function_decrypt.total_time,
function_outputs_latency = final_function_decrypt.latency,
function_outputs_max_parallelism = %final_function_decrypt.max_parallelism,
"estimated DiamondIO PRF final-FE 3-way breakdown"
);
let final_summary =
sequential_summaries(&[final_prg.clone(), final_output_decrypt.clone()]);
let compute_without_refresh_decoder =
Expand Down
81 changes: 75 additions & 6 deletions tests/test_gpu_diamond_we.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,22 @@ use keccak_asm::Keccak256;
use mxx::{
bench_estimator::{
BggEncodingBenchEstimator, BggPublicKeyBenchEstimator, BggPublicKeyBenchSamples,
PublicLutSampleAuxBenchEstimator,
},
bgg::{
public_key::BggPublicKey,
sampler::{BGGEncodingSampler, BGGPublicKeySampler},
},
bgg::{public_key::BggPublicKey, sampler::BGGPublicKeySampler},
circuit::{PolyCircuit, gate::GateId},
element::PolyElem,
func_enc::NoCircuitEvaluator,
input_injector::DiamondInjector,
io::diamond_io::GpuDCRTPolyMatrixNativeBenchEstimator,
lookup::{PublicLut, lwe::LWEBGGPubKeyPltEvaluator},
matrix::{dcrt_poly::DCRTPolyMatrix, gpu_dcrt_poly::GpuDCRTPolyMatrix},
lookup::{
PltEvaluator, PublicLut,
lwe::{LWEBGGEncodingPltEvaluator, LWEBGGPubKeyPltEvaluator},
},
matrix::{PolyMatrix, dcrt_poly::DCRTPolyMatrix, gpu_dcrt_poly::GpuDCRTPolyMatrix},
poly::{
Poly, PolyParams,
dcrt::{
Expand All @@ -22,14 +29,14 @@ use mxx::{
},
},
sampler::{
PolyTrapdoorSampler,
DistType, PolyTrapdoorSampler, PolyUniformSampler,
gpu::{GpuDCRTPolyHashSampler, GpuDCRTPolyUniformSampler},
hash::DCRTPolyHashSampler,
trapdoor::{DCRTPolyTrapdoorSampler, GpuDCRTPolyTrapdoorSampler},
uniform::DCRTPolyUniformSampler,
},
slot_transfer::bgg_pubkey::BggPublicKeySTEvaluator,
storage::write::init_storage_system,
storage::write::{init_storage_system, wait_for_all_writes},
utils::bigdecimal_bits_ceil,
we::{
DiamondWE, WitnessEnc,
Expand Down Expand Up @@ -95,6 +102,7 @@ type GpuDiamondWE = DiamondWE<
>;
type GpuPubKeyPltEvaluator =
LWEBGGPubKeyPltEvaluator<GpuMatrix, GpuHashSampler, GpuDCRTPolyTrapdoorSampler>;
type GpuEncodingPltEvaluator = LWEBGGEncodingPltEvaluator<GpuMatrix, GpuHashSampler>;
type GpuPubKeySlotEvaluator = BggPublicKeySTEvaluator<
GpuMatrix,
GpuDCRTPolyUniformSampler,
Expand Down Expand Up @@ -541,6 +549,67 @@ fn build_public_key_bench_estimator(
)
}

async fn build_encoding_bench_estimator(
params: &GpuDCRTPolyParams,
cfg: &DiamondWEGpuBenchConfig,
bench_dir: PathBuf,
) -> BggEncodingBenchEstimator<GpuMatrix> {
if bench_dir.exists() {
fs::remove_dir_all(&bench_dir).expect("failed to clear DiamondWE encoding bench directory");
}
ensure_dir(&bench_dir);

@chatgpt-codex-connector chatgpt-codex-connector Bot Jun 28, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Initialize storage for encoding aux files

When this helper runs, the global LUT writer is still initialized to final_dir from the caller, so write_dummy_aux_for_poly_encode_bench writes the batch files there while wait_for_all_writes(bench_dir) writes the index under final_dir/encoding_bench; the following LWEBGGEncodingPltEvaluator then reads from bench_dir and cannot find the referenced batch files. Initialize the storage system with bench_dir.clone() before enqueueing these dummy aux writes, as the other new encoding benchmark helpers do.

Useful? React with 👍 / 👎.

init_storage_system(bench_dir.clone());

let (public_lut, public_lut_id, public_lut_gate_id) =
build_benchmark_public_lookup_gate(params);
let public_key_sampler =
BGGPublicKeySampler::<_, GpuHashSampler>::new(DEFAULT_HASH_KEY, cfg.d_secret);
let uniform_sampler = GpuDCRTPolyUniformSampler::new();
let secrets =
uniform_sampler.sample_uniform(params, 1, cfg.d_secret, DistType::TernaryDist).get_row(0);
let plaintexts = vec![GpuDCRTPoly::from_usize_to_constant(params, 2)];
let public_keys =
public_key_sampler.sample(params, b"test_gpu_diamond_we_encoding_bench", &[true]);
let encoding_sampler =
BGGEncodingSampler::<GpuDCRTPolyUniformSampler>::new(params, &secrets, None);
let encodings = encoding_sampler.sample(params, &public_keys, &plaintexts);

let public_lut_aux_writer = build_lwe_scalar_pubkey_plt_evaluator(
params,
DEFAULT_HASH_KEY,
cfg.d_secret,
bench_dir.clone(),
);
let lookup_base_matrix = Arc::clone(&public_lut_aux_writer.pub_matrix);
public_lut_aux_writer.write_dummy_aux_for_poly_encode_bench(
params,
&public_lut,
&[plaintexts[0].const_coeff_u64()],
public_lut_id,
public_lut_gate_id,
cfg.error_sigma,
);
wait_for_all_writes(bench_dir.clone())
.await
.expect("DiamondWE encoding public lookup benchmark aux writes must flush");

let secret_vec = GpuMatrix::from_poly_vec_row(params, secrets);
let c_b = secret_vec * lookup_base_matrix.as_ref();
let encoding_evaluator = GpuEncodingPltEvaluator::new(DEFAULT_HASH_KEY, bench_dir, c_b);
let one = encodings[0].clone();
let input = encodings[1].clone();
BggEncodingBenchEstimator::<GpuMatrix>::benchmark(params, cfg.bench_iterations, || {
encoding_evaluator.public_lookup(
params,
&public_lut,
&one,
&input,
public_lut_gate_id,
public_lut_id,
)
})
}

#[tokio::test]
#[sequential_test::sequential]
async fn test_gpu_diamond_we_error_search_bench_estimate_and_round_trip() {
Expand Down Expand Up @@ -646,7 +715,7 @@ async fn test_gpu_diamond_we_error_search_bench_estimate_and_round_trip() {
build_public_key_bench_estimator(&gpu_params, &cfg, final_dir.join("pubkey_bench"));
info!("starting DiamondWE GPU encoding bench estimator construction");
let encoding_estimator =
BggEncodingBenchEstimator::<GpuMatrix>::benchmark(&gpu_params, cfg.bench_iterations, || ());
build_encoding_bench_estimator(&gpu_params, &cfg, final_dir.join("encoding_bench")).await;
info!("starting DiamondWE GPU native bench estimator construction");
let native_estimator =
GpuDCRTPolyMatrixNativeBenchEstimator::new(gpu_params.clone(), cfg.bench_iterations);
Expand Down
118 changes: 114 additions & 4 deletions tests/test_gpu_lwe_nested_rns_modq_arith.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use mxx::{
DEFAULT_MAX_UNREDUCED_MULS, NestedRnsPoly, NestedRnsPolyContext, encode_nested_rns_poly,
},
lookup::{
PublicLut,
PltEvaluator, PublicLut,
lwe::{LWEBGGEncodingPltEvaluator, LWEBGGPubKeyPltEvaluator},
poly::PolyPltEvaluator,
},
Expand All @@ -38,7 +38,11 @@ use mxx::{
gpu::{GpuDCRTPolyHashSampler, GpuDCRTPolyUniformSampler},
trapdoor::GpuDCRTPolyTrapdoorSampler,
},
simulator::{SimulatorContext, error_norm::NormPltLWEEvaluator},
simulator::{
SimulatorContext,
error_norm::NormPltLWEEvaluator,
lattice_estimator::{Distribution, run_lattice_estimator_cli_with_timeout},
},
slot_transfer::bgg_pubkey::BggPublicKeySTEvaluator,
storage::write::{init_storage_system, wait_for_all_writes},
utils::bigdecimal_bits_ceil,
Expand All @@ -48,7 +52,7 @@ use std::{
env, fs,
path::{Path, PathBuf},
sync::Arc,
time::Instant,
time::{Duration, Instant},
};
use tracing::{debug, info};

Expand All @@ -66,6 +70,7 @@ const DEFAULT_BENCH_ITERATIONS: usize = 1;
const DEFAULT_BENCH_SEED: [u8; 32] = [0u8; 32];
const TRAPDOOR_SIGMA: f64 = 4.578;
const ERROR_SIM_ACTIVE_LEVELS: usize = 1;
const LATTICE_ESTIMATOR_TIMEOUT: Duration = Duration::from_secs(60 * 60);
const ACTUAL_HASH_KEY_CHECKPOINT_FILE: &str = "actual_lwe_hash_key.bin";
const ACTUAL_TRAPDOOR_CHECKPOINT_FILE: &str = "actual_lwe_trapdoor.bin";
const ACTUAL_PUB_MATRIX_CHECKPOINT_FILE: &str = "actual_lwe_pub_matrix.bin";
Expand All @@ -74,6 +79,7 @@ type GpuMatrix = GpuDCRTPolyMatrix;
type GpuHashSampler = GpuDCRTPolyHashSampler<Keccak256>;
type GpuScalarPubKeyPltEvaluator =
LWEBGGPubKeyPltEvaluator<GpuMatrix, GpuHashSampler, GpuDCRTPolyTrapdoorSampler>;
type GpuEncodingPltEvaluator = LWEBGGEncodingPltEvaluator<GpuMatrix, GpuHashSampler>;
type GpuPubKeySlotEvaluator = BggPublicKeySTEvaluator<
GpuMatrix,
GpuDCRTPolyUniformSampler,
Expand Down Expand Up @@ -487,6 +493,46 @@ fn find_crt_depth_for_modq_arith(
);
}

fn log_lattice_estimator_security_for_selected_crt_depth(
params: &DCRTPolyParams,
cfg: &ModqArithConfig,
crt_depth: usize,
) {
let q = params.modulus();
let ring_dim = BigUint::from(params.ring_dimension());
let s_dist = Distribution::Ternary;
let e_dist =
Distribution::DiscreteGaussian { stddev: cfg.error_sigma.to_string(), mean: None, n: None };
info!(
crt_depth,
ring_dim = params.ring_dimension(),
modulus_bits = q.bits(),
error_sigma = cfg.error_sigma,
"running lattice-estimator security check for selected nested-RNS modq arithmetic parameters"
);
let achieved_secpar_for_gauss = run_lattice_estimator_cli_with_timeout(
&ring_dim,
q.as_ref(),
&s_dist,
&e_dist,
None,
false,
LATTICE_ESTIMATOR_TIMEOUT,
)
.unwrap_or_else(|err| {
panic!(
"nested-RNS modq arithmetic benchmark requires a working lattice-estimator-cli on PATH to log the selected CRT-depth security level; failed: {err}"
)
});
info!(
crt_depth,
ring_dim = params.ring_dimension(),
modulus_bits = q.bits(),
achieved_secpar_for_gauss,
"nested-RNS modq arithmetic selected parameter lattice-estimator security"
);
}

fn build_lwe_scalar_pubkey_plt_evaluator(
params: &GpuDCRTPolyParams,
hash_key: [u8; 32],
Expand Down Expand Up @@ -565,6 +611,69 @@ fn constant_and_shared_benchmark_pubkeys(
(&pubkeys[0], &pubkeys[1])
}

async fn build_encoding_bench_estimator(
params: &GpuDCRTPolyParams,
cfg: &ModqArithConfig,
bench_dir: PathBuf,
) -> BggEncodingBenchEstimator<GpuMatrix> {
if bench_dir.exists() {
fs::remove_dir_all(&bench_dir)
.expect("failed to clear modq arithmetic encoding bench directory");
}
fs::create_dir_all(&bench_dir)
.expect("failed to create modq arithmetic encoding bench directory");
init_storage_system(bench_dir.clone());

let (public_lut, public_lut_id, public_lut_gate_id) =
build_benchmark_public_lookup_gate(params);
let public_key_sampler =
BGGPublicKeySampler::<_, GpuHashSampler>::new(DEFAULT_BENCH_SEED, cfg.d_secret);
let uniform_sampler = GpuDCRTPolyUniformSampler::new();
let secrets =
uniform_sampler.sample_uniform(params, 1, cfg.d_secret, DistType::TernaryDist).get_row(0);
let plaintexts = vec![GpuDCRTPoly::from_usize_to_constant(params, 2)];
let public_keys =
public_key_sampler.sample(params, b"LWE_NESTED_RNS_MODQ_ARITH_ENCODING_BENCH", &[true]);
let encoding_sampler =
BGGEncodingSampler::<GpuDCRTPolyUniformSampler>::new(params, &secrets, None);
let encodings = encoding_sampler.sample(params, &public_keys, &plaintexts);

let public_lut_aux_writer = build_lwe_scalar_pubkey_plt_evaluator(
params,
DEFAULT_BENCH_SEED,
cfg.d_secret,
bench_dir.clone(),
);
let lookup_base_matrix = Arc::clone(&public_lut_aux_writer.pub_matrix);
public_lut_aux_writer.write_dummy_aux_for_poly_encode_bench(
params,
&public_lut,
&[plaintexts[0].const_coeff_u64()],
public_lut_id,
public_lut_gate_id,
cfg.error_sigma,
);
wait_for_all_writes(bench_dir.clone())
.await
.expect("modq arithmetic encoding public lookup benchmark aux writes must flush");

let secret_vec = GpuMatrix::from_poly_vec_row(params, secrets);
let c_b = secret_vec * lookup_base_matrix.as_ref();
let encoding_evaluator = GpuEncodingPltEvaluator::new(DEFAULT_BENCH_SEED, bench_dir, c_b);
let one = encodings[0].clone();
let input = encodings[1].clone();
BggEncodingBenchEstimator::<GpuMatrix>::benchmark(params, cfg.bench_iterations, || {
encoding_evaluator.public_lookup(
params,
&public_lut,
&one,
&input,
public_lut_gate_id,
public_lut_id,
)
})
}

fn round_div_biguint(value: &BigUint, divisor: &BigUint) -> BigUint {
let half = divisor / BigUint::from(2u64);
(value + half) / divisor
Expand All @@ -582,6 +691,7 @@ async fn test_gpu_lwe_nested_rns_modq_arith() {
.checked_shl(cfg.height as u32)
.expect("LWE_NESTED_RNS_MODQ_ARITH_HEIGHT is too large");
let (crt_depth, cpu_params) = find_crt_depth_for_modq_arith(&cfg, q_level);
log_lattice_estimator_security_for_selected_crt_depth(&cpu_params, &cfg, crt_depth);
let (moduli, _, _) = cpu_params.to_crt();
let detected_gpu_ids = detected_gpu_device_ids();
let detected_gpu_count = detected_gpu_device_count();
Expand Down Expand Up @@ -752,7 +862,7 @@ async fn test_gpu_lwe_nested_rns_modq_arith() {
);

let encoding_scalar_bench =
BggEncodingBenchEstimator::<GpuMatrix>::benchmark(&params, cfg.bench_iterations, || ());
build_encoding_bench_estimator(&params, &cfg, dir.join("encoding_bench")).await;
let encoding_circuit_bench = encoding_scalar_bench.estimate_circuit_bench(&circuit);
info!(
"modq_arith scalar bgg encoding circuit bench estimate: total_time={:.6} latency={:.6} max_parallelism={} peak_vram={}",
Expand Down
Loading
Loading