Skip to content

Simplify original arenas #3438

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

Open
Schaeff wants to merge 2 commits into
base: main
Choose a base branch
from
+34 −59
Open

Simplify original arenas #3438

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2ef6701
simplify original arenas
Schaeff Nov 12, 2025
26fe2c9
fix gpu
Schaeff Nov 13, 2025
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
57 changes: 14 additions & 43 deletions openvm/src/powdr_extension/executor/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,51 +69,20 @@ impl<A: Arena> OriginalArenas<A> {
apc_call_count_estimate: impl Fn() -> usize,
original_airs: &OriginalAirs<BabyBear>,
apc: &Arc<Apc<BabyBear, Instr<BabyBear>>>,
) {
) -> &mut InitializedOriginalArenas<A> {
match self {
OriginalArenas::Uninitialized => {
*self = OriginalArenas::Initialized(InitializedOriginalArenas::new(
apc_call_count_estimate(),
original_airs,
apc,
));
match self {
OriginalArenas::Initialized(i) => i,
_ => unreachable!(),
}
}
OriginalArenas::Initialized(_) => {}
}
}

/// Returns a mutable reference to the arenas.
/// - Panics if the arenas are not initialized.
pub fn arenas_mut(&mut self) -> &mut HashMap<String, A> {
match self {
OriginalArenas::Uninitialized => panic!("original arenas are uninitialized"),
OriginalArenas::Initialized(initialized) => &mut initialized.arenas,
}
}

/// Returns a reference to the arenas.
/// - Panics if the arenas are not initialized.
pub fn arenas(&self) -> &HashMap<String, A> {
match self {
OriginalArenas::Uninitialized => panic!("original arenas are uninitialized"),
OriginalArenas::Initialized(initialized) => &initialized.arenas,
}
}

/// Returns a mutable reference to the number of calls.
/// - Panics if the arenas are not initialized.
pub fn number_of_calls_mut(&mut self) -> &mut usize {
match self {
OriginalArenas::Uninitialized => panic!("original arenas are uninitialized"),
OriginalArenas::Initialized(initialized) => &mut initialized.number_of_calls,
}
}

/// Returns the number of calls. If not initialized, `Preflight::execute` is never called, and thus return 0.
pub fn number_of_calls(&self) -> usize {
match self {
OriginalArenas::Uninitialized => 0,
OriginalArenas::Initialized(initialized) => initialized.number_of_calls,
OriginalArenas::Initialized(i) => i,
}
}
}
Expand Down Expand Up @@ -425,9 +394,10 @@ impl PreflightExecutor<BabyBear, MatrixRecordArena<BabyBear>> for PowdrExecutor
// Recover an estimate of how many times the APC is called in this segment based on the current ctx height and width
let apc_call_count = || ctx.trace_buffer.len() / ctx.width;

original_arenas.ensure_initialized(apc_call_count, &self.air_by_opcode_id, &self.apc);
let original_arenas =
original_arenas.ensure_initialized(apc_call_count, &self.air_by_opcode_id, &self.apc);

let arenas = original_arenas.arenas_mut();
let arenas = &mut original_arenas.arenas;
Copy link
Collaborator Author

@Schaeff Schaeff Nov 12, 2025

Choose a reason for hiding this comment

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

we use to match here


// execute the original instructions one by one
for instruction in self.apc.instructions() {
Expand Down Expand Up @@ -458,7 +428,7 @@ impl PreflightExecutor<BabyBear, MatrixRecordArena<BabyBear>> for PowdrExecutor
}

// Update the real number of calls to the APC
*original_arenas.number_of_calls_mut() += 1;
original_arenas.number_of_calls += 1;
Copy link
Collaborator Author

@Schaeff Schaeff Nov 12, 2025

Choose a reason for hiding this comment

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

we use to match here


Ok(())
}
Expand Down Expand Up @@ -500,9 +470,10 @@ impl PreflightExecutor<BabyBear, DenseRecordArena> for PowdrExecutor {
buf.len() / bytes_per_row
};

original_arenas.ensure_initialized(apc_call_count, &self.air_by_opcode_id, &self.apc);
let original_arenas =
original_arenas.ensure_initialized(apc_call_count, &self.air_by_opcode_id, &self.apc);

let arenas = original_arenas.arenas_mut();
let arenas = &mut original_arenas.arenas;

// execute the original instructions one by one
for instruction in self.apc.instructions() {
Expand Down Expand Up @@ -533,7 +504,7 @@ impl PreflightExecutor<BabyBear, DenseRecordArena> for PowdrExecutor {
}

// Update the real number of calls to the APC
*original_arenas.number_of_calls_mut() += 1;
original_arenas.number_of_calls += 1;

Ok(())
}
Expand Down
19 changes: 11 additions & 8 deletions openvm/src/powdr_extension/trace_generator/cpu/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,16 +94,19 @@ impl PowdrTraceGeneratorCpu {

pub fn generate_witness(
&self,
mut original_arenas: OriginalArenas<MatrixRecordArena<BabyBear>>,
original_arenas: OriginalArenas<MatrixRecordArena<BabyBear>>,
) -> DenseMatrix<BabyBear> {
use powdr_autoprecompiles::trace_handler::{generate_trace, TraceData};

let num_apc_calls = original_arenas.number_of_calls();
if num_apc_calls == 0 {
// If the APC isn't called, early return with an empty trace.
let width = self.apc.machine().main_columns().count();
return RowMajorMatrix::new(vec![], width);
}
let width = self.apc.machine().main_columns().count();

let original_arenas = match original_arenas {
OriginalArenas::Initialized(arenas) => arenas,
OriginalArenas::Uninitialized => {
// if the arenas are uninitialized, the apc was not called, so we return an empty trace
return RowMajorMatrix::new(vec![], width);
}
};

let chip_inventory = {
let airs: AirInventory<BabyBearSC> =
Expand All @@ -119,7 +122,7 @@ impl PowdrTraceGeneratorCpu {
.inventory
};

let arenas = original_arenas.arenas_mut();
let (mut arenas, num_apc_calls) = (original_arenas.arenas, original_arenas.number_of_calls);

let dummy_trace_by_air_name: HashMap<String, SharedCpuTrace<BabyBear>> = chip_inventory
.chips()
Expand Down
17 changes: 9 additions & 8 deletions openvm/src/powdr_extension/trace_generator/cuda/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,14 +196,15 @@ impl PowdrTraceGeneratorGpu {

fn try_generate_witness(
&self,
mut original_arenas: OriginalArenas<DenseRecordArena>,
original_arenas: OriginalArenas<DenseRecordArena>,
) -> Option<DeviceMatrix<BabyBear>> {
let num_apc_calls = original_arenas.number_of_calls();

if num_apc_calls == 0 {
// If the APC isn't called, early return with an empty trace.
return None;
}
let original_arenas = match original_arenas {
OriginalArenas::Initialized(arenas) => arenas,
OriginalArenas::Uninitialized => {
// if the arenas are uninitialized, the apc was not called, so we return early
return None;
}
};

let chip_inventory = {
let airs: AirInventory<BabyBearSC> =
Expand All @@ -219,7 +220,7 @@ impl PowdrTraceGeneratorGpu {
.inventory
};

let arenas = original_arenas.arenas_mut();
let (mut arenas, num_apc_calls) = (original_arenas.arenas, original_arenas.number_of_calls);

let dummy_trace_by_air_name: HashMap<String, DeviceMatrix<BabyBear>> = chip_inventory
.chips()
Expand Down
Loading