markdown TVM使用内联和数学函数

Posted 2021-05-05

# TVM使用内联和数学函数

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

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

```python
from __future__ import absolute_import, print_function

import tvm
import numpy as np
```

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

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

```python
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())
```

输出：

```c
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`来执行指数运算。

```python
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())
```

输出：

```c
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数据类型。

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

输出：

```c
__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规则。

```python
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`。

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

输出（不是__expf函数）：

```c
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`。

```python
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())
```

输出：

```c
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`
- 可以通过定义自己的规则来自定义内联行为