TVM Compiler中文教程：TVM使用内联和数学函数

TVM使用内联和数学函数

TVM支持基础算术运算操作，但是在很多情况下我们需要更复杂的內建函数。例如exp指数函数。

这些內建函数取决于目标系统，在不同的平台可能有不同的名字。这个教程中，我们将学习调用目标特定的內建函数，和怎么能够通过TVM内联API统一接口。

from __future__ import absolute_import, print_function

import tvm
import numpy as np

直接声明外部数学函数调用

调用目标特定函数最直接的方法是通过TVM外部（extrern）函数调用并构造。在下面例子中，我们使用tvm.call_pure_extern调用CUDA的__expf函数。

n = tvm.var("n")
A = tvm.placeholder((n,) name='A')
#调用特定平台的函数
B = tvm.compute(A.shape,lambda i: tvm.call_pure_extern("float32","__expf",A[i]),name='B')
s = tvm.create_schedule(B.op)
num_thread = 64
#分裂成内外循环便于多线程计算
bx, tx = s[B].split(B.op.axis[0], factor=64)
#绑定一维块和线程
s[B].bind(bx, tvm.thread_axis("blockIdx.x"))
s[B].bind(tx, tvm.thread_axis("threadIdx.x"))
f = tvm.build(s, [A,B], "cuda", name='myexp')
#打印生成的myexp内核函数
print(f.import_modules[0].get_source())

输出：

extern "C" __global__ void myexp_kernel0( float* __restrict__ B,  float* __restrict__ A, int n) {
  if (((int)blockIdx.x) < (n / 64)) {
    B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = __expf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
  } else {
    if ((((int)blockIdx.x) * 64) < (n - ((int)threadIdx.x))) {
      B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = __expf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
    }
  }
}

统一内联函数调用

上面代码验证了，直接外部调用设备特定函数。不管怎样，上面的方法只工作在CUDA浮点类型，我们通常想写通用于任何设备和数据类型的代码。

TVM内部为用户提供实现此目的的机制，这是我们推荐的方法。下面的代码使用tvm.exp，它创建一个内联调用tvm.exp来执行指数运算。

n = tvm.var("n")
A = tvm.placeholder((n,), name='A')
#这行替换为tvm.exp
B = tvm.compute(A.shape, lambda i: tvm.exp(A[i]), name="B")
s = tvm.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, tvm.thread_axis("blockIdx.x"))
s[B].bind(tx, tvm.thread_axis("threadIdx.x"))
fcuda = tvm.build(s, [A, B], "cuda", name="myexp")
print(fcuda.imported_modules[0].get_source())

输出：

extern "C" __global__ void myexp_kernel0( float* __restrict__ B,  float* __restrict__ A, int n) {
  if (((int)blockIdx.x) < (n / 64)) {
    B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = __expf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
  } else {
    if ((((int)blockIdx.x) * 64) < (n - ((int)threadIdx.x))) {
      B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = __expf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
    }
  }
}

我们可以发现这个代码同时适用于CUDA和opencl。同样tvm.exp也能用于float64数据类型。

fopnencl = tvm.build(s, [A,B], "opencl", name='myexp')
print(fopencl.import_modules[0].get_source())

输出：

__kernel void myexp_kernel0(__global float* restrict B, __global float* restrict A, int n) {
  if (((int)get_group_id(0)) < (n / 64)) {
    B[((((int)get_group_id(0)) * 64) + ((int)get_local_id(0)))] = exp(A[((((int)get_group_id(0)) * 64) + ((int)get_local_id(0)))]);
  } else {
    if ((((int)get_group_id(0)) * 64) < (n - ((int)get_local_id(0)))) {
      B[((((int)get_group_id(0)) * 64) + ((int)get_local_id(0)))] = exp(A[((((int)get_group_id(0)) * 64) + ((int)get_local_id(0)))]);
    }
  }
}

内部Lowering规则

当tvm.exp被调用，TVM创建一个内部调用表示，TVM使用变换规则去转换内部调用去设备特定外部调用。

TVM也允许用户在运行时自定义规则。下面例子为exp展示自定义CUDA lowering规则。

def my_cuda_math_rule(op):
    #自定义CUDA内部lowering规则
    assert isinstance(op, tvm.expr.Call)
    if op.dtype == "float32":
        #调用浮点函数
        return tvm.call_pure_extern("float32", "%sf" % op.name, op.args[0])
    elif op.dtype == "float64":
        #调用双精浮点函数
        return tvm.call_pure_extern("float64", op.name, op.args[0])
    else:
        return op
tvm.register_intrin_rule("cuda", "exp", my_cuda_math_rule, override=True)

使用覆盖选项将规则注册到TVM去覆盖现有规则。注意与前一个打印代码之间的区别：我们的新规则使用数学函数expf而不是快速数学版本__expf。

fcuda = tvm.build(s, [A, B], "cuda", name="myexp")
print(fcuda.imported_modules[0].get_source())

输出（不是__expf函数）：

extern "C" __global__ void myexp_kernel0( float* __restrict__ B,  float* __restrict__ A, int n) {
  if (((int)blockIdx.x) < (n / 64)) {
    B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = expf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
  } else {
    if ((((int)blockIdx.x) * 64) < (n - ((int)threadIdx.x))) {
      B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = expf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
    }
  }
}

如果存在TVM未提供的内联函数。用户可以使用内联函数规则系统轻松添加新的内联函数。以下示例向系统添加内联函数mylog。

def mylog(x):
    """customized log intrinsic function"""
    return tvm.call_pure_intrin(x.dtype, "mylog", x)

def my_cuda_mylog_rule(op):
    """CUDA lowering rule for log"""
    if op.dtype == "float32":
        return tvm.call_pure_extern("float32", "logf", op.args[0])
    elif op.dtype == "float64":
        return tvm.call_pure_extern("float64", "log", op.args[0])
    else:
        return op
tvm.register_intrin_rule("cuda", "mylog", my_cuda_mylog_rule, override=True)

n = tvm.var("n")
A = tvm.placeholder((n,), name='A')
B = tvm.compute(A.shape, lambda i: mylog(A[i]), name="B")
s = tvm.create_schedule(B.op)
num_thread = 64
bx, tx = s[B].split(B.op.axis[0], factor=num_thread)
s[B].bind(bx, tvm.thread_axis("blockIdx.x"))
s[B].bind(tx, tvm.thread_axis("threadIdx.x"))
fcuda = tvm.build(s, [A, B], "cuda", name="mylog")
print(fcuda.imported_modules[0].get_source())

输出：

extern "C" __global__ void mylog_kernel0( float* __restrict__ B,  float* __restrict__ A, int n) {
  if (((int)blockIdx.x) < (n / 64)) {
    B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = logf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
  } else {
    if ((((int)blockIdx.x) * 64) < (n - ((int)threadIdx.x))) {
      B[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))] = logf(A[((((int)blockIdx.x) * 64) + ((int)threadIdx.x))]);
    }
  }
}

总结

• TVM能调用外部特定平台的数学函数
• 使用intrinsic内联去为函数定义一个统一接口
• 有关TVM中更多可用内联函数，参考tvm.intrin
• 可以通过定义自己的规则来自定义内联行为