Apache TVMは、CPU、GPU、そして様々な機械学習アクセラレータチップに適した、エンドツーエンドのディープラーニング構築フレームワークです。中国語版のTVMドキュメントは、→ https://tvm.hyper.ai/ をご覧ください。
著者：Tianqi Chen
この記事では、TVMにおける縮約の方法について説明します。結合縮約演算子（sum/max/minなど）は、線形代数演算の典型的な構成要素です。

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

説明行の合計

NumPy 構文では、行の合計を計算することはB = numpy.sum(A, axis=1)
次の行は行の合計演算を記述しています。縮約式を作成するには、縮約対象の範囲を受け取るte.reduce_axisを使用して縮約軸を宣言します。`te.sum` te.sum縮約対象の式と縮約軸を受け取り、宣言された範囲内のすべての k 値の合計を計算します。
同等の C コードは次のとおりです。

 for (int i = 0; i < n; ++i) { B[i] = 0; for (int k = 0; k < m; ++k) { B[i] = B[i] + A[i][k]; } }

 n = te.var("n") m = te.var("m") A = te.placeholder((n, m), name="A") k = te.reduce_axis((0, m), "k") B = te.compute((n,), lambda i: te.sum(A[i, k], axis=k), name="B")

スケジュール仕様

スケジュールを短縮する方法はいくつかあります。まず、デフォルトのスケジュールのIRコードを印刷します。

 s = te.create_schedule(B.op) 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, m: int32], [stride, stride_2: int32], type="auto"), B_1: B_3: Buffer(B_2, float32, [n], [stride_1], type="auto")} { for (i: int32, 0, n) { B[(i*stride_1)] = 0f32 for (k: int32, 0, m) { B[(i*stride_1)] = (B[(i*stride_1)] + A[((i*stride) + (k*stride_2))]) } } }

IR コードは C コードと非常に似ており、縮小軸は通常の軸と似ており、分割できます。
次のコードは、B の行軸と軸を異なる係数で分割し、ネストされた削減を実現します。

 ko, ki = s[B].split(B.op.reduce_axis[0], factor=16) 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*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, m: int32], [stride, stride_2: int32], type="auto"), B_1: B_3: Buffer(B_2, float32, [n], [stride_1], type="auto")} { for (i.outer: int32, 0, floordiv((n + 31), 32)) { for (i.inner: int32, 0, 32) { if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) { B[(((i.outer*32) + i.inner)*stride_1)] = 0f32 } if @tir.likely((((i.outer*32) + i.inner) < n), dtype=bool) { for (k.outer: int32, 0, floordiv((m + 15), 16)) { for (k.inner: int32, 0, 16) { if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) { let cse_var_1: int32 = ((i.outer*32) + i.inner) B[(cse_var_1*stride_1)] = (B[(cse_var_1*stride_1)] + A[((cse_var_1*stride) + (((k.outer*16) + k.inner)*stride_2))]) } } } } } } }

行 B を GPU スレッドにバインドして GPU カーネルを構築します。

 s[B].bind(xo, te.thread_axis("blockIdx.x")) s[B].bind(xi, 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, m: int32], [stride, stride_2: int32], 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 + 31), 32); attr [IterVar(threadIdx.x: int32, (nullptr), "ThreadIndex", "threadIdx.x")] "thread_extent" = 32 { if @tir.likely((((blockIdx.x*32) + threadIdx.x) < n), dtype=bool) { B[(((blockIdx.x*32) + threadIdx.x)*stride_1)] = 0f32 } for (k.outer: int32, 0, floordiv((m + 15), 16)) { for (k.inner: int32, 0, 16) { if @tir.likely((((blockIdx.x*32) + threadIdx.x) < n), dtype=bool) { if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) { B[(((blockIdx.x*32) + threadIdx.x)*stride_1)] = (B[(((blockIdx.x*32) + threadIdx.x)*stride_1)] + A[((((blockIdx.x*32) + threadIdx.x)*stride) + (((k.outer*16) + k.inner)*stride_2))]) } } } } } }

因数分解と並列化の削減

