外部张量函数#

原作者: Tianqi Chen

虽然 TVM 支持透明的代码生成，但有时将手动编写的代码合并到 pipeline 中也是有帮助的。例如，可能想要使用 cuDNN 来处理一些卷积核，并定义其余的阶段。

TVM本身就支持这些黑盒函数调用。具体来说，TVM 支持所有与 DLPack 兼容的张量函数。这意味着可以调用任何带有 POD 类型（pointer、int、float）或指向 DLTensor 的指针作为参数的函数。

备注

这里需要设定 cmake：

# Whether use BLAS, choices: openblas, atlas, apple
set(USE_BLAS atlas)

import env # 加载 TVM 环境
import tvm
from tvm import te
import numpy as np
from tvm.contrib import cblas
import tvm.testing

if not tvm.get_global_func("tvm.contrib.cblas.matmul", allow_missing=True):
    raise Exception("Not compiled with cblas support; can't build this tutorial")

使用 Extern 张量函数#

在下面的例子中，使用 te.extern 添加外部数组函数调用。在 extern 调用中，声明输出张量的形状。在第二个参数中，提供了输入列表。

用户需要提供描述如何计算结果的函数。compute 函数接受输入的符号占位符列表和输出的符号占位符列表，并返回正在执行的语句。

在本例中，只需调用已注册的 TVM 函数，该函数调用 CBLAS 回调。TVM 不控制 extern 数组函数的内部，并将其视为黑盒。可以进一步混合可调度的 TVM 调用，为结果添加 bias 项。

n = 1024
l = 128
m = 235
bias = te.var("bias", dtype="float32")
A = te.placeholder((n, l), name="A")
B = te.placeholder((l, m), name="B")
C = te.extern(
    (n, m),
    [A, B],
    lambda ins, outs: tvm.tir.call_packed(
        "tvm.contrib.cblas.matmul", ins[0], ins[1], outs[0], False, False
    ),
    name="C",
)
D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D")
s = te.create_schedule(D.op)

验证结果#

可以验证结果是否符合期望。

dev = tvm.cpu(0)
f = tvm.build(s, [A, B, D, bias], "llvm")
a = tvm.nd.array(np.random.uniform(size=(n, l)).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=(l, m)).astype(B.dtype), dev)
d = tvm.nd.array(np.zeros((n, m), dtype=D.dtype), dev)
bb = 10.0
f(a, b, d, bb)
tvm.testing.assert_allclose(d.numpy(), np.dot(a.numpy(), b.numpy()) + 10, rtol=1e-5)

封装 Extern Contrib#

TVM 还为有用的 extern 调用提供了 extern contrib 封装器，下面一行与前面的示例等价。

from tvm.contrib import cblas

C = cblas.matmul(A, B)
D = te.compute(C.shape, lambda i, j: C[i, j] + bias, name="D")
s = te.create_schedule(D.op)

Hook Python 函数作为 Extern#

因为可以在 TVM 中调用任何 PackedFunc。所以可以使用 extern 函数回调到 python。

下面的例子注册了 python 函数到 TVM 运行时系统，并使用它来完成计算的一个阶段。这使得 TVM 更加灵活。例如，可以插入前端回调来检查中间结果，或者将定制代码与 TVM 混合使用。

@tvm.register_func("tvm.contrib.my_tvm_addone")
def my_tvm_addone(x, y):
    print(f"my_tvm_addone signatures: {type(x)}, {type(y)}")
    tvm.nd.array(x.numpy() + 1).copyto(y)


A = te.placeholder((n,), name="A")
B = te.extern(
    A.shape,
    [A],
    lambda ins, outs: tvm.tir.call_packed("tvm.contrib.my_tvm_addone", ins[0], outs[0]),
    name="C",
)
s = te.create_schedule(B.op)
f = tvm.build(s, [A, B], "llvm")
a = tvm.nd.array(np.random.uniform(size=(n,)).astype(A.dtype), dev)
b = tvm.nd.array(np.random.uniform(size=(n,)).astype(B.dtype), dev)
f(a, b)
tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1, rtol=1e-5)

my_tvm_addone signatures: <class 'tvm.runtime.ndarray.NDArray'>, <class 'tvm.runtime.ndarray.NDArray'>

小结#

  • TVM 通过 te.extern 调用外部张量函数

  • 使用 contrib 包装的外部张量调用的简短语法糖。

  • 可以将前端函数 hook 为外部张量回调函数。