Apache TVM 是一个端到端的深度学习编译框架，适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 → https://tvm.hyper.ai/
作者：Tianqi Chen
尽管 TVM 支持基本的算术运算，但很多时候，也需要复杂的内置函数，例如 exp 取指函数。
这些函数是依赖 target 系统的，并且在不同 target 平台中可能具有不同的名称。本教程会学习到如何调用这些 target-specific 函数，以及如何通过 TVM 内联 API 统一接口。
from __future__ import absolute_import, print_function
import numpy as np
import tvm
from tvm import te
from tvm.ir import register_op_attr, register_intrin_lowering直接声明外部数学调用
调用 target-specific 函数最直接方法，就是通过 TVM 中的 extern 函数调用构造。以下示例用 tvm.tir.call_pure_extern 来调用 __expf 函数（仅在 CUDA 下可用）。
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: tvm.tir.call_pure_extern("float32", "__expf", A[i]), name="B")
s = te.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
f = tvm.build(s, [A, B], "cuda", name="myexp")
print(f.imported_modules[0].get_source())输出结果：
#ifdef _WIN32
using uint = unsigned int;
using uchar = unsigned char;
using ushort = unsigned short;
using int64_t = long long;
using uint64_t = unsigned long long;
#else
#define uint unsigned int
#define uchar unsigned char
#define ushort unsigned short
#define int64_t long long
#define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) myexp_kernel0(float* __restrict__ B, float* __restrict__ A, int n, int stride, int stride1) {
if (((int)blockIdx.x) < (n >> 6)) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
}
}
}统一内联调用
以上代码验证了直接外部调用可用于 device-specific 的函数。但上述方式仅适用于带有浮点类型的 CUDA target。理想情况下，我们希望写一套代码，即可适用于任何设备以及任何数据类型。
TVM 内联函数为用户提供了实现机制，且推荐用这个方法来解决问题。以下代码用的是 te.exp，它创建了一个内联调用 tvm.te.exp() 来做指数。
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: te.exp(A[i]), name="B")
s = te.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
fcuda = tvm.build(s, [A, B], "cuda", name="myexp")
print(fcuda.imported_modules[0].get_source())输出结果：
#ifdef _WIN32
using uint = unsigned int;
using uchar = unsigned char;
using ushort = unsigned short;
using int64_t = long long;
using uint64_t = unsigned long long;
#else
#define uint unsigned int
#define uchar unsigned char
#define ushort unsigned short
#define int64_t long long
#define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) myexp_kernel0(float* __restrict__ B, float* __restrict__ A, int n, int stride, int stride1) {
if (((int)blockIdx.x) < (n >> 6)) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
}
}
}该代码适用于 CUDA 和 opencl，相同的 te.exp 也可用于 float64 数据类型。
fopencl = tvm.build(s, [A, B], "opencl", name="myexp")
print(fopencl.imported_modules[0].get_source())输出结果：
// Function: myexp_kernel0
__kernel void myexp_kernel0(__global float* restrict B, __global float* restrict A, int n, int stride, int stride1) {
if (((int)get_group_id(0)) < (n >> 6)) {
B[(((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride)] = exp(A[(((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride1)]);
} else {
if (((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) < n) {
B[(((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride)] = exp(A[(((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride1)]);
}
}
}内联函数降级规则
当调用 tvm.te.exp() 时，TVM 会创建一个 intrinsic Call Expr。TVM 使用转换规则（transformation rules），将内联调用（intrinsic call）转换为特定设备的外部调用（extern calls）。
TVM 支持在运行时自定义规则，以下示例为 exp 自定义 CUDA 降级规则。
def my_cuda_math_rule(op):
"""自定义 CUDA 内联函数降级规则"""
assert isinstance(op, tvm.tir.Call)
name = op.op.name
assert name.startswith("tir.")
dispatch_name = name[4:]
if op.dtype == "float32":
# 调用浮点函数
return tvm.tir.call_pure_extern("float32", "%sf" % dispatch_name, op.args[0])
elif op.dtype == "float64":
# 调用双精度函数
return tvm.tir.call_pure_extern("float32", dispatch_name, op.args[0])
else:
# 不能转换，返回自身。
return op
register_intrin_lowering("tir.exp", target="cuda", f=my_cuda_math_rule, level=99)输出结果：
<function my_cuda_math_rule at 0x7f7017159dd0>用选项覆盖现有规则，从而将规则注册到 TVM。注意，打印代码与之前代码的区别：新规则用数学函数 expf，而不是快速数学版本 __expf。
fcuda = tvm.build(s, [A, B], "cuda", name="myexp")
print(fcuda.imported_modules[0].get_source())输出结果：
#ifdef _WIN32
using uint = unsigned int;
using uchar = unsigned char;
using ushort = unsigned short;
using int64_t = long long;
using uint64_t = unsigned long long;
#else
#define uint unsigned int
#define uchar unsigned char
#define ushort unsigned short
#define int64_t long long
#define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) myexp_kernel0(float* __restrict__ B, float* __restrict__ A, int n, int stride, int stride1) {
if (((int)blockIdx.x) < (n >> 6)) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
}
}
}添加内联函数
对于 TVM 未提供的内联函数，用户可以借助内联规则系统，添加新的内联函数。以下是将内联函数 mylog 添加到系统的示例：
def mylog(x):
"""自定义日志内联函数"""
return tvm.tir.call_intrin(x.dtype, "tir.mylog", x)
def my_cuda_mylog_rule(op):
"""CUDA 降级日志的规则"""
if op.dtype == "float32":
return tvm.tir.call_pure_extern("float32", "logf", op.args[0])
elif op.dtype == "float64":
return tvm.tir.call_pure_extern("float64", "log", op.args[0])
else:
return op
# 新的注册操作是通过注册操作的属性来触发的
register_op_attr("tir.mylog", "TCallEffectKind", tvm.tir.CallEffectKind.Pure)
register_intrin_lowering("tir.mylog", target="cuda", f=my_cuda_mylog_rule, level=99)
n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.compute(A.shape, lambda i: mylog(A[i]), name="B")
s = te.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
fcuda = tvm.build(s, [A, B], "cuda", name="mylog")
print(fcuda.imported_modules[0].get_source())输出结果：
#ifdef _WIN32
using uint = unsigned int;
using uchar = unsigned char;
using ushort = unsigned short;
using int64_t = long long;
using uint64_t = unsigned long long;
#else
#define uint unsigned int
#define uchar unsigned char
#define ushort unsigned short
#define int64_t long long
#define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) mylog_kernel0(float* __restrict__ B, float* __restrict__ A, int n, int stride, int stride1) {
if (((int)blockIdx.x) < (n >> 6)) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = logf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = logf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
}
}
}总结
- TVM 能调用依赖 target 的外部数学函数。
- 用内联函数为函数定义统一的接口。
- 有关 TVM 中更多可用的内联函数，查看 tvm.tir。
- 通过自定义规则，从而自定义内联行为。
**粗体** _斜体_ [链接](http://example.com) `代码` - 列表 > 引用。你还可以使用@来通知其他用户。