Closed jcasas00 closed 1 year ago
def test_if_elif_scope(): hcl.init() rshape = (2,) def kernel(): r = hcl.compute(rshape, lambda _:0, dtype=hcl.Int(32)) a = hcl.scalar(2, "a", dtype='uint32') b = hcl.scalar(1, "b", dtype='uint32') r[1] = 4 with hcl.if_(a.v == 0): r[0] = 1 with hcl.elif_(a.v == 1): r[0] = 2 with hcl.if_(b.v == 1): r[0] = 3 with hcl.else_(): # this else gets scoped with if_(b.v==1) r[0] = 4 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() assert np_res[0] == np_res[1]
generates the following IR:
%c1_i32_0 = arith.constant 1 : i32 %5 = arith.cmpi eq, %3, %c1_i32_0 : i32 scf.if %4 { affine.store %c1_i32_0, %0[0] {to = "compute_0"} : memref<2xi32> } else { scf.if %5 { %c2_i32_1 = arith.constant 2 : i32 affine.store %c2_i32_1, %0[0] {to = "compute_0"} : memref<2xi32> %6 = affine.load %2[0] {from = "b", moved, unsigned} : memref<1xi32> %c1_i32_2 = arith.constant {moved} 1 : i32 %7 = arith.cmpi eq, %6, %c1_i32_2 {moved} : i32 scf.if %7 { %c3_i32 = arith.constant 3 : i32 affine.store %c3_i32, %0[0] {to = "compute_0"} : memref<2xi32> } else { affine.store %c4_i32, %0[0] {to = "compute_0"} : memref<2xi32> } } }
Note that the final else block (that assigns %c4_i32 to %0) is scoped with scf.if %7 instead of scl.if %5.
Scope issues are fixed by https://github.com/cornell-zhang/heterocl/pull/473 Test case added by cornell-zhang/heterocl@4e3a5905549703a524b4f10e4acd6ade04e80e81
generates the following IR:
Note that the final else block (that assigns %c4_i32 to %0) is scoped with scf.if %7 instead of scl.if %5.