Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
48 changes: 48 additions & 0 deletions circuit-std-rs/src/logup.rs
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,35 @@ impl LogUpSingleKeyTable {

assert_eq_rational(builder, &v_table, &v_query);
}

pub fn final_check_with_query_count<C: Config, B: RootAPI<C>>(
&mut self,
builder: &mut B,
query_count: &[Variable],
) {
if self.table.is_empty() || self.query_keys.is_empty() {
panic!("empty table or empty query");
Comment on lines +337 to +338
Copy link
Contributor

Choose a reason for hiding this comment

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

medium

Using panic! for empty table or query inputs can lead to an abrupt program termination. Consider returning a Result type to allow callers to handle this error gracefully, or at least logging a more informative error before panicking, especially in a library context where consumers might expect more robust error handling.

Suggested change
if self.table.is_empty() || self.query_keys.is_empty() {
panic!("empty table or empty query");
if self.table.is_empty() || self.query_keys.is_empty() {
// TODO: Consider returning a Result<_, Error> instead of panicking
// or logging a more specific error message.
panic!("empty table or empty query");
}

}

let value_len = self.table[0].len();

let alpha = builder.get_random_value();
let randomness = get_column_randomness(builder, value_len);

let table_combined = combine_columns(builder, &self.table, &randomness);
let v_table = logup_poly_val(builder, &table_combined, query_count, &alpha);

let query_combined = combine_columns(builder, &self.query_results, &randomness);
let one = builder.constant(1);
let v_query = logup_poly_val(
builder,
&query_combined,
&vec![one; query_combined.len()],
&alpha,
);

assert_eq_rational(builder, &v_table, &v_query);
}
}

pub struct LogUpRangeProofTable {
Expand Down Expand Up @@ -455,6 +484,25 @@ impl LogUpRangeProofTable {
);
assert_eq_rational(builder, &v_table, &v_query);
}

pub fn final_check_with_query_count<C: Config, B: RootAPI<C>>(
&mut self,
builder: &mut B,
query_count: &[Variable],
) {
let alpha = builder.get_random_value();

let v_table = logup_poly_val(builder, &self.table_keys, query_count, &alpha);

let one = builder.constant(1);
let v_query = logup_poly_val(
builder,
&self.query_keys,
&vec![one; self.query_keys.len()],
&alpha,
);
assert_eq_rational(builder, &v_table, &v_query);
}
}

pub fn query_count_hint<F: Field>(inputs: &[F], outputs: &mut [F]) -> Result<(), Error> {
Expand Down
11 changes: 11 additions & 0 deletions expander_compiler/src/zkcuda/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -888,6 +888,17 @@ impl<C: Config, H: HintCaller<CircuitField<C>>> Context<C, H> {
ContextState::WitnessDone,
"Please finish computation graph and witness solving before exporting device memories."
);
self.export_device_memories_impl()
}

/// Export device memories without checking the context state.
/// Use this when you need to export memories outside the normal workflow,
/// e.g., for memory optimization where you want to export and then drop the context.
pub fn export_device_memories_unchecked(&self) -> Vec<Vec<SIMDField<C>>> {
self.export_device_memories_impl()
}

fn export_device_memories_impl(&self) -> Vec<Vec<SIMDField<C>>> {
self.device_memories
.iter()
.map(|dm| {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,3 +73,15 @@ where
wait_async(ClientHttpHelper::request_exit())
}
}

impl<ZC: ZKCudaConfig> ExpanderNoOverSubscribe<ZC>
where
<GetPCS<ZC> as ExpanderPCS<GetFieldConfig<ZC>>>::Commitment:
AsRef<<GetPCS<ZC> as ExpanderPCS<GetFieldConfig<ZC>>>::Commitment>,
{
/// Lightweight prove that doesn't require computation_graph or prover_setup.
/// Use this after setup() to allow releasing those large data structures before proving.
pub fn prove_lightweight(device_memories: Vec<Vec<SIMDField<ZC::ECCConfig>>>) {
client_send_witness_and_prove::<ZC::GKRConfig, ZC::ECCConfig>(device_memories);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,11 @@ where
let mpi_size = if allow_oversubscribe {
max_parallel_count
} else {
let num_cpus = prev_power_of_two(num_cpus::get_physical());
let num_cpus = std::env::var("ZKML_NUM_CPUS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or_else(num_cpus::get_physical);
let num_cpus = prev_power_of_two(num_cpus);
if max_parallel_count > num_cpus {
num_cpus
} else {
Expand Down