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

著者: ジャン ジヘン

TVM はカーネルを効率的に構築するためのドメイン固有言語です。

このチュートリアルでは、TVM が提供するさまざまなプリミティブを使用して計算をスケジュールする方法を説明します。

 from __future__ import absolute_import, print_function import tvm from tvm import te import numpy as np

同じ結果を計算する方法は数多くありますが、方法が異なると局所性やパフォーマンスが変化する可能性があります。そのため、TVMでは、ユーザーはスケジュールに従って計算を実行する必要があります。

スケジュールは、プログラム内のループ計算を変換するために使用できる計算変換のセットです。

 # 声明变量，供之后使用n = te.var("n") m = te.var("m")

スケジュールは演算子のリストから作成でき、デフォルトでは行優先順にテンソルを順番に計算します。

 # 声明一个矩阵元素乘法A = te.placeholder((m, n), name="A") B = te.placeholder((m, n), name="B") C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], name="C") s = te.create_schedule([C.op]) # lower 会将计算从定义转换为实际可调用的函数。 # 使用参数`simple_mode=True` 会返回一个可读的类C 的语句，这里用它来打印schedule 结果。 print(tvm.lower(s, [A, B, C], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle, C_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto"), C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B, C_1: C} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m, n: int32], [stride, stride_3: int32], type="auto"), B_1: B_3: Buffer(B_2, float32, [m, n], [stride_1, stride_4: int32], type="auto"), C_1: C_3: Buffer(C_2, float32, [m, n], [stride_2, stride_5: int32], type="auto")} { for (i: int32, 0, m) { for (j: int32, 0, n) { C[((i*stride_2) + (j*stride_5))] = (A[((i*stride) + (j*stride_3))]*B[((i*stride_1) + (j*stride_4))]) } } }

スケジュールは複数のステージで構成され、ステージは1つの操作のスケジュールを表します。各ステージのスケジュール設定には複数の方法があります。

スプリット

split factorに基づいて指定された軸を 2 つの軸に分割できます。

 A = te.placeholder((m,), name="A") B = te.compute((m,), lambda i: A[i] * 2, name="B") s = te.create_schedule(B.op) xo, xi = s[B].split(B.op.axis[0], factor=32) print(tvm.lower(s, [A, B], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m], [stride], type="auto"), B_1: B_3: Buffer(B_2, float32, [m], [stride_1], type="auto")} { for (i.outer: int32, 0, floordiv((m + 31), 32)) { for (i.inner: int32, 0, 32) { if @tir.likely((((i.outer*32) + i.inner) < m), dtype=bool) { let cse_var_1: int32 = ((i.outer*32) + i.inner) B[(cse_var_1*stride_1)] = (A[(cse_var_1*stride)]*2f32) } } } }

npartsを使用して軸を分割することもできます。これは、 factorとは逆の方法で軸を分割します。

 A = te.placeholder((m,), name="A") B = te.compute((m,), lambda i: A[i], name="B") s = te.create_schedule(B.op) bx, tx = s[B].split(B.op.axis[0], nparts=32) print(tvm.lower(s, [A, B], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m], [stride], type="auto"), B_1: B_3: Buffer(B_2, float32, [m], [stride_1], type="auto")} { for (i.outer: int32, 0, 32) { for (i.inner: int32, 0, floordiv((m + 31), 32)) { if @tir.likely(((i.inner + (i.outer*floordiv((m + 31), 32))) < m), dtype=bool) { B[((i.inner + (i.outer*floordiv((m + 31), 32)))*stride_1)] = A[((i.inner + (i.outer*floordiv((m + 31), 32)))*stride)] } } } }

タイル

tile両方の軸でブロックごとに計算を実行できます。

 A = te.placeholder((m, n), name="A") B = te.compute((m, n), lambda i, j: A[i, j], name="B") s = te.create_schedule(B.op) xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5) print(tvm.lower(s, [A, B], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m, n: int32], [stride, stride_2: int32], type="auto"), B_1: B_3: Buffer(B_2, float32, [m, n], [stride_1, stride_3: int32], type="auto")} { for (i.outer: int32, 0, floordiv((m + 9), 10)) { for (j.outer: int32, 0, floordiv((n + 4), 5)) { for (i.inner: int32, 0, 10) { if @tir.likely((((i.outer*10) + i.inner) < m), dtype=bool) { for (j.inner: int32, 0, 5) { if @tir.likely((((j.outer*5) + j.inner) < n), dtype=bool) { let cse_var_2: int32 = ((j.outer*5) + j.inner) let cse_var_1: int32 = ((i.outer*10) + i.inner) B[((cse_var_1*stride_1) + (cse_var_2*stride_3))] = A[((cse_var_1*stride) + (cse_var_2*stride_2))] } } } } } } }

