Why we should clean build result for IfOp's condition

[zzzDavid]

This thread shows why we should clean the condition's building result for IfOp. This discussion is relevant to the frontend implementation with AST, where we build a HeteroCL AST during create_schedule, and built MLIR IR in a second pass with IRBuilder.

How do we know the MLIR operation is built for an AST node in the second pass? Each AST node (Operation, Expr) has these two members: .result, .ir_op. When an MLIR operation is built for that node, the ir_op is set to the MLIR operation object, and .result is set to the result of the MLIR operation, which is an OpResult object. So we can check whether .ir_op is None for any AST node to determine if its MLIR operation has been built, and access the OpResult object with .result.

I want to show that we should clean the previous build result of if operation's condition expression because we may run into dominance issue if we use previous build results. Here is an example:

HeteroCL Code

def test_if_struct_access():
    hcl.init()

    def kernel():
        # tag = hcl.scalar(0, "tag", dtype='uint32')
        res = hcl.compute((1,), lambda _: 0, name="res", dtype="uint32")
        stype = hcl.Struct({"x": hcl.UInt(8), "y": hcl.UInt(8)})
        s_xy = hcl.scalar(0x1234, "xy", dtype=stype)
        xy = s_xy.v
        x = xy.x
        y = xy.y
        # use x and y
        with hcl.if_(x == 1):
            res[0] = 1
        with hcl.elif_(y == 2):
            res[0] = 2
        # use get bit and get slice
        with hcl.if_(x[0] == 1):
            res[0] = 8
        with hcl.elif_(y[1] == 1):
            res[0] = 9
        return res

    s = hcl.create_schedule([], kernel)
    hcl.lower(s)

IR that didn't pass MLIR verifier

Full IR

