Single step with branching #2274

chriseth · 2024-12-20T12:36:40Z

Infers single-step update code including branching.

TODO:

perform constant propagation after branching (using the evaluator in only_concrete_known()-setting)
extend the "can process fully"-interface to allow answers like "yes, I am authorized to process, but it will never succeed given these range constraints". This way we can remove conflicting combinations of e.g. instruction flags. Another way would be to directly return range constraints on variables with the "can process" call, that way we could save branching. Maybe it's also fine to just implement this for the fixed machine. In any case, we should double-check that we do not create another machine call for an identity we already solved.

…ching

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: `b58afc1`	Previous: `c83472f`	Ratio
`evaluator-benchmark/std::math::ff::reduce`	`748` ns/iter (`± 1`)	`588` ns/iter (`± 1`)	`1.27`

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

Co-authored-by: Georg Wiese <[email protected]>

…r into single_step_simple

…ching

georgwiese

Wild!

I'm a bit worried generating code for every single combination of instruction flags. How would that not be infeasible for realistic examples, such as RISC-V? Also, I'm not sure if it is always possible to derive a unique witness if multiple instructions are active. Their constraints could be conflicting, right? In that case, code gen would fail?

It seems to me like we need more information which we can only get by running the PC lookup at compile time. We can figure out that the PC determines all other values in the PC lookup, and there is only a small number of possible combinations. These should be the cases we branch on, right? In FixedLookup::can_process_fully, we do have the list of all possible values.

executor/src/witgen/jit/single_step_processor.rs

georgwiese · 2025-01-03T23:39:01Z

executor/src/witgen/jit/single_step_processor.rs

-        let mut complete = HashSet::new();
-        for iteration in 0.. {
-            let mut progress = false;
+            // TODO propagate known.


Not 100% sure what this means! Is that about propagating concrete values?

georgwiese · 2025-01-03T23:53:46Z

executor/src/witgen/jit/single_step_processor.rs

+                .iter()
+                .filter_map(|var| witgen.range_constraint(var).map(|rc| (var, rc)))
+                .filter(|(_, rc)| rc.try_to_single_value().is_none())
+                .sorted()


I guess this has the effect that when in two range widths are equal we pick the column with the smaller ID. Is this so that the generated code is deterministic? I think this deserves a comment then :)

georgwiese · 2025-01-04T00:17:20Z

executor/src/witgen/jit/witgen_inference.rs

@@ -327,7 +363,7 @@ impl<'a, T: FieldElement, FixedEval: FixedEvaluator<T>> 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.
-    fn range_constraint(&self, variable: &Variable) -> Option<RangeConstraint<T>> {
+    pub fn range_constraint(&self, variable: &Variable) -> Option<RangeConstraint<T>> {


I think it would be better to instead have a public method that selects the variable to branch on. Then the processor does not need to care about range constraints at all.

We could use some information specific to the processor, for example if we know it's a single step processor in a VM, we can branch on the first variable in the longest lookup that is binary constrained (it will likely be an instruction flag).

georgwiese · 2025-01-04T00:21:01Z

executor/src/witgen/jit/single_step_processor.rs

+                .filter_map(|var| witgen.range_constraint(var).map(|rc| (var, rc)))
+                .filter(|(_, rc)| rc.try_to_single_value().is_none())
+                .sorted()
+                .min_by_key(|(_, rc)| rc.range_width())


Should we have some limit on the width? In practice, this will bisect until it's a concrete value, right? Which would take forever for very large ranges?

I would say we'll tune it later and see where we get with this.

georgwiese · 2025-01-04T14:15:44Z

executor/src/witgen/jit/single_step_processor.rs

+VM::B[1] = VM::B[0];
+if (VM::instr_add[0] == 1) {
+    if (VM::instr_mul[0] == 1) {
+        VM::A[1] = -((-(VM::A[0] + VM::B[0]) + -(VM::A[0] * VM::B[0])) + VM::A[0]);


Interesting. It doesn't know that instruction flags are exclusive. Does that mean that for $n$ instructions it generates $2^n$ branches? (For the Keccak RISC-V example, $n = 36$.)

georgwiese · 2025-01-04T14:21:13Z

executor/src/witgen/jit/witgen_inference.rs

 enum VariableOrValue<T, V> {
    Variable(V),
    Value(T),
 }

-pub trait FixedEvaluator<T: FieldElement> {
+pub trait FixedEvaluator<T: FieldElement>: Clone {


In practice, we just pass a reference, which implements Clone, right? I guess the advantage of this is that we can also pass something like Rc<_>, and it can automatically free the memory?

I guess it's the same for CanProcessCall below, which is not restricted to be Clone as well, but it is in practice. I feel like accepting references and not requiring Clone would reduce verbosity, but no strong opinion...

I'd like to avoid the lifetimes...

executor/src/witgen/jit/compiler.rs

georgwiese · 2025-01-04T14:40:42Z

executor/src/witgen/jit/compiler.rs

+            },
+            vec![assignment(&y, symbol(&x) + number(1))],
+            vec![assignment(&y, symbol(&x) + number(2))],
+        )];


Suggested change

)];

)];

