cornell-zhang/hcl-dialect

[Op][SetBitOp] All zero tensor after SetBitOp

Closed this issue · 0 comments

Test case:

def test_set_bit_tensor():

    hcl.init()

    def kernel(A, B):
        with hcl.for_(0, 10) as i:
            B[i][0] = A[i]

    A = hcl.placeholder((10,), dtype=hcl.Int(1))
    B = hcl.placeholder((10,))
    s = hcl.create_schedule([A, B], kernel)
    f = hcl.build(s)
    print(hcl.lower(s))

    np_A = np.random.randint(1, size=(10,))
    np_B = np.random.randint(10, size=(10,))
    golden = (np_B & 0b1110) | np_A
    hcl_A = hcl.asarray(np_A)
    hcl_B = hcl.asarray(np_B)

    f(hcl_A, hcl_B)

    ret = hcl_B.asnumpy()
    print(ret)
    print(golden)
    assert np.array_equal(golden, ret)

IR:

module {
  func @top(%arg0: memref<10xi1>, %arg1: memref<10xi32>) attributes {bit, extra_itypes = "ss", extra_otypes = "", llvm.emit_c_interface, top} {
    affine.for %arg2 = 0 to 10 {
      %0 = affine.load %arg0[%arg2] {from = "compute_0"} : memref<10xi1>
      %1 = affine.load %arg1[%arg2] {from = "compute_1"} : memref<10xi32>
      %c0 = arith.constant 0 : index
      hcl.set_bit(%1 : i32, %c0, %0 : i1)
      affine.store %1, %arg1[%arg2] {to = "compute_1"} : memref<10xi32>
    } {loop_name = "loop_0"}
    return
  }
}