TensorIR初始化部分没有生成T.float32(0)而是0

[DzAvril]

最后一题我生成的TensorIR基本和TargetModule结构类似，唯一的区别是我在使用了decompose_reduction之后，初始化部分变成了：

with tir.block("Y_init"):
      vn, vi = tir.axis.remap("SS", [n, i])
      vj = tir.axis.spatial(128, j_0 * 8 + j_1_init)
      tir.reads()
      tir.writes(Y[vn, vi, vj])
      Y[vn, vi, vj] = 0

TargetModule里的这一部分是：

 with T.block("Y_init"):
        n, i = T.axis.remap("SS", [i0, i1])
        j = T.axis.spatial(128, i2_0 * 8 + ax0_init)
        Y[n, i, j] = T.float32(0)

区别在于我的结果中Y[vn, vi, vj] = 0，而Target里是Y[n, i, j] = T.float32(0)，导致在assert_structural_equal的时候通不过。但是time_evaluator可以正常执行，结果也正确。

不知道是否因为mlc的版本问题?

我的解答是这样的：

sch = tvm.tir.Schedule(MyBmmRelu)

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

b, i, j, k = sch.get_loops(Y)
sch.parallel(b)

j0, j1 = sch.split(j, factors=[None, 8])
sch.reorder(b, i, j0, k, j1)

Y_init = sch.decompose_reduction(Y, k)
k0, k1 = sch.split(k, factors=[None, 4])

_, _, _, j1_init = sch.get_loops(Y_init)
sch.vectorize(j1_init)

sch.reverse_compute_at(C, j0)

_, _, _, i2_c = sch.get_loops(C)
sch.vectorize(i2_c)

sch.unroll(k1)

print(sch.mod.script())