diff --git a/executor/Cargo.toml b/executor/Cargo.toml index bb7080b17..3c5c0c8e6 100644 --- a/executor/Cargo.toml +++ b/executor/Cargo.toml @@ -12,6 +12,7 @@ powdr-ast.workspace = true powdr-number.workspace = true powdr-parser-util.workspace = true powdr-pil-analyzer.workspace = true +powdr-jit-compiler.workspace = true itertools = "0.13" log = { version = "0.4.17" } diff --git a/executor/src/constant_evaluator/interpreter.rs b/executor/src/constant_evaluator/interpreter.rs new file mode 100644 index 000000000..410f16ee1 --- /dev/null +++ b/executor/src/constant_evaluator/interpreter.rs @@ -0,0 +1,158 @@ +use std::{ + collections::{BTreeMap, HashMap}, + sync::{Arc, RwLock}, +}; + +use powdr_ast::{ + analyzed::{Analyzed, Expression, FunctionValueDefinition, Symbol, TypedExpression}, + parsed::{ + types::{ArrayType, Type}, + IndexAccess, + }, +}; +use powdr_number::{BigInt, BigUint, DegreeType, FieldElement}; +use powdr_pil_analyzer::evaluator::{self, Definitions, SymbolLookup, Value}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; + +/// Evaluates the fixed polynomial `name` on all values from 0 to `degree - 1` +/// using an interpreter. +/// If `index` is `Some(i)`, evaluates the `i`-th element of the array. +pub fn generate_values( + analyzed: &Analyzed, + degree: DegreeType, + name: &str, + body: &FunctionValueDefinition, + index: Option, +) -> Vec { + let symbols = CachedSymbols { + symbols: &analyzed.definitions, + solved_impls: &analyzed.solved_impls, + cache: Arc::new(RwLock::new(Default::default())), + degree, + }; + let result = match body { + FunctionValueDefinition::Expression(TypedExpression { e, type_scheme }) => { + if let Some(type_scheme) = type_scheme { + assert!(type_scheme.vars.is_empty()); + let ty = &type_scheme.ty; + if ty == &Type::Col { + assert!(index.is_none()); + } else if let Type::Array(ArrayType { base, length: _ }) = ty { + assert!(index.is_some()); + assert_eq!(base.as_ref(), &Type::Col); + } else { + panic!("Invalid fixed column type: {ty}"); + } + }; + let index_expr; + let e = if let Some(index) = index { + index_expr = IndexAccess { + array: e.clone().into(), + index: Box::new(BigUint::from(index).into()), + } + .into(); + &index_expr + } else { + e + }; + let fun = evaluator::evaluate(e, &mut symbols.clone()).unwrap(); + (0..degree) + .into_par_iter() + .map(|i| { + evaluator::evaluate_function_call( + fun.clone(), + vec![Arc::new(Value::Integer(BigInt::from(i)))], + &mut symbols.clone(), + )? + .try_to_field_element() + }) + .collect::, _>>() + } + FunctionValueDefinition::Array(values) => { + assert!(index.is_none()); + values + .to_repeated_arrays(degree) + .map(|elements| { + let items = elements + .pattern() + .iter() + .map(|v| { + let mut symbols = symbols.clone(); + evaluator::evaluate(v, &mut symbols) + .and_then(|v| v.try_to_field_element()) + }) + .collect::, _>>()?; + + Ok(items + .into_iter() + .cycle() + .take(elements.size() as usize) + .collect::>()) + }) + .collect::, _>>() + .map(|values| { + let values: Vec = values.into_iter().flatten().collect(); + assert_eq!(values.len(), degree as usize); + values + }) + } + FunctionValueDefinition::TypeDeclaration(_) + | FunctionValueDefinition::TypeConstructor(_, _) + | FunctionValueDefinition::TraitDeclaration(_) + | FunctionValueDefinition::TraitFunction(_, _) => panic!(), + }; + match result { + Err(err) => { + eprintln!("Error evaluating fixed polynomial {name}{body}:\n{err}"); + panic!("{err}"); + } + Ok(v) => v, + } +} + +type SymbolCache<'a, T> = HashMap>, Arc>>>; + +#[derive(Clone)] +pub struct CachedSymbols<'a, T> { + symbols: &'a HashMap)>, + solved_impls: &'a HashMap, Arc>>, + cache: Arc>>, + degree: DegreeType, +} + +impl<'a, T: FieldElement> SymbolLookup<'a, T> for CachedSymbols<'a, T> { + fn lookup( + &mut self, + name: &'a str, + type_args: &Option>, + ) -> Result>, evaluator::EvalError> { + if let Some(v) = self + .cache + .read() + .unwrap() + .get(name) + .and_then(|map| map.get(type_args)) + { + return Ok(v.clone()); + } + let result = Definitions::lookup_with_symbols( + self.symbols, + self.solved_impls, + name, + type_args, + self, + )?; + self.cache + .write() + .unwrap() + .entry(name.to_string()) + .or_default() + .entry(type_args.clone()) + .or_insert_with(|| result.clone()); + Ok(result) + } + + fn degree(&self) -> Result>, evaluator::EvalError> { + Ok(Value::Integer(self.degree.into()).into()) + } +} diff --git a/executor/src/constant_evaluator/jit_compiler.rs b/executor/src/constant_evaluator/jit_compiler.rs new file mode 100644 index 000000000..b8d7e1290 --- /dev/null +++ b/executor/src/constant_evaluator/jit_compiler.rs @@ -0,0 +1,63 @@ +use std::collections::HashMap; + +use powdr_ast::analyzed::{Analyzed, PolyID}; +use powdr_number::FieldElement; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; + +use super::VariablySizedColumn; + +/// Tries to just-in-time compile all fixed columns in `analyzed` +/// and then evaluates those functions to return `VariablySizedColumn`s. +/// Ignoreds all columns where the compilation fails. +pub fn generate_values( + analyzed: &Analyzed, +) -> HashMap<(String, PolyID), VariablySizedColumn> { + let fun_map = match powdr_jit_compiler::compile(analyzed, &symbols_to_compile(analyzed)) { + Err(err) => { + log::error!("Failed to compile some constant columns: {}", err); + return HashMap::new(); + } + Ok(fun_map) => fun_map, + }; + + analyzed + .constant_polys_in_source_order() + .into_iter() + .filter_map(|(symbol, _)| { + let fun = fun_map.get(symbol.absolute_name.as_str())?; + Some((symbol, fun)) + }) + .map(|(symbol, fun)| { + let column_values: Vec> = symbol + .degree + .unwrap() + .iter() + .map(|degree| { + let values = (0..degree) + .into_par_iter() + .map(|i| { + let result = fun.call(i as u64); + T::from(result) + }) + .collect(); + values + }) + .collect(); + + ( + (symbol.absolute_name.clone(), symbol.into()), + column_values.into(), + ) + }) + .collect() +} + +fn symbols_to_compile(analyzed: &Analyzed) -> Vec<&str> { + analyzed + .constant_polys_in_source_order() + .into_iter() + .filter_map(|(symbol, value)| { + (!symbol.is_array() && value.is_some()).then_some(symbol.absolute_name.as_str()) + }) + .collect() +} diff --git a/executor/src/constant_evaluator/mod.rs b/executor/src/constant_evaluator/mod.rs index 47f0327c1..c24be88eb 100644 --- a/executor/src/constant_evaluator/mod.rs +++ b/executor/src/constant_evaluator/mod.rs @@ -1,22 +1,11 @@ -use std::{ - collections::{BTreeMap, HashMap}, - sync::{Arc, RwLock}, -}; - pub use data_structures::{get_uniquely_sized, get_uniquely_sized_cloned, VariablySizedColumn}; use itertools::Itertools; -use powdr_ast::{ - analyzed::{Analyzed, Expression, FunctionValueDefinition, Symbol, TypedExpression}, - parsed::{ - types::{ArrayType, Type}, - IndexAccess, - }, -}; -use powdr_number::{BigInt, BigUint, DegreeType, FieldElement}; -use powdr_pil_analyzer::evaluator::{self, Definitions, SymbolLookup, Value}; -use rayon::prelude::{IntoParallelIterator, ParallelIterator}; +use powdr_ast::analyzed::Analyzed; +use powdr_number::FieldElement; mod data_structures; +mod interpreter; +mod jit_compiler; /// Generates the fixed column values for all fixed columns that are defined /// (and not just declared). @@ -24,171 +13,35 @@ mod data_structures; /// Arrays of columns are flattened, the name of the `i`th array element /// is `name[i]`. pub fn generate(analyzed: &Analyzed) -> Vec<(String, VariablySizedColumn)> { - let mut fixed_cols = HashMap::new(); + let mut fixed_cols = jit_compiler::generate_values(analyzed); for (poly, value) in analyzed.constant_polys_in_source_order() { if let Some(value) = value { // For arrays, generate values for each index, // for non-arrays, set index to None. for (index, (name, id)) in poly.array_elements().enumerate() { - let index = poly.is_array().then_some(index as u64); - let range = poly.degree.unwrap(); - let values = range - .iter() - .map(|degree| generate_values(analyzed, degree, &name, value, index)) - .collect::>() - .into(); - assert!(fixed_cols.insert(name, (id, values)).is_none()); + if !fixed_cols.contains_key(&(name.clone(), id)) { + let index = poly.is_array().then_some(index as u64); + let range = poly.degree.unwrap(); + let values = range + .iter() + .map(|degree| { + interpreter::generate_values(analyzed, degree, &name, value, index) + }) + .collect::>() + .into(); + fixed_cols.insert((name, id), values); + } } } } fixed_cols .into_iter() - .sorted_by_key(|(_, (id, _))| *id) - .map(|(name, (_, values))| (name, values)) + .sorted_by_key(|((_, id), _)| *id) + .map(|((name, _), values)| (name, values)) .collect() } -fn generate_values( - analyzed: &Analyzed, - degree: DegreeType, - name: &str, - body: &FunctionValueDefinition, - index: Option, -) -> Vec { - let symbols = CachedSymbols { - symbols: &analyzed.definitions, - solved_impls: &analyzed.solved_impls, - cache: Arc::new(RwLock::new(Default::default())), - degree, - }; - let result = match body { - FunctionValueDefinition::Expression(TypedExpression { e, type_scheme }) => { - if let Some(type_scheme) = type_scheme { - assert!(type_scheme.vars.is_empty()); - let ty = &type_scheme.ty; - if ty == &Type::Col { - assert!(index.is_none()); - } else if let Type::Array(ArrayType { base, length: _ }) = ty { - assert!(index.is_some()); - assert_eq!(base.as_ref(), &Type::Col); - } else { - panic!("Invalid fixed column type: {ty}"); - } - }; - let index_expr; - let e = if let Some(index) = index { - index_expr = IndexAccess { - array: e.clone().into(), - index: Box::new(BigUint::from(index).into()), - } - .into(); - &index_expr - } else { - e - }; - let fun = evaluator::evaluate(e, &mut symbols.clone()).unwrap(); - (0..degree) - .into_par_iter() - .map(|i| { - evaluator::evaluate_function_call( - fun.clone(), - vec![Arc::new(Value::Integer(BigInt::from(i)))], - &mut symbols.clone(), - )? - .try_to_field_element() - }) - .collect::, _>>() - } - FunctionValueDefinition::Array(values) => { - assert!(index.is_none()); - values - .to_repeated_arrays(degree) - .map(|elements| { - let items = elements - .pattern() - .iter() - .map(|v| { - let mut symbols = symbols.clone(); - evaluator::evaluate(v, &mut symbols) - .and_then(|v| v.try_to_field_element()) - }) - .collect::, _>>()?; - - Ok(items - .into_iter() - .cycle() - .take(elements.size() as usize) - .collect::>()) - }) - .collect::, _>>() - .map(|values| { - let values: Vec = values.into_iter().flatten().collect(); - assert_eq!(values.len(), degree as usize); - values - }) - } - FunctionValueDefinition::TypeDeclaration(_) - | FunctionValueDefinition::TypeConstructor(_, _) - | FunctionValueDefinition::TraitDeclaration(_) - | FunctionValueDefinition::TraitFunction(_, _) => panic!(), - }; - match result { - Err(err) => { - eprintln!("Error evaluating fixed polynomial {name}{body}:\n{err}"); - panic!("{err}"); - } - Ok(v) => v, - } -} - -type SymbolCache<'a, T> = HashMap>, Arc>>>; - -#[derive(Clone)] -pub struct CachedSymbols<'a, T> { - symbols: &'a HashMap)>, - solved_impls: &'a HashMap, Arc>>, - cache: Arc>>, - degree: DegreeType, -} - -impl<'a, T: FieldElement> SymbolLookup<'a, T> for CachedSymbols<'a, T> { - fn lookup( - &mut self, - name: &'a str, - type_args: &Option>, - ) -> Result>, evaluator::EvalError> { - if let Some(v) = self - .cache - .read() - .unwrap() - .get(name) - .and_then(|map| map.get(type_args)) - { - return Ok(v.clone()); - } - let result = Definitions::lookup_with_symbols( - self.symbols, - self.solved_impls, - name, - type_args, - self, - )?; - self.cache - .write() - .unwrap() - .entry(name.to_string()) - .or_default() - .entry(type_args.clone()) - .or_insert_with(|| result.clone()); - Ok(result) - } - - fn degree(&self) -> Result>, evaluator::EvalError> { - Ok(Value::Integer(self.degree.into()).into()) - } -} - #[cfg(test)] mod test { use itertools::Itertools;