import tvm

n = 1024

dtype = "float32"

A = tvm.placeholder((n, n), dtype=dtype, name='A')

k = tvm.reduce_axis((0, n), name='k')

B = tvm.compute((n,), lambda i: tvm.sum(A[i, k], axis=k), name='B')

s = tvm.create_schedule(B.op)

print(tvm.lower(s, [A, B], simple_mode=True))

print("---------cutting line---------")

BW = s.cache_write(B, "local")

print(tvm.lower(s, [A, B], simple_mode=True))