ヒューズ

fuse計算で 2 つの連続した軸を組み合わせることができます。

 A = te.placeholder((m, n), name="A") B = te.compute((m, n), lambda i, j: A[i, j], name="B") s = te.create_schedule(B.op) # 首先调用tile 平铺到四个axis: (i.outer, j.outer, i.inner, j.inner) xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5) # 然后将(i.inner, j.inner) 融合成一个轴： (i.inner.j.inner.fused) fused = s[B].fuse(xi, yi) print(tvm.lower(s, [A, B], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m, n: int32], [stride, stride_2: int32], type="auto"), B_1: B_3: Buffer(B_2, float32, [m, n], [stride_1, stride_3: int32], type="auto")} { for (i.outer: int32, 0, floordiv((m + 9), 10)) { for (j.outer: int32, 0, floordiv((n + 4), 5)) { for (i.inner.j.inner.fused: int32, 0, 50) { if @tir.likely((((i.outer*10) + floordiv(i.inner.j.inner.fused, 5)) < m), dtype=bool) { if @tir.likely((((j.outer*5) + floormod(i.inner.j.inner.fused, 5)) < n), dtype=bool) { let cse_var_2: int32 = ((j.outer*5) + floormod(i.inner.j.inner.fused, 5)) let cse_var_1: int32 = ((i.outer*10) + floordiv(i.inner.j.inner.fused, 5)) B[((cse_var_1*stride_1) + (cse_var_2*stride_3))] = A[((cse_var_1*stride) + (cse_var_2*stride_2))] } } } } } }

並べ替え

reorder関数は、指定された順序で軸を並べ替えることができます。

 A = te.placeholder((m, n), name="A") B = te.compute((m, n), lambda i, j: A[i, j], name="B") s = te.create_schedule(B.op) # 首先调用tile 平铺到四个轴: (i.outer, j.outer, i.inner, j.inner) xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5) # 然后将axis 重新排序：（i.inner，j.outer，i.outer，j.inner） s[B].reorder(xi, yo, xo, yi) print(tvm.lower(s, [A, B], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m, n: int32], [stride, stride_2: int32], type="auto"), B_1: B_3: Buffer(B_2, float32, [m, n], [stride_1, stride_3: int32], type="auto")} { for (i.inner: int32, 0, 10) { for (j.outer: int32, 0, floordiv((n + 4), 5)) { for (i.outer: int32, 0, floordiv((m + 9), 10)) { if @tir.likely((((i.outer*10) + i.inner) < m), dtype=bool) { for (j.inner: int32, 0, 5) { if @tir.likely((((j.outer*5) + j.inner) < n), dtype=bool) { let cse_var_2: int32 = ((j.outer*5) + j.inner) let cse_var_1: int32 = ((i.outer*10) + i.inner) B[((cse_var_1*stride_1) + (cse_var_2*stride_3))] = A[((cse_var_1*stride) + (cse_var_2*stride_2))] } } } } } } }

バインド

bind指定された軸をスレッド軸にバインドすることができ、GPU プログラミングでよく使用されます。

 A = te.placeholder((n,), name="A") B = te.compute(A.shape, lambda i: A[i] * 2, name="B") s = te.create_schedule(B.op) bx, tx = s[B].split(B.op.axis[0], factor=64) s[B].bind(bx, te.thread_axis("blockIdx.x")) s[B].bind(tx, te.thread_axis("threadIdx.x")) print(tvm.lower(s, [A, B], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*n: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*n)], [], type="auto")} buffer_map = {A_1: A, B_1: B} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [n], [stride], type="auto"), B_1: B_3: Buffer(B_2, float32, [n], [stride_1], type="auto")} { attr [IterVar(blockIdx.x: int32, (nullptr), "ThreadIndex", "blockIdx.x")] "thread_extent" = floordiv((n + 63), 64); attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 64; if @tir.likely((((blockIdx.x*64) + threadIdx.x) < n), dtype=bool) { B[(((blockIdx.x*64) + threadIdx.x)*stride_1)] = (A[(((blockIdx.x*64) + threadIdx.x)*stride)]*2f32) } }

計算場所

複数の演算子を含むスケジュールの場合、TVM はデフォルトでルートのテンソルを個別に計算します。

 A = te.placeholder((m,), name="A") B = te.compute((m,), lambda i: A[i] + 1, name="B") C = te.compute((m,), lambda i: B[i] * 2, name="C") s = te.create_schedule(C.op) print(tvm.lower(s, [A, B, C], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle, C_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto"), C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B, C_1: C} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m], [stride], type="auto"), B_1: B_3: Buffer(B_2, float32, [m], [stride_1], type="auto"), C_1: C_3: Buffer(C_2, float32, [m], [stride_2], type="auto")} { for (i: int32, 0, m) { B[(i*stride_1)] = (A[(i*stride)] + 1f32) } for (i_1: int32, 0, m) { C[(i_1*stride_2)] = (B[(i_1*stride_1)]*2f32) } }

compute_at 、B の計算を C の計算の最初の軸に移動できます。

 A = te.placeholder((m,), name="A") B = te.compute((m,), lambda i: A[i] + 1, name="B") C = te.compute((m,), lambda i: B[i] * 2, name="C") s = te.create_schedule(C.op) s[B].compute_at(s[C], C.op.axis[0]) print(tvm.lower(s, [A, B, C], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle, C_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto"), C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B, C_1: C} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m], [stride], type="auto"), B_1: B_3: Buffer(B_2, float32, [m], [stride_1], type="auto"), C_1: C_3: Buffer(C_2, float32, [m], [stride_2], type="auto")} { for (i: int32, 0, m) { B[(i*stride_1)] = (A[(i*stride)] + 1f32) C[(i*stride_2)] = (B[(i*stride_1)]*2f32) } }

インライン計算

compute_inlineステージをインラインとしてマークし、計算ボリュームを拡張して、テンソルが必要なアドレスに挿入します。

 A = te.placeholder((m,), name="A") B = te.compute((m,), lambda i: A[i] + 1, name="B") C = te.compute((m,), lambda i: B[i] * 2, name="C") s = te.create_schedule(C.op) s[B].compute_inline() print(tvm.lower(s, [A, B, C], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle, C_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto"), C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B, C_1: C} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m], [stride], type="auto"), B_1: B_3: Buffer(B_2, float32, [m], [stride_1], type="auto"), C_1: C_3: Buffer(C_2, float32, [m], [stride_2], type="auto")} { for (i: int32, 0, m) { C[(i*stride_2)] = ((A[(i*stride)] + 1f32)*2f32) } }

計算ルート

compute_rootステージの計算をルートに移動できます。

 A = te.placeholder((m,), name="A") B = te.compute((m,), lambda i: A[i] + 1, name="B") C = te.compute((m,), lambda i: B[i] * 2, name="C") s = te.create_schedule(C.op) s[B].compute_at(s[C], C.op.axis[0]) s[B].compute_root() print(tvm.lower(s, [A, B, C], simple_mode=True))

出力結果:

 @main = primfn(A_1: handle, B_1: handle, C_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {A: Buffer(A_2: Pointer(float32), float32, [(stride: int32*m: int32)], [], type="auto"), B: Buffer(B_2: Pointer(float32), float32, [(stride_1: int32*m)], [], type="auto"), C: Buffer(C_2: Pointer(float32), float32, [(stride_2: int32*m)], [], type="auto")} buffer_map = {A_1: A, B_1: B, C_1: C} preflattened_buffer_map = {A_1: A_3: Buffer(A_2, float32, [m], [stride], type="auto"), B_1: B_3: Buffer(B_2, float32, [m], [stride_1], type="auto"), C_1: C_3: Buffer(C_2, float32, [m], [stride_2], type="auto")} { for (i: int32, 0, m) { B[(i*stride_1)] = (A[(i*stride)] + 1f32) } for (i_1: int32, 0, m) { C[(i_1*stride_2)] = (B[(i_1*stride_1)]*2f32) } }

まとめ

このチュートリアルでは、TVM のスケジュール プリミティブ (ユーザーが計算を簡単かつ柔軟にスケジュールできるようにする) を紹介します。

カーネルのパフォーマンスを向上させるための一般的なワークフローは次のとおりです。

• 一連の操作を通じて計算を説明します。
• プリミティブを使用して計算をスケジュールします。
• コードをコンパイルして実行し、パフォーマンスの違いを確認します。
• 結果に基づいてスケジュールを調整します。

Pythonソースコードをダウンロードする: schedule_primitives.py

Jupyter Notebook をダウンロード: schedule_primitives.ipynb