chore: array_set_window_optimization

[asterite]

Description

Problem

Trying out a new optimization on top of #11659 to see if it solves the regression.

Summary

If this works (it removes the regression) I'll open a separate PR with just this new SSA pass so we could merge it before #11659.

Additional Context

I wrote this with the help of Claude. I reviewed the initial code but then I kept adding requirements so I need to review the final code, but I wanted to push this to see if it fixes the regression: it seems it doesn't. At one point I think it did but then I kept refining the logic so maybe then it regressed again as it can't be optimized as much as I thought.

User Documentation

Check one:

• No user documentation needed.
• Changes in docs/ included in this PR.
• [For Experimental Features] Changes in docs/ 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.

[aakoshh]

I looked at why the tests time out.

The ski_calculus seems easy to replicate, for example the test_identity took over 35s on my laptop and the next test_logic I had to abort.

cargo run -q -p nargo_cli -- test -Zenums test_identity --benchmark-codegen --exact

On master this shows:

Flattening (1) (step 33): 877 ms
Mem2Reg (4) (step 34): 807 ms
Inlining (2) (step 35): 992 ms
ArraySet optimization (5) (step 36): 447 ms
ArrayGet optimization (6) (step 37): 538 ms

while on this branch:

Flattening (1) (step 33): 987 ms
ArraySet Window optimization (1) (step 34): 5570 ms
Mem2Reg (4) (step 35): 10208 ms
Inlining (2) (step 36): 5229 ms
ArraySet optimization (5) (step 37): 1965 ms
ArrayGet optimization (6) (step 38): 565 ms

I'm guessing that the way this pass replaces 1 array set with N-1 array_get and a make_array makes a lot of work for mem2reg. The arrays in this test have ~1000 elements.

The name_shadowing test didn't seem to have changed, it takes the same amount of time on both branches.

[aakoshh]

I tried it this way and then the test stays fast, but at the same time all replaces_* unit tests fail, because none of them work on eligible make_array instructions.

self.simple_optimization(|context| {
            let inst_id = context.instruction_id;
            if !candidates.contains(&inst_id) {
                return;
            }

            let Instruction::ArraySet { index, value, .. } = *context.instruction() else {
                unreachable!("candidate must be an ArraySet instruction");
            };

            let Some(const_index) =
                context.dfg.get_numeric_constant(index).and_then(|v| v.try_to_u32())
            else {
                unreachable!("candidate ArraySet index must be a constant u32");
            };

            // Only handle cases where we know the elements at compile time. Otherwise we would have to insert `array_get`
            // for every other element to re-pack them as a a `make_array`, which could blow up the SSA, making subsequent
            // passes work exponentially harder.
            let Some((mut elements, array_type)) = context.dfg.get_array_constant(value) else {
                return;
            };

            // Only replace safe index.
            if elements.len() <= const_index as usize {
                return;
            }

            // Remove the array_set; we will emit replacement instructions instead.
            context.remove_current_instruction();

            elements[const_index as usize] = value;
            let make_array = Instruction::MakeArray { elements, typ: array_type.clone() };
            let new_result = context.insert_instruction(make_array, Some(vec![array_type])).first();

            let [old_result] = context.dfg.instruction_result(inst_id);
            context.replace_value(old_result, new_result);
        });

[asterite]

@aakoshh Thank you for investigating this! I changed the optimization to run on arrays up to 64 elements. Maybe with 16 it would work just as well... maybe it's only needed for this poseidon case, not sure, in which case I think we'd only need to look at arrays of 4 elements? We could even only apply this optimization on chains of array_set + poseidon but I don't know if it would be too specific then and for other programs it would regress.

Let's see if the circuits are still optimized and CI passes now...

[asterite]

That seems to have worked.

[aakoshh]

Nice, the benchmarks show that it does rein back the 3% opcode increase incurred on the rollup contracts in aztec-packages 🎉