Interpreter for witgen effects (#2301)

Implementation of an interpreter for jit witgen effects.
Performance seems to be around 0.20x of the compiled jit code (poseidon
benchmark).
- Disabled by default, can be enabled setting `POWDR_JIT_INTERPRETER=1`.
- Also introduces a way to disable JIT in general, with
`POWDR_JIT_DISABLE=1`.

executor/src/witgen/data_structures/finalizable_data.rs

@@ -91,7 +91,7 @@ impl<T: FieldElement> CompactData<T> {
     }
 
     /// Sets an entire row at the given index
-    pub fn set(&mut self, row: usize, new_row: Row<T>) {
+    pub fn set_row(&mut self, row: usize, new_row: Row<T>) {
         let idx = row * self.column_count;
         for (i, col_id) in self.column_ids().enumerate() {
             if let Some(v) = new_row.value(&col_id) {
@@ -109,6 +109,7 @@ impl<T: FieldElement> CompactData<T> {
         self.known_cells.append_empty_rows(count);
     }
 
+    #[inline]
     fn index(&self, row: usize, col: u64) -> usize {
         let col = col - self.first_column_id;
         row * self.column_count + col as usize
@@ -121,6 +122,14 @@ impl<T: FieldElement> CompactData<T> {
         })
     }
 
+    /// Sets a single cell
+    pub fn set(&mut self, row: usize, col: u64, value: T) {
+        let idx = self.index(row, col);
+        self.data[idx] = value;
+        let relative_col = col - self.first_column_id;
+        self.known_cells.set(row, relative_col, true);
+    }
+
     pub fn get(&self, row: usize, col: u64) -> (T, bool) {
         let idx = self.index(row, col);
         let relative_col = col - self.first_column_id;
@@ -146,8 +155,8 @@ impl<T: FieldElement> CompactData<T> {
 /// only for a certain block of rows, starting from row index zero.
 /// It allows negative row indices as well.
 pub struct CompactDataRef<'a, T> {
-    data: &'a mut CompactData<T>,
-    row_offset: usize,
+    pub data: &'a mut CompactData<T>,
+    pub row_offset: usize,
 }
 
 impl<'a, T: FieldElement> CompactDataRef<'a, T> {
@@ -160,10 +169,6 @@ impl<'a, T: FieldElement> CompactDataRef<'a, T> {
     pub fn as_mut_slices(&mut self) -> (&mut [T], &mut [u32]) {
         self.data.as_mut_slices()
     }
-
-    pub fn row_offset(&self) -> usize {
-        self.row_offset
-    }
 }
 
 /// A data structure that stores witness data.
@@ -413,7 +418,7 @@ impl<'a, T: FieldElement> FinalizableData<'a, T> {
     pub fn set(&mut self, i: usize, row: Row<T>) {
         match self.location_of_row(i) {
             Location::Finalized(local) => {
-                self.finalized_data.set(local, row);
+                self.finalized_data.set_row(local, row);
             }
             Location::PostFinalized(local) => self.post_finalized_data[local] = row,
         }

executor/src/witgen/jit/compiler.rs

@@ -1,4 +1,4 @@
-use std::{cmp::Ordering, ffi::c_void, iter, mem, sync::Arc};
+use std::{cmp::Ordering, ffi::c_void, mem, sync::Arc};
 
 use itertools::Itertools;
 use libloading::Library;
@@ -20,7 +20,7 @@ use super::{
     variable::Variable,
 };
 
-pub struct WitgenFunction<T> {
+pub struct WitgenFunction<T: FieldElement> {
     // TODO We might want to pass arguments as direct function parameters
     // (instead of a struct), so that
     // they are stored in registers instead of the stack. Should be checked.
@@ -39,7 +39,7 @@ impl<T: FieldElement> WitgenFunction<T> {
         params: &mut [LookupCell<T>],
         mut data: CompactDataRef<'_, T>,
     ) {
-        let row_offset = data.row_offset().try_into().unwrap();
+        let row_offset = data.row_offset.try_into().unwrap();
         let (data, known) = data.as_mut_slices();
         (self.function)(WitgenFunctionParams {
             data: data.into(),
@@ -160,7 +160,7 @@ fn witgen_code<T: FieldElement>(
     let main_code = format_effects(effects);
     let vars_known = effects
         .iter()
-        .flat_map(written_vars_in_effect)
+        .flat_map(Effect::written_vars)
         .map(|(var, _)| var)
         .collect_vec();
     let store_values = vars_known
@@ -224,30 +224,6 @@ extern "C" fn witgen(
     )
 }
 
-/// Returns an iterator over all variables written to in the effect.
-/// The flag indicates if the variable is the return value of a machine call and thus needs
-/// to be declared mutable.
-fn written_vars_in_effect<T: FieldElement>(
-    effect: &Effect<T, Variable>,
-) -> Box<dyn Iterator<Item = (&Variable, bool)> + '_> {
-    match effect {
-        Effect::Assignment(var, _) => Box::new(iter::once((var, false))),
-        Effect::RangeConstraint(..) => unreachable!(),
-        Effect::Assertion(..) => Box::new(iter::empty()),
-        Effect::MachineCall(_, known, vars) => Box::new(
-            vars.iter()
-                .zip_eq(known)
-                .flat_map(|(v, known)| (!known).then_some((v, true))),
-        ),
-        Effect::Branch(_, first, second) => Box::new(
-            first
-                .iter()
-                .chain(second)
-                .flat_map(|e| written_vars_in_effect(e)),
-        ),
-    }
-}
-
 pub fn format_effects<T: FieldElement>(effects: &[Effect<T, Variable>]) -> String {
     format_effects_inner(effects, true)
 }
@@ -321,7 +297,7 @@ fn format_effect<T: FieldElement>(effect: &Effect<T, Variable>, is_top_level: bo
                 first
                     .iter()
                     .chain(second)
-                    .flat_map(|e| written_vars_in_effect(e))
+                    .flat_map(|e| e.written_vars())
                     .sorted()
                     .dedup()
                     .map(|(v, needs_mut)| {

executor/src/witgen/jit/effect.rs

@@ -27,6 +27,27 @@ pub enum Effect<T: FieldElement, V> {
     Branch(BranchCondition<T, V>, Vec<Effect<T, V>>, Vec<Effect<T, V>>),
 }
 
+impl<T: FieldElement> Effect<T, Variable> {
+    /// Returns an iterator over all variables written to in the effect.
+    /// The flag indicates if the variable is the return value of a machine call and thus needs
+    /// to be declared mutable.
+    pub fn written_vars(&self) -> Box<dyn Iterator<Item = (&Variable, bool)> + '_> {
+        match self {
+            Effect::Assignment(var, _) => Box::new(iter::once((var, false))),
+            Effect::RangeConstraint(..) => unreachable!(),
+            Effect::Assertion(..) => Box::new(iter::empty()),
+            Effect::MachineCall(_, known, vars) => Box::new(
+                vars.iter()
+                    .zip_eq(known)
+                    .flat_map(|(v, known)| (!known).then_some((v, true))),
+            ),
+            Effect::Branch(_, first, second) => {
+                Box::new(first.iter().chain(second).flat_map(|e| e.written_vars()))
+            }
+        }
+    }
+}
+
 impl<T: FieldElement, V: Hash + Eq> Effect<T, V> {
     pub fn referenced_variables(&self) -> impl Iterator<Item = &V> {
         let iter: Box<dyn Iterator<Item = &V>> = match self {