关于课程三课后习题2.5.2.2

[cjx0709]

我的原始代码：

@tvm.script.ir_module
class MyBmmRelu:
    @T.prim_func
    def bmm_relu(A: T.Buffer[(16, 128, 128), "float32"], B: T.Buffer[(16, 128, 128), "float32"], C: T.Buffer[(16, 128, 128), "float32"]) -> None:
        T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
        Y = T.alloc_buffer([16, 128, 128], dtype="float32")
        for b, i, j, k in T.grid(16, 128, 128, 128):
            with T.block("Y"):
                vb, vi, vj, vk = T.axis.remap("SSSR", [b, i, j, k])
                with T.init():
                    Y[vb, vi, vj] = T.float32(0)
                Y[vb, vi, vj] = Y[vb, vi, vj] + A[vb, vi, vk] * B[vb, vk, vj]
        for b, i, j in T.grid(16, 128, 128):
            with T.block("C"):
                vb, vi, vj = T.axis.remap("SSS", [b, i, j])
                C[vb, vi, vj] = T.max(Y[vb, vi, vj], T.float32(0))

我的transform过程：

block_Y = sch.get_block("Y", func_name="bmm_relu")
block_C = sch.get_block("C", "bmm_relu")
b, i, j, k = sch.get_loops(block_Y)
i2_0, ax0 = sch.split(j, [16, None])
ax1_0, ax1_1 = sch.split(k, [None, 4])
sch.reorder(ax1_0, ax1_1, ax0)
sch.parallel(b)
sch.unroll(ax1_1)
sch.vectorize(ax0)
sch.reverse_compute_at(block_C, i2_0)
sch.decompose_reduction(block_Y, ax1_0)

我的print结果：

@tvm.script.ir_module
class Module:
    @T.prim_func
    def bmm_relu(A: T.Buffer[(16, 128, 128), "float32"], B: T.Buffer[(16, 128, 128), "float32"], C: T.Buffer[(16, 128, 128), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
        # body
        # with T.block("root")
        Y = T.alloc_buffer([16, 128, 128], dtype="float32")
        for b in T.parallel(16):
            for i, j_0 in T.grid(128, 16):
                for j_1_init in T.vectorized(8):
                    with T.block("Y_init"):
                        vb, vi = T.axis.remap("SS", [b, i])
                        vj = T.axis.spatial(128, j_0 * 8 + j_1_init)
                        T.reads()
                        T.writes(Y[vb, vi, vj])
                        Y[vb, vi, vj] = T.float32(0)
                for k_0 in T.serial(32):
                    for k_1 in T.unroll(4):
                        for j_1 in T.vectorized(8):
                            with T.block("Y_update"):
                                vb, vi = T.axis.remap("SS", [b, i])
                                vj = T.axis.spatial(128, j_0 * 8 + j_1)
                                vk = T.axis.reduce(128, k_0 * 4 + k_1)
                                T.reads(Y[vb, vi, vj], A[vb, vi, vk], B[vb, vk, vj])
                                T.writes(Y[vb, vi, vj])
                                Y[vb, vi, vj] = Y[vb, vi, vj] + A[vb, vi, vk] * B[vb, vk, vj]
                for ax0 in T.serial(8):
                    with T.block("C"):
                        vb, vi = T.axis.remap("SS", [b, i])
                        vj = T.axis.spatial(128, j_0 * 8 + ax0)
                        T.reads(Y[vb, vi, vj])
                        T.writes(C[vb, vi, vj])
                        C[vb, vi, vj] = T.max(Y[vb, vi, vj], T.float32(0))

可以看到，和最终题目提供的结果相比，有几个点不一样：
1.在中间循环的j_1部分变成vectorized了；
2.最后relu部分，ax0没有被vectoried。
产生这个问题的原因在于我不知道如何去控制原来初始化和relu和mm合并后如何单独去控制vectorize，求大佬解释下或者告知下正确答案我参考学习下，多谢啦。

[cjx0709]

我自己好像解决这个问题了：
1.关于block C部分，可以在compute at后，直接去用get_loops去找到对应的循环变量，再进行操作；
2.关于block Y_init部分，在decompose_reduction之后，会自动生成一个Y_init的block，因此，操作同对于block_C部分的操作。
修改后的tansform部分如下：

sch = tvm.tir.Schedule(MyBmmRelu)

block_Y = sch.get_block("Y", func_name="bmm_relu")
block_C = sch.get_block("C", "bmm_relu")

b, i, j, k = sch.get_loops(block_Y)
i2_0, ax0 = sch.split(j, [16, None])
ax1_0, ax1_1 = sch.split(k, [None, 4])
sch.reorder(ax1_0, ax1_1, ax0)
sch.parallel(b)
sch.unroll(ax1_1)
sch.reverse_compute_at(block_C, i2_0)
sch.decompose_reduction(block_Y, ax1_0)

i0, i1, i2_0, i2_1 = sch.get_loops(block_C)
sch.vectorize(i2_1)

block_Y_init = sch.get_block("Y_init", "bmm_relu")

b, i, j0, j_1_init = sch.get_loops(block_Y_init)
sch.vectorize(j_1_init)