feat(ssa): Unroll small loops in brillig

[aakoshh]

Description

Problem*

Resolves #6470

Summary*

• Look at the runtime during unrolling to preserve ref count instructions, which ACIR did not need and are currently eliminated, but Brilling would need to keep using
• If the runtime is Brillig, enable unrolling the loops if the following check passes:
  • Identify the loop header, the back edge and the pre-header (exists in the module already)
  • Find the lower and upper bounds of the loop, which are expected to be turned into constants by inlining by the time we reach the Ssa::unroll_loops_iteratively pass.
  • Count all instructions in the loop
  • Detect the presence of jumps that look like break or continue that we couldn't unroll
  • If there are no blocks with terminators coming from break or continue instructions:
    • Collect all references allocated before the pre-header, for the detection of boilerplate load/stores in the loop body
    • Collect all load/store pairs of pre-header variables in the loop between header and back-edge, to estimate which instructions we can get rid of
    • Compute loop_instructions = all_instructions_in_loop - 2 * num_load_store_pairs - const_loop_jmp_overhead
    • Compare loop_instructions * num_iterations vs all_instructions_in_loop to find out if it's worth unrolling

I added some extra docstrings to existing functions to illustrate what the code does, and also refactored functions into methods where they already had an obvious parameter corresponding to a self; I found it made it easier to orient myself. Sorry about that code churn, here's a diff to make it easier to review, or we can break it up into two PRs.

Additional Context

Take for example this code:

fn main(sum: u32) {
    assert(loop(10, 14) == sum);
}

fn loop(from: u32, to: u32) -> u32 {
    let mut sum = 0;
    for i in from..to {
        sum = sum + i;
    }
    sum
}

Here's the SSA before and after unrolling.

After `static_assert` and `assert_constant`:
acir(inline) fn main f0 {
  b0(v0: u32):                         // Pre-header
    v2 = allocate -> u32
    store u32 0 at v2                  // This happens even if we do 4 additions in `main`
    jmp b1(u32 10)                     // -: The lower bound is a constant
  b1(v1: u32):                         // Loop header
    v6 = lt v1, u32 14                 // -: the upper bound is a constant
    jmpif v6 then: b3, else: b2        // -: loop one more time?
  b3():
    v9 = load v2 -> u32                // -: Loading v2 from the pre-header
    v10 = add v9, v1                   // +: sum = sum + 1
    store v10 at v2                    // -: Storing modified v2 for next loop
    v12 = add v1, u32 1                // -: Increment loop counter
    jmp b1(v12)                        // -: Back edge
  b2():
    v7 = load v2 -> u32
    v8 = eq v7, v0
    constrain v7 == v0
    return
}

After Unrolling:
acir(inline) fn main f0 {
  b0(v0: u32):
    v1 = allocate -> u32
    store u32 0 at v1
    v3 = load v1 -> u32
    v5 = add v3, u32 10
    store v5 at v1
    v6 = load v1 -> u32
    v8 = add v6, u32 11
    store v8 at v1
    v9 = load v1 -> u32
    v11 = add v9, u32 12
    store v11 at v1
    v12 = load v1 -> u32
    v14 = add v12, u32 13
    store v14 at v1
    jmp b1()
  b1():
    v15 = load v1 -> u32
    v16 = eq v15, v0
    constrain v15 == v0
    return
}

Lines marked // -: can be removed if we unroll the loops, and then running mem2reg gets rid of the redundant store and load operations:

After Mem2Reg (2nd):
acir(inline) fn main f0 {
  b0(v0: u32):
    v1 = allocate -> u32
    v3 = eq u32 46, v0
    constrain u32 46 == v0
    return
}

Documentation*

Check one:

• No documentation needed.
• Documentation included in this PR.
• [For Experimental Features] Documentation to be submitted in a separate PR.

PR Checklist*

• I have tested the changes locally.
• I have formatted the changes with Prettier and/or cargo fmt on default settings.

[github-actions]

Changes to Brillig bytecode sizes

Generated at commit: 0a3f9fff799e3d28957298939cf07a4022b613ba, compared to commit: 5f730e82f959cff9a06afea8a7bf89403850d687

🧾 Summary (10% most significant diffs)

Program	Brillig opcodes (+/-)	%
generics	-111 ✅	-53.37%
conditional_regression_547	-21 ✅	-53.85%
loop	-58 ✅	-56.86%
import	-88 ✅	-59.06%

---

Full diff report 👇

