张量形状符号

接下来以张量加法为例探讨使用张量表达式定义不定形状的张量。

不定形状的张量算子

%cd ../..
import set_env

import numpy as np
import tvm
from tvm import te
from tvm.ir.module import IRModule

声明用于记录张量形状的符号变量（var()）：

N = te.var(name='N')
print(type(N), N.dtype)
N

<class 'tvm.tir.expr.Var'> int32

N

声明两个张量占位符：

A = te.placeholder((N,), name='A')
B = te.placeholder((N,), name='B')
A

Tensor(shape=[N], op.name=A)

定义张量加法运算：

C = te.compute(A.shape, lambda i: A[i] + B[i], name='C')

构建张量原语函数：

te_func = te.create_prim_func([A, B, C])
te_func.show()

# from tvm.script import tir as T
@T.prim_func
def func(var_A: T.handle, var_B: T.handle, var_C: T.handle):
    # function attr dict
    T.func_attr({"global_symbol": "main", "tir.noalias": True})
    N = T.var("int32")
    A = T.match_buffer(var_A, [N], dtype="float32")
    B = T.match_buffer(var_B, [N], dtype="float32")
    C = T.match_buffer(var_C, [N], dtype="float32")
    # body
    # with T.block("root")
    for i0 in T.serial(N):
        with T.block("C"):
            i = T.axis.spatial(N, i0)
            T.reads(A[i], B[i])
            T.writes(C[i])
            C[i] = A[i] + B[i]

将张量原语函数变换为模块：

Module = IRModule({"add": te_func})
Module.show()

# from tvm.script import tir as T
@tvm.script.ir_module
class Module:
    @T.prim_func
    def add(var_A: T.handle, var_B: T.handle, var_C: T.handle):
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        N = T.var("int32")
        A = T.match_buffer(var_A, [N], dtype="float32")
        B = T.match_buffer(var_B, [N], dtype="float32")
        C = T.match_buffer(var_C, [N], dtype="float32")
        # body
        # with T.block("root")
        for i0 in T.serial(N):
            with T.block("C"):
                i = T.axis.spatial(N, i0)
                T.reads(A[i], B[i])
                T.writes(C[i])
                C[i] = A[i] + B[i]
    

测试一致性

构建运行时库：

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

Module(llvm, 55cdd2c145a8)

定义数据：

a_np = np.arange(16, dtype="float32")
b_np = np.arange(16, 0, -1, dtype="float32")
c_np = a_np + b_np # 基准结果
c_np

array([16., 16., 16., 16., 16., 16., 16., 16., 16., 16., 16., 16., 16.,
       16., 16., 16.], dtype=float32)

将 numpy 数组变换为 tvm 数组：

a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.empty((16,), dtype="float32")
type(c_nd)

tvm.runtime.ndarray.NDArray

验证结果：

rt_lib(a_nd, b_nd, c_nd)
c_nd

<tvm.nd.NDArray shape=(16,), cpu(0)>
array([16., 16., 16., 16., 16., 16., 16., 16., 16., 16., 16., 16., 16.,
       16., 16., 16.], dtype=float32)

可以对不同形状张量计算：

a_np = np.arange(10, dtype="float32")
b_np = np.arange(10, 0, -1, dtype="float32")
c_np = a_np + b_np # 基准结果
a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.empty((10,), dtype="float32")
rt_lib(a_nd, b_nd, c_nd)
c_nd

<tvm.nd.NDArray shape=(10,), cpu(0)>
array([10., 10., 10., 10., 10., 10., 10., 10., 10., 10.], dtype=float32)