【TVM 教程】外部张量函数

Apache TVM 是一个端到端的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 → tvm.hyper.ai/

作者:Tianqi Chen

虽然 TVM 支持透明代码生成,但有时也需将手写的代码合并到流水线,例如对一些卷积核使用 cuDNN,并定义其余阶段。

原生 TVM 就支持黑盒函数调用。具体来说,TVM 支持所有与 DLPack 兼容的张量函数。这意味着可以使用 POD 类型(指针、整数、浮点数),或者将指向 DLTensor 的指针作为参数,调用任何函数。

c 复制代码
from __future__ import absolute_import, print_function

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")

使用外部张量函数

以下示例用 te.extern 来添加一个外部数组函数调用。外部调用声明了输出张量的 shape,第二个参数给出了输入列表。

用户需要提供一个描述如何对结果进行计算的函数。计算函数获取输入和输出的符号占位符列表,并返回执行语句。

这种情况只需调用一个注册的 TVM 函数,它会调用 CBLAS。TVM 不控制外部数组函数的内部,将其视为黑盒。可以进一步混合可调度的 TVM 函数,为结果添加偏差项。

c 复制代码
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)

验证结果

验证结果是否符合预期。

c 复制代码
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)

外部 Contrib Wrappers

TVM 为外部调用提供了外部contrib Wrappers,以下代码与前面的示例等效。

c 复制代码
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)

将 Python 函数 Hook 为 Extern

由于可以调用 TVM 中的任何 PackedFunc,所以可以用外部函数回调到 Python 中。

以下示例将一个 Python 函数注册到 TVM runtime 系统,并用它来完成一个阶段的计算,这使得 TVM 更加灵活。例如,可通过插入前端回调来检查中间结果,或将自定义代码与 TVM 混合。

c 复制代码
@tvm.register_func("tvm.contrib.my_tvm_addone")
def my_tvm_addone(x, y):
    print("my_tvm_addone signatures: %s, %s" % (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)

输出结果:

c 复制代码
my_tvm_addone signatures: <class 'tvm.runtime.ndarray.NDArray'>, <class 'tvm.runtime.ndarray.NDArray'>

总结

  • TVM 通过 te.extern 调用外部张量函数。
  • 对外部张量调用使用 contrib wrappers。
  • 将前端函数 hook 为外部张量的回调。

下载 Python 源代码:extern_op.py

下载 Jupyter Notebook:extern_op.ipynb

相关推荐
聚客AI7 分钟前
⭐精准率暴跌50%?RAG开发者必避的十大认知误区
人工智能·llm·agent
codeGoogle15 分钟前
大厂研发之谜:千亿投入砸出利润大缩水
前端·人工智能·后端
豆浩宇32 分钟前
Conda环境隔离和PyCharm配置,完美同时运行PaddlePaddle和PyTorch
人工智能·pytorch·算法·计算机视觉·pycharm·conda·paddlepaddle
京东云开发者42 分钟前
DeepSeek冲击(含本地化部署实践)
人工智能
@国境以南,太阳以西1 小时前
基于Grad-CAM(Gradient-weighted Class Activation Mapping)的可解释性分析
人工智能·深度学习
AI人工智能+1 小时前
表格识别技术:通过计算机视觉和OCR,实现非结构化表格向结构化数据的转换,推动数字化转型。
人工智能·计算机视觉·ocr
小宁爱Python2 小时前
基于 Django+Vue3 的 AI 海报生成平台开发(海报模块专项)
人工智能·python·django
破烂儿2 小时前
基于机器学习的缓存准入策略研究
人工智能·机器学习·缓存
算法打盹中2 小时前
SimLingo:纯视觉框架下的自动驾驶视觉 - 语言 - 动作融合模型
人工智能·机器学习·计算机视觉·语言模型·自动驾驶
大嘴带你水论文2 小时前
震惊!仅用10张照片就能随意编辑3D人脸?韩国KAIST最新黑科技FFaceNeRF解析!
论文阅读·人工智能·python·科技·计算机视觉·3d·transformer