let witgen_code = witgen_code(&[x.clone()], &effects);

assert_eq!(

witgen_code,

"

#[no_mangle]

extern \"C\" fn witgen(

WitgenFunctionParams{

data,

known,

row_offset,

params,

mutable_state,

call_machine

}: WitgenFunctionParams<FieldElement>,

) {

let known = known_to_slice(known, data.len);

let data = data.to_mut_slice();

let params = params.to_mut_slice();

let p_0 = get_param(params, 0);

let p_1;

if 7 <= IntType::from(p_0) && IntType::from(p_0) <= 20 {

p_1 = (p_0 + FieldElement::from(1));} else {

p_1 = (p_0 + FieldElement::from(2));}

set_param(params, 1, p_1);

set_param(params, 1, p_1);

}

"

);

At least for me, it helps to see a concrete example in the tests.

I'll instead add unit tests for the branch code itself. I don't really like these long auto-generated code tests because they need a lot of changes if we decide to do something in a different order (i.e. they don't really test the core that needs to be tested).

executor/src/witgen/jit/single_step_processor.rs

georgwiese · 2025-01-04T15:03:59Z

BTW, we'll need branching in block machines too, right? For stuff like operation_id * (a + b - c) + (1 - operation_id) * (a * b - c). So hopefully, we can pull some stuff out of the SingleStepProcessor in the future.

georgwiese · 2025-01-04T16:00:00Z

Ah, now I read & (somewhat) understood the issue description ^^

Another way would be to directly return range constraints on variables with the "can process" call, that way we could save branching.

I think this is the way to go, because the number of combinations is exponential and I don't think we want to ask the sub-machine that many times if it can succeed.

So with this solution what would happen in the example from the test is:

Branch on instr_add:
- instr_add == 0: Call can_process_call_fully, learn that instr_mul must be 1 => uniquely determines the full witness.
- instr_add == 1: Call can_process_call_fully, learn that instr_mul must be 0 => uniquely determines the full witness.

Sounds nice!

Also, I think returning range constraints might be a good idea anyway, because it would give us a "transfer of range constraints" ([a, b] in [c, d] => any range constraint on d transfers to b) also in cases where the constraint is not global (because e.g. it depends on the operation ID).

chriseth · 2025-01-06T21:03:48Z

Indeed, and it is already implemented in this PR: #2300

Co-authored-by: Georg Wiese <[email protected]>

…abs/powdr into single_step_with_branching

Co-authored-by: Georg Wiese <[email protected]>

…abs/powdr into single_step_with_branching

georgwiese

Vamos!

georgwiese · 2025-01-07T12:30:59Z

executor/src/witgen/jit/witgen_inference.rs

@@ -62,6 +62,16 @@ impl<T: Display> Display for Value<T> {
    }
 }

+/// Return type of the `branch_on` method.
+pub struct BranchResult<'a, T: FieldElement, FixedEval> {


chriseth added 17 commits December 19, 2024 23:15

Bisect range constraints.

3710c19

fix mask.

200cfa9

remove commented-out code.

00f53da

Single step processor.

cce4375

Negative test.

c54c0dd

branching test.

71666c0

loop

96b3660

c

405bdc1

compute branch var

2bdbe81

Merge branch 'bisect' into single_step_with_branching

140e7d4

almost branch.

79dd351

Merge remote-tracking branch 'origin/main' into single_step_with_bran…

e5b9846

…ching

Add "can process"

e41976a

Fix test.

b346f62

Fix test again.

e271ef7

fix test.

0f8ac79

swap branches

b58afc1

github-actions bot reviewed Dec 20, 2024

View reviewed changes

chriseth and others added 7 commits January 2, 2025 12:06

Merge remote-tracking branch 'origin/main' into single_step_simple

a40a133

Apply suggestions from code review

eb245ac

Co-authored-by: Georg Wiese <[email protected]>

Merge branch 'single_step_simple' of ssh://github.com/powdr-labs/powd…

1530cae

…r into single_step_simple

adjust tests.

1588eb2

Remove test function.

cc70c19

Merge branch 'remove_test_function' into single_step_simple

b13df2d

Fix testing functions.

96d5506

