Take induction variables into account during loop invariant code motion

[vezenovm]

Problem

When we have a nested loop as such:

fn loop(upper_bound: u32, x: u32) {
    let arr = [2 as u32; 5];
    for i in 0..upper_bound {
        for _ in 0..upper_bound {
            assert_eq(arr[i], x);
        }
    }
}

We generate the following SSA:

brillig(inline) fn main f0 {
  b0(v0: u32, v1: u32):
    v5 = make_array [u32 2, u32 2, u32 2, u32 2, u32 2] : [u32; 5]
    jmp b1(u32 0)
  b1(v2: u32):
    v8 = lt v2, u32 4
    jmpif v8 then: b3, else: b2
  b3():
    jmp b4(u32 0)
  b4(v3: u32):
    v9 = lt v3, u32 4
    jmpif v9 then: b6, else: b5
  b6():
    v12 = array_get v5, index v2 -> u32
    constrain v12 == v0
    v13 = add v3, u32 1
    jmp b4(v13)
  b5():
    v11 = add v2, u32 1
    jmp b1(v11)
  b2():
    return
}

The instruction v12 = array_get v5, index v2 -> u32 is executed inside of the inner loop body in b6, when it can be executed in the preheader to the inner loop b3.

Happy Case

Looking at the example above, even though v2 comes from a block parameter, we know it is an induction variable on a loop with a constant bound. From this information, we can determine whether array_get is a safe index to move out of the loop.

We should be able to treat basic induction variables from parent loops as loop invariants when we are performing loop invariant code motion for inner loops.

Workaround

Yes

Workaround Description

Manually moving the loop invariant.

e.g. for the snippet above:

fn loop(upper_bound: u32, x: u32) {
    let arr = [2 as u32; 5];
    for i in 0..upper_bound {
        let arr_i = arr[i];
        for _ in 0..upper_bound {
            assert_eq(arr_i, x);
        }
    }
}

To better note the benefit the optimization provides:
Writing out the above snippet manually for an upper_bound of 4 actually reduced the number of opcodes executed from 332 to 104.

Additional Context

No response

Project Impact

None

Blocker Context

No response

Would you like to submit a PR for this Issue?

None

Support Needs

No response