Interpreter for witgen effects #2301

pacheco · 2025-01-02T19:34:22Z

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.

chriseth · 2025-01-02T21:25:10Z

executor/src/witgen/jit/interpreter.rs

+    }
+
+    // Execute the machine effects for the given the parameters
+    pub fn call<Q: QueryCallback<T>>(&self, params: WitgenFunctionParams<'_, T>) {


Is it possible to use a higher-level interface than WitgenFunctionParams? They require unsafe code to be used and are tailored towards an FFI.

I think CompactDataRef<'_, T>, would be the right thing to pass

chriseth · 2025-01-02T21:25:46Z

executor/src/witgen/jit/interpreter.rs

+            idx
+        };
+
+        // load known inputs


Sounds like this code would be perfect inside a function called load_known_inputs ;)

chriseth · 2025-01-02T21:29:21Z

executor/src/witgen/jit/symbolic_expression.rs

+        RPNExpression { elems }
+    }
+
+    fn to_rpn_inner(&self, elems: &mut Vec<RPNExpressionElem<T, V>>) {


Are you worried about machine stack size or is the RPN transformation a performance optimization?

from some profiling, the "evaluate expression" code was a hot spot... RPN was quite a bit faster.
I'm not sure if its just due to avoiding recursion or because the whole expression is kept in a vec (avoiding pointers).

Yeah, thought, about the same, locality might be a thing.

Could you move this function to the interpreter, though? I think SymbolicExpression should not care that there is a thing called RPNExpressionElem.

executor/src/witgen/jit/interpreter.rs

Schaeff · 2025-01-03T10:49:13Z

executor/src/witgen/jit/interpreter.rs

+    }
+}
+
+// the following functions come from the interface.rs file also included in the compiled jit code


Does this mean it's duplicated? Any way it could be reused?

yep, I was lazy because the include was not part of the module tree, let me try and improve this

I think we should just use CompactDataRef

Schaeff

Just some mostly non-blocking comments

chriseth · 2025-01-03T11:25:30Z

executor/src/witgen/jit/function_cache.rs

-                )
-                .unwrap()
+                let compiled_jit =
+                    !matches!(std::env::var("POWDR_JIT_INTERPRETER"), Ok(val) if val == "1");


I don't particularly like the use of environment variables. Could we move them as far up as possible? In my opinion, it should just be a configuration setting in the constructor of FuncitonCache.

i removed the other env var (for fully disabling jit), can we just keep this one as is for the moment (while interpreter is fully integrated) for easy testing?

To be honest, I would prefer to remove the interpreter from function_cache for the time being...

ok, removed it completely

chriseth · 2025-01-03T11:27:51Z

While this is a great piece of software, I think we should actually just not at all integrate it for now. The FunctionCache is still specialized for the BlockMachineProcessor (i.e. it does not support dynamic machines), and it will be even harder to disentangle if we also have to deal with the interpreter as well.

chriseth · 2025-01-03T11:31:22Z

executor/src/witgen/jit/interpreter.rs

