UnusedYield rewrite rule is too aggressive

[johnzl-777]

Given the following program:

@structural(fold=False, typeinfer=True)
def mwe():

    result = 0
    start = ilist.IList([0])
    stop = ilist.IList([1])
    for _ in range(10):
        result = start[0] + stop[0]
        start = stop
        stop = ilist.IList([result])

    return result

If you run it as-is you get the correct output of 89 but this falls apart if you apply: rewrite.Walk(scf.trim.UnusedYield()).rewrite(mwe.code)

Which now gives the incorrect output of 1.

Looking at the SSA output, it appears that the necessary things to be yielded are completely dropped, which leaves the following:

func.func @mwe() -> Literal(89,int) {
  ^0(%mwe_self):
  │   %result = py.constant.constant 0 : !py.int
  │        %0 = py.constant.constant 0 : !py.int
  │    %start = py.ilist.new(values=(%0)){elem_type=!py.int} : !py.IList[!py.int, Literal(1,int)]
  │        %1 = py.constant.constant 1 : !py.int
  │     %stop = py.ilist.new(values=(%1)){elem_type=!py.int} : !py.IList[!py.int, Literal(1,int)]
  │        %2 = py.constant.constant 0 : !py.int
  │        %3 = py.constant.constant 10 : !py.int
  │        %4 = py.constant.constant 1 : !py.int
  │        %5 = py.ilist.range(start=%2, stop=%3, step=%4) : !py.IList[!py.int, Literal(10,int)]
  │ %result_1 = %_ : !py.int in %5 -> !py.int
  │             │ iter_args(%result_2 : !py.int = %result) {
  │             │        %6 = py.constant.constant 0 : !py.int
  │             │        %7 = py.indexing.getitem(%start, %6) : !py.int
  │             │        %8 = py.constant.constant 0 : !py.int
  │             │        %9 = py.indexing.getitem(%stop, %8) : !py.int
  │             │ %result_3 = py.binop.add(%7, %9) : !py.int
  │             │  %start_1 = py.assign.alias start = %stop
  │             │   %stop_1 = py.ilist.new(values=(%result_3)){elem_type=~T} : !py.IList[!py.int, Literal(1,int)]
  │             │             scf.yield %result_3
  │             } -> purity=True
  │             func.return %result_1

@weinbe58 first identified this while helping me debug some problems with my attempts to support stim's REPEAT from the squin dialect in bloqade-circuit. This isn't blocking that work (the workaround for now is to not use the trim UnusedYield rewrite rule but I don't know if that bodes well for the IfToStim pass).

[weinbe58]

OK so I think I understand the issue with TrimUnused the problem is that it assumes that if the result of the For statement is not being used by later statements in the program, than you safely remove that result and also remove that same equivilant argument from inside the body of the loop.

In the MWE this is clearly not the case because you need the values of start and stop to remain a part of the loop variables even though they are not being used after the loop because these variables are mutating during the loop.

[weinbe58]

The current rewrite works only for IfElse because in that statement the body for the then and else blocks are only executed once when executing the IfElse statement.

What we need to do is to split the logic to handle IfElse and For in slightly different ways. I think for the For loop you need to also check to make sure the block arguments are not being used also inside the body of the loop as well and if it is being used then you need to keep that argument even if the result value isn't used outside of the loop.