超神经HyperAI · 10月22日 · 天津

【TVM 教程】内联及数学函数

微信图片_20241016191925.png
Apache TVM 是一个端到端的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 → https://tvm.hyper.ai/

作者:Tianqi Chen

尽管 TVM 支持基本的算术运算,但很多时候,也需要复杂的内置函数,例如 exp 取指函数。
这些函数是依赖 target 系统的,并且在不同 target 平台中可能具有不同的名称。本教程会学习到如何调用这些 target-specific 函数,以及如何通过 TVM 内联 API 统一接口。

from __future__ import absolute_import, print_function

import numpy as np
import tvm
from tvm import te
from tvm.ir import register_op_attr, register_intrin_lowering

直接声明外部数学调用

调用 target-specific 函数最直接方法,就是通过 TVM 中的 extern 函数调用构造。以下示例用 tvm.tir.call_pure_extern 来调用 __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())

输出结果:

#ifdef _WIN32
  using uint = unsigned int;
  using uchar = unsigned char;
  using ushort = unsigned short;
  using int64_t = long long;
  using uint64_t = unsigned long long;
#else
  #define uint unsigned int
  #define uchar unsigned char
  #define ushort unsigned short
  #define int64_t long long
  #define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) 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)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
  } else {
    if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
      B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
    }
  }
}

统一内联调用

以上代码验证了直接外部调用可用于 device-specific 的函数。但上述方式仅适用于带有浮点类型的 CUDA target。理想情况下,我们希望写一套代码,即可适用于任何设备以及任何数据类型。

TVM 内联函数为用户提供了实现机制,且推荐用这个方法来解决问题。以下代码用的是 te.exp,它创建了一个内联调用 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())

输出结果:

#ifdef _WIN32
  using uint = unsigned int;
  using uchar = unsigned char;
  using ushort = unsigned short;
  using int64_t = long long;
  using uint64_t = unsigned long long;
#else
  #define uint unsigned int
  #define uchar unsigned char
  #define ushort unsigned short
  #define int64_t long long
  #define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) 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)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
  } else {
    if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
      B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = __expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
    }
  }
}

该代码适用于 CUDA 和 opencl,相同的 te.exp 也可用于 float64 数据类型。

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

输出结果:

// Function: myexp_kernel0
__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))) * stride)] = exp(A[(((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride1)]);
  } 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))) * stride)] = exp(A[(((((int)get_group_id(0)) * 64) + ((int)get_local_id(0))) * stride1)]);
    }
  }
}

内联函数降级规则

当调用 tvm.te.exp() 时,TVM 会创建一个 intrinsic Call Expr。TVM 使用转换规则(transformation rules),将内联调用(intrinsic call)转换为特定设备的外部调用(extern calls)。

TVM 支持在运行时自定义规则,以下示例为 exp 自定义 CUDA 降级规则。

def my_cuda_math_rule(op):
    """自定义 CUDA 内联函数降级规则"""
    assert isinstance(op, tvm.tir.Call)
    name = op.op.name
    assert name.startswith("tir.")
    dispatch_name = name[4:]
    if op.dtype == "float32":
        # 调用浮点函数
        return tvm.tir.call_pure_extern("float32", "%sf" % dispatch_name, op.args[0])
    elif op.dtype == "float64":
        # 调用双精度函数
        return tvm.tir.call_pure_extern("float32", dispatch_name, op.args[0])
    else:
        # 不能转换,返回自身。
        return op

register_intrin_lowering("tir.exp", target="cuda", f=my_cuda_math_rule, level=99)

输出结果:

<function my_cuda_math_rule at 0x7f7017159dd0>

用选项覆盖现有规则,从而将规则注册到 TVM。注意,打印代码与之前代码的区别:新规则用数学函数 expf,而不是快速数学版本 __expf。

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

输出结果:

#ifdef _WIN32
  using uint = unsigned int;
  using uchar = unsigned char;
  using ushort = unsigned short;
  using int64_t = long long;
  using uint64_t = unsigned long long;
#else
  #define uint unsigned int
  #define uchar unsigned char
  #define ushort unsigned short
  #define int64_t long long
  #define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) 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)) * stride)] = expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
  } else {
    if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
      B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = expf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
    }
  }
}

添加内联函数

对于 TVM 未提供的内联函数,用户可以借助内联规则系统,添加新的内联函数。以下是将内联函数 mylog 添加到系统的示例:

def mylog(x):
    """自定义日志内联函数"""
    return tvm.tir.call_intrin(x.dtype, "tir.mylog", x)

def my_cuda_mylog_rule(op):
    """CUDA 降级日志的规则"""
    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

# 新的注册操作是通过注册操作的属性来触发的
register_op_attr("tir.mylog", "TCallEffectKind", tvm.tir.CallEffectKind.Pure)
register_intrin_lowering("tir.mylog", target="cuda", f=my_cuda_mylog_rule, level=99)

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

输出结果:

#ifdef _WIN32
  using uint = unsigned int;
  using uchar = unsigned char;
  using ushort = unsigned short;
  using int64_t = long long;
  using uint64_t = unsigned long long;
#else
  #define uint unsigned int
  #define uchar unsigned char
  #define ushort unsigned short
  #define int64_t long long
  #define uint64_t unsigned long long
#endif
extern "C" __global__ void __launch_bounds__(64) 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)) * stride)] = logf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
  } else {
    if (((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) < n) {
      B[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride)] = logf(A[(((((int)blockIdx.x) * 64) + ((int)threadIdx.x)) * stride1)]);
    }
  }
}

总结

  • TVM 能调用依赖 target 的外部数学函数。
  • 用内联函数为函数定义统一的接口。
  • 有关 TVM 中更多可用的内联函数,查看 tvm.tir。
  • 通过自定义规则,从而自定义内联行为。

下载 Python 源代码:intrin_math.py

下载 Jupyter Notebook:intrin_math.ipynb

推荐阅读
关注数
666
内容数
225
链接人工智能新场景
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息