Question about reshape of schedule

Questions
#1

If I wann do reshape + reorder, result will not meet expectation. Split can’t work In these demo while do reorder.

src_shape = (1, 16, 16000, 21)  # 16449
dst_shape = (1, 256000, 21)  # 263184 263169 256000
fp16 = "float16"
src_tensor = tvm.placeholder(src_shape, dtype=fp16, name='ph_1')
dst_tensor = topi.reshape(src_tensor, (1, 256000, 21))
# dst_tensor = tvm.compute(dst_shape, lambda *i: src_tensor[i[0], i[1] // 16000, i[1] % 16000, i[2]], name="dst")
sch = tvm.create_schedule(dst_tensor.op)

crt_tensor = sch.cache_read(src_tensor, "local.UB", [dst_tensor])

sch[crt_tensor].storage_align(crt_tensor.op.axis[1], 16, 0)
sch[crt_tensor].set_buffer_size(60000)

"""
dst_tensor shape:
(1, 256000, 21) -> (1, 16, 16000, 21)
                        |----| h_o, h_i
(1, 16, 16000, 21) -> (1, 16, 400, 40, 21)
                               |----| ub_outer, ub_inner
assume (1,16,40,21) in ub, 400 is outer of ub,
if i reorder (1, 16, 400, 40, 21) -> (400, 1, 16, 40, 21), split is not work in crt_tensor.
"""
h_o, h_i = sch[dst_tensor].split(dst_tensor.op.axis[1], nparts=16)
ub_outer, ub_inner = sch[dst_tensor].split(h_i, factor=40)

# _reorder = [dst_tensor.op.axis[0], ub_outer, h_o, ub_inner, dst_tensor.op.axis[-1]]
# sch[dst_tensor].reorder(*_reorder)

sch[crt_tensor].compute_at(sch[dst_tensor], ub_outer)
#2

Is there anyone meet the same problem or can help me figure out it. I would appreciate it very much~

#3

The root cause is InferBound get unexpected value.

1.The complex expression may infer wrong bound.

2.You split 2 axis, if you compute at “ub_outer”, the “ub_inner” range will not be updated which cause the parent’s range of “ho/hi” not be updated.

I hope that was helpful.