リダクションを構築する際、単純にリダクション軸上で並列化することはできません。リダクションを分割し、ローカルなリダクション結果を配列に格納し、その後、一時配列をリダクションする必要があります。
rfactorプリミティブは上記のように計算を書き換えます。次のスケジュールでは、Bの結果が一時結果B.rfに書き込まれ、分解された次元がB.rfの最初の次元になります。

 s = te.create_schedule(B.op) ko, ki = s[B].split(B.op.reduce_axis[0], factor=16) BF = s.rfactor(B, ki) 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, m: int32], [stride, stride_2: int32], type="auto"), B_1: B_3: Buffer(B_2, float32, [n], [stride_1], type="auto")} { allocate(B.rf: Pointer(global float32), float32, [(n*16)]), storage_scope = global { for (k.inner: int32, 0, 16) { for (i: int32, 0, n) { B.rf_1: Buffer(B.rf, float32, [(16*n)], [])[((k.inner*n) + i)] = 0f32 for (k.outer: int32, 0, floordiv((m + 15), 16)) { if @tir.likely((((k.outer*16) + k.inner) < m), dtype=bool) { B.rf_1[((k.inner*n) + i)] = (B.rf_1[((k.inner*n) + i)] + A[((i*stride) + (((k.outer*16) + k.inner)*stride_2))]) } } } } for (ax0: int32, 0, n) { B[(ax0*stride_1)] = 0f32 for (k.inner.v: int32, 0, 16) { B[(ax0*stride_1)] = (B[(ax0*stride_1)] + B.rf_1[((k.inner.v*n) + ax0)]) } } } }

B のスケジューリング演算子は、最初の軸上の Bf の削減結果の合計として書き換えられます。

 print(s[B].op.body)

出力結果:

 [reduce(combiner=comm_reducer(result=[(x + y)], lhs=[x], rhs=[y], identity_element=[0f]), source=[B.rf[k.inner.v, ax0]], init=[], axis=[iter_var(k.inner.v, range(min=0, ext=16))], where=(bool)1, value_index=0)]

クロススレッド規約

次に、因子軸上で並列化を実行します。この場合、B の縮約軸はスレッドとしてマークされます。TVM では、デバイス内のスレッド間で縮約できる縮約軸が 1 つだけである場合、縮約軸をスレッドとしてマークできます。
あるいは、ブロック折り畳み（BF）を縮約軸上で直接計算することもできます。最終的に生成されるカーネルは、行をblockIdx.xで割り、threadIdx.y列をthreadIdx.xで割り、最後にthreadIdx.xに対してクロススレッド縮約を実行します。

 xo, xi = s[B].split(s[B].op.axis[0], factor=32) s[B].bind(xo, te.thread_axis("blockIdx.x")) s[B].bind(xi, te.thread_axis("threadIdx.y")) tx = te.thread_axis("threadIdx.x") s[B].bind(s[B].op.reduce_axis[0], tx) s[BF].compute_at(s[B], s[B].op.reduce_axis[0]) s[B].set_store_predicate(tx.var.equal(0)) fcuda = tvm.build(s, [A, B], "cuda") print(fcuda.imported_modules[0].get_source())

出力結果:

 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 700) #define __shfl_sync(mask, var, lane, width) \ __shfl((var), (lane), (width)) #define __shfl_down_sync(mask, var, offset, width) \ __shfl_down((var), (offset), (width)) #define __shfl_up_sync(mask, var, offset, width) \ __shfl_up((var), (offset), (width)) #endif #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__(512) default_function_kernel0(float* __restrict__ A, float* __restrict__ B, int m, int n, int stride, int stride1, int stride2) { float B_rf[1]; float red_buf0[1]; B_rf[0] = 0.000000e+00f; for (int k_outer = 0; k_outer < (m >> 4); ++k_outer) { if (((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) < n) { B_rf[0] = (B_rf[0] + A[((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride) + (((k_outer * 16) + ((int)threadIdx.x)) * stride1))]); } } for (int k_outer1 = 0; k_outer1 < (((m & 15) + 15) >> 4); ++k_outer1) { if (((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) < n) { if (((((m >> 4) * 16) + (k_outer1 * 16)) + ((int)threadIdx.x)) < m) { B_rf[0] = (B_rf[0] + A[((((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride) + (((((m >> 4) * 16) + (k_outer1 * 16)) + ((int)threadIdx.x)) * stride1))]); } } } uint mask[1]; float t0[1]; red_buf0[0] = B_rf[0]; mask[0] = (__activemask() & ((uint)(65535 << (((int)threadIdx.y) * 16)))); t0[0] = __shfl_down_sync(mask[0], red_buf0[0], 8, 32); red_buf0[0] = (red_buf0[0] + t0[0]); t0[0] = __shfl_down_sync(mask[0], red_buf0[0], 4, 32); red_buf0[0] = (red_buf0[0] + t0[0]); t0[0] = __shfl_down_sync(mask[0], red_buf0[0], 2, 32); red_buf0[0] = (red_buf0[0] + t0[0]); t0[0] = __shfl_down_sync(mask[0], red_buf0[0], 1, 32); red_buf0[0] = (red_buf0[0] + t0[0]); red_buf0[0] = __shfl_sync(mask[0], red_buf0[0], (((int)threadIdx.y) * 16), 32); if (((int)threadIdx.x) == 0) { B[(((((int)blockIdx.x) * 32) + ((int)threadIdx.y)) * stride2)] = red_buf0[0]; } }

結果のカーネルは NumPy と比較され、その正確性が検証されます。

 nn = 128 dev = tvm.cuda(0) a = tvm.nd.array(np.random.uniform(size=(nn, nn)).astype(A.dtype), dev) b = tvm.nd.array(np.zeros(nn, dtype=B.dtype), dev) fcuda(a, b) tvm.testing.assert_allclose(b.numpy(), np.sum(a.numpy(), axis=1), rtol=1e-4)

2次元縮約を用いた畳み込みの記述

TVMでは、畳み込みは単純な2次元縮小（フィルタサイズ= [3, 3]、ストライド= [1, 1]）を使用して記述されます。

 n = te.var("n") Input = te.placeholder((n, n), name="Input") Filter = te.placeholder((3, 3), name="Filter") di = te.reduce_axis((0, 3), name="di") dj = te.reduce_axis((0, 3), name="dj") Output = te.compute( (n - 2, n - 2), lambda i, j: te.sum(Input[i + di, j + dj] * Filter[di, dj], axis=[di, dj]), name="Output", ) s = te.create_schedule(Output.op) print(tvm.lower(s, [Input, Filter, Output], simple_mode=True))

出力結果:

 @main = primfn(Input_1: handle, Filter_1: handle, Output_1: handle) -> () attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True} buffers = {Input: Buffer(Input_2: Pointer(float32), float32, [(stride: int32*n: int32)], [], type="auto"), Filter: Buffer(Filter_2: Pointer(float32), float32, [9], []), Output: Buffer(Output_2: Pointer(float32), float32, [((n - 2)*(n - 2))], [])} buffer_map = {Input_1: Input, Filter_1: Filter, Output_1: Output} preflattened_buffer_map = {Input_1: Input_3: Buffer(Input_2, float32, [n, n], [stride, stride_1: int32], type="auto"), Filter_1: Filter_3: Buffer(Filter_2, float32, [3, 3], []), Output_1: Output_3: Buffer(Output_2, float32, [(n - 2), (n - 2)], [])} { for (i: int32, 0, (n - 2)) { for (j: int32, 0, (n - 2)) { Output[((i*(n - 2)) + j)] = 0f32 for (di: int32, 0, 3) { for (dj: int32, 0, 3) { Output[((i*(n - 2)) + j)] = (Output[((i*(n - 2)) + j)] + (Input[(((i + di)*stride) + ((j + dj)*stride_1))]*Filter[((di*3) + dj)])) } } } } }

一般的な可換縮約演算を定義します。

te.sum 、 tvm.te.min 、 tvm.te.maxなどの組み込み削減操作に加えて、 te.comm_reducerを通じて交換削減操作を定義することもできます。

 n = te.var("n") m = te.var("m") product = te.comm_reducer(lambda x, y: x * y, lambda t: tvm.tir.const(1, dtype=t), name="product") A = te.placeholder((n, m), name="A") k = te.reduce_axis((0, m), name="k") B = te.compute((n,), lambda i: product(A[i, k], axis=k), name="B")

まとめ

このチュートリアルでは、スケジュールを設定する方法を説明します。
削減軸は削減を記述するために使用されます。
並列処理の場合は、rfactor を使用して軸を分解します。
te.comm_reducerを使用して新しい削減操作を定義します。