Failed to prefetch on splitted axis

import tvm
from tvm import te

M = 64
K = 256
N = 256

dtype = "float32"

k = te.reduce_axis((0, K), name='k')
A = te.placeholder((M, K), dtype=dtype, name='A')
B = te.placeholder((N, K), dtype=dtype, name='B')
C = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[j, k], axis=k), name='C')

s = te.create_schedule(C.op)
y, x = s[C].op.axis
k = s[C].op.reduce_axis[0]
y2, y1 = s[C].split(y, factor=4)
y3, y2 = s[C].split(y2, factor=4)
x2, x1 = s[C].split(x, factor=8)
x3, x2 = s[C].split(x2, factor=8)
k2, k1 = s[C].split(k, factor=2)
k3, k2 = s[C].split(k2, factor=2)
s[C].reorder(y3, x3, k3, y2, x2, k2, y1, x1, k1)
s[C].prefetch(A, y1, 1)

target = "llvm"

lib = tvm.build(s, [A, B, C], target)

Error message:

/tvm/src/tir/transforms/storage_flatten.cc", line 1023
Check failed: e.buffer->shape.size() == op->bounds.size() (2 vs. 0) : Prefetch dim should be the same as buffer dim

hi @Mirage, may I know have you solved this issue, due to I met similar prefetch issue when I attempt to split the axis:

Check failed: e.buffer->shape.size() == op->bounds.size() (4 vs. 0) : Prefetch dim should be the same as buffer dim

@were, could you help with this prefetch issue? thanks~

hi @huangzhiyuan no I havent.