算子本质与数学函数
TVM支持基本的算术运算。在许多情况下，通常需要更复杂的内置函数。例如exp取函数的指数。
这些函数依赖于目标系统，可能具有不同目标平台的不同名称。本文将学习如何调用这些特定于目标的函数，以及如何通过tvm的内在API统一接口。
from __future__ import absolute_import, print_function
import tvm
from tvm import te
import numpy as np
Direct Declare Extern Math Call
调用特定于目标函数的最直接的方法是通过tvm中的extern函数调用构造。在下面的示例中，使用call __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())
Out:
extern "C" __global__ void 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)) * stride1))] = __expf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1))] = __expf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
}
}
}
Unified Intrinsic Call
上面的代码验证了直接外部调用是否可以用于调用特定于设备的函数。但是，上述方法只适用于浮点型的CUDA目标。理想情况下，希望为任何设备和任何数据类型编写相同的代码。
TVM内在机制为用户提供了实现这一点的机制，这是解决问题的推荐方法。以下代码使用te.exp公司而是创建一个call :py: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())
Out:
extern "C" __global__ void 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)) * stride1))] = __expf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1))] = __expf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
}
}
}
可以发现代码对CUDA和opencl都有效。同样的te.exp也可用于float64数据类型。
fopencl = tvm.build(s, [A, B], "opencl", name="myexp")
print(fopencl.imported_modules[0].get_source())
Out:
__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))) * stride1))] = exp(A[((((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride))]);
} 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))) * stride1))] = exp(A[((((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride))]);
}
}
}
Intrinsic Lowering Rule
什么时候tvm.te.exp()调用时，TVM将创建一个内部调用表达式。TVM使用转换规则将内部调用转换为特定于设备的外部调用。
TVM还允许用户在运行时自定义规则。以下示例为exp自定义CUDA降低规则。
def my_cuda_math_rule(op):
"""Customized CUDA intrinsic lowering rule"""
assert isinstance(op, tvm.tir.Call)
name = op.op.name
assert name.startswith("tir.")
dispatch_name = name[4:]
if op.dtype == "float32":
# call float function
return tvm.tir.call_pure_extern("float32", "%sf" % dispatch_name, op.args[0])
elif op.dtype == "float64":
# call double function
return tvm.tir.call_pure_extern("float32", dispatch_name, op.args[0])
else:
# cannot do translation, return self.
return op
tvm.target.register_intrin_rule("cuda", "exp", my_cuda_math_rule, override=True)
使用override选项将规则注册到TVM以覆盖现有规则。注意打印的代码与以前的代码之间的区别：新规则使用数学函数expf，而不是version __expf快速版本。
fcuda = tvm.build(s, [A, B], "cuda", name="myexp")
print(fcuda.imported_modules[0].get_source())
Out:
extern "C" __global__ void 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)) * stride1))] = expf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1))] = expf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
}
}
}
Add Your Own Intrinsic
如果存在TVM未提供的内在特性。用户可以通过使用内在规则系统轻松地添加新的内在规则。下面的示例向系统添加一个内部mylog。
def mylog(x):
"""customized log intrinsic function"""
return tvm.tir.call_intrin(x.dtype, "tir.mylog", x)
def my_cuda_mylog_rule(op):
"""CUDA lowering rule for log"""
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
# new op registration is triggered by registering an attribute of the op
tvm.ir.register_op_attr("tir.mylog", "TCallEffectKind", tvm.tir.CallEffectKind.Pure)
tvm.target.register_intrin_rule("cuda", "mylog", my_cuda_mylog_rule, override=True)
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())
Out:
extern "C" __global__ void 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)) * stride1))] = logf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
} else {
if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
B[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1))] = logf(A[((((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride))]);
}
}
}
Summary
TVM可以调用外部目标相关的数学函数。
为统一接口定义的函数。
有关tvm中可用的更多内部函数，请查看tvm.tir
可以通过定义自己的规则来定制内部行为。
下载Python源代码：intrin_math.py
下载Jupyter笔记本：intrin_math.ipynb