// Verification failed, printing generic form
#map0 = affine_map<(d0) -> (d0)>
#map1 = affine_map<() -> (0)>
"builtin.module"() ({
  "func.func"() ({
    %0 = "arith.constant"() {value = 0 : index} : () -> index
    %1 = "arith.constant"() {value = 0 : index} : () -> index
    %2 = "memref.alloc"() {name = "res", operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<1xi32>
    %3 = "arith.constant"() {value = 0 : i32} : () -> i32
    "affine.store"(%3, %2, %1) {map = #map0, to = "res", unsigned} : (i32, memref<1xi32>, index) -> ()
    %4 = "memref.alloc"() {name = "xy", operand_segment_sizes = dense<0> : vector<2xi32>} : () -> memref<1x!hcl.struct<i8, i8>>
    %5 = "arith.constant"() {unsigned, value = 52 : i8} : () -> i8
    %6 = "arith.constant"() {unsigned, value = 18 : i8} : () -> i8
    %7 = "hcl.struct_construct"(%5, %6) : (i8, i8) -> !hcl.struct<i8, i8>
    "affine.store"(%7, %4, %0) {map = #map0, to = "xy"} : (!hcl.struct<i8, i8>, memref<1x!hcl.struct<i8, i8>>, index) -> ()
    %8 = "affine.load"(%4) {from = "xy", map = #map1} : (memref<1x!hcl.struct<i8, i8>>) -> !hcl.struct<i8, i8>
    %9 = "hcl.struct_get"(%8) {index = 0 : i64} : (!hcl.struct<i8, i8>) -> i8
    %10 = "arith.constant"() {value = 1 : i32} : () -> i32
    %11 = "arith.extui"(%9) : (i8) -> i32
    %12 = "arith.cmpi"(%11, %10) {predicate = 0 : i64} : (i32, i32) -> i1
    "scf.if"(%12) ({
      %18 = "arith.constant"() {value = 1 : i32} : () -> i32
      "affine.store"(%18, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
      "scf.yield"() : () -> ()
    }, {
      %18 = "hcl.struct_get"(%8) {index = 1 : i64} : (!hcl.struct<i8, i8>) -> i8
      %19 = "arith.constant"() {value = 2 : i32} : () -> i32
      %20 = "arith.extui"(%18) : (i8) -> i32
      %21 = "arith.cmpi"(%20, %19) {predicate = 0 : i64} : (i32, i32) -> i1
      "scf.if"(%21) ({
        %22 = "arith.constant"() {value = 2 : i32} : () -> i32
        "affine.store"(%22, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
        "scf.yield"() : () -> ()
      }, {
      }) : (i1) -> ()
      "scf.yield"() : () -> ()
    }) : (i1) -> ()
    %13 = "arith.constant"() {unsigned, value = 0 : index} : () -> index
    %14 = "hcl.get_bit"(%9, %13) : (i8, index) -> i1
    %15 = "arith.constant"() {value = 1 : i32} : () -> i32
    %16 = "arith.extui"(%14) : (i1) -> i32
    %17 = "arith.cmpi"(%16, %15) {predicate = 0 : i64} : (i32, i32) -> i1
    "scf.if"(%17) ({
      %18 = "arith.constant"() {value = 8 : i32} : () -> i32
      "affine.store"(%18, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
      "scf.yield"() : () -> ()
    }, {
      %18 = "arith.constant"() {unsigned, value = 1 : index} : () -> index
      %19 = "hcl.get_bit"(%18, %18) : (i8, index) -> i1
      %20 = "arith.constant"() {value = 1 : i32} : () -> i32
      %21 = "arith.extui"(%19) : (i1) -> i32
      %22 = "arith.cmpi"(%21, %20) {predicate = 0 : i64} : (i32, i32) -> i1
      "scf.if"(%22) ({
        %23 = "arith.constant"() {value = 9 : i32} : () -> i32
        "affine.store"(%23, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
        "scf.yield"() : () -> ()
      }, {
      }) : (i1) -> ()
      "scf.yield"() : () -> ()
    }) : (i1) -> ()
    "func.return"(%2) : (memref<1xi32>) -> ()
  }) {bit, function_type = () -> memref<1xi32>, itypes = "", otypes = "u", sym_name = "top"} : () -> ()
}) : () -> ()

The important bit:

   "scf.if"(%12) ({
      %18 = "arith.constant"() {value = 1 : i32} : () -> i32
      "affine.store"(%18, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
      "scf.yield"() : () -> ()
    }, {
      %18 = "hcl.struct_get"(%8) {index = 1 : i64} : (!hcl.struct<i8, i8>) -> i8
      %19 = "arith.constant"() {value = 2 : i32} : () -> i32
      %20 = "arith.extui"(%18) : (i8) -> i32
      %21 = "arith.cmpi"(%20, %19) {predicate = 0 : i64} : (i32, i32) -> i1
      "scf.if"(%21) ({
        %22 = "arith.constant"() {value = 2 : i32} : () -> i32
        "affine.store"(%22, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
        "scf.yield"() : () -> ()
      }, {
      }) : (i1) -> ()
      "scf.yield"() : () -> ()
    }) : (i1) -> ()
    %13 = "arith.constant"() {unsigned, value = 0 : index} : () -> index
    %14 = "hcl.get_bit"(%9, %13) : (i8, index) -> i1
    %15 = "arith.constant"() {value = 1 : i32} : () -> i32
    %16 = "arith.extui"(%14) : (i1) -> i32
    %17 = "arith.cmpi"(%16, %15) {predicate = 0 : i64} : (i32, i32) -> i1
    "scf.if"(%17) ({
      %18 = "arith.constant"() {value = 8 : i32} : () -> i32
      "affine.store"(%18, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
      "scf.yield"() : () -> ()
    }, {
      %18 = "arith.constant"() {unsigned, value = 1 : index} : () -> index
      %19 = "hcl.get_bit"(%18, %18) : (i8, index) -> i1  // PROBLEM: the first  %18 is defined in a block that does not dominate this usage
      %20 = "arith.constant"() {value = 1 : i32} : () -> i32
      %21 = "arith.extui"(%19) : (i1) -> i32
      %22 = "arith.cmpi"(%21, %20) {predicate = 0 : i64} : (i32, i32) -> i1
      "scf.if"(%22) ({
        %23 = "arith.constant"() {value = 9 : i32} : () -> i32
        "affine.store"(%23, %2) {map = #map1, to = "res", unsigned} : (i32, memref<1xi32>) -> ()
        "scf.yield"() : () -> ()
      }, {
      }) : (i1) -> ()
      "scf.yield"() : () -> ()
    }) : (i1) -> ()

Note that because %18 is defined in the first scf.if's body scope:

%18 = "hcl.struct_get"(%8) {index = 1 : i64} : (!hcl.struct<i8, i8>) -> i8

The use of its result is not dominated by its definition:

%19 = "hcl.get_bit"(%18, %18) : (i8, index)

This is caused by the build result of this condition with hcl.elif_(y == 2) being reused when this condition with hcl.elif_(y[1] == 1): is built.

[zzzDavid]

So we need to add an "ast cleaner" to remove the previous build result and build the condition again at the current scope.
This is how we build the IR for IfOp in IRBuilder at the time of this discussion:

    def build_if_op(self, op: ast.IfOp, ip):
        """Build IfOp"""
        # build condition
        # clear condition build result from previous build
        self.cleaner.visit(op.cond)
        self.build_visitor(op.cond, ip)
        has_else = op.else_branch_valid
        if_op = scf_d.IfOp(op.cond.result, hasElse=has_else, results_=[], ip=ip)
        # build then body
        ip = InsertionPoint(if_op.then_block)
        for body_op in op.body:
            self.build_visitor(body_op, ip)
        scf_d.YieldOp([], ip=ip)
        # build else body
        if has_else:
            ip = InsertionPoint(if_op.else_block)
            for body_op in op.else_body:
                self.build_visitor(body_op, ip)
            scf_d.YieldOp([], ip=ip)
        op.ir_op = i

[zzzDavid]

With the condition being freshly built at every scf.if's condition scope, the IR becomes:

module {
  func.func @top() -> memref<1xi32> attributes {bit, itypes = "", otypes = "u"} {
    %c0 = arith.constant 0 : index
    %0 = memref.alloc() {name = "res"} : memref<1xi32>
    %c0_i32 = arith.constant 0 : i32
    affine.store %c0_i32, %0[%c0] {to = "res", unsigned} : memref<1xi32>
    %1 = memref.alloc() {name = "xy"} : memref<1x!hcl.struct<i8, i8>>
    %c52_i8 = arith.constant {unsigned} 52 : i8
    %c18_i8 = arith.constant {unsigned} 18 : i8
    %2 = hcl.struct_construct(%c52_i8, %c18_i8) : i8, i8 -> <i8, i8>
    affine.store %2, %1[%c0] {to = "xy"} : memref<1x!hcl.struct<i8, i8>>
    %3 = affine.load %1[0] {from = "xy"} : memref<1x!hcl.struct<i8, i8>>
    %4 = hcl.struct_get %3[0] : <i8, i8> -> i8
    %c1_i32 = arith.constant 1 : i32
    %5 = arith.extui %4 : i8 to i32
    %6 = arith.cmpi eq, %5, %c1_i32 : i32
    scf.if %6 {
      affine.store %c1_i32, %0[0] {to = "res", unsigned} : memref<1xi32>
    } else {
      %12 = affine.load %1[0] {from = "xy"} : memref<1x!hcl.struct<i8, i8>>
      %13 = hcl.struct_get %12[1] : <i8, i8> -> i8
      %c2_i32 = arith.constant 2 : i32
      %14 = arith.extui %13 : i8 to i32
      %15 = arith.cmpi eq, %14, %c2_i32 : i32
      scf.if %15 {
        affine.store %c2_i32, %0[0] {to = "res", unsigned} : memref<1xi32>
      }
    }
    %7 = affine.load %1[0] {from = "xy"} : memref<1x!hcl.struct<i8, i8>>
    %8 = hcl.struct_get %7[0] : <i8, i8> -> i8
    %c0_0 = arith.constant {unsigned} 0 : index
    %9 = hcl.get_bit(%8 : i8, %c0_0) -> i1
    %10 = arith.extui %9 : i1 to i32
    %11 = arith.cmpi eq, %10, %c1_i32 : i32
    scf.if %11 {
      %c8_i32 = arith.constant 8 : i32
      affine.store %c8_i32, %0[0] {to = "res", unsigned} : memref<1xi32>
    } else {
      %12 = affine.load %1[0] {from = "xy"} : memref<1x!hcl.struct<i8, i8>>
      %13 = hcl.struct_get %12[1] : <i8, i8> -> i8
      %c1 = arith.constant {unsigned} 1 : index
      %14 = hcl.get_bit(%13 : i8, %c1) -> i1
      %15 = arith.extui %14 : i1 to i32
      %16 = arith.cmpi eq, %15, %c1_i32 : i32
      scf.if %16 {
        %c9_i32 = arith.constant 9 : i32
        affine.store %c9_i32, %0[0] {to = "res", unsigned} : memref<1xi32>
      }
    }
    return %0 : memref<1xi32>
  }
}

Note that there is a new struct_get operation before each scf.if. This means the struct field is freshly accessed from the struct when the scf.if condition evaluates, and therefore, no more "definition does not dominate usage" issue.