From 178cba3c972cc787c465c50b5d6c22987ae78ec1 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 21 Jan 2025 14:13:38 +0000 Subject: [PATCH 1/5] Access to nonconstant fixed columns. --- .../src/witgen/jit/block_machine_processor.rs | 20 +-- executor/src/witgen/jit/compiler.rs | 135 ++++++++++++++++-- executor/src/witgen/jit/function_cache.rs | 4 +- executor/src/witgen/jit/includes/interface.rs | 2 + executor/src/witgen/jit/interpreter.rs | 4 +- .../src/witgen/jit/single_step_processor.rs | 25 +++- executor/src/witgen/jit/variable.rs | 24 +++- executor/src/witgen/jit/witgen_inference.rs | 46 +++--- 8 files changed, 215 insertions(+), 45 deletions(-) diff --git a/executor/src/witgen/jit/block_machine_processor.rs b/executor/src/witgen/jit/block_machine_processor.rs index 443f0d4169..7ffe687f83 100644 --- a/executor/src/witgen/jit/block_machine_processor.rs +++ b/executor/src/witgen/jit/block_machine_processor.rs @@ -2,14 +2,14 @@ use std::collections::HashSet; use bit_vec::BitVec; use itertools::Itertools; -use powdr_ast::analyzed::AlgebraicReference; +use powdr_ast::analyzed::{PolyID, PolynomialType}; use powdr_number::FieldElement; use crate::witgen::{jit::processor::Processor, machines::MachineParts, FixedData}; use super::{ effect::Effect, - variable::Variable, + variable::{Cell, Variable}, witgen_inference::{CanProcessCall, FixedEvaluator, WitgenInference}, }; @@ -123,17 +123,19 @@ impl<'a, T: FieldElement> BlockMachineProcessor<'a, T> { } impl FixedEvaluator for &BlockMachineProcessor<'_, T> { - fn evaluate(&self, var: &AlgebraicReference, row_offset: i32) -> Option { - assert!(var.is_fixed()); - let values = self.fixed_data.fixed_cols[&var.poly_id].values_max_size(); + fn evaluate(&self, fixed_cell: &Cell) -> Option { + let poly_id = PolyID { + id: fixed_cell.id, + ptype: PolynomialType::Constant, + }; + let values = self.fixed_data.fixed_cols[&poly_id].values_max_size(); // By assumption of the block machine, all fixed columns are cyclic with a period of . // An exception might be the first and last row. - assert!(row_offset >= -1); + assert!(fixed_cell.row_offset >= -1); assert!(self.block_size >= 1); - // The current row is guaranteed to be at least 1. - let current_row = (2 * self.block_size as i32 + row_offset) as usize; - let row = current_row + var.next as usize; + // The row is guaranteed to be at least 1. + let row = (2 * self.block_size as i32 + fixed_cell.row_offset) as usize; assert!(values.len() >= self.block_size * 4); diff --git a/executor/src/witgen/jit/compiler.rs b/executor/src/witgen/jit/compiler.rs index 05bc55feda..978bf7250f 100644 --- a/executor/src/witgen/jit/compiler.rs +++ b/executor/src/witgen/jit/compiler.rs @@ -2,7 +2,10 @@ use std::{cmp::Ordering, ffi::c_void, mem, sync::Arc}; use itertools::Itertools; use libloading::Library; -use powdr_ast::indent; +use powdr_ast::{ + analyzed::{PolyID, PolynomialType}, + indent, +}; use powdr_number::{FieldElement, KnownField}; use crate::witgen::{ @@ -11,7 +14,7 @@ use crate::witgen::{ profiling::{record_end, record_start}, LookupCell, }, - QueryCallback, + FixedData, QueryCallback, }; use super::{ @@ -35,6 +38,7 @@ impl WitgenFunction { /// This function always succeeds (unless it panics). pub fn call>( &self, + fixed_data: &FixedData<'_, T>, mutable_state: &MutableState<'_, T, Q>, params: &mut [LookupCell], mut data: CompactDataRef<'_, T>, @@ -48,10 +52,26 @@ impl WitgenFunction { params: params.into(), mutable_state: mutable_state as *const _ as *const c_void, call_machine: call_machine::, + fixed_data: fixed_data as *const _ as *const c_void, + get_fixed_value: get_fixed_value::, }); } } +extern "C" fn get_fixed_value( + fixed_data: *const c_void, + column: u64, + row: u64, +) -> T { + let fixed_data = unsafe { &*(fixed_data as *const FixedData<'_, T>) }; + let poly_id = PolyID { + id: column, + ptype: PolynomialType::Constant, + }; + // TODO which size? + fixed_data.fixed_cols[&poly_id].values_max_size()[row as usize] +} + extern "C" fn call_machine>( mutable_state: *const c_void, identity_id: u64, @@ -105,6 +125,11 @@ struct WitgenFunctionParams<'a, T: 'a> { mutable_state: *const c_void, /// A callback to call submachines. call_machine: extern "C" fn(*const c_void, u64, MutSlice>) -> bool, + /// A pointer to the "fixed data". + fixed_data: *const c_void, + /// A callback to retrieve values from fixed columns. + /// The parameters are: fixed data pointer, fixed column id, row number. + get_fixed_value: extern "C" fn(*const c_void, u64, u64) -> T, } #[repr(C)] @@ -150,6 +175,7 @@ fn witgen_code( format!("get(data, row_offset, {}, {})", c.row_offset, c.id) } Variable::Param(i) => format!("get_param(params, {i})"), + Variable::FixedColumn(_) => panic!("Fixed columns should not be known inputs."), Variable::MachineCallParam(_) => { unreachable!("Machine call variables should not be pre-known.") } @@ -157,6 +183,27 @@ fn witgen_code( format!(" let {var_name} = {value};") }) .format("\n"); + + // Pre-load all the fixed columns so that we can treat them as + // plain variables later. + let load_fixed = effects + .iter() + .flat_map(|e| e.referenced_variables()) + .filter_map(|v| match v { + Variable::FixedColumn(c) => Some((v, c)), + _ => None, + }) + .unique() + .map(|(var, cell)| { + format!( + " let {} = get_fixed_value(fixed_data, {}, (row_offset + {}));", + variable_to_string(var), + cell.id, + cell.row_offset, + ) + }) + .format("\n"); + let main_code = format_effects(effects); let vars_known = effects .iter() @@ -173,6 +220,7 @@ fn witgen_code( cell.row_offset, cell.id, )), Variable::Param(i) => Some(format!(" set_param(params, {i}, {value});")), + Variable::FixedColumn(_) => panic!("Fixed columns should not be written to."), Variable::MachineCallParam(_) => { // This is just an internal variable. None @@ -186,7 +234,7 @@ fn witgen_code( .iter() .filter_map(|var| match var { Variable::Cell(cell) => Some(cell), - Variable::Param(_) | Variable::MachineCallParam(_) => None, + Variable::Param(_) | Variable::FixedColumn(_) | Variable::MachineCallParam(_) => None, }) .map(|cell| { format!( @@ -205,19 +253,28 @@ extern "C" fn witgen( row_offset, params, mutable_state, - call_machine + call_machine, + fixed_data, + get_fixed_value, }}: WitgenFunctionParams, ) {{ let known = known_to_slice(known, data.len); let data = data.to_mut_slice(); let params = params.to_mut_slice(); + // Pre-load fixed column values into local variables +{load_fixed} + + // Load all known inputs into local variables {load_known_inputs} + // Perform the main computations {main_code} + // Store the newly derived witness cell values {store_values} + // Store the "known" flags {store_known} }} "# @@ -373,6 +430,14 @@ fn variable_to_string(v: &Variable) -> String { format_row_offset(cell.row_offset) ), Variable::Param(i) => format!("p_{i}"), + Variable::FixedColumn(cell) => { + format!( + "f_{}_{}_{}", + escape_column_name(&cell.column_name), + cell.id, + cell.row_offset + ) + } Variable::MachineCallParam(call_var) => { format!( "call_var_{}_{}_{}", @@ -461,6 +526,8 @@ fn util_code(first_column_id: u64, column_count: usize) -> Resu #[cfg(test)] mod tests { + use std::ptr::null; + use pretty_assertions::assert_eq; use test_log::test; @@ -546,8 +613,8 @@ mod tests { let known_inputs = vec![a0.clone()]; let code = witgen_code(&known_inputs, &effects); assert_eq!( - code, - " + code, + " #[no_mangle] extern \"C\" fn witgen( WitgenFunctionParams{ @@ -556,15 +623,22 @@ extern \"C\" fn witgen( row_offset, params, mutable_state, - call_machine + call_machine, + fixed_data, + get_fixed_value, }: WitgenFunctionParams, ) { let known = known_to_slice(known, data.len); let data = data.to_mut_slice(); let params = params.to_mut_slice(); + // Pre-load fixed column values into local variables + + + // Load all known inputs into local variables let c_a_2_0 = get(data, row_offset, 0, 2); + // Perform the main computations let c_x_0_0 = (FieldElement::from(7) * c_a_2_0); let call_var_7_1_0 = c_x_0_0; let mut call_var_7_1_1 = FieldElement::default(); @@ -573,16 +647,18 @@ extern \"C\" fn witgen( let c_y_1_1 = (c_y_1_m1 + c_x_0_0); assert!(c_y_1_m1 == c_x_0_0); + // Store the newly derived witness cell values set(data, row_offset, 0, 0, c_x_0_0); set(data, row_offset, -1, 1, c_y_1_m1); set(data, row_offset, 1, 1, c_y_1_1); + // Store the \"known\" flags set_known(known, row_offset, 0, 0); set_known(known, row_offset, -1, 1); set_known(known, row_offset, 1, 1); } " - ); + ); } extern "C" fn no_call_machine( @@ -593,6 +669,10 @@ extern \"C\" fn witgen( false } + extern "C" fn get_fixed_data_test(_: *const c_void, col_id: u64, row: u64) -> GoldilocksField { + GoldilocksField::from(col_id * 2000 + row) + } + fn witgen_fun_params<'a>( data: &mut [GoldilocksField], known: &mut [u32], @@ -604,6 +684,8 @@ extern \"C\" fn witgen( params: Default::default(), mutable_state: std::ptr::null(), call_machine: no_call_machine, + fixed_data: null(), + get_fixed_value: get_fixed_data_test, } } @@ -655,6 +737,8 @@ extern \"C\" fn witgen( params: Default::default(), mutable_state: std::ptr::null(), call_machine: no_call_machine, + fixed_data: null(), + get_fixed_value: get_fixed_data_test, }; (f2.function)(params2); assert_eq!(data[0], GoldilocksField::from(7)); @@ -746,6 +830,8 @@ extern \"C\" fn witgen( params: params.as_mut_slice().into(), mutable_state: std::ptr::null(), call_machine: no_call_machine, + fixed_data: null(), + get_fixed_value: get_fixed_data_test, }; (f.function)(params); assert_eq!(y_val, GoldilocksField::from(7 * 2)); @@ -768,6 +854,33 @@ extern \"C\" fn witgen( assert!(code.contains(&format!("let c_x_1_0 = (c_a_0_0 & {large_num:#x});"))); } + #[test] + fn fixed_column_access() { + let a = cell("a", 0, 0); + let x = Variable::FixedColumn(Cell { + column_name: "X".to_string(), + id: 15, + row_offset: 6, + }); + let effects = vec![assignment(&a, symbol(&x))]; + let f = compile_effects(0, 1, &[], &effects).unwrap(); + let mut data = vec![7.into()]; + let mut known = vec![0]; + let mut params = vec![]; + let params = WitgenFunctionParams { + data: data.as_mut_slice().into(), + known: known.as_mut_ptr(), + row_offset: 0, + params: params.as_mut_slice().into(), + mutable_state: std::ptr::null(), + call_machine: no_call_machine, + fixed_data: null(), + get_fixed_value: get_fixed_data_test, + }; + (f.function)(params); + assert_eq!(data[0], GoldilocksField::from(30006)); + } + extern "C" fn mock_call_machine( _: *const c_void, id: u64, @@ -820,6 +933,8 @@ extern \"C\" fn witgen( params: Default::default(), mutable_state: std::ptr::null(), call_machine: mock_call_machine, + fixed_data: null(), + get_fixed_value: get_fixed_data_test, }; (f.function)(params); assert_eq!(data[0], GoldilocksField::from(9)); @@ -854,6 +969,8 @@ extern \"C\" fn witgen( params: params.as_mut_slice().into(), mutable_state: std::ptr::null(), call_machine: no_call_machine, + fixed_data: null(), + get_fixed_value: get_fixed_data_test, }; (f.function)(params); assert_eq!(y_val, GoldilocksField::from(8)); @@ -867,6 +984,8 @@ extern \"C\" fn witgen( params: params.as_mut_slice().into(), mutable_state: std::ptr::null(), call_machine: no_call_machine, + fixed_data: null(), + get_fixed_value: get_fixed_data_test, }; (f.function)(params); assert_eq!(y_val, GoldilocksField::from(4)); diff --git a/executor/src/witgen/jit/function_cache.rs b/executor/src/witgen/jit/function_cache.rs index ab5c13dc4e..955149581b 100644 --- a/executor/src/witgen/jit/function_cache.rs +++ b/executor/src/witgen/jit/function_cache.rs @@ -23,6 +23,7 @@ struct CacheKey { } pub struct FunctionCache<'a, T: FieldElement> { + fixed_data: &'a FixedData<'a, T>, /// The processor that generates the JIT code processor: BlockMachineProcessor<'a, T>, /// The cache of JIT functions. If the entry is None, we attempted to generate the function @@ -47,6 +48,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { BlockMachineProcessor::new(fixed_data, parts.clone(), block_size, latch_row); FunctionCache { + fixed_data, processor, column_layout: metadata, witgen_functions: HashMap::new(), @@ -167,7 +169,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { .expect("Need to call compile_cached() first!") .as_ref() .expect("compile_cached() returned false!"); - f.call(mutable_state, &mut values, data); + f.call(self.fixed_data, mutable_state, &mut values, data); Ok(true) } diff --git a/executor/src/witgen/jit/includes/interface.rs b/executor/src/witgen/jit/includes/interface.rs index dcdc599bca..c0610395cd 100644 --- a/executor/src/witgen/jit/includes/interface.rs +++ b/executor/src/witgen/jit/includes/interface.rs @@ -102,4 +102,6 @@ pub struct WitgenFunctionParams<'a, T: 'a> { params: MutSlice>, mutable_state: *const std::ffi::c_void, call_machine: extern "C" fn(*const std::ffi::c_void, u64, MutSlice>) -> bool, + fixed_data: *const std::ffi::c_void, + get_fixed_value: extern "C" fn(*const std::ffi::c_void, u64, u64) -> T, } diff --git a/executor/src/witgen/jit/interpreter.rs b/executor/src/witgen/jit/interpreter.rs index 556f13bef6..8eef39816d 100644 --- a/executor/src/witgen/jit/interpreter.rs +++ b/executor/src/witgen/jit/interpreter.rs @@ -69,7 +69,8 @@ impl EffectsInterpreter { let idx = var_mapper.map_var(var); InterpreterAction::ReadParam(idx, *i) } - Variable::MachineCallParam(_) => unreachable!(), + // TODO maybe load it to avoid multiple access? + Variable::FixedColumn(_) | Variable::MachineCallParam(_) => unreachable!(), })); } @@ -136,6 +137,7 @@ impl EffectsInterpreter { let idx = var_mapper.get_var(var).unwrap(); actions.push(InterpreterAction::WriteParam(idx, *i)); } + Variable::FixedColumn(_) => panic!("Should not write to fixed column."), Variable::MachineCallParam(_) => { // This is just an internal variable. } diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index f137f27969..42c5b01519 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -1,7 +1,9 @@ #![allow(dead_code)] use itertools::Itertools; -use powdr_ast::analyzed::{AlgebraicExpression as Expression, AlgebraicReference, PolyID}; +use powdr_ast::analyzed::{ + AlgebraicExpression as Expression, AlgebraicReference, PolyID, PolynomialType, +}; use powdr_number::FieldElement; use crate::witgen::{machines::MachineParts, FixedData}; @@ -111,9 +113,12 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { /// Evaluator for fixed columns which are constant except for the first and last row. impl FixedEvaluator for &SingleStepProcessor<'_, T> { - fn evaluate(&self, var: &AlgebraicReference, _row_offset: i32) -> Option { - assert!(var.is_fixed()); - self.fixed_data.fixed_cols[&var.poly_id].has_constant_inner_value() + fn evaluate(&self, fixed_cell: &Cell) -> Option { + let poly_id = PolyID { + id: fixed_cell.id, + ptype: PolynomialType::Constant, + }; + self.fixed_data.fixed_cols[&poly_id].has_constant_inner_value() } } @@ -292,4 +297,16 @@ machine_call(2, [Known(call_var(2, 1, 0)), Known(call_var(2, 1, 1)), Unknown(cal VM::instr_mul[1] = 1;" ); } + + #[test] + fn nonconstant_fixed_columns() { + let input = " + namespace VM(256); + let STEP: col = |i| i; + let w; + w = STEP; + "; + let code = generate_single_step(input, "Main").unwrap(); + assert_eq!(format_code(&code), "VM::w[1] = VM::STEP[1];"); + } } diff --git a/executor/src/witgen/jit/variable.rs b/executor/src/witgen/jit/variable.rs index 34ac92c36d..33a6ebd1db 100644 --- a/executor/src/witgen/jit/variable.rs +++ b/executor/src/witgen/jit/variable.rs @@ -16,6 +16,8 @@ pub enum Variable { /// An input or output value of a machine call on a certain /// identity on a certain row offset. MachineCallParam(MachineCallVariable), + /// A fixed column cell. + FixedColumn(Cell), } impl Display for Variable { @@ -30,6 +32,7 @@ impl Display for Variable { ret.identity_id, ret.row_offset, ret.index ) } + Variable::FixedColumn(cell) => write!(f, "{cell}"), } } } @@ -37,21 +40,29 @@ impl Display for Variable { impl Variable { /// Create a variable from an algebraic reference. pub fn from_reference(r: &AlgebraicReference, row_offset: i32) -> Self { - assert!(r.is_witness()); - Self::Cell(Cell { + let cell = Cell { column_name: r.name.clone(), id: r.poly_id.id, row_offset: r.next as i32 + row_offset, - }) + }; + match r.poly_id.ptype { + PolynomialType::Committed => Self::Cell(cell), + PolynomialType::Constant => Self::FixedColumn(cell), + _ => panic!(), + } } - /// If this variable corresponds to a witness cell, return the corresponding polynomial ID. - pub fn try_to_witness_poly_id(&self) -> Option { + /// If this variable corresponds to a fixed or witness cell, return the corresponding polynomial ID. + pub fn try_to_poly_id(&self) -> Option { match self { Variable::Cell(cell) => Some(PolyID { id: cell.id, ptype: PolynomialType::Committed, }), + Variable::FixedColumn(cell) => Some(PolyID { + id: cell.id, + ptype: PolynomialType::Constant, + }), Variable::Param(_) | Variable::MachineCallParam(_) => None, } } @@ -64,7 +75,8 @@ pub struct MachineCallVariable { pub index: usize, } -/// The identifier of a witness cell in the trace table. +/// The identifier of a witness cell in the trace table +/// or a fixed column cell. /// The `row_offset` is relative to a certain "zero row" defined /// by the component that uses this data structure. #[derive(Debug, Clone, Eq)] diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 79248b86d2..09672b5bef 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -21,7 +21,7 @@ use crate::witgen::{ use super::{ affine_symbolic_expression::{AffineSymbolicExpression, Error, ProcessResult}, effect::{BranchCondition, Effect}, - variable::{MachineCallVariable, Variable}, + variable::{Cell, MachineCallVariable, Variable}, }; /// This component can generate code that solves identities. @@ -108,7 +108,11 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F } pub fn is_known(&self, variable: &Variable) -> bool { - self.known_variables.contains(variable) + if let Variable::FixedColumn(_) = variable { + true + } else { + self.known_variables.contains(variable) + } } pub fn is_complete(&self, identity: &Identity, row_offset: i32) -> bool { @@ -458,9 +462,16 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F /// Returns the current best-known range constraint on the given variable /// combining global range constraints and newly derived local range constraints. + /// For fixed columns, it also invokes the fixed evaluator. pub fn range_constraint(&self, variable: &Variable) -> RangeConstraint { + if let Variable::FixedColumn(fixed_cell) = variable { + if let Some(v) = self.fixed_evaluator.evaluate(fixed_cell) { + // TODO store these in derived range constraints for performance reasons? + return RangeConstraint::from_value(v); + } + } variable - .try_to_witness_poly_id() + .try_to_poly_id() .and_then(|poly_id| { self.fixed_data .global_range_constraints @@ -518,14 +529,8 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Evaluator<'a, T, FixedEv ) -> Option> { Some(match expr { Expression::Reference(r) => match r.poly_id.ptype { - PolynomialType::Constant => self - .witgen_inference - .fixed_evaluator - .evaluate(r, offset)? - .into(), - PolynomialType::Committed => { - let variable = Variable::from_reference(r, offset); - self.evaluate_variable(variable) + PolynomialType::Constant | PolynomialType::Committed => { + self.evaluate_variable(Variable::from_reference(r, offset)) } PolynomialType::Intermediate => { let definition = @@ -608,7 +613,12 @@ enum VariableOrValue { } pub trait FixedEvaluator: Clone { - fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option { + /// Evaluate a fixed column cell if that cell has a constant value, + /// otherwise return None. + /// If this function returns `None`, the value of the fixed column will + /// be treated as symbolically known but not compile-time constant + /// (i.e. it depends on the row). + fn evaluate(&self, _fixed_cell: &Cell) -> Option { None } } @@ -640,6 +650,7 @@ impl> CanProcessCall for &MutableState<' #[cfg(test)] mod test { + use powdr_ast::analyzed::PolyID; use powdr_number::GoldilocksField; use pretty_assertions::assert_eq; use test_log::test; @@ -656,10 +667,13 @@ mod test { #[derive(Clone)] pub struct FixedEvaluatorForFixedData<'a, T: FieldElement>(pub &'a FixedData<'a, T>); impl FixedEvaluator for FixedEvaluatorForFixedData<'_, T> { - fn evaluate(&self, var: &AlgebraicReference, row_offset: i32) -> Option { - assert!(var.is_fixed()); - let values = self.0.fixed_cols[&var.poly_id].values_max_size(); - let row = (row_offset + var.next as i32 + values.len() as i32) as usize % values.len(); + fn evaluate(&self, fixed_cell: &Cell) -> Option { + let poly_id = PolyID { + id: fixed_cell.id, + ptype: PolynomialType::Constant, + }; + let values = self.0.fixed_cols[&poly_id].values_max_size(); + let row = (fixed_cell.row_offset + values.len() as i32) as usize % values.len(); Some(values[row]) } } From bbed563ea493dd05e0f64d31b2a1eea3de562520 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 21 Jan 2025 16:44:47 +0000 Subject: [PATCH 2/5] Interpreter. --- executor/src/witgen/jit/interpreter.rs | 47 ++++++++++++++++++++++---- 1 file changed, 40 insertions(+), 7 deletions(-) diff --git a/executor/src/witgen/jit/interpreter.rs b/executor/src/witgen/jit/interpreter.rs index 8eef39816d..c03deca955 100644 --- a/executor/src/witgen/jit/interpreter.rs +++ b/executor/src/witgen/jit/interpreter.rs @@ -7,7 +7,9 @@ use super::variable::{Cell, Variable}; use crate::witgen::data_structures::finalizable_data::CompactDataRef; use crate::witgen::data_structures::mutable_state::MutableState; use crate::witgen::machines::LookupCell; -use crate::witgen::QueryCallback; +use crate::witgen::{FixedData, QueryCallback}; +use itertools::Itertools; +use powdr_ast::analyzed::{PolyID, PolynomialType}; use powdr_number::FieldElement; use std::collections::{BTreeSet, HashMap}; @@ -23,6 +25,7 @@ pub struct EffectsInterpreter { enum InterpreterAction { ReadCell(usize, Cell), ReadParam(usize, usize), + ReadFixedColumn(usize, Cell), AssignExpression(usize, RPNExpression), WriteCell(usize, Cell), WriteParam(usize, usize), @@ -43,6 +46,7 @@ impl EffectsInterpreter { let mut actions = vec![]; let mut var_mapper = VariableMapper::new(); + Self::load_fixed_column_values(&mut var_mapper, &mut actions, effects); Self::load_known_inputs(&mut var_mapper, &mut actions, known_inputs); Self::process_effects(&mut var_mapper, &mut actions, effects); Self::write_data(&mut var_mapper, &mut actions, effects); @@ -55,6 +59,27 @@ impl EffectsInterpreter { ret } + fn load_fixed_column_values( + var_mapper: &mut VariableMapper, + actions: &mut Vec>, + effects: &[Effect], + ) { + actions.extend( + effects + .iter() + .flat_map(|e| e.referenced_variables()) + .filter_map(|v| match v { + Variable::FixedColumn(c) => Some((v, c)), + _ => None, + }) + .unique() + .map(|(var, cell)| { + let idx = var_mapper.map_var(var); + InterpreterAction::ReadFixedColumn(idx, cell.clone()) + }), + ) + } + fn load_known_inputs( var_mapper: &mut VariableMapper, actions: &mut Vec>, @@ -69,7 +94,6 @@ impl EffectsInterpreter { let idx = var_mapper.map_var(var); InterpreterAction::ReadParam(idx, *i) } - // TODO maybe load it to avoid multiple access? Variable::FixedColumn(_) | Variable::MachineCallParam(_) => unreachable!(), })); } @@ -166,12 +190,14 @@ impl EffectsInterpreter { /// Execute the machine effects for the given the parameters pub fn call>( &self, + fixed_data: &FixedData<'_, T>, mutable_state: &MutableState<'_, T, Q>, params: &mut [LookupCell], data: CompactDataRef<'_, T>, ) { let mut vars = vec![T::zero(); self.var_count]; + let row_offset: i64 = data.row_offset.try_into().unwrap(); let mut eval_stack = vec![]; for action in &self.actions { match action { @@ -180,19 +206,25 @@ impl EffectsInterpreter { vars[*idx] = val; } InterpreterAction::ReadCell(idx, c) => { - let row_offset: i32 = data.row_offset.try_into().unwrap(); vars[*idx] = data .data - .get((row_offset + c.row_offset).try_into().unwrap(), c.id) + .get((row_offset + c.row_offset as i64).try_into().unwrap(), c.id) .0; } + InterpreterAction::ReadFixedColumn(idx, c) => { + let poly_id = PolyID { + id: c.id, + ptype: PolynomialType::Constant, + }; + vars[*idx] = fixed_data.fixed_cols[&poly_id].values_max_size() + [usize::try_from(row_offset + c.row_offset as i64).unwrap()]; + } InterpreterAction::ReadParam(idx, i) => { vars[*idx] = get_param(params, *i); } InterpreterAction::WriteCell(idx, c) => { - let row_offset: i32 = data.row_offset.try_into().unwrap(); data.data.set( - (row_offset + c.row_offset).try_into().unwrap(), + (row_offset + c.row_offset as i64).try_into().unwrap(), c.id, vars[*idx], ); @@ -306,6 +338,7 @@ impl VariableMapper { self.count } + /// Returns the index of the variable, allocates it if it does not exist. pub fn map_var(&mut self, var: &Variable) -> usize { let idx = *self.var_idx.entry(var.clone()).or_insert_with(|| { self.count += 1; @@ -526,7 +559,7 @@ mod test { let mut data = CompactData::new(poly_ids.iter()); data.append_new_rows(31); let data_ref = CompactDataRef::new(&mut data, 0); - interpreter.call(&mutable_state, &mut param_lookups, data_ref); + interpreter.call(&fixed_data, &mutable_state, &mut param_lookups, data_ref); assert_eq!( ¶ms, From 45338783f60ff13451a9f3174f55a4d7c0d22dd8 Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 21 Jan 2025 16:57:43 +0000 Subject: [PATCH 3/5] Clean up known. --- executor/src/witgen/jit/witgen_inference.rs | 24 ++++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 09672b5bef..58074128f9 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -138,7 +138,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F pub fn branch_on(mut self, variable: &Variable) -> BranchResult<'a, T, FixedEval> { // The variable needs to be known, we need to have a range constraint but // it cannot be a single value. - assert!(self.known_variables.contains(variable)); + assert!(self.is_known(variable)); let rc = self.range_constraint(variable); assert!(rc.try_to_single_value().is_none()); @@ -395,7 +395,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F if self.add_range_constraint(variable.clone(), assignment.range_constraint()) { updated_variables.push(variable.clone()); } - if self.known_variables.insert(variable.clone()) { + if self.record_known(variable.clone()) { log::trace!("{variable} := {assignment}"); updated_variables.push(variable.clone()); self.code.push(e); @@ -416,7 +416,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F assert!(process_result.complete); for v in vars { // Inputs are already known, but it does not hurt to add all of them. - if self.known_variables.insert(v.clone()) { + if self.record_known(v.clone()) { updated_variables.push(v.clone()); } } @@ -448,10 +448,9 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F if rc == old_rc { return false; } - if !self.known_variables.contains(&variable) { - if let Some(v) = rc.try_to_single_value() { - // Special case: Variable is fixed to a constant by range constraints only. - self.known_variables.insert(variable.clone()); + if let Some(v) = rc.try_to_single_value() { + // Special case: Variable is fixed to a constant by range constraints only. + if self.record_known(variable.clone()) { self.code .push(Effect::Assignment(variable.clone(), v.into())); } @@ -460,13 +459,22 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F true } + /// Record a variable as known. Return true if it was not known before. + fn record_known(&mut self, variable: Variable) -> bool { + // We do not record fixed columns as known. + if matches!(variable, Variable::FixedColumn(_)) { + false + } else { + self.known_variables.insert(variable) + } + } + /// Returns the current best-known range constraint on the given variable /// combining global range constraints and newly derived local range constraints. /// For fixed columns, it also invokes the fixed evaluator. pub fn range_constraint(&self, variable: &Variable) -> RangeConstraint { if let Variable::FixedColumn(fixed_cell) = variable { if let Some(v) = self.fixed_evaluator.evaluate(fixed_cell) { - // TODO store these in derived range constraints for performance reasons? return RangeConstraint::from_value(v); } } From 8fe94ac77d907bab4ffe1ef7190ecc47e675f88a Mon Sep 17 00:00:00 2001 From: chriseth Date: Tue, 21 Jan 2025 16:59:14 +0000 Subject: [PATCH 4/5] clarify comment. --- executor/src/witgen/jit/witgen_inference.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 58074128f9..0ee31b1874 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -621,8 +621,8 @@ enum VariableOrValue { } pub trait FixedEvaluator: Clone { - /// Evaluate a fixed column cell if that cell has a constant value, - /// otherwise return None. + /// Evaluate a fixed column cell and returns its value if it is + /// compile-time constant, otherwise return None. /// If this function returns `None`, the value of the fixed column will /// be treated as symbolically known but not compile-time constant /// (i.e. it depends on the row). From 8bb1602ff7d0e0f292a53aea62f5939e3758d81a Mon Sep 17 00:00:00 2001 From: chriseth Date: Thu, 23 Jan 2025 08:46:00 +0000 Subject: [PATCH 5/5] review --- executor/src/witgen/jit/compiler.rs | 23 +++++++++---------- executor/src/witgen/jit/function_cache.rs | 2 +- executor/src/witgen/jit/interpreter.rs | 10 ++++---- executor/src/witgen/jit/processor.rs | 6 ++--- .../src/witgen/jit/single_step_processor.rs | 2 +- executor/src/witgen/jit/variable.rs | 18 +++++++-------- executor/src/witgen/jit/witgen_inference.rs | 8 +++---- 7 files changed, 34 insertions(+), 35 deletions(-) diff --git a/executor/src/witgen/jit/compiler.rs b/executor/src/witgen/jit/compiler.rs index 978bf7250f..9b9e414de4 100644 --- a/executor/src/witgen/jit/compiler.rs +++ b/executor/src/witgen/jit/compiler.rs @@ -68,7 +68,6 @@ extern "C" fn get_fixed_value( id: column, ptype: PolynomialType::Constant, }; - // TODO which size? fixed_data.fixed_cols[&poly_id].values_max_size()[row as usize] } @@ -171,11 +170,11 @@ fn witgen_code( .map(|v| { let var_name = variable_to_string(v); let value = match v { - Variable::Cell(c) => { + Variable::WitnessCell(c) => { format!("get(data, row_offset, {}, {})", c.row_offset, c.id) } Variable::Param(i) => format!("get_param(params, {i})"), - Variable::FixedColumn(_) => panic!("Fixed columns should not be known inputs."), + Variable::FixedCell(_) => panic!("Fixed columns should not be known inputs."), Variable::MachineCallParam(_) => { unreachable!("Machine call variables should not be pre-known.") } @@ -190,7 +189,7 @@ fn witgen_code( .iter() .flat_map(|e| e.referenced_variables()) .filter_map(|v| match v { - Variable::FixedColumn(c) => Some((v, c)), + Variable::FixedCell(c) => Some((v, c)), _ => None, }) .unique() @@ -215,12 +214,12 @@ fn witgen_code( .filter_map(|var| { let value = variable_to_string(var); match var { - Variable::Cell(cell) => Some(format!( + Variable::WitnessCell(cell) => Some(format!( " set(data, row_offset, {}, {}, {value});", cell.row_offset, cell.id, )), Variable::Param(i) => Some(format!(" set_param(params, {i}, {value});")), - Variable::FixedColumn(_) => panic!("Fixed columns should not be written to."), + Variable::FixedCell(_) => panic!("Fixed columns should not be written to."), Variable::MachineCallParam(_) => { // This is just an internal variable. None @@ -233,8 +232,8 @@ fn witgen_code( let store_known = vars_known .iter() .filter_map(|var| match var { - Variable::Cell(cell) => Some(cell), - Variable::Param(_) | Variable::FixedColumn(_) | Variable::MachineCallParam(_) => None, + Variable::WitnessCell(cell) => Some(cell), + Variable::Param(_) | Variable::FixedCell(_) | Variable::MachineCallParam(_) => None, }) .map(|cell| { format!( @@ -423,14 +422,14 @@ fn format_condition(condition: &BranchCondition) - /// Returns the name of a local (stack) variable for the given expression variable. fn variable_to_string(v: &Variable) -> String { match v { - Variable::Cell(cell) => format!( + Variable::WitnessCell(cell) => format!( "c_{}_{}_{}", escape_column_name(&cell.column_name), cell.id, format_row_offset(cell.row_offset) ), Variable::Param(i) => format!("p_{i}"), - Variable::FixedColumn(cell) => { + Variable::FixedCell(cell) => { format!( "f_{}_{}_{}", escape_column_name(&cell.column_name), @@ -552,7 +551,7 @@ mod tests { // } fn cell(column_name: &str, id: u64, row_offset: i32) -> Variable { - Variable::Cell(Cell { + Variable::WitnessCell(Cell { column_name: column_name.to_string(), row_offset, id, @@ -857,7 +856,7 @@ extern \"C\" fn witgen( #[test] fn fixed_column_access() { let a = cell("a", 0, 0); - let x = Variable::FixedColumn(Cell { + let x = Variable::FixedCell(Cell { column_name: "X".to_string(), id: 15, row_offset: 6, diff --git a/executor/src/witgen/jit/function_cache.rs b/executor/src/witgen/jit/function_cache.rs index 9602085049..dd5d75b54b 100644 --- a/executor/src/witgen/jit/function_cache.rs +++ b/executor/src/witgen/jit/function_cache.rs @@ -120,7 +120,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> { .iter() .flat_map(|effect| effect.referenced_variables()) .filter_map(|var| match var { - Variable::Cell(cell) => Some(cell.row_offset), + Variable::WitnessCell(cell) => Some(cell.row_offset), _ => None, }) .all(|row_offset| row_offset >= -1 && row_offset < self.block_size as i32); diff --git a/executor/src/witgen/jit/interpreter.rs b/executor/src/witgen/jit/interpreter.rs index c03deca955..f36a175514 100644 --- a/executor/src/witgen/jit/interpreter.rs +++ b/executor/src/witgen/jit/interpreter.rs @@ -69,7 +69,7 @@ impl EffectsInterpreter { .iter() .flat_map(|e| e.referenced_variables()) .filter_map(|v| match v { - Variable::FixedColumn(c) => Some((v, c)), + Variable::FixedCell(c) => Some((v, c)), _ => None, }) .unique() @@ -86,7 +86,7 @@ impl EffectsInterpreter { known_inputs: &[Variable], ) { actions.extend(known_inputs.iter().map(|var| match var { - Variable::Cell(c) => { + Variable::WitnessCell(c) => { let idx = var_mapper.map_var(var); InterpreterAction::ReadCell(idx, c.clone()) } @@ -94,7 +94,7 @@ impl EffectsInterpreter { let idx = var_mapper.map_var(var); InterpreterAction::ReadParam(idx, *i) } - Variable::FixedColumn(_) | Variable::MachineCallParam(_) => unreachable!(), + Variable::FixedCell(_) | Variable::MachineCallParam(_) => unreachable!(), })); } @@ -153,7 +153,7 @@ impl EffectsInterpreter { .flat_map(Effect::written_vars) .for_each(|(var, _mutable)| { match var { - Variable::Cell(cell) => { + Variable::WitnessCell(cell) => { let idx = var_mapper.get_var(var).unwrap(); actions.push(InterpreterAction::WriteCell(idx, cell.clone())); } @@ -161,7 +161,7 @@ impl EffectsInterpreter { let idx = var_mapper.get_var(var).unwrap(); actions.push(InterpreterAction::WriteParam(idx, *i)); } - Variable::FixedColumn(_) => panic!("Should not write to fixed column."), + Variable::FixedCell(_) => panic!("Should not write to fixed column."), Variable::MachineCallParam(_) => { // This is just an internal variable. } diff --git a/executor/src/witgen/jit/processor.rs b/executor/src/witgen/jit/processor.rs index 5aaeb84728..0583159d6f 100644 --- a/executor/src/witgen/jit/processor.rs +++ b/executor/src/witgen/jit/processor.rs @@ -278,7 +278,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv .known_variables() .iter() .filter_map(|var| match var { - Variable::Cell(cell) => Some(cell.id), + Variable::WitnessCell(cell) => Some(cell.id), _ => None, }) .collect(); @@ -287,7 +287,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv .known_variables() .iter() .filter_map(|var| match var { - Variable::Cell(cell) if cell.id == column_id => Some(cell.row_offset), + Variable::WitnessCell(cell) if cell.id == column_id => Some(cell.row_offset), _ => None, }) .collect::>(); @@ -301,7 +301,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> Processor<'a, T, FixedEv _ => false, }; let cell_var = |row_offset| { - Variable::Cell(Cell { + Variable::WitnessCell(Cell { // Column name does not matter. column_name: "".to_string(), id: column_id, diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs index 42c5b01519..4acf702443 100644 --- a/executor/src/witgen/jit/single_step_processor.rs +++ b/executor/src/witgen/jit/single_step_processor.rs @@ -95,7 +95,7 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> { } fn cell(&self, id: PolyID, row_offset: i32) -> Variable { - Variable::Cell(Cell { + Variable::WitnessCell(Cell { column_name: self.fixed_data.column_name(&id).to_string(), id: id.id, row_offset, diff --git a/executor/src/witgen/jit/variable.rs b/executor/src/witgen/jit/variable.rs index 33a6ebd1db..b7bb25465a 100644 --- a/executor/src/witgen/jit/variable.rs +++ b/executor/src/witgen/jit/variable.rs @@ -8,8 +8,8 @@ use powdr_ast::analyzed::{AlgebraicReference, PolyID, PolynomialType}; #[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Debug)] /// A variable that can be used in the inference engine. pub enum Variable { - /// A witness cell in the current block machine. - Cell(Cell), + /// A witness cell in the current machine. + WitnessCell(Cell), /// A parameter (input or output) of the machine. #[allow(dead_code)] Param(usize), @@ -17,13 +17,13 @@ pub enum Variable { /// identity on a certain row offset. MachineCallParam(MachineCallVariable), /// A fixed column cell. - FixedColumn(Cell), + FixedCell(Cell), } impl Display for Variable { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Variable::Cell(cell) => write!(f, "{cell}"), + Variable::WitnessCell(cell) => write!(f, "{cell}"), Variable::Param(i) => write!(f, "params[{i}]"), Variable::MachineCallParam(ret) => { write!( @@ -32,7 +32,7 @@ impl Display for Variable { ret.identity_id, ret.row_offset, ret.index ) } - Variable::FixedColumn(cell) => write!(f, "{cell}"), + Variable::FixedCell(cell) => write!(f, "{cell}"), } } } @@ -46,8 +46,8 @@ impl Variable { row_offset: r.next as i32 + row_offset, }; match r.poly_id.ptype { - PolynomialType::Committed => Self::Cell(cell), - PolynomialType::Constant => Self::FixedColumn(cell), + PolynomialType::Committed => Self::WitnessCell(cell), + PolynomialType::Constant => Self::FixedCell(cell), _ => panic!(), } } @@ -55,11 +55,11 @@ impl Variable { /// If this variable corresponds to a fixed or witness cell, return the corresponding polynomial ID. pub fn try_to_poly_id(&self) -> Option { match self { - Variable::Cell(cell) => Some(PolyID { + Variable::WitnessCell(cell) => Some(PolyID { id: cell.id, ptype: PolynomialType::Committed, }), - Variable::FixedColumn(cell) => Some(PolyID { + Variable::FixedCell(cell) => Some(PolyID { id: cell.id, ptype: PolynomialType::Constant, }), diff --git a/executor/src/witgen/jit/witgen_inference.rs b/executor/src/witgen/jit/witgen_inference.rs index 0ee31b1874..d4bbd8e5a9 100644 --- a/executor/src/witgen/jit/witgen_inference.rs +++ b/executor/src/witgen/jit/witgen_inference.rs @@ -108,7 +108,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F } pub fn is_known(&self, variable: &Variable) -> bool { - if let Variable::FixedColumn(_) = variable { + if let Variable::FixedCell(_) = variable { true } else { self.known_variables.contains(variable) @@ -462,7 +462,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F /// Record a variable as known. Return true if it was not known before. fn record_known(&mut self, variable: Variable) -> bool { // We do not record fixed columns as known. - if matches!(variable, Variable::FixedColumn(_)) { + if matches!(variable, Variable::FixedCell(_)) { false } else { self.known_variables.insert(variable) @@ -473,7 +473,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator> WitgenInference<'a, T, F /// combining global range constraints and newly derived local range constraints. /// For fixed columns, it also invokes the fixed evaluator. pub fn range_constraint(&self, variable: &Variable) -> RangeConstraint { - if let Variable::FixedColumn(fixed_cell) = variable { + if let Variable::FixedCell(fixed_cell) = variable { if let Some(v) = self.fixed_evaluator.evaluate(fixed_cell) { return RangeConstraint::from_value(v); } @@ -708,7 +708,7 @@ mod test { let known_cells = known_cells.iter().map(|(name, row_offset)| { let id = fixed_data.try_column_by_name(name).unwrap().id; - Variable::Cell(Cell { + Variable::WitnessCell(Cell { column_name: name.to_string(), id, row_offset: *row_offset,