Skip to content

Commit f8a14e7

Browse files
SchaeffSchaeff
committed
clean up
1 parent ab93fee commit f8a14e7

File tree

2 files changed

+73
-66
lines changed

2 files changed

+73
-66
lines changed

‎openvm/src/powdr_extension/trace_generator/cpu/mod.rs‎

Lines changed: 52 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,20 @@ use openvm_circuit::{
1010
utils::next_power_of_two_or_zero,
1111
};
1212
use openvm_stark_backend::{
13-
Chip, p3_field::FieldAlgebra, p3_matrix::dense::{DenseMatrix, RowMajorMatrix}, prover::{
14-
hal::ProverBackend, types::{AirProvingContext, AirProvingContexts, Rejected}
15-
}
13+
p3_field::FieldAlgebra,
14+
p3_matrix::dense::RowMajorMatrix,
15+
prover::{
16+
hal::ProverBackend,
17+
types::{AirProvingContext, AirProvingContexts, Rejected},
18+
},
19+
Chip,
1620
};
1721
use openvm_stark_sdk::p3_baby_bear::BabyBear;
18-
use powdr_autoprecompiles::{expression::RowEvaluator, trace_handler::TraceTrait, Apc};
22+
use powdr_autoprecompiles::{
23+
expression::{AlgebraicEvaluator, ConcreteBusInteraction, RowEvaluator},
24+
trace_handler::TraceTrait,
25+
Apc,
26+
};
1927
use powdr_constraint_solver::constraint_system::ComputationMethod;
2028

2129
use crate::{
@@ -27,7 +35,6 @@ use crate::{
2735
},
2836
BabyBearSC, Instr,
2937
};
30-
use openvm_stark_backend::p3_field::PrimeField32;
3138

3239
use openvm_stark_backend::p3_field::Field;
3340

@@ -61,22 +68,18 @@ impl<F> From<Arc<RowMajorMatrix<F>>> for SharedCpuTrace<F> {
6168
}
6269
}
6370

64-
impl<R, PB: ProverBackend<Matrix = Arc<RowMajorMatrix<BabyBear>>>> Chip<R, PB> for PowdrChipCpu {
71+
impl<R, PB: ProverBackend<Val = BabyBear, Matrix = Arc<RowMajorMatrix<BabyBear>>>> Chip<R, PB>
72+
for PowdrChipCpu
73+
{
6574
fn generate_proving_ctx(&self, _records: R) -> AirProvingContext<PB> {
6675
unreachable!()
6776
}
6877

6978
fn generate_proving_ctxs(&self, _: R) -> AirProvingContexts<PB> {
7079
tracing::trace!("Generating air proof input for PowdrChip {}", self.name);
7180

72-
let (trace, rejected) = self
73-
.trace_generator
74-
.generate_witness(self.record_arena_by_air_name.take());
75-
76-
AirProvingContexts {
77-
main: AirProvingContext::simple(trace, vec![]),
78-
rejected
79-
}
81+
self.trace_generator
82+
.generate_witness::<PB>(self.record_arena_by_air_name.take())
8083
}
8184
}
8285

@@ -102,13 +105,12 @@ impl PowdrTraceGeneratorCpu {
102105
}
103106
}
104107

