算子本质与数学函数
算子本質與數學函數
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
總結
