cornell-zhang/hcl-dialect

if/elif scoping not correct

Closed this issue · 1 comments

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 cornell-zhang/heterocl#473
Test case added by cornell-zhang/heterocl@4e3a590