[提问] 03课中提到的axis_range和loop_range关系

[Aguin]

课程视频21:50处以及英文笔记提到如果把axis_range改为和loop_range不一致会报错. 我尝试了以下code:

import numpy as np
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T


@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def mm_relu(A: T.Buffer[(128, 128), "float32"],
                B: T.Buffer[(128, 128), "float32"],
                C: T.Buffer[(128, 128), "float32"]):
        T.func_attr({"global_symbol": "mm_relu", "Tnoalias": True})
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i, j, k in T.grid(127, 127, 127):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                vk = T.axis.reduce(128, k)
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(127, 127):
            with T.block("C"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

rt_lib = tvm.build(MyModule, target="llvm")

dtype = "float32"
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)
a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.empty((128, 128), dtype="float32")

func_mm_relu = rt_lib["mm_relu"]
func_mm_relu(a_nd, b_nd, c_nd)

发现可以顺利执行完成, 但是若把127改为129则会报错:

free(): corrupted unsorted chunks
[1] 13786 abort python3 test_axis_range.py

请问这个现象符合预期吗? 能否简单解释一下原理呢? 感谢助教.

[Hzfengsy]

感谢提问，这个问题的本质分为几个部分：

1. 如果axis_range 和 loop_range不一致是合法的吗？
   答：显然是不合法的
2. 如果不合法会出现什么情况？
   答：目前对于程序正确性验证这一块确实代码还不完善。理论上应该报错，但实际上由于没有严格检查会产生undefined behavior。该问题我们也会尽快修复。
3. 为什么127 和129行为不一致？
   答：undefined behavior会产生不可控的结果。从具体实现上来看，129实际上超过了数组本身的定义空间所以报错了；127没有越界。但这都是”非法行为“不予以鼓励。