Apache TVMは、CPU、GPU、そして様々な機械学習アクセラレーションチップに適した、エンドツーエンドのディープラーニング構築フレームワークです。中国語版のTVMドキュメントは、→ https://tvm.hyper.ai/ をご覧ください。

著者：Tianqi Chen

TVM は基本的な算術演算をサポートしていますが、exp フェッチ関数などの複雑な組み込み関数が必要になることがよくあります。
これらの関数はターゲット固有であり、ターゲットプラットフォームによって名前が異なる場合があります。このチュートリアルでは、これらのターゲット固有の関数の呼び出し方法と、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

外部の数学的な呼び出しを直接宣言する

ターゲット固有の関数を呼び出す最も直接的な方法は、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)]); } } }

統合インラインコール

上記のコードは、デバイス固有の関数に対して直接外部呼び出しが使用できることを確認しています。ただし、このアプローチは浮動小数点型を扱うCUDAターゲットにのみ適用されます。理想的には、あらゆるデバイスとあらゆるデータ型で動作する単一のコードセットを記述したいと考えています。

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 は組み込み呼び出し expr を作成します。TVM は変換ルールを使用して、組み込み呼び出しをデバイス固有の外部呼び出しに変換します。

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 は、ターゲットに依存する外部の数学関数を呼び出すことができます。
• インライン関数を使用して、関数の統一されたインターフェースを定義します。
• TVM で利用可能なインライン関数の詳細については、tvm.tir を参照してください。
• カスタム ルールを定義してインライン動作をカスタマイズします。

Pythonソースコードをダウンロード: inrin_math.py

Jupyter ノートブックをダウンロード: inrin_math.ipynb