Unable to partition tensors inside a function

[chhzh123]

A simple example is shown below (modified from the tutorial),

def test():
    hcl.init()
    A = hcl.placeholder((10, 10), "A")
    def kernel(A):
        B = hcl.compute(A.shape, lambda x, y: A[x][y] + 1, "B")
        C = hcl.compute(A.shape, lambda x, y: B[x][y] + 1, "C") # add this line
        return C
    s = hcl.create_schedule(A, kernel)
    s.partition(kernel.B)
    print(hcl.lower(s))

which causes Runtime Error.

Traceback (most recent call last):
  File "partition.py", line 16, in <module>
    test()
  File "partition.py", line 13, in test
    print(hcl.lower(s))
  File "/home/chz/heterocl/python/heterocl/api.py", line 276, in lower
    return _lower(schedule.sch, new_inputs, simple_mode=True)
  File "/home/chz/heterocl/python/heterocl/tvm/build_module.py", line 349, in lower
    stmt = ir_pass.LiftAllocateAttrs(stmt)
  File "/home/chz/heterocl/python/heterocl/tvm/_ffi/function.py", line 280, in my_api_func
    return flocal(*args)
  File "/home/chz/heterocl/python/heterocl/tvm/_ffi/_ctypes/function.py", line 183, in __call__
    ctypes.byref(ret_val), ctypes.byref(ret_tcode)))
  File "/home/chz/heterocl/python/heterocl/tvm/_ffi/base.py", line 66, in check_call
    raise TVMError(py_str(_LIB.TVMGetLastError()))
heterocl.tvm._ffi.base.TVMError: [12:29:02] src/ir/IR.cpp:445: Check failed: first.defined() Block of undefined

It's able to run when partitioning the A array or the C array, thus only inner tensors cause the problem. I have also tried s.partition(kernel.B._op), but it cannot work either.

[seanlatias]

This should be fixed with #206, which we will merge soon. Thanks for filing the issue.

[hecmay]

The error is caused by the IR pass lift_allocate_attrs creating an invalid block statement. But the root cause occurs before that.

The bug has been fixed in #206. I am able to run your example. But it is not fully as expected... I can still see an allocate statement for the partitioned buffer:

// attr [_top] storage_scope = "global"
allocate _top[int32 * 1]
produce _top {
  // attr [0] extern_scope = 0
  // attr [B] storage_scope = "global"
  allocate B[int32 * 10 * 10]
  array partition variable=B complete factor=0 dim=0
  produce B {
    // attr [0] extern_scope = 0
    // attr [B.partitioned] storage_scope = "global"
    allocate B.partitioned[int32 * 1]
    for (x, 0, 10) {
      for (y, 0, 10) {
        B[(y + (x*10))] = (A[(y + (x*10))] + 1)
      }
    }
  }
  produce C {
    // attr [0] extern_scope = 0
    for (x, 0, 10) {
      for (y, 0, 10) {
        C[(y + (x*10))] = (B[(y + (x*10))] + 1)
      }
    }
  }
}

Here is the test case I added before: https://github.com/Hecmay/heterocl/blob/stream_to/tests/test_schedule_stream.py#L531

[chhzh123]

Yes, it's weird... I see many xxx_partitioned variables before each loop.

[hecmay]

I need to look into this pass: https://github.com/cornell-zhang/heterocl/blob/master/tvm/src/pass/lift_allocate_attrs.cc

It seems that some assumptions are not satisfied...