Fold pass removes if / elif when it shouldn't

[david-pl]

MWE:

from kirin.prelude import structural_no_opt
from kirin.passes import Fold

@structural_no_opt
def main(n: int):
    x = 0

    if x == n:
        x += 1
    elif x == n+1:
        x += 3
    
    return x

main.print()

Fold(main.dialects)(main)

main.print()

Output:

func.func main(!py.int) -> !Any {
  ^0(%main_self, %n):
  │ %x = py.constant.constant 0 : !py.int
  │      func.return %x
} // func.func main

Things work correctly if you replace the elif by another if, but it's not technically identical since x is updated in-place which might lead to the second if to be True..

[david-pl]

Oh, I think the issue is actually not with the Fold pass. This is working as intended, kind of.

The issue is that as soon as there's an elif, the SSA value doesn't get updated properly anymore.

@structural_no_opt
def main(n: int):
    x = 0
    if x == n:
        x = 1
    elif x == n+1:
        x = 3
    return x

main.print()

Output:

func.func @main(n : !py.int) -> !Any {
  ^0(%main_self, %n):
  │ %x = py.constant.constant 0 : !py.int
  │ %0 = py.cmp.eq(lhs=%x, rhs=%n : !py.int) : !py.bool
  │      scf.if %0 {
  │      │ ^1(%1):
  │      │ │ %x_1 = py.constant.constant 1 : !py.int
  │      │ │        scf.yield 
  │      } else {
  │      │ ^2(%5):
  │      │ │ %2 = py.constant.constant 1 : !py.int
  │      │ │ %3 = py.binop.add(%n : !py.int, %2) : ~T
  │      │ │ %4 = py.cmp.eq(lhs=%x, rhs=%3) : !py.bool
  │      │ │      scf.if %4 {
  │      │ │      │ ^3(%6):
  │      │ │      │ │ %x_2 = py.constant.constant 3 : !py.int
  │      │ │      │ │        scf.yield 
  │      │ │      } -> purity=False
  │      │ │      scf.yield 
  │      } -> purity=False
  │      func.return %x
} // func.func main

[david-pl]

Slowly digging further into this: the bug is actually triggered whenever there is an assignment to a variable from the outer scope in only one of the branches. Consider, for example

@ir.dialect_group(python_basic.union([func, scf, lowering.func]))
def kernel(self):
    def run_pass(method):
        pass

    return run_pass

@kernel
def main(n: int):
    x = 0

    if x == n:
        x = 1
    else:# x == n + 1:
        y = 2

    return x

main.print()

The output again shows the missing assignment.

As for the elif: it's parsed as an else with a nested if and the assignment in the nested if isn't picked up.

[david-pl]

The issue is, I think, right here:

kirin/src/kirin/dialects/scf/lowering.py

Line 33 in 1e2b257

for name in body_frame.defs.keys():

When there's an else body, we loop over all definitions in the then body and check whether it's also present in the else. So, if there's a then body that redefines a variable from the outer scope, but it's not redefined in the else body, we miss it.