关于reverse_compute_at后，无法控制剩下的ax0轴的问题

[AmosChenYQ]

在做第一次练习的时候，发现无论如何都无法控制最后一个ax轴的vectorize化。

@tvm.script.ir_module
class MyBmmRelu:
  @T.prim_func
  def bmm_relu(A: T.Buffer[(16, 128, 128), "int32"], B: T.Buffer[(16, 128, 128), "int32"], C: T.Buffer[(16, 128, 128), "int32"]):
    T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
    Y = T.alloc_buffer([16, 128, 128], dtype="int32")
    for n, i, j, k in T.grid(16, 128, 128, 128):
      with T.block("Y"):
        vn = T.axis.spatial(16, n)
        vi = T.axis.spatial(128, i)
        vj = T.axis.spatial(128, j)
        vk = T.axis.reduce(128, k)
        with T.init():
          Y[vn, vi, vj] = T.int64(0)
        Y[vn, vi, vj] = Y[vn, vi, vj] + A[vn, vi, vk] * B[vn, vk, vj]
    for n, i, j in T.grid(16, 128, 128):
      with T.block("C"):
        vn = T.axis.spatial(16, n)
        vi = T.axis.spatial(128, i)
        vj = T.axis.spatial(128, j)
        C[vn, vi, vj] = T.max(Y[vn, vi, vj], T.int32(0))
      

以下是我的transformation

sch = tvm.tir.Schedule(MyBmmRelu)
# Step 1. Get blocks
Y = sch.get_block("Y", func_name="bmm_relu")
C = sch.get_block("C", func_name="bmm_relu")
# Step 2. Get loops
update_n, update_i, update_j, update_k = sch.get_loops(Y)
sch.parallel(update_n)
update_j_0, update_j_1 = sch.split(update_j, factors=[16, 8])
update_k_0, update_k_1 = sch.split(update_k, factors=[32, 4])
c_n, c_i, c_j = sch.get_loops(C)
# **getting vectorize at here doesn't work**
c_j_0, c_j_1 = sch.split(c_j, factors=[16, 8])
sch.vectorize(c_j_1)

# Step 3. Organize the loops
sch.unroll(update_k_1)
sch.reorder(update_k_0, update_k_1, update_j_1)
# **trying to change reverse_compuate_at position, get error or the same unexpected results**
sch.reverse_compute_at(C, update_j_0)
# Step 4. decompose reduction
Y_init = sch.decompose_reduction(Y, update_k_0)
init_n, init_i, init_j_0, init_j_1 = sch.get_loops(Y_init)
sch.vectorize(init_j_1)
# Step 5. vectorize / parallel / unroll
# **nor does getting vectorize at here**
# c_j_0, c_j_1 = sch.split(c_j, factors=[16, 8])
# sch.vectorize(c_j_1)

IPython.display.Code(sch.mod.script(), language="python")

我查阅课上的文档以及看了下reverse_compute_at的typehint似乎也没有返回值，那么请问应该如何去控制截图中的ax0轴呢

[AmosChenYQ]

在英文论坛里找到了一种方法
mlc-ai/mlc-en#17

需要在reverse_compuate_at后重新拿一次切分后的轴