Program	Brillig opcodes (+/-)	%
keccak256	2,276 (+482)	+26.87%
poseidon_bn254_hash_width_3	5,673 (-17)	-0.30%
poseidon_bn254_hash	5,673 (-17)	-0.30%
sha256_regression	7,176 (-102)	-1.40%
6_array	425 (-8)	-1.85%
higher_order_functions	711 (-30)	-4.05%
sha256_var_size_regression	1,902 (-146)	-7.13%
array_dynamic_blackbox_input	1,078 (-103)	-8.72%
slices	2,081 (-237)	-10.22%
sha2_byte	3,083 (-368)	-10.66%
fold_complex_outputs	552 (-66)	-10.68%
hash_to_field	132 (-17)	-11.41%
brillig_rc_regression_6123	175 (-24)	-12.06%
strings	904 (-138)	-13.24%
slice_regex	2,179 (-398)	-15.44%
struct_inputs	256 (-48)	-15.79%
uhashmap	14,464 (-2,735)	-15.90%
array_to_slice	844 (-209)	-19.85%
simple_2d_array	110 (-30)	-21.43%
slice_loop	273 (-76)	-21.78%
pedersen_check	600 (-172)	-22.28%
brillig_pedersen	600 (-172)	-22.28%
hashmap	22,204 (-6,482)	-22.60%
brillig_constant_reference_regression	102 (-30)	-22.73%
references	204 (-110)	-35.03%
brillig_loop_size_regression	40 (-34)	-45.95%
reference_only_used_as_alias	171 (-146)	-46.06%
generics	97 (-111)	-53.37%
conditional_regression_547	18 (-21)	-53.85%
loop	44 (-58)	-56.86%
import	61 (-88)	-59.06%

[aakoshh]

Here's the complete example without unrolling:

brillig(inline) fn main f0 {
  b0(v0: Field, v1: [u8; 32]):
    v11 = truncate v0 to 8 bits, max_bit_size: 254
    v12 = cast v11 as u8
    inc_rc [u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0] of u8
    v15 = allocate -> &mut [u8; 136]
    // Empty array at v15
    store [u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0, u8 0] of u8 at v15
    v17 = call is_unconstrained() -> u1
    jmp b1(u32 0)                                 // loop from 0
  b1(v2: u32):                                    // Single iteration loop header
    v19 = eq v2, u32 0                            // loop to 0, ie. 0..len where len==1
    jmpif v19 then: b30, else: b2
  b30():                                          // loop body
    v119 = load v15 -> [u8; 136]                  // load the empty array
    v121 = array_get [v12] of u8, index v2 -> u8  // get the n-th byte from the input, in our case just the single byte
    v122 = array_set v119, index v2, value v121   // set the input byte in the empty array
    v123 = add v2, u32 1                          
    store v122 at v15                             // store the modified array
    v124 = add v2, u32 1                          // increment induction variable
    jmp b1(v124)                                  // back edge

This b1-to-b30 loop gets unrolled and then instead of working with array references it becomes full array constants.

[aakoshh]

FWIW I tried it like this:

        for (_, function) in ssa.functions.iter_mut() {
            let mut orig_function = Function::clone_with_id(function.id(), &function);
            ...
            if function.runtime().is_brillig() {
                function.mem2reg();
                let old_size = instruction_size(&orig_function);                
                let new_size = instruction_size(&function);
                if old_size < new_size {
                    std::mem::swap(function, &mut orig_function);
                }
            }
        }

I tried two approaches for instruction_size:

• Generate Brillig opcodes with crate::brillig::brillig_gen::convert_ssa_function: this doesn't work on this state of the SSA
• Just print the instructions and see how many characters it takes in total, as an approximation for data size

The latter shows the increase if the function.mem2reg(); call is present. We would of course implement it properly with recursion and take both the number of instructions as well as the size of the constants into account.

Do you think some a strategy like this would make sense to avoid increasing the byte code at all costs?

[jfecher]

Do you think some a strategy like this would make sense to avoid increasing the byte code at all costs?

I don't think it's a goal to avoid increasing byte code size at all costs but I see two directions we can go:

• Use heuristics to determine whether we should unroll
• Use this exact check plus maybe an optimization level multiplier to determine whether we should unroll.
  • Unrolling will always lead to fewer instructions executed since we guarantee we eliminate the increment & jump instructions, but may lead to a larger size. So this multiplier would be to adjust the balance between execution & size we want. E.g. we may say "unrolling is fine, up until it is 2x the size of the original loop, then we just want to keep the original loop"
  • Downside of this option of course is that it is significantly slower. Maybe this should be only enabled in a release mode of sorts

[aakoshh]

Thanks for suggesting the heuristics; do you think it can be in a follow up PR?

[TomAFrench]

I think we can have a follow-up PR in order to set the exact trade-off. Looking at the benchmarks we're only regressing in keccak256 atm by 25% which is higher than I'd like but manageable.