105-
pub fn generate_witness(
108+
pub fn generate_witness<
109+
PB: ProverBackend<Val = BabyBear, Matrix = Arc<RowMajorMatrix<BabyBear>>>,
110+
>(
106111
&self,
107112
mut original_arenas: OriginalArenas<MatrixRecordArena<BabyBear>>,
108-
) -> (
109-
Arc<DenseMatrix<BabyBear>>,
110-
Rejected<Arc<DenseMatrix<BabyBear>>>,
111-
) {
113+
) -> AirProvingContexts<PB> {
112114
let mut rejected_pcs = vec![];
113115

114116
use powdr_autoprecompiles::trace_handler::{generate_trace, TraceData};
@@ -117,10 +119,8 @@ impl PowdrTraceGeneratorCpu {
117119
if num_apc_calls == 0 {
118120
// If the APC isn't called, early return with an empty trace.
119121
let width = self.apc.machine().main_columns().count();
120-
return (
121-
Arc::new(RowMajorMatrix::new(vec![], width)),
122-
Default::default()
123-
);
122+
return AirProvingContext::simple_no_pis(Arc::new(RowMajorMatrix::new(vec![], width)))
123+
.into();
124124
}
125125

126126
let chip_inventory = {
@@ -227,27 +227,12 @@ impl PowdrTraceGeneratorCpu {
227227
let evaluator = MappingRowEvaluator::new(row_slice, &apc_poly_id_to_index);
228228

229229
// check the constraints and bus interactions
230+
// let row_is_valid = unimplemented!("evaluate constraints and bus interactions, or just specialization constraints? For now we reject the first call, just for testing.");
230231

231232
if row_is_valid {
232233
// replay the side effects of this row on the main periphery
233-
self.apc
234-
.machine()
235-
.bus_interactions
236-
.iter()
237-
.for_each(|interaction| {
238-
use powdr_autoprecompiles::expression::{
239-
AlgebraicEvaluator, ConcreteBusInteraction,
240-
};
241-
242-
let ConcreteBusInteraction { id, mult, args } =
243-
evaluator.eval_bus_interaction(interaction);
244-
self.periphery.real.apply(
245-
id as u16,
246-
mult.as_canonical_u32(),
247-
args.map(|arg| arg.as_canonical_u32()),
248-
&self.periphery.bus_ids,
249-
);
250-
});
234+
self.periphery
235+
.replay_bus_interactions(self.apc.machine(), &evaluator);
251236
} else {
252237
// set the whole row to zero
253238
// TODO: this generates a gap in the table. Instead, reuse the row in the next iteration.
@@ -257,29 +242,29 @@ impl PowdrTraceGeneratorCpu {
257242

258243
// for each original row
259244
for original_row_reference in dummy_values {
260-
// replay the side effects of this row on the real periphery
245+
// build an evaluator over the row
261246
let original_row_data = &original_row_reference.data[original_row_reference
262247
.start()
263248
..original_row_reference.start() + original_row_reference.length];
264249
let evaluator = RowEvaluator::new(original_row_data);
265250
let (machine, _) =
266251
&self.original_airs.air_name_to_machine[original_row_reference.air_id];
267252

268-
for interaction in &machine.bus_interactions {
269-
use powdr_autoprecompiles::expression::{
270-
AlgebraicEvaluator, ConcreteBusInteraction,
271-
};
272-
273-
let ConcreteBusInteraction { id, mult, args } =
274-
evaluator.eval_bus_interaction(interaction);
275-
276-
rejected_pcs.extend(self.periphery.real.apply(
277-
id as u16,
278-
mult.as_canonical_u32(),
279-
args.map(|arg| arg.as_canonical_u32()),
280-
&self.periphery.bus_ids,
281-
));
282-
}
253+
// replay the side effects of this row on the real periphery
254+
self.periphery.replay_bus_interactions(machine, &evaluator);
255+
256+
// find the concrete value of the received pc
257+
rejected_pcs.push(
258+
machine
259+
.bus_interactions
260+
.iter()
261+
.find_map(|interaction| {
262+
let ConcreteBusInteraction { id, mut args, .. } =
263+
evaluator.eval_bus_interaction(interaction);
264+
(id == 2).then(|| args.next().unwrap())
265+
})
266+
.unwrap(),
267+
);
283268

284269
// add the row index to the rejected set
285270
rejected_rows_per_air
@@ -295,18 +280,19 @@ impl PowdrTraceGeneratorCpu {
295280
pcs: rejected_pcs,
296281
rows_per_air: rejected_rows_per_air
297282
.into_iter()
283+
// if this original table contains any rejected rows
284+
.filter(|(_, indices)| !indices.is_empty())
285+
// return the table with its rejected rows
298286
.map(|(name, indices)| {
299-
(
300-
name.clone(),
301-
(
302-
dummy_trace_by_air_name.remove(&name).unwrap().matrix,
303-
indices,
304-
),
305-
)
287+
let original_trace = dummy_trace_by_air_name.remove(&name).unwrap().matrix;
288+
(name, (original_trace, indices))
306289
})
307290
.collect(),
308-
};
291+
};
309292

310-
(Arc::new(RowMajorMatrix::new(values, width)), rejected)
293+
AirProvingContexts {
294+
main: AirProvingContext::simple_no_pis(Arc::new(RowMajorMatrix::new(values, width))),
295+
rejected,
296+
}
311297
}
312298
}

‎openvm/src/powdr_extension/trace_generator/cpu/periphery.rs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@ use openvm_circuit::arch::RowMajorMatrixArena;
1616
use openvm_stark_backend::config::Val;
1717
use openvm_stark_backend::engine::StarkEngine;
1818
use openvm_stark_backend::prover::cpu::{CpuBackend, CpuDevice};
19+
use powdr_autoprecompiles::{
20+
expression::{AlgebraicEvaluator, ConcreteBusInteraction},
21+
SymbolicMachine,
22+
};
1923

2024
use crate::powdr_extension::trace_generator::common::DummyExecutor;
2125
use crate::PeripheryBusIds;
@@ -70,6 +74,23 @@ impl PowdrPeripheryInstancesCpu {
7074
bus_ids,
7175
}
7276
}
77+
78+
pub fn replay_bus_interactions<T: PrimeField32>(
79+
&self,
80+
machine: &SymbolicMachine<T>,
81+
evaluator: &impl AlgebraicEvaluator<T, T>,
82+
) {
83+
machine.bus_interactions.iter().for_each(|interaction| {
84+
let ConcreteBusInteraction { id, mult, args } =
85+
evaluator.eval_bus_interaction(interaction);
86+
self.real.apply(
87+
id as u16,
88+
mult.as_canonical_u32(),
89+
args.map(|arg| arg.as_canonical_u32()),
90+
&self.bus_ids,
91+
);
92+
});
93+
}
7394
}
7495

7596
impl<F: PrimeField32> VmExecutionExtension<F> for SharedPeripheryChipsCpu {

0 commit comments

Comments
 (0)