cornell-zhang / heterocl

HeteroCL: A Multi-Paradigm Programming Infrastructure for Software-Defined Heterogeneous Computing
https://cornell-zhang.github.io/heterocl/
Apache License 2.0
322 stars 92 forks source link

[Fix] Halide IR Simplify pass wrongly optimizes away if condition #456

Closed hecmay closed 2 years ago

hecmay commented 2 years ago

For fixing issues

Fixed issue: #450

Detailed description:

Minimal test case: in the top-level else branch, the var has been updated before going in another if statement. Namely, the second with hcl.if_(var == 0) is not always false, and should not be optimized away.

def two_stage(A):
    var = hcl.scalar(0, "v", dtype=hcl.UInt(32))
    var.v = 1
    with hcl.if_(var == 0):
        hcl.print((),"A\n")
    with hcl.else_():
        var.v = var - 1
        with hcl.if_(var == 0):
            hcl.print((),"B\n")
    return A

This is the buggy IR generated by HCL before the ifx.

// generated IR
produce _top {
  // attr [0] extern_scope = 0
  // attr [v] storage_scope = "global"
  allocate v[uint32 * 1]
  produce v {
    // attr [0] extern_scope = 0
    for "stage_name"="v" (x, 0, 1) {
      v[x] = (uint32)0
    }
  }
  v[0] = (uint32)1
  if ((v[0] == (uint32)0)) {
    print:
  } else {
    v[0] = (v[0] - (uint32)1)
  }
}

Here is the expected IR that should be generated after applying the fix.

// Expected IR
produce _top {                                                                                                                                                                         [76/1883]
  // attr [0] extern_scope = 0
  // attr [v] storage_scope = "global"
  allocate v[uint32 * 1]
  produce v {
    // attr [0] extern_scope = 0
    for "stage_name"="v" (x, 0, 1) {
      v[x] = (uint32)0
    }
  }
  v[0] = (uint32)1
  if ((v[0] == (uint32)0)) {
    print:
  } else {
    v[0] = (v[0] - (uint32)1)
    if ((v[0] == (uint32)0)) {
      print:
    }
  }
}

Link to the tests: tests/issues/test_issue_450.py