Combine util functions for jit. (#2430)

The JIT compiler used for computing fixed column and the new code for
witgen used two different "util" code snippets. Since the field
implementation for the new witgen code is more advanced, it is now moved
to the generic jit compiler and used by both. It also adds some types
that will later be used by prover functions (PilVec, Callable, etc).
Those types were present already in the fixed column jit code.

Author: chriseth

executor/src/witgen/jit/compiler.rs

@@ -1,15 +1,19 @@
-use std::{cmp::Ordering, ffi::c_void, mem, sync::Arc};
+use std::{cmp::Ordering, ffi::c_void, sync::Arc};
 
 use itertools::Itertools;
 use libloading::Library;
 use powdr_ast::{
     analyzed::{PolyID, PolynomialType},
     indent,
 };
-use powdr_number::{FieldElement, KnownField};
+use powdr_jit_compiler::util_code::util_code;
+use powdr_number::FieldElement;
 
 use crate::witgen::{
-    data_structures::{finalizable_data::CompactDataRef, mutable_state::MutableState},
+    data_structures::{
+        finalizable_data::{ColumnLayout, CompactDataRef},
+        mutable_state::MutableState,
+    },
     machines::{
         profiling::{record_end, record_start},
         LookupCell,
@@ -84,14 +88,20 @@ extern "C" fn call_machine<T: FieldElement, Q: QueryCallback<T>>(
 
 /// Compile the given inferred effects into machine code and load it.
 pub fn compile_effects<T: FieldElement>(
-    first_column_id: u64,
-    column_count: usize,
+    column_layout: ColumnLayout,
     known_inputs: &[Variable],
     effects: &[Effect<T, Variable>],
 ) -> Result<WitgenFunction<T>, String> {
-    let utils = util_code::<T>(first_column_id, column_count)?;
+    let utils = util_code::<T>()?;
+    let interface = interface_code(column_layout);
     let witgen_code = witgen_code(known_inputs, effects);
-    let code = format!("{utils}\n//-------------------------------\n{witgen_code}");
+    let code = format!(
+        "{utils}\n\
+        //-------------------------------\n\
+        {interface}\n\
+        //-------------------------------\n\
+        {witgen_code}"
+    );
 
     record_start("JIT-compilation");
     let start = std::time::Instant::now();
@@ -496,62 +506,20 @@ fn format_row_offset(row_offset: i32) -> String {
     }
 }
 
-/// Returns the rust code containing utility functions given a first column id and a column count
-/// that is used to store the column table.
-fn util_code<T: FieldElement>(first_column_id: u64, column_count: usize) -> Result<String, String> {
-    if !(T::has_direct_repr() && (mem::size_of::<T>() == 8 || mem::size_of::<T>() == 4)) {
-        return Err(format!(
-            "Field {}not supported",
-            T::known_field()
-                .map(|f| format!("{f} "))
-                .unwrap_or_default()
-        ));
-    }
-
-    let field_impl = match T::known_field() {
-        Some(KnownField::GoldilocksField) => {
-            include_str!("includes/field_goldilocks.rs").to_string()
-        }
-        _ => {
-            let int_type = if mem::size_of::<T>() == 8 {
-                "u64"
-            } else {
-                "u32"
-            };
-            let double_int_type = if mem::size_of::<T>() == 8 {
-                "u128"
-            } else {
-                "u64"
-            };
-            let modulus = T::modulus();
-
-            format!(
-                "\
-                #[derive(Clone, Copy, Default)]\n\
-                #[repr(transparent)]\n\
-                struct FieldElement({int_type});\n\
-                \n\
-                type IntType = {int_type};\n\
-                type DoubleIntType = {double_int_type};\n\
-                const MODULUS: IntType = {modulus}_{int_type};\n\
-                {}\
-                ",
-                include_str!("includes/field_generic_up_to_64.rs")
-            )
-        }
-    };
-
-    let interface = format!(
+/// Returns the rust code containing functions and data structures used to
+/// interface with witgen functions given the layout of the trace table.
+fn interface_code(column_layout: ColumnLayout) -> String {
+    let ColumnLayout {
+        column_count,
+        first_column_id,
+    } = column_layout;
+    format!(
         "\
         const column_count: u64 = {column_count};\n\
         const first_column_id: u64 = {first_column_id};\n\
         {}",
         include_str!("includes/interface.rs")
-    );
-
-    Ok(format!(
-        "#![allow(non_snake_case, unused_parens, unused_variables)]\n{field_impl}\n{interface}"
-    ))
+    )
 }
 
 #[cfg(test)]
@@ -570,9 +538,24 @@ mod tests {
 
     use super::*;
 
+    fn compile_effects(
+        column_count: usize,
+        known_inputs: &[Variable],
+        effects: &[Effect<GoldilocksField, Variable>],
+    ) -> Result<WitgenFunction<GoldilocksField>, String> {
+        super::compile_effects(
+            ColumnLayout {
+                column_count,
+                first_column_id: 0,
+            },
+            known_inputs,
+            effects,
+        )
+    }
+
     #[test]
     fn compile_util_code_goldilocks() {
-        compile_effects::<GoldilocksField>(0, 2, &[], &[]).unwrap();
+        compile_effects(2, &[], &[]).unwrap();
     }
 
     // We would like to test the generic field implementation, but
@@ -728,7 +711,7 @@ extern \"C\" fn witgen(
             assignment(&x, number(7)),
             assignment(&y, symbol(&x) + number(2)),
         ];
-        let f = compile_effects(0, 1, &[], &effects).unwrap();
+        let f = compile_effects(1, &[], &effects).unwrap();
         let mut data = vec![GoldilocksField::from(0); 2];
         let mut known = vec![0; 1];
         (f.function)(witgen_fun_params(&mut data, &mut known));
@@ -751,8 +734,8 @@ extern \"C\" fn witgen(
         let row_count = 2;
         let column_count = 2;
         let data_len = column_count * row_count;
-        let f1 = compile_effects(0, column_count, &[], &effects1).unwrap();
-        let f2 = compile_effects(0, column_count, &[], &effects2).unwrap();
+        let f1 = compile_effects(column_count, &[], &effects1).unwrap();
+        let f2 = compile_effects(column_count, &[], &effects2).unwrap();
         let mut data = vec![GoldilocksField::from(0); data_len];
         let mut known = vec![0; row_count];
         (f1.function)(witgen_fun_params(&mut data, &mut known));
@@ -788,7 +771,7 @@ extern \"C\" fn witgen(
             assignment(&cell("x", 0, 3), number(8).field_div(&-number(2))),
             assignment(&cell("x", 0, 4), (-number(8)).field_div(&-number(2))),
         ];
-        let f = compile_effects(0, 1, &[], &effects).unwrap();
+        let f = compile_effects(1, &[], &effects).unwrap();
         let mut data = vec![GoldilocksField::from(0); 5];
         let mut known = vec![0; 5];
         (f.function)(witgen_fun_params(&mut data, &mut known));
@@ -809,7 +792,7 @@ extern \"C\" fn witgen(
         let z = cell("z", 2, 0);
         let effects = vec![assignment(&x, symbol(&y) * symbol(&z))];
         let known_inputs = vec![y.clone(), z.clone()];
-        let f = compile_effects(0, 3, &known_inputs, &effects).unwrap();
+        let f = compile_effects(3, &known_inputs, &effects).unwrap();
         let mut data = vec![
             GoldilocksField::from(0),
             GoldilocksField::from(3),
@@ -830,7 +813,7 @@ extern \"C\" fn witgen(
             assignment(&z, symbol(&x).integer_div(&-number(10))),
         ];
         let known_inputs = vec![x.clone()];
-        let f = compile_effects(0, 3, &known_inputs, &effects).unwrap();
+        let f = compile_effects(3, &known_inputs, &effects).unwrap();
         let mut data = vec![
             GoldilocksField::from(23),
             GoldilocksField::from(0),
@@ -850,7 +833,7 @@ extern \"C\" fn witgen(
         let x_val: GoldilocksField = 7.into();
         let mut y_val: GoldilocksField = 9.into();
         let effects = vec![assignment(&y, symbol(&x) + number(7))];
-        let f = compile_effects(0, 1, &[x], &effects).unwrap();
+        let f = compile_effects(1, &[x], &effects).unwrap();
         let mut data = vec![];
         let mut known = vec![];
         let mut params = vec![LookupCell::Input(&x_val), LookupCell::Output(&mut y_val)];
@@ -894,7 +877,7 @@ extern \"C\" fn witgen(
             row_offset: 6,
         });
         let effects = vec![assignment(&a, symbol(&x))];
-        let f = compile_effects(0, 1, &[], &effects).unwrap();
+        let f = compile_effects(1, &[], &effects).unwrap();
         let mut data = vec![7.into()];
         let mut known = vec![0];
         let mut params = vec![];
@@ -954,7 +937,7 @@ extern \"C\" fn witgen(
             Effect::Assignment(y.clone(), symbol(&r2)),
         ];
         let known_inputs = vec![];
-        let f = compile_effects(0, 3, &known_inputs, &effects).unwrap();
+        let f = compile_effects(3, &known_inputs, &effects).unwrap();
         let mut data = vec![GoldilocksField::from(0); 3];
         let mut known = vec![0; 1];
         let params = WitgenFunctionParams {
@@ -987,7 +970,7 @@ extern \"C\" fn witgen(
             vec![assignment(&y, symbol(&x) + number(1))],
             vec![assignment(&y, symbol(&x) + number(2))],
         )];
-        let f = compile_effects(0, 1, &[x], &effects).unwrap();
+        let f = compile_effects(1, &[x], &effects).unwrap();
         let mut data = vec![];
         let mut known = vec![];
 

executor/src/witgen/jit/function_cache.rs

@@ -158,13 +158,7 @@ impl<'a, T: FieldElement> FunctionCache<'a, T> {
             .collect::<Vec<_>>();
 
         log::trace!("Compiling effects...");
-        let function = compile_effects(
-            self.column_layout.first_column_id,
-            self.column_layout.column_count,
-            &known_inputs,
-            &code,
-        )
-        .unwrap();
+        let function = compile_effects(self.column_layout.clone(), &known_inputs, &code).unwrap();
         log::trace!("Compilation done.");
 
         Some(CacheEntry {

jit-compiler/src/compiler.rs

@@ -13,20 +13,16 @@ use powdr_ast::{
         types::{FunctionType, Type, TypeScheme},
     },
 };
-use powdr_number::{FieldElement, LargeInt};
+use powdr_number::FieldElement;
 
-use crate::{codegen::escape_symbol, CompiledPIL, FixedColFunction};
+use crate::{codegen::escape_symbol, util_code::util_code, CompiledPIL, FixedColFunction};
 
 pub fn generate_glue_code<T: FieldElement>(
     symbols: &[(&str, String)],
     analyzed: &Analyzed<T>,
 ) -> Result<String, String> {
-    if T::BITS > 64 {
-        return Err(format!(
-            "Fields with more than 64 bits not supported, requested {}",
-            T::BITS,
-        ));
-    }
+    let utils = util_code::<T>()?;
+
     let mut glue = String::new();
     let int_int_fun: TypeScheme = Type::Function(FunctionType {
         params: vec![Type::Int],
@@ -54,128 +50,7 @@ pub fn generate_glue_code<T: FieldElement>(
         ));
     }
 
-    Ok(format!(
-        "{PREAMBLE}\n{}\n{glue}\n",
-        field_specific_preamble::<T>()
-    ))
-}
-
-const PREAMBLE: &str = r#"
-#![allow(unused_parens, unused_variables)]
-
-static DEGREE: std::sync::RwLock<Option<ibig::IBig>> = std::sync::RwLock::new(None);
-
-#[no_mangle]
-pub extern "C" fn __set_degree(degree: u64) {
-    *DEGREE.write().unwrap() = Some(ibig::IBig::from(degree));
-}
-
-#[derive(Clone, Copy)]
-struct FieldElement(u64);
-impl From<FieldElement> for u64 {
-    fn from(x: FieldElement) -> u64 {
-        x.0
-    }
-}
-impl From<ibig::IBig> for FieldElement {
-    fn from(x: ibig::IBig) -> Self {
-        FieldElement(u64::try_from(x).unwrap())
-    }
-}
-impl From<FieldElement> for ibig::IBig {
-    fn from(x: FieldElement) -> Self {
-        ibig::IBig::from(x.0)
-    }
-}
-
-#[derive(Clone)]
-enum Callable<Args, Ret> {
-    Fn(fn(Args) -> Ret),
-    Closure(std::sync::Arc<dyn Fn(Args) -> Ret + Send + Sync>),
-}
-impl<Args, Ret> Callable<Args, Ret> {
-    #[inline(always)]
-    fn call(&self, args: Args) -> Ret {
-        match self {
-            Callable::Fn(f) => f(args),
-            Callable::Closure(f) => f(args),
-        }
-    }
-}
-
-#[derive(Clone)]
-struct PilVec<T>(std::sync::Arc<Vec<T>>);
-
-impl<T> PilVec<T> {
-    fn len(&self) -> usize {
-        self.0.len()
-    }
-}
-impl<T> From<Vec<T>> for PilVec<T> {
-    fn from(v: Vec<T>) -> Self {
-        PilVec(std::sync::Arc::new(v))
-    }
-}
-impl<T> std::ops::Index<usize> for PilVec<T> {
-    type Output = T;
-
-    #[inline]
-    fn index(&self, index: usize) -> &T {
-        &self.0[index]
-    }
-}
-    
-
-trait Add {
-    fn add(a: Self, b: Self) -> Self;
-}
-
-impl Add for ibig::IBig {
-    fn add(a: Self, b: Self) -> Self { a + b }
-}
-
-impl<T: Clone> Add for PilVec<T> {
-    fn add(a: Self, b: Self) -> Self {
-        // TODO for a regular "push" or array::map this is very slow.
-        // We could optimize this by sharing a larger backing vector
-        // across prefix instances, allowing to extend the backing vector if
-        // our view is the full vector.
-        PilVec(std::sync::Arc::new(
-            a.0.as_ref().iter().chain(b.0.as_ref()).cloned().collect::<Vec<_>>()
-        ))
-    }
-}
-
-trait FromLiteral {
-    fn from_u64(x: u64) -> Self;
-}
-impl FromLiteral for ibig::IBig {
-    fn from_u64(x: u64) -> Self { ibig::IBig::from(x) }
-}
-impl FromLiteral for FieldElement {
-    fn from_u64(x: u64) -> Self { FieldElement::from(x) }
-}
-
-"#;
-
-fn field_specific_preamble<T: FieldElement>() -> String {
-    let modulus = u64::try_from(T::modulus().to_arbitrary_integer()).unwrap();
-    format!(
-        r#"
-        impl From<u64> for FieldElement {{
-            fn from(x: u64) -> Self {{
-                // TODO this is inefficient.
-                FieldElement(x % {modulus}_u64)
-            }}
-        }}
-        impl Add for FieldElement {{
-            fn add(a: Self, b: Self) -> Self {{
-                // TODO this is inefficient.
-                Self(u64::try_from((u128::from(a.0) + u128::from(b.0)) % u128::from({modulus}_u64)).unwrap())
-            }}
-        }}
-        "#
-    )
+    Ok(format!("{utils}\n{glue}\n",))
 }
 
 const CARGO_TOML: &str = r#"