loop_invariant_elimination: require the body to return the loop var itself

[LeSingh1]

loop_invariant_elimination decides that loop var i is invariant if either of these holds:

vx_out = block.outputs[i]
return_input_as_output = vx_in == vx_out
output_from_outside_of_block = enclosing_block.is_var_visible_in_block(vx_out, upto_op=while_op)
if return_input_as_output or output_from_outside_of_block:
    loop_invariant_ids.append(i)

The first test is sound. The second is not: Block.is_var_visible_in_block is true of any var in an enclosing scope — a function input, an _internal_var, or the output of any op before the while_loop. It says nothing about whether that var is while_op.loop_vars[i].

When the body returns a different outer var, the loop var is not invariant. It is the seed on the first iteration and the other var from the second on. The pass nevertheless drops the loop var and rewrites the while_loop output to the seed.

Repro

@mb.program(input_specs=[mb.TensorSpec(shape=(1, 2))] * 3)
def prog(a, b, c):
    def body(ax, bx):
        return mb.add(x=ax, y=np.float32([[1.0, 1.0]])), c   # not loop_vars[1], which is b
    def cond(ax, bx):
        return mb.less(x=mb.reduce_mean(x=ax, axes=[0, 1]), y=np.float32(3.0))
    return mb.while_loop(_cond=cond, _body=body, loop_vars=(a, b))

The loop runs three times, so output 1 should be c. Converting on main and predicting with a = 0, b = 10, c = 20:

{'while_loop_1_0': [[3.0, 3.0]], 'while_loop_1_1': [[10.0, 10.0]]}

i.e. b. The pass reduced the loop to one loop var and emitted %while_loop_0_1 = identity(x=%b). Nothing warns; op.enclosing_block.validate() at the end of the pass only checks that the program is structurally well formed.

The existing test_loop_invariant_elimination2 is titled "Invariant pattern: Block outputs var from outside of the block", and there the body does return an outer var — but it returns b, which is loop_vars[1]. That coincidence is the precondition the code never checks.

Fix

Test the body output against while_op.loop_vars[i] instead of against visibility. The existing "outputs a var from outside of the block" case still qualifies, because there the outer var is the loop var. loop_vars[i] is an operand of the while_loop, so it is visible in the enclosing block by construction and the visibility query adds nothing.

The pass also has to give such a body output a definition inside the body, because a block output has to be produced in its own block. Left alone, the program reaches the backend with -> (%add_0, %c) and Core ML segfaults while loading the compiled model (in MLModel._get_proxy_and_spec). Re-emitting the value with an identity inside the body block converts and runs correctly, returning [[20.0, 20.0]].

Also dropped a stale comment: block here is blocks[1], the body, whose outputs are the loop vars, not the cond var.

Testing

Two new tests in TestLoopInvariantElimination, both failing on main:

• test_loop_invariant_elimination_other_outer_var runs the pass and checks the loop var survives and the body output is now defined in the body.
• test_loop_invariant_elimination_other_outer_var_prediction runs the same program through run_compare_builder and checks output 1 is c. Skipped on the neuralnetwork backend, which has no while_loop.

The two existing tests are unchanged and still pass, as does the rest of test_cleanup_passes.py.