markdowntvm使用内联和数学函数(代码片段)

author author     2022-12-13     621

关键词:

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

markdowntvm调用外部张量函数(代码片段)

查看详情

内联函数和函数重载(代码片段)

...中函数调用的效率问题。函数重载允许程序员定义函数时使用相同的函数名,以不同数据类型的数据作为形参,实现相似功能的函数。内联函数在调用函数时,系统要将程序当前的一些状态信息、断点信息保存到堆栈中,同时转到被... 查看详情

inline内联函数(代码片段)

...替代宏代码片段就可以避免宏的副作用!) C++中推荐使用内联函数替代宏代码片段 C++中使用inline关键字声明内联函数  内联函数声明时inline关键字必须和函数定义结合在一起,否则编译器会直接忽略内联请求。&nb... 查看详情

markdowntvm使用autotvm调优nvidiagpu上的高性能卷积(代码片段)

查看详情

内联宏与普通函数(代码片段)

...普通函数与宏的特点与区别。1、普通函数的特点:函数使用可以避免相同代码重复写几次的麻烦,同时减少可执行程序的体积,但同样会造成程序运行时间的开销。如果一个函数内部语句较少,执行时间非常短,函数调用产生... 查看详情

c++内联函数的使用(代码片段)

...,修改函数要比找出并修改每一处等价表达式容易得多③使用函数可以确保统一的行为,每个测试都保证以相同的方式实现④函 查看详情

markdowntvm基准实验(代码片段)

查看详情

c++——引用内联函数(代码片段)

...什么?1.1引用的特性1.2常引用(对const的引用)1.3引用的使用场景1.4传值、传引用效率比较1.5引用和指针的区别二、内联函数1.什么是内联函数2.内联函数的特性三、auto关键字(C++11)auto的使用细则一、引用是什么?引用... 查看详情

类中的静态内联函数

...不可能的话,我希望它们分组在一个类中。我只想像这样使用我的功能:MyMath::Pow(2,2);MayMath::PI;所 查看详情

markdowntvm工作流(代码片段)

查看详情

c++入门篇之引用,内联函数,auto和范围遍历(代码片段)

...用效率比较内联函数1.概念2.特性auto关键字1.概念2.auto的使用细则3.auto不能推导的场景基于范围的for循环(C++11)使用条件指针空值nullptr前言承接上文入门篇1,博主这次将会继续更新以下内容:extern,引用,内联, 查看详情

markdowntvm调度原语(scheduleprimitives)(代码片段)

查看详情

markdowntvm部署ssd模型(代码片段)

查看详情

markdowntvm如何优化gpu卷积(代码片段)

查看详情

mysql函数-数学函数,控制函数(代码片段)

...值函数、三角函数、对数函数和随机函数等。注意:使用数学函数时,如有错误发生,该函数将返回null示例:--使用ceil(x)和ceiling(x)返回不小于x的最小整数。selectceil(2),ce 查看详情

复习笔记——c++内联函数(代码片段)

...定义inline关键字必须用于函数体定义,在函数声明前使用无效inlineintadd(intx,inty);//声明处使用无效intadd(intx,inty);//正确声明方式inlineintadd(intx,inty)//定义处使用有效returnx+y;内联函数优劣优点:①函数内联可以避免了频繁... 查看详情

内联汇编的使用(代码片段)

...会像Solidity一样解析注释,字面量和标识符。所以你可以使用//和/**/的方式注释。内联编译的在Solidity中的语法是包裹在assembly...,下面是可用的语法,后续有更详细的内容。字面量。如0x123,42或abc(字符串最多是32个字符)操作... 查看详情

markdowntvm如何优化cpugemm(矩阵乘法)(代码片段)

查看详情