[Op][SetBitOp] All zero tensor after SetBitOp
Closed this issue · 0 comments
zzzDavid commented
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
}
}