+    ) -> Self {
+        let mut actions = vec![];
+        let mut var_idx = HashMap::new();
+        let mut vars = 0;


var_counter?

logic moved to struct

chriseth · 2025-01-03T11:45:33Z

executor/src/witgen/jit/interpreter.rs

+                    let idx = map_var_idx(var);
+                    actions.push(InterpreterAction::AssignExpression(
+                        idx,
+                        e.map_variables(&mut map_var_idx).to_rpn(),


Would it be possible to perform variable mapping and rpn conversion in a single struct function where the variable mapping is internal state of the struct?

chriseth · 2025-01-03T11:48:20Z

executor/src/witgen/jit/interpreter.rs

+            match action {
+                InterpreterAction::AssignExpression(idx, e) => {
+                    let val = self.evaluate_expression(&mut eval_stack, &vars, e);
+                    assert!(vars[*idx].replace(val).is_none());


Instead of this runtime assertion, we can actually check in the conversion inside the new function if each variable is written to exactly once (and only read after it has been written).

chriseth · 2025-01-03T11:56:17Z

executor/src/witgen/jit/interpreter.rs

+                    let mut arg_values: Vec<_> = arguments
+                        .iter()
+                        .map(|a| match a {
+                            MachineCallArgument::Unknown(v) => (Some(v), Default::default()),


Is this Some(v) used anywhere?

Do you think there is a way to directly store a &mut of the variable inside args instead of having to construct this second arg_values vector?

Then we also don't need the second "write output variables" step. Of course you need to first store a Some(0) - but maybe we could also just initialize all the variables with 0 once we checked consistency (i.e. all variables written exactly once and only read after written)?

changed things around a bit here: the knowns/unknowns are transformed into variable indexes (known arguments generate a previous assignment action)

for the second part your comment, i managed to do it, but the compiler didn't like it (needed a little unsafe)

chriseth · 2025-01-03T11:58:34Z

executor/src/witgen/jit/interpreter.rs

+        }
+    }
+
+    // evaluates an expression using the provided variable values


Suggested change

// evaluates an expression using the provided variable values

/// Evaluates an expression using the provided variable values

chriseth · 2025-01-03T11:59:11Z

executor/src/witgen/jit/interpreter.rs

+
+    // evaluates an expression using the provided variable values
+    fn evaluate_expression(
+        &self,


self is unused as far as I can see. can this be a function on RPNExpresssion?

chriseth · 2025-01-03T12:01:51Z

executor/src/witgen/jit/interpreter.rs

+                        BinaryOperator::Mul => left * right,
+                        BinaryOperator::Div => left / right,
+                        BinaryOperator::IntegerDiv => T::from(
+                            left.to_integer().try_into_u64().unwrap()


I think we have to turn this into arbitrary integer. Or (maybe better) we require division on FieldElement::Integer.

I'm not sure where's the right place to put the Div bound on FieldElement::Integer, if that's what you mean

chriseth · 2025-01-03T12:05:36Z

executor/src/witgen/jit/symbolic_expression.rs

 };

 use num_traits::Zero;
 use powdr_number::FieldElement;

 use crate::witgen::range_constraints::RangeConstraint;

+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RPNExpressionElem<T: FieldElement, S> {


Can we move this to the interpreter?

chriseth · 2025-01-03T12:06:18Z

executor/src/witgen/jit/symbolic_expression.rs

 /// A value that is known at run-time, defined through a complex expression
 /// involving known cells or variables and compile-time constants.
 /// Each of the sub-expressions can have its own range constraint.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]


PartialEq is a recursive operation - where do you need it? Maybe it can be done based on pointers?

Although maybe the PartialEq operation of Arc shortcuts if the pointers are the same...

mm I don't recall why I added it, ill see if its needed

chriseth · 2025-01-09T09:59:42Z

executor/src/witgen/jit/interpreter.rs

+            InterpreterAction::ReadCell(idx, _) => {
+                set.insert(*idx);
+            }
+            InterpreterAction::ReadParam(idx, _) => {
+                set.insert(*idx);
+            }
+            InterpreterAction::AssignExpression(idx, _) => {
+                set.insert(*idx);
+            }


Suggested change

InterpreterAction::ReadCell(idx, _) => {

set.insert(*idx);

}

InterpreterAction::ReadParam(idx, _) => {

set.insert(*idx);

}

InterpreterAction::AssignExpression(idx, _) => {

set.insert(*idx);

}

InterpreterAction::ReadCell(idx, _) | InterpreterAction::ReadParam(idx, _) | InterpreterAction::AssignExpression(idx, _) => {

set.insert(*idx);

}

github-actions

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'Benchmarks'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 1.20.

Benchmark suite	Current: `6ff230c`	Previous: `0b25a80`	Ratio
`jit-benchmark/sqrt_879882356`	`18085` ns/iter (`± 39`)	`2603` ns/iter (`± 1`)	`6.95`
`jit-benchmark/sqrt_1882356`	`13943` ns/iter (`± 28`)	`2086` ns/iter (`± 1`)	`6.68`
`jit-benchmark/sqrt_1187956`	`13613` ns/iter (`± 17`)	`2063` ns/iter (`± 1`)	`6.60`
`jit-benchmark/sqrt_56`	`7049` ns/iter (`± 6`)	`1229` ns/iter (`± 2`)	`5.74`
`jit-benchmark/sort_33`	`386485` ns/iter (`± 466`)	`71906` ns/iter (`± 99`)	`5.37`
`jit-benchmark/sort_100`	`1507571` ns/iter (`± 1204`)	`269179` ns/iter (`± 387`)	`5.60`
`jit-benchmark/sort_300`	`6501239` ns/iter (`± 4108`)	`1035405` ns/iter (`± 1164`)	`6.28`
`jit-benchmark/sort_900`	`34087991` ns/iter (`± 26144`)	`4374283` ns/iter (`± 7843`)	`7.79`
`jit-benchmark/sort_2700`	`222377682` ns/iter (`± 104891`)	`20974189` ns/iter (`± 70438`)	`10.60`

This comment was automatically generated by workflow using github-action-benchmark.

chriseth · 2025-01-09T15:03:19Z

executor/src/witgen/jit/compiler.rs

    log::trace!("Calling cargo...");
    let r = powdr_jit_compiler::call_cargo(&code);
-    log::trace!("Done compiling, took {:.2}s", start.elapsed().as_secs_f32());


Why did you remove this?

oh, must've been an oversight

chriseth · 2025-01-09T15:04:02Z

executor/src/witgen/jit/compiler.rs

@@ -227,7 +225,7 @@ 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>(
+pub fn written_vars_in_effect<T: FieldElement>(


If this is going to be pub, then maybe turn it into a function of Effect?

chriseth · 2025-01-09T15:05:39Z

executor/src/witgen/jit/interpreter.rs

+use std::collections::{BTreeSet, HashMap};
+
+/// Witgen effects compiled into interpreter instructions.
+pub struct EffectsInterpreter<T: FieldElement> {


InterpreterEffects? InterpreterCode?

I am actually thinking about splitting up Effects into Effects (used at inference time) and Code (used for the final code), since they do not really overlap that much any more. Just as a side note...

ah sorry, this is the actual interpreter. all good :)

chriseth · 2025-01-09T15:12:34Z

executor/src/witgen/jit/interpreter.rs

+        params: &mut [LookupCell<T>],
+        data: CompactDataRef<'_, T>,
+    ) {
+        let mut vars = vec![T::zero(); self.var_count];


We could actually keep that one in the state - but can be done in another PR. There is also no need to reset it to zero between calls, since we checked that nothing is read before it is written.

yep, i've done exactly this in the other branch I have

chriseth · 2025-01-09T15:13:59Z

executor/src/witgen/jit/interpreter.rs

+                    vars[*idx] = val;
+                }
+                InterpreterAction::ReadCell(idx, c) => {
+                    let cell_offset: usize = c.row_offset.try_into().unwrap();


I think it would be better to convert data.row_offset to i32 instead, do the addition in i32 and then convert back to usize.

chriseth · 2025-01-09T15:14:43Z

executor/src/witgen/jit/interpreter.rs

+                    vars[*idx] = get_param(params, *i);
+                }
+                InterpreterAction::WriteCell(idx, c) => {
+                    let cell_offset: usize = c.row_offset.try_into().unwrap();


chriseth · 2025-01-09T15:25:57Z

Only some cosmetic comments left. It looks great!
Could you add some tests to the interpreter? Something like running the block machine processor on "../test_data/pil/binary.pil" and executing one block?

chriseth · 2025-01-09T15:26:19Z

Or poseidon for that matter.

chriseth · 2025-01-09T15:26:41Z

although it would be nice to have a submachine call in it.

chriseth · 2025-01-10T12:38:07Z

executor/src/witgen/jit/effect.rs

@@ -60,6 +60,27 @@ pub struct Assertion<T: FieldElement, V> {
    pub expected_equal: bool,
 }

+impl<T: FieldElement> Effect<T, Variable> {


This way it is between struct Assertion and its impl. Can you move it up to effect?

chriseth · 2025-01-10T12:43:29Z

executor/src/witgen/jit/interpreter.rs

+        data.append_new_rows(31);
+        let data_ref = CompactDataRef::new(&mut data, 0);
+        interpreter.call(&mutable_state, &mut param_lookups, data_ref);
+    }


Can you add an assertion about the outputs?

pacheco added 5 commits January 2, 2025 11:10

POWDR_JIT_DISABLE flag

b5af9e2

jit interpreter

49dd47f

comment

78d4743

bring todo back

a3a0cdf

move interpreter to own file

db4aea4

chriseth reviewed Jan 2, 2025

View reviewed changes

Schaeff reviewed Jan 3, 2025

View reviewed changes

executor/src/witgen/jit/interpreter.rs Outdated Show resolved Hide resolved

Schaeff reviewed Jan 3, 2025

View reviewed changes

executor/src/witgen/jit/interpreter.rs Outdated Show resolved Hide resolved

Schaeff reviewed Jan 3, 2025

View reviewed changes

executor/src/witgen/jit/interpreter.rs Outdated Show resolved Hide resolved

Schaeff reviewed Jan 3, 2025

View reviewed changes

executor/src/witgen/jit/interpreter.rs Outdated Show resolved Hide resolved

Schaeff reviewed Jan 3, 2025

View reviewed changes

executor/src/witgen/jit/interpreter.rs Outdated Show resolved Hide resolved

Schaeff reviewed Jan 3, 2025

View reviewed changes

executor/src/witgen/jit/interpreter.rs Outdated Show resolved Hide resolved

Schaeff reviewed Jan 3, 2025

View reviewed changes

chriseth reviewed Jan 3, 2025

View reviewed changes

pacheco added 3 commits January 3, 2025 10:59

VariableMapper and moving RPN stuff to interpreter

1cb10fb

move evaluate to RPNExpression

14cc011

extract some functions

f9af6e0

chriseth reviewed Jan 9, 2025

View reviewed changes

pacheco added 7 commits January 9, 2025 09:00

make members of CompactDataRef public

597faf1

comments

7e0245f

avoid expression cloning when converting to RPN

ef800aa

comment

79330d6

comment

f5742be

collapse match arms

82cf204

remove interpreter integration

6ff230c

github-actions bot reviewed Jan 9, 2025

View reviewed changes

chriseth reviewed Jan 9, 2025

View reviewed changes

pacheco added 4 commits January 9, 2025 13:34

bring back code removed by mistake

6cc5c32

Effect::written_vars

762a723

test

504306f

Merge remote-tracking branch 'origin/main' into jit-interpreter

cb6b06e

chriseth reviewed Jan 10, 2025

View reviewed changes

review

d2dd90b

chriseth approved these changes Jan 13, 2025

View reviewed changes

chriseth added this pull request to the merge queue Jan 13, 2025

Merged via the queue into main with commit 3eba5c4 Jan 13, 2025
16 checks passed

chriseth deleted the jit-interpreter branch January 13, 2025 15:47

	// evaluates an expression using the provided variable values
	/// Evaluates an expression using the provided variable values

Interpreter for witgen effects #2301

Interpreter for witgen effects #2301

Conversation

pacheco commented Jan 2, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Schaeff left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chriseth commented Jan 3, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chriseth Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

pacheco Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

⚠️ Performance Alert ⚠️

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chriseth commented Jan 9, 2025

chriseth commented Jan 9, 2025

chriseth commented Jan 9, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chriseth Jan 3, 2025 •

edited

Loading

pacheco Jan 3, 2025 •

edited

Loading