chriseth changed the base branch from main to single_step_simple January 2, 2025 14:40

Merge remote-tracking branch 'origin/main' into single_step_with_bran…

b6a2ad9

…ching

chriseth changed the base branch from single_step_simple to main January 2, 2025 14:43

chriseth added 2 commits January 2, 2025 14:49

Merge branch 'single_step_simple' into single_step_with_branching

b799c35

Fix merge

819d794

chriseth added 4 commits January 2, 2025 15:23

Fix written vars.

250253c

Fix test expectation.

93373a3

Code generation for branches.

ce159f6

clippy

b45dd7f

Base automatically changed from single_step_simple to main January 3, 2025 12:03

chriseth added 4 commits January 3, 2025 15:11

Merge branch 'single_step_simple'

4fc4a2e

Merge commit '50fd4580b' into single_step_with_branching

3517df3

Merge branch 'main' into single_step_with_branching

8fab6ae

Fix formatting and tests.

cf9ea9f

chriseth marked this pull request as ready for review January 3, 2025 15:30

georgwiese reviewed Jan 4, 2025

View reviewed changes

chriseth and others added 9 commits January 6, 2025 22:05

Update executor/src/witgen/jit/single_step_processor.rs

352d63f

Co-authored-by: Georg Wiese <[email protected]>

Extend comment

232b0d8

Merge branch 'single_step_with_branching' of ssh://github.com/powdr-l…

671546f

…abs/powdr into single_step_with_branching

Return both copies in on_branch.

7d69841

Improve error message.

df85d7f

Apply suggestions from code review

0242bfe

Co-authored-by: Georg Wiese <[email protected]>

Update executor/src/witgen/jit/compiler.rs

0ab316e

Co-authored-by: Georg Wiese <[email protected]>

Test codegen.

747ff3b

Merge branch 'single_step_with_branching' of ssh://github.com/powdr-l…

c4cbc05

…abs/powdr into single_step_with_branching

georgwiese approved these changes Jan 7, 2025

View reviewed changes

georgwiese added this pull request to the merge queue Jan 7, 2025

Merged via the queue into main with commit dc4c5cd Jan 7, 2025
16 checks passed

georgwiese deleted the single_step_with_branching branch January 7, 2025 13:08

This was referenced Jan 21, 2025

Runtime-branchnig for jit-witgen #2237

Closed

Auto-generated prover functions #2066

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Single step with branching #2274

Single step with branching #2274

chriseth commented Dec 20, 2024 •

edited

Loading

github-actions bot left a comment

georgwiese left a comment •

edited

Loading

georgwiese Jan 3, 2025

georgwiese Jan 3, 2025

georgwiese Jan 4, 2025

chriseth Jan 6, 2025

georgwiese Jan 4, 2025

chriseth Jan 6, 2025

georgwiese Jan 4, 2025

georgwiese Jan 4, 2025

chriseth Jan 6, 2025

georgwiese Jan 4, 2025

georgwiese Jan 4, 2025

chriseth Jan 6, 2025

georgwiese commented Jan 4, 2025

georgwiese commented Jan 4, 2025

chriseth commented Jan 6, 2025

georgwiese left a comment

georgwiese Jan 7, 2025

-        )];
+        )];
+        let witgen_code = witgen_code(&[x.clone()], &effects);
+        assert_eq!(
+            witgen_code,
+            "
+#[no_mangle]
+extern \"C\" fn witgen(
+    WitgenFunctionParams{
+        data,
+        known,
+        row_offset,
+        params,
+        mutable_state,
+        call_machine
+    }: WitgenFunctionParams<FieldElement>,
+) {
+    let known = known_to_slice(known, data.len);
+    let data = data.to_mut_slice();
+    let params = params.to_mut_slice();
+    let p_0 = get_param(params, 0);
+    let p_1;
+    if 7 <= IntType::from(p_0) && IntType::from(p_0) <= 20 {
+        p_1 = (p_0 + FieldElement::from(1));} else {
+        p_1 = (p_0 + FieldElement::from(2));}
+    set_param(params, 1, p_1);
+    set_param(params, 1, p_1);
+}
+"
+        );

Single step with branching #2274

Single step with branching #2274

Conversation

chriseth commented Dec 20, 2024 • edited Loading

github-actions bot left a comment

Choose a reason for hiding this comment

⚠️ Performance Alert ⚠️

georgwiese left a comment • 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

georgwiese commented Jan 4, 2025

georgwiese commented Jan 4, 2025

chriseth commented Jan 6, 2025

georgwiese left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chriseth commented Dec 20, 2024 •

edited

Loading

georgwiese left a comment •

edited

Loading