if/else scoping not correct if body is empty (pass)
jcasas00 opened this issue · 8 comments
def test_if_else_scope():
hcl.init()
rshape = (1,)
def kernel():
f = hcl.scalar(0, "f", dtype='uint32')
xy = hcl.scalar(0x12340001, "foo", dtype='uint32')
with hcl.if_(xy.v == 0):
hcl.print((),"0")
with hcl.if_(f.v == 0):
pass
with hcl.else_(): # should be scoped as the else clause of if_(xy.y == 0)
hcl.print((),"1")
#
s = hcl.create_schedule([], kernel)
print(hcl.lower(s))
generates the following IR:
module {
func.func @top() attributes {itypes = "", otypes = ""} {
%0 = memref.alloc() {name = "f", unsigned} : memref<1xi32>
%c0_i32 = arith.constant {unsigned} 0 : i32
affine.store %c0_i32, %0[0] {to = "f", unsigned} : memref<1xi32>
%1 = memref.alloc() {name = "foo", unsigned} : memref<1xi32>
%c305397761_i32 = arith.constant {unsigned} 305397761 : i32
affine.store %c305397761_i32, %1[0] {to = "foo", unsigned} : memref<1xi32>
%2 = affine.load %1[0] {from = "foo", moved, unsigned} : memref<1xi32>
%c0_i32_0 = arith.constant {moved} 0 : i32
%3 = arith.cmpi eq, %2, %c0_i32_0 {moved} : i32
scf.if %3 {
%c0_i32_1 = arith.constant 0 : i32
hcl.print(%c0_i32_1) {format = "0", signedness = "_"} : i32
%4 = affine.load %0[0] {from = "f", moved, unsigned} : memref<1xi32>
%5 = arith.cmpi eq, %4, %c0_i32_0 {moved} : i32
scf.if %5 {
} else {
hcl.print(%c0_i32_1) {format = "1", signedness = "_"} : i32
}
}
return
}
}
The hcl.else_() body got scoped with the hcl.if_(f.v==0) test instead of the hcl.if_(xy.v == 0) test. If instead of pass, there was an hcl statement, it gets scoped correctly.
This case came about because the code we have builds the hcl model based on python variables for configuration and one combination ended up with no hcl statements under the hcl.if_ test (i.e., just like using pass).
This looks like an issue with the insertion point. The insertion point should be reset when the program exits the with hcl.if_ context. However, it seems that the _exit_cb function is not executed when the body is empty.
I found that it is not caused by _exit_cb is not called, but Schedule._IfThenStack not popped. We use Schedule._IfThenStack to keep track of which if operation should the else branch be inserted to. The current implementation only pops Schedule._IfThenStack in hcl.else_. So, in the above case with hcl.if_(f.v == 0): is not popped from stack because it is not closed by any hcl.else_.
I added logic to pop the stack in _exit_cb function if an If op is not closed by hcl.else_: cornell-zhang/heterocl@494d39b
Still seeing an issue with if/else scoping (or at least looks like it).
def test_if_else_multiple():
hcl.init()
rshape = (1,)
stype0 = hcl.Struct({"x": hcl.UInt(8), "y": hcl.UInt(64-8)})
stype1 = stype0
stype2 = stype0
def kernel():
inst_enc = hcl.scalar(0, "inst_enc", dtype='uint64')
inst_id = hcl.scalar(0, "inst_id", dtype='uint16')
inst_enc_v = inst_enc.v
inst_lat = hcl.scalar(0, "inst_lat", dtype='uint16')
tmp = hcl.scalar(0, "tmp", dtype=hcl.UInt(1))
with hcl.if_(inst_lat.v <= 1):
tmp.v = 1
with hcl.if_(tmp.v != 0):
exe_inst = hcl.scalar(inst_enc_v, "exe_inst", dtype='uint64')
with hcl.if_(exe_inst.v == 0):
inste = hcl.scalar(exe_inst.v, f"inste0", dtype=stype0).v
with hcl.else_():
with hcl.if_(exe_inst.v == 1):
inste = hcl.scalar(exe_inst.v, f"inste1", dtype=stype1).v
r = hcl.compute(rshape, lambda _:0, dtype=hcl.Int(32))
return r
#
s = hcl.create_schedule([], kernel)
print(hcl.lower(s))
hcl_res = hcl.asarray(np.zeros(rshape, dtype=np.uint32), dtype=hcl.UInt(32))
f = hcl.build(s)
f(hcl_res)
np_res = hcl_res.asnumpy()
golden = np.zeros(rshape, dtype=np.int32)
assert np.array_equal(golden, np_res)
Generates the following error:
%24 = "affine.load"(%24) {from = "exe_inst", map = #map0, moved, unsigned} : (memref<1xi64>) -> i64
loc("-":60:7): error: definition of SSA value '%24#0' has type 'i64'
This test snippet was derived from a bigger program (took a while to isolate). Both the hcl.if_(inst_lat.v <= 1) test and 2nd level hcl.if_(exec_inst.v == 1) tests are needed to sensitize this issue.
Common to the error in the bigger program, the argument to affine.load is not correct and gets flagged (thankfully) as a type error.
Feel free to open another issue if this is a different issue.
Latest urgent issue as there's no temporary workaround
Added the following lines in dsl.py:if_ seems to fix the issue. Please check if it makes sense to you. If okay, something similar needs to be done for elseif_() .
def if_(cond):
"""Construct an IF branch."""
hcl_mlir.enable_build_inplace()
if isinstance(cond, hcl_mlir.ExprOp):
if_op = hcl_mlir.make_if(cond, ip=hcl_mlir.GlobalInsertionPoint.get())
else:
raise RuntimeError("Not implemented")
hcl_mlir.GlobalInsertionPoint.save(if_op.then_block.operations[0])
+ if len(Schedule._IfElseStack) > Schedule._CurrentIf:
+ Schedule._IfElseStack.pop()
Schedule._IfElseStack.append(if_op)
Schedule._CurrentIf += 1
+ assert len(Schedule._IfElseStack) == Schedule._CurrentIf
The issue is that for code like:
with hcl.if_(...):
...
with hcl.if_(...):
...
The second if_ effectively "closes" the first if and pop the stack (meant for an else clause of the first if_).
Hi Jeremy, sorry for the late reply, I am pretty occupied these few days by other things. You are right, the purpose of keep a stack of if operations is to find which if the next elif or else should be inserted to. When a if is followed by another if, the previous if can no longer have else or elif branches, so it should be popped out from the stack. I'll add this fix to the HeteroCL repo, thanks!
No problem. Removed urgent label since this if/else seems to work so far.
The frontend is revamped by cornell-zhang/heterocl#473, and if/elseif/else scope are no longer implemented with a stack. This issue is resolved and a test case is added by cornell-zhang/heterocl@0c8a1c0