如何TensorIR替代TE来定义算子的实现？

[DzAvril]

当前的relay+te的框架下模型的转换大致流程为：front model -> relay -> te -> tir
社区现在主推的relax+TensorIR的流程我的理解是:front model -> relax -> TnesorIR -> tir
其中当前算子的te实现部分将来是要被TensorIR替代。

如果我上述理解是正确的话，那么我的问题就来了。
目前为止我看到的tvm.script写的tir.prim_fun是这种：

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # We exchange data between function by handles, which are similar to pointer.
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # Create buffer from handles.
        A = T.match_buffer(a, (8,), dtype="float32")
        B = T.match_buffer(b, (8,), dtype="float32")
        for i in range(8):
            # A block is an abstraction for computation.
            with T.block("B"):
                # Define a spatial block iterator and bind it to value i.
                vi = T.axis.spatial(8, i)
                B[vi] = A[vi] + 1.0

或是这种：

    def relu0(X: T.Buffer[(1, 128), "float32"],
              Y: T.Buffer[(1, 128), "float32"]):
        # function attr dict
        T.func_attr({"global_symbol": "relu0", "tir.noalias": True})
        for i, j in T.grid(1, 128):
            with T.block("Y"):
                vi, vj = T.axis.remap("SS", [i, j])
                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))

这两种写法中，传入到tir.prim_func的shape都是固定的。比如图一中的A = T.match_buffer(a, (8,), dtype="float32")及图二中的X: T.Buffer[(1, 784), "float32"],。

如果是写一个算子，我希望的是传入一个tensor类，在算子内可以通过该类获取shape大小，以达到如下效果:

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.tensor, b: T.tensor):
        # We exchange data between function by handles, which are similar to pointer.
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # Create buffer from handles.
        A = T.match_buffer(a, a.shape, dtype=a.dtype)
        B = T.match_buffer(b, b.shape, dtype=b.dtype)
        for i in range(a.shape[0]):
            # A block is an abstraction for computation.
            with T.block("B"):
                # Define a spatial block iterator and bind it to value i.
                vi = T.axis.spatial(b.shape[0], i)
                B[vi] = A[vi] + 1.0

不知道我理解的TensorIR的作用是否是替代te来写算子？如果是的话可否提供一个用TensorIR写算子的例子？

[Hzfengsy]

感谢提问。

首先需要指出一个点，我们的流程并不是front model -> relax -> TnesorIR -> tir，准确应该是front model -> IRModule(relax + TnesorIR) -> Runtime Module。

目前来说我们还不具备这种动态shape的表达方式，关于更灵活的Script写法可以关注RFC apache/tvm-rfcs#79

[DzAvril]

感谢指正。

如果还不能支持这种动态表达方式，那目前就没法使用TensorIR来替代TE来实现算子对吧？

那TensorIR在从前端模型到最终可部署的模型.so文件这个流程中扮演了一个什么样的角色？

[yzh119]

@DzAvril 维度已知的话，可以在TIR script内定义variable

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.tensor, b: T.tensor):
        # We exchange data between function by handles, which are similar to pointer.
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        n = T.var('int32')
        # Create buffer from handles.
        A = T.match_buffer(a, (n,), dtype=a.dtype)
        B = T.match_buffer(b, (n,), dtype=b.dtype)
        for i in range(a.shape[0]):
            # A block is an abstraction for computation.
            with T.block("B"):
                # Define a spatial block iterator and bind it to value i.
                vi = T.axis.spatial(b.shape[0], i)
                B[vi] = A[vi] + 1.0

[DzAvril]

感谢回答。

我想要的是用TensorIR写一个通用的算子，在解析前端模型时，将前端算子转成tvm的算子，比如torch.add变成tvm.add。而tvm.add调用了我用TensorIR写的算子，那么传入到该算子里的tensor的shape只有在模型解析的时候才能获取。在定义这个算子的时候模型传入tensor的shape是未知的，我希望可以在算子内部通过传入的tensor获取其shape，然后根据shape做一些运算。

[tqchen]

我们的目标并非完全去掉te。本身relax也会支持通过直接和te交互来快速构建TensorIR。针对te的schedule本身会被TensorIR变换替代。

课程的核心是让大家强化本身IRModule的概念，我们也鼓励用te或者其他方式本身可以用来meta programming（可以把这些东西看成是generic生成IRModule的一种方式）， 也会继续做相关的整合。

而这些方式生成的目标IRModule本身会包含TensorIR和relax，并且大家可以对它进行进一步的变换。

[DzAvril]

之前是在relay+te的框架下写算子，最近遇到难表达的算子，看到社区介绍的tvm script构建TensorIR更具灵活性，所以想着用tvm script来写算子，由此遇到了本贴提到的问题。

可能我之前理解有误，认为tvm script是用来替代te的。现在看来它最大的作用还是替换现在的te schedule。

感谢回答，我也会持续关注relax+TensorIR的进展