search

cornell-zhang / hcl-dialect

HeteroCL-MLIR dialect for accelerator design
https://cornell-zhang.github.io/heterocl/index.html
Other
38 stars 17 forks source link

[API][Pass] Support `.reshape()` primitive #46

Closed chhzh123 closed 2 years ago

chhzh123 commented 2 years ago

The .reshape() primitive operates on the schedule, and should directly generate new memref for tensors.

def test_index_split_reshape():
    hcl.init()
    A = hcl.placeholder((10, 10), "A")
    B = hcl.compute(A.shape, lambda y, x: A[y][x], "B")
    s = hcl.create_schedule([A, B])
    s[B].split(B.axis[0], 5)
    s.reshape(B, (2, 5, 10))