Insertion Point Issue for Imperative Code

[zzzDavid]

There are quite a few recent issues caused by insertion point management for imperative code, so I would like to open a thread to explain and discuss the problem.

HeteroCL code example

This is an example with an if and an else if branch, and two conditions.

def kernel():
    cond_a = hcl.scalar(0, "cond_a", dtype=hcl.UInt(16))
    cond_b = hcl.scalar(1, "cond_b", dtype=hcl.UInt(16))
    with hcl.if_(cond_a.v == 0):
        pass
    with hcl.elif_(cond_b.v == 0):
        pass
    return

We implement such else if branch as a nested if in the first if's else branch, so the IR should look like this:

  func.func @top()  {
    %0 = memref.alloc() {name = "cond_a", unsigned} : memref<1xi16>
    %c0 = arith.constant 0 : index
    %c0_i16 = arith.constant {unsigned} 0 : i16
    affine.store %c0_i16, %0[0] {to = "cond_a", unsigned} : memref<1xi16>
    %1 = memref.alloc() {name = "cond_b", unsigned} : memref<1xi16>
    %c0_0 = arith.constant 0 : index
    %c1_i16 = arith.constant {unsigned} 1 : i16
    affine.store %c1_i16, %1[0] {to = "cond_b", unsigned} : memref<1xi16>
    // evaluating first if's condition
    %2 = affine.load %0[0] {from = "cond_a", unsigned} : memref<1xi16>
    %c0_i32 = arith.constant 0 : i32
    %3 = arith.extui %2 : i16 to i32
    %4 = arith.cmpi eq, %3, %c0_i32 : i32
    scf.if %4 { // FIRST IF
    
    } else {
      // evaluating second if's condition
      %5 = affine.load %1[0] {from = "cond_b", unsigned} : memref<1xi16>
      %6 = arith.extui %5 : i16 to i32
      %c0_i32_1 = arith.constant 0 : i32
      %7 = arith.cmpi eq, %6, %c0_i32_1 : i32
      
      scf.if %7 { // SECOND IF, generated by elif
      }
    }
    return
  }

Note that the operations that evaluate the second if's condition should be inserted inside the else block.

However, because of the insertion point issue, we get IR like this:

  func.func @top()  {
    %0 = memref.alloc() {name = "cond_a", unsigned} : memref<1xi16>
    %c0 = arith.constant 0 : index
    %c0_i16 = arith.constant {unsigned} 0 : i16
    affine.store %c0_i16, %0[0] {to = "cond_a", unsigned} : memref<1xi16>
    %1 = memref.alloc() {name = "cond_b", unsigned} : memref<1xi16>
    %c0_0 = arith.constant 0 : index
    %c1_i16 = arith.constant {unsigned} 1 : i16
    affine.store %c1_i16, %1[0] {to = "cond_b", unsigned} : memref<1xi16>
    %2 = affine.load %0[0] {from = "cond_a", unsigned} : memref<1xi16>
    %c0_i32 = arith.constant 0 : i32
    %3 = arith.extui %2 : i16 to i32
    %4 = arith.cmpi eq, %3, %c0_i32 : i32
    scf.if %4 { // FIRST IF
     
    } else {      
      scf.if %7 { // SECOND IF, generated by elif
      }
    }
    %5 = affine.load %1[0] {from = "cond_b", unsigned} : memref<1xi16>
    %6 = arith.extui %5 : i16 to i32
    %c0_i32_1 = arith.constant 0 : i32
    %7 = arith.cmpi eq, %6, %c0_i32_1 : i32
    return
  }

Note that in this case, the condition operations are inserted after the scf.if operation. The condition of second if %7 is used before defined, which causes an error.

The insertion point issue

Concisely, this issue is caused by resetting the operation insertion point when the program exits with hcl.if_() scope, and the condition expression is evaluated before the program entering the with hcl.elif() scope. An insertion point is a position we insert operations in the IR, for example, at the beginning or end of a block, or after an operation. When we build a if operation with a with hcl.if_() scope, we set the insertion point to the then block of the scf.if operation. When we exit the with hcl.if_() scope, we need to reset the insertion point back to after the scf.if operation. Here is an example to illustrate this point:

with hcl.if(cond):
    hcl.print(...) # insertion point should be set inside `scf.if`'s `then` block
hcl.compute(...) # insertion point should be set after the `scf.if` operation, so these operations don't get inserted to `scf.if`

However, due to the evaluation order of Python, when we have a if + elif statement like this:

with hcl.if(cond1):
    pass
with hcl.elif(cond2):
    pass

cond2 is evaluated before the program jumps to with hcl.elif. At this point when cond2 is evaluated, the insertion point is set to "after scf.if operation", hence the wrong IR above.