Deep Learning 向けコンパイラ TVM の紹介 11/10 コンパイラ勉強会 増田正博

自己紹介 • 名前: 増田正博 • コンピュータビジョン・グラフィックスや, 周辺の技術に興味 • Fixstars でインターン経験あり • AMDGPU バックエンドの開発から, TVM にコントリビュートするように • TVM 関係の記事もいくつか書きました https://qiita.com/masahi オレが始めた

2017 ～ Deep Learning コンパイラ? Deep Learning コンパイラ 2015 ～ 現在

前置き • Deep Learning における, 学習済みモデルの推論のみを扱います • 入力として有向グラフと多次元配列 (画像など) が与えられる • ノード: なんらかの演算を表す (畳み込み, 行列積, シンプルな要素ごとの演算 etc) • エッジ: 演算間のデータフローを表す. 多次元配列 (Tensor) が流れていくイメージ GoogleNet, from Going Deeper with Convolutions 計算グラフ ( Computational Graph)

本日のテーマ Root から各ノードの演算を実行し, 出力を得る. これをとにかく速くやりたい DL 用途に特化した (広い意味での) 最適化コンパイラ → • 内部でどのような処理をしているかをざっくり紹介 • TVM をどのように使うか, などの説明は最小限 高速化のポイント ・各ノードに対応する演算 (オペレータ) の最適化 ・不必要なノードを刈る, 隣接するノードを一つにまとめる, などのグラフ構造の最適化 TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

Contents • TVM の概要 • オペレータレベルの最適化 • グラフレベルの最適化 • パフォーマンス • デプロイ用のツールとしての側面 • オープンソースプロジェクトとしての側面

TVM とは? • DL フレームワークで学習したモデルをデプロイしたい • 各ハードウェアバックエンド上で高速に推論実行したい そのための ドメイン特化言語 コンパイラ ランタイムライブラリ https://tvm.ai/about

TVM と NNVM CUDA ONNX MXNet Tensorflow Keras グラフレベル + オペレータレベル 最適化 LLVM OpenCL x86 ARM AMDGPU NNVM グラフ全体を扱う TVM オペレータ(各ノード)を扱う

LLVM とのアナロジー • フロントエンド → IRの最適化 → バックエンドコード生成 の流れは同じ From “The Architecture of Open Source Applications: LLVM” by Chris Lattner

Halide の影響 • “アルゴリズムとスケジュールの分離” を踏襲 • Halide と同じ IR を使用している. メモリ上の表現はほぼ同じ Func sample(Func A, Func B){ Func f; Var x, y; f(x, y) = A(x, y) + B(x, y); // schedule f.parallel(y).vectorize(x, 8); return f; } def sample(A, B): f = tvm.compute(A.shape, lambda y, x: A[y, x] + B[y, x]) # schedule sch = tvm.create_schedule(f.op) y, x = sch[f].op.axis sch[f].parallel(y) xo, xi = sch[f].split(x, factor=8) sch[f].vectorize(xi) return f TVM による記述 Halide による記述 詳しくは Halide による画像処理プログラミング入門 Prelude to Halide

核となるアイデア • cuDNN などの “手で最適化された” コードは使わない. 自分でコード生成する • 出力されるコードのひな形を登録 • 登録されたひな形から, 各バックエンド向けに Auto Tuning + コード生成 def matmul(A, B): M, K = A.shape K, N = B.shape k = tvm.reduce_axis((0, K), name='k’) C = tvm.compute((M, N), lambda i, j: tvm.sum(A[i, k] * B[k, j], axis=k), name="C") return C def schedule_matmul(C, s): y, x = C.op.axis k, = s[C].op.reduce_axis yo, xo, yi, xi = s[C].tile(y, x, 16, 16) ko, ki = s[C].split(k, factor=8) s[C].reorder(yo, xo, ko, yi, ki, xi) fused = s[C].fuse(yo, xo) s[C].parallel(fused) s[C].vectorize(xi) 計算の定義 (バックエンド非依存) CPU 向けスケジュール

Contents • TVM の概要 • オペレータレベルの最適化 • グラフレベルの最適化 • パフォーマンス • デプロイ用のツールとしての側面 • オープンソースプロジェクトとしての側面

オペレータレベルの最適化 • オペレータ = 1つのレイヤー, 計算グラフのノードに相当 • Convolution, Pooling, Batch normalization など スケジューリング 計算の定義 Lowering コード生成 TVM IR → LLVM IR LLVM による最適化 CUDA Kernel 生成 より抽象度の低い IR に変換 ループ範囲・バッ ファサイズの推論 ループ構造など を指定 各ターゲットごと に記述 Auto Tuning 純粋関数 限られた記述力 ターゲット非依存

計算の定義 def matmul(A, B): M, K = A.shape K, N = B.shape k = tvm.reduce_axis((0, K), name='k’) C = tvm.compute((M, N), lambda i, j: tvm.sum(A[i, k] * B[k, j], axis=k), name="C") return C • 行列積の TVM による定義 • CPU, GPU バックエンドともに同じ定義を使用 • 記述できるものは, 行列の演算や畳み込みなどに限られる • DL 用途としてはこれで十分

TVM IR を見てみる M = N = K = 512 A = tvm.placeholder((M, K), name='A') B = tvm.placeholder((K, N) , name='B') C = matmul(A, B) s = tvm.create_schedule(C.op) print(tvm.lower(s, [A, B, C], simple_mode=True)) produce C { for (i, 0, 512) { for (j, 0, 512) { C[((i*512) + j)] = 0.000000f for (k, 0, 512) { C[((i*512) + j)] = (C[((i*512) + j)] + (A[((i*512) + k)]*B[(j + (k*512))])) } } } } デフォルトのスケジュール 計算の定義からループのネスト 構造を自動生成

スケジューリング • 定義した計算の意味を変えずに, ループ構造・種類を変更 tile_size = 16 y, x = C.op.axis k, = s[C].op.reduce_axis yo, xo, yi, xi = s[C].tile(y, x, tile_size, tile_size) ko, ki = s[C].split(k, factor=8) s[C].reorder(yo, xo, ko, yi, ki, xi) fused = s[C].fuse(yo, xo) s[C].parallel(fused) s[C].vectorize(xi) tile, split, reorder, fuse, vectorize, parallel は Halide 由来 16 x 16 のブロックに分割 内積をとる軸を分割し, ループの順序を変更 外側ブロック数についてのループを並列化 最内ループをベクトル化

スケジューリング適用後 produce C { parallel (i.outer.j.outer.fused, 0, 1024) { for (i.inner.init, 0, 16) { C[ramp(((((i.outer.j.outer.fused/32)*512) + ((i.inner.init*32) + (i.outer.j.outer.fused % 32)))*16), 1, 16)] = x16(0.000000f) } for (k.outer, 0, 64) { for (i.inner, 0, 16) { for (k.inner, 0, 8) { C[ramp(((((i.outer.j.outer.fused/32)*512) + ((i.inner*32) + (i.outer.j.outer.fused % 32)))*16), 1, 16)] = (C[ramp(((((i.outer.j.outer.fused/32)*512) + ((i.inner*32) + (i.outer.j.outer.fused % 32)))*16), 1, 16)] + (x16(A[((((((i.outer.j.outer.fused/32)*1024) + k.outer) + (i.inner*64))*8) + k.inner)])*B[ramp(((((k.inner*32) + (i.outer.j.outer.fused % 32)) + (k.outer*256))*16), 1, 16)])) } } } } }

GPU へのマッピング tile_size = 16 yo, xo, yi, xi = s[C].tile(y, x, tile_size, tile_size) s[C].bind(yo, tvm.thread_axis("blockIdx.y")) s[C].bind(xo, tvm.thread_axis("blockIdx.x")) s[C].bind(yi, tvm.thread_axis("threadIdx.y")) s[C].bind(xi, tvm.thread_axis("threadIdx.x")) produce C { // attr [iter_var(blockIdx.y, , blockIdx.y)] thread_extent = 32 // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 32 // attr [iter_var(threadIdx.y, , threadIdx.y)] thread_extent = 16 // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 16 C[(((((blockIdx.y*512) + blockIdx.x) + (threadIdx.y*32))*16) + threadIdx.x)] = 0.000000f for (k, 0, 512) { C[(((((blockIdx.y*512) + blockIdx.x) + (threadIdx.y*32))*16) + threadIdx.x)] = (C[(((((blockIdx.y*512) + blockIdx.x) + (threadIdx.y*32))*16) + threadIdx.x)] + (A[((((blockIdx.y*16) + threadIdx.y)*512) + k)]*B[(((blockIdx.x*16) + threadIdx.x) + (k*512))])) } }

共有メモリの利用 num_thread = 16 k, = s[C].op.reduce_axis ko, ki = s[C].split(k, factor=num_thread) A, B = s[C].op.input_tensors AA = s.cache_read(A, "shared", [C]) BB = s.cache_read(B, "shared", [C]) s[AA].compute_at(s[C], ko) s[BB].compute_at(s[C], ko) # load one value into shared mem per thread y, x = s[AA].op.axis _, _, ty, tx = s[AA].tile(y, x, num_thread, num_thread) s[AA].bind(ty, tvm.thread_axis("threadIdx.y")) s[AA].bind(tx, tvm.thread_axis("threadIdx.x")) y, x = s[BB].op.axis _, _, ty, tx = s[BB].tile(y, x, num_thread, num_thread) s[BB].bind(ty, tvm.thread_axis("threadIdx.y")) s[BB].bind(tx, tvm.thread_axis("threadIdx.x")) http://www.es.ele.tue.nl/~mwijtvliet/5KK73/?page=mmcuda)

共有メモリの利用 produce C { // attr [iter_var(blockIdx.y, , blockIdx.y)] thread_extent = 64 // attr [A.shared] storage_scope = "shared" allocate A.shared[float32 * 16 * 16] // attr [B.shared] storage_scope = “shared” allocate B.shared[float32 * 16 * 16] // attr [iter_var(blockIdx.x, , blockIdx.x)] thread_extent = 64 // attr [iter_var(threadIdx.y, , threadIdx.y)] thread_extent = 16 // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 16 C[(((((blockIdx.y*1024) + blockIdx.x) + (threadIdx.y*64))*16) + threadIdx.x)] = 0.000000f for (k.outer, 0, 64) { produce A.shared { // attr [iter_var(threadIdx.y, , threadIdx.y)] thread_extent = 16 // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 16 A.shared[((threadIdx.y*16) + threadIdx.x)] = A[(((((blockIdx.y*1024) + k.outer) + (threadIdx.y*64))*16) + threadIdx.x)] } produce B.shared { // attr [iter_var(threadIdx.y, , threadIdx.y)] thread_extent = 16 // attr [iter_var(threadIdx.x, , threadIdx.x)] thread_extent = 16 B.shared[((threadIdx.y*16) + threadIdx.x)] = B[((((blockIdx.x + (k.outer*1024)) + (threadIdx.y*64))*16) + threadIdx.x)] } for (k.inner, 0, 16) { C[(((((blockIdx.y*1024) + blockIdx.x) + (threadIdx.y*64))*16) + threadIdx.x)] = (C[(((((blockIdx.y*1024) + blockIdx.x) + (threadIdx.y*64))*16) + threadIdx.x)] + (A.shared[((threadIdx.y*16) + k.inner)]*B.shared[(threadIdx.x + (k.inner*16))])) } } }

Lowering からコード生成へ ループ範囲・バッファサイズの推論 ループのベクトル化, アンローリング 再利用できるバッファの検出 共有メモリ, ローカルメモリの割り当て ダブルバッファリング スレッド同期命令の挿入 Lowering コード生成 NVRTC LLVM TVM IR TVM IR から LLVM IR へ変換 または, 直接 CUDA カーネル生成

コード生成: AVX 2 M = N = K = 1024 A = tvm.placeholder((M, K), name='A') B = tvm.placeholder((K, N) , name='B') C = matmul(A, B) s = schedule(C) target = "llvm -mcpu=core-avx2" func = tvm.build(s, [A, B, C], target=target) print(func.get_source("asm")) .LBB2_4: leal (%rdx,%rsi), %edi movslq %edi, %rdi vbroadcastss -28(%rbx), %ymm4 vmovaps 32(%rcx,%rdi,4), %ymm5 vfmadd231ps 64(%rsp), %ymm4, %ymm5 vmovups 32(%rsp), %ymm6 vfmadd213ps (%rcx,%rdi,4), %ymm6, %ymm4 vbroadcastss -24(%rbx), %ymm6 vfmadd231ps %ymm7, %ymm6, %ymm4 vfmadd132ps (%rsp), %ymm5, %ymm6 vbroadcastss -20(%rbx), %ymm5 vfmadd231ps %ymm8, %ymm5, %ymm6 vfmadd213ps %ymm4, %ymm9, %ymm5 vbroadcastss -16(%rbx), %ymm4 vfmadd231ps %ymm11, %ymm4, %ymm5 vfmadd213ps %ymm6, %ymm10, %ymm4 vbroadcastss -12(%rbx), %ymm6 vfmadd231ps %ymm12, %ymm6, %ymm4 vfmadd213ps %ymm5, %ymm13, %ymm6 vbroadcastss -8(%rbx), %ymm5 vfmadd231ps %ymm15, %ymm5, %ymm6 vfmadd213ps %ymm4, %ymm14, %ymm5 vbroadcastss -4(%rbx), %ymm4 vfmadd231ps %ymm1, %ymm4, %ymm5 vfmadd213ps %ymm6, %ymm2, %ymm4 vbroadcastss (%rbx), %ymm6 vfmadd231ps %ymm3, %ymm6, %ymm4 vfmadd213ps %ymm5, %ymm0, %ymm6 vmovaps %ymm6, 32(%rcx,%rdi,4) vmovaps %ymm4, (%rcx,%rdi,4) addq $1024, %rsi addq $4096, %rbx cmpq $16384, %rsi jne .LBB2_4

コード生成: NEON target = tvm.target.arm_cpu(model="rasp3b") func = tvm.build(s, [A, B, C], target=target) print(func.get_source("asm")) vld1.32 {d18, d19}, [r3:128]! vld1.32 {d26[], d27[]}, [r2:32]! vld1.64 {d22, d23}, [r0:128] vld1.64 {d24, d25}, [r6:128] vmla.f32 q9, q13, q15 vldmia r4, {d30, d31} vld1.64 {d20, d21}, [r3:128] vld1.32 {d28[], d29[]}, [r2:32] vmla.f32 q11, q13, q7 vmla.f32 q12, q13, q8 add r2, r1, #8 add r4, sp, #32 vmla.f32 q10, q13, q15 vld1.32 {d26[], d27[]}, [r2:32] vmla.f32 q12, q14, q4 vmla.f32 q9, q14, q6 vmla.f32 q11, q14, q3 add r2, r1, #12 vmla.f32 q12, q13, q0 vmla.f32 q10, q14, q5 vldmia r4, {d28, d29} add r4, sp, #96 vmla.f32 q9, q13, q2 vmla.f32 q10, q13, q1 vmla.f32 q11, q13, q14 vld1.32 {d28[], d29[]}, [r2:32] vldmia r4, {d26, d27} add r4, sp, #80 add r2, r1, #16 vmla.f32 q11, q14, q13 vldmia r4, {d26, d27} ..... ターゲットを変えるだけで, 全く違った アーキテクチャ用のコード生成が可能

コード生成: CUDA s = schedule_shared_mem(C) target = "cuda" func = tvm.build(s, [A, B, C], target=target) print(func.imported_modules[0].get_source()) extern "C" __global__ void default_function_kernel0( float* __restrict__ C, float* __restrict__ A, float* __restrict__ B) { __shared__ float A_shared[256]; __shared__ float B_shared[256]; C[((((((int)blockIdx.y) * 16384) + (((int)blockIdx.x) * 16)) + (((int)threadIdx.y) * 1024)) + ((int)threadIdx.x))] = 0.000000e+00f; for (int k_outer = 0; k_outer < 64; ++k_outer) { __syncthreads(); A_shared[((((int)threadIdx.y) * 16) + ((int)threadIdx.x))] = A[((((((int)blockIdx.y) * 16384) + (k_outer * 16)) + (((int)threadIdx.y) * 1024)) + ((int)threadIdx.x))]; B_shared[((((int)threadIdx.y) * 16) + ((int)threadIdx.x))] = B[((((((int)blockIdx.x) * 16) + (k_outer * 16384)) + (((int)threadIdx.y) * 1024)) + ((int)threadIdx.x))]; __syncthreads(); for (int k_inner = 0; k_inner < 16; ++k_inner) { C[((((((int)blockIdx.y) * 16384) + (((int)blockIdx.x) * 16)) + (((int)threadIdx.y) * 1024)) + ((int)threadIdx.x))] = (C[((((((int)blockIdx.y) * 16384) + (((int)blockIdx.x) * 16)) + (((int)threadIdx.y) * 1024)) + ((int)threadIdx.x))] + (A_shared[((((int)threadIdx.y) * 16) + k_inner)] * B_shared[(((int)threadIdx.x) + (k_inner * 16))])); } } }

コード生成: AMDGPU for_body: ; preds = %for_body, %entry %.promoted = phi float [ 0.000000e+00, %entry ], [ %147, %for_body ] %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for_body ] tail call void @llvm.amdgcn.s.barrier() %84 = trunc i64 %indvars.iv to i32 %85 = add i32 %9, %84 %86 = shl i32 %85, 4 %87 = add nsw i32 %86, %6 %88 = sext i32 %87 to i64 %89 = getelementptr inbounds float, float addrspace(1)* %1, i64 %88 %90 = bitcast float addrspace(1)* %89 to i32 addrspace(1)* %91 = load i32, i32 addrspace(1)* %90, align 4, !tbaa !8 store i32 %91, i32 addrspace(3)* %18, align 4, !tbaa !11 %indvars.iv.tr = trunc i64 %indvars.iv to i32 %92 = shl i32 %indvars.iv.tr, 10 %93 = add i32 %19, %92 %94 = shl i32 %93, 4 %95 = add nsw i32 %94, %6 %96 = sext i32 %95 to i64 %97 = getelementptr inbounds float, float addrspace(1)* %2, i64 %96 %98 = bitcast float addrspace(1)* %97 to i32 addrspace(1)* %99 = load i32, i32 addrspace(1)* %98, align 4, !tbaa !14 store i32 %99, i32 addrspace(3)* %21, align 4, !tbaa !17 tail call void @llvm.amdgcn.s.barrier() %100 = load float, float addrspace(3)* %22, align 16, !tbaa !11 %101 = load float, float addrspace(3)* %23, align 4, !tbaa !17 %102 = tail call float @llvm.fmuladd.f32(float %100, float %101, float %.promoted) %103 = load float, float addrspace(3)* %25, align 4, !tbaa !11 %104 = load float, float addrspace(3)* %27, align 4, !tbaa !17 %105 = tail call float @llvm.fmuladd.f32(float %103, float %104, float %102) %106 = load float, float addrspace(3)* %29, align 8, !tbaa !11 %107 = load float, float addrspace(3)* %31, align 4, !tbaa !17 %108 = tail call float @llvm.fmuladd.f32(float %106, float %107, float %105) ..... BB0_1: v_add_u32_e32 v11, vcc, s6, v0 v_ashrrev_i32_e32 v12, 31, v11 v_ashrrev_i32_e32 v10, 31, v9 v_lshlrev_b64 v[11:12], 2, v[11:12] v_lshlrev_b64 v[13:14], 2, v[9:10] v_mov_b32_e32 v1, s1 v_add_u32_e32 v10, vcc, s0, v11 v_addc_u32_e32 v11, vcc, v1, v12, vcc v_mov_b32_e32 v15, s3 v_add_u32_e32 v12, vcc, s2, v13 v_addc_u32_e32 v13, vcc, v15, v14, vcc flat_load_dword v1, v[10:11] flat_load_dword v10, v[12:13] s_waitcnt vmcnt(0) lgkmcnt(0) s_barrier s_add_u32 s6, s6, 16 s_addc_u32 s7, s7, 0 v_add_u32_e32 v9, vcc, 0x4000, v9 s_cmp_lg_u64 s[6:7], s[4:5] ds_write_b32 v6, v1 ds_write_b32 v5, v10 s_waitcnt lgkmcnt(0) s_barrier ds_read2_b32 v[18:19], v8 offset1:16 ds_read2_b64 v[10:13], v7 offset1:1 ds_read2_b64 v[14:17], v7 offset0:2 offset1:3 s_waitcnt lgkmcnt(1) v_mac_f32_e32 v4, v10, v18 v_mac_f32_e32 v4, v11, v19 ds_read2_b32 v[10:11], v8 offset0:32 offset1:48 s_waitcnt lgkmcnt(0) v_mac_f32_e32 v4, v12, v10 v_mac_f32_e32 v4, v13, v11 ds_read2_b32 v[10:11], v8 offset0:64 offset1:80 ds_read2_b32 v[12:13], v8 offset0:96 offset1:112 ds_read2_b32 v[18:19], v8 offset0:128 offset1:144 s_waitcnt lgkmcnt(2) v_mac_f32_e32 v4, v14, v10 v_mac_f32_e32 v4, v15, v11 s_waitcnt lgkmcnt(1) .....

AutoTVM – 学習ベースの Auto Tuning フレームワーク • 以前の TVM • 特定のネットワーク向けに手でパラメータチューニング • チューニング対象外のネットワークでは遅かった • AutoTVM によって, 高速なスケジュールを自動で生成できるようになった ブロック行列積のタイルの大きさをチューニング可能にする例 from tvm import autotvm cfg = autotvm.get_config() cfg.define_knob("tile_size", [4, 8, 16, 32]) yo, xo, yi, xi = s[C].tile(y, x, cfg['tile_size'].val, cfg['tile_size'].val) tile_size = 16 yo, xo, yi, xi = s[C].tile(y, x, tile_size, tile_size)

AutoTVM 使用例 • CUDA バックエンド Direct Convolution スケジュール • Thread block の大きさなどが チューニングされる tvm/topi/python/topi/cuda/conv2d_direct.py def schedule_direct_cuda(cfg, s, conv): n, f, y, x = s[conv].op.axis rc, ry, rx = s[conv].op.reduce_axis cfg.define_split("tile_f", f, num_outputs=4) cfg.define_split("tile_y", y, num_outputs=4) cfg.define_split("tile_x", x, num_outputs=4) cfg.define_split("tile_rc", rc, num_outputs=2) cfg.define_split("tile_ry", ry, num_outputs=2) cfg.define_split("tile_rx", rx, num_outputs=2) cfg.define_knob("auto_unroll_max_step", [0, 512, 1500]) cfg.define_knob("unroll_explicit", [0, 1]) ..... bf, vf, tf, fi = cfg["tile_f"].apply(s, output, f) by, vy, ty, yi = cfg["tile_y"].apply(s, output, y) bx, vx, tx, xi = cfg["tile_x"].apply(s, output, x) bf = s[output].fuse(n, bf) s[output].bind(bf, tvm.thread_axis("blockIdx.z")) s[output].bind(by, tvm.thread_axis("blockIdx.y")) s[output].bind(bx, tvm.thread_axis("blockIdx.x")) ..... rco, rci = cfg['tile_rc'].apply(s, OL, rc) ryo, ryi = cfg['tile_rx'].apply(s, OL, ry) rxo, rxi = cfg['tile_ry'].apply(s, OL, rx) ..... for load in [AA, WW]: n, f, y, x = s[load].op.axis fused = s[load].fuse(n, f, y, x) tz, fused = s[load].split(fused, nparts=cfg["tile_f"].size[2]) ty, fused = s[load].split(fused, nparts=cfg["tile_y"].size[2]) tx, fused = s[load].split(fused, nparts=cfg["tile_x"].size[2]) ..... s[output].pragma(kernel_scope, 'auto_unroll_max_step', cfg['auto_unroll_max_step'].val) s[output].pragma(kernel_scope, 'unroll_explicit', cfg['unroll_explicit'].val) AutoTVM についての詳細は, チュートリアル 及び 論文を参照

TVM による高速化 ケーススタディ • CPU 行列積 • CUDA Direct Convolution • CUDA Winograd Convolution

Contents • TVM の概要 • オペレータレベルの最適化 • グラフレベルの最適化 • パフォーマンス • デプロイ用のツールとしての側面 • オープンソースプロジェクトとしての側面

グラフレベルの最適化 • TVM の上位レイヤーにあたる, NNVM が担当 • WebDNN の発表資料もおすすめ • グラフ最適化がとても詳しく説明されている Graph Transformation モデルインポート 各オペレータを TVM でコンパイル ここまでの話 ここからの話

グラフコンパイル ワークフロー • DL フレームワークから直接, または ONNX 経由でモデルをインポート • ターゲット, 入力のサイズを指定して, グラフをコンパイル import mxnet as mx sym_mx, params_mx = mx.load_checkpoint(...) # load a trained model import nnvm net, params = nnvm.frontend.from_mxnet(sym_mx, params_mx) target = "cuda" x = np.zeros((1, 3, 224, 224)).astype(np.float32) shape_dict = {'data': x.shape} graph, lib, params = nnvm.compiler.build(net, target, shape_dict, params=params)

推論実行 ctx = tvm.context(target, 0) module = runtime.create(graph, lib, ctx) module.set_input(**params) module.set_input("data", x) module.run() out = module.get_output(0) • 生成されたカーネルとデバイスごとのランタイムを繋ぎ合わせて実行

NNVM によるネットワーク構造の定義 • 通常は, 直接 NNVM API でモデルを定義することはない • TensorFlow などと同様の, データフロー記述 import nnvm.symbol as sym def get_net(filters=8): data = sym.Variable(name="data") data = sym.conv2d(data, kernel_size=(3,3), channels=filters) data = sym.batch_norm(data) data = sym.relu(data) data = sym.max_pool2d(data, pool_size=(2, 2), strides=(2, 2)) data = sym.flatten(data) data = sym.dense(data, units=1000) data = sym.softmax(data, axis=1) return data nnvm.sym.conv2d に対して, TVM による ・計算の定義 ・各バックエンド用のスケジュール が紐づけられている 他のシンボルについても同様

Graph(%data, %conv2d0_weight, %conv2d0_bias, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, %dense0_weight, %dense0_bias) { %5 = tvm_op(%data, %conv2d0_weight, %conv2d0_bias, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, flatten_data='0', func_name='fuse_conv2d_broadcast_mul_broadcast_add_relu', num_inputs='5', num_outputs='1’) %6 = tvm_op(%5, flatten_data='0', func_name='fuse_max_pool2d', num_inputs='1', num_outputs='1’) %7 = tvm_op(%6, flatten_data='0', func_name='fuse_flatten', num_inputs='1', num_outputs='1’) %10 = tvm_op(%7, %dense0_weight, %dense0_bias, flatten_data='0', func_name='fuse_dense', num_inputs='3', num_outputs='1’) %11 = tvm_op(%10, flatten_data='0', func_name='fuse_softmax', num_inputs='1', num_outputs='1’) ret %11 } Graph IR • コンパイルしたグラフの IR をダンプ • Convolution, Batch norm, ReLU が フューズされている net = get_net() net, params = utils.create_workload(net, 1, (3, 224, 224)) graph, lib, params = nnvm.compiler.build(net, target, shape_dict, params=params) print(graph.ir())

最適化レベルを上げる • opt_level を 3 に (デフォルトは2) • Batch norm のスケーリングが消え, Convolution のパラメータが代わりに スケールされる Graph(%data, %conv2d0_weight_sc, %conv2d0_bias_sc, %batch_norm0_add_beta_expand, %dense0_weight, %dense0_bias) { %4 = tvm_op(%data, %conv2d0_weight_sc, %conv2d0_bias_sc, %batch_norm0_add_beta_expand, num_outputs='1', num_inputs='4', flatten_data='0', func_name='fuse_conv2d_broadcast_add_relu’) %5 = tvm_op(%4, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_max_pool2d’) %6 = tvm_op(%5, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_flatten’) %9 = tvm_op(%6, %dense0_weight, %dense0_bias, num_outputs='1', num_inputs='3', flatten_data='0', func_name='fuse_dense’) %10 = tvm_op(%9, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_softmax’) ret %10 } with nnvm.compiler.build_config(opt_level=3): graph, lib, params = nnvm.compiler.build(net, target, shape_dict, params=params)

Winograd アルゴリズム を使う • 入力サイズとフィルター数を, Winograd Convolution に適したものに変更 • Direct Convolution が Winograd Convolution に置き換わる Graph(%data, %transpose0, %conv2d0_bias, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, %dense0_weight,%dense0_bias) { %5 = tvm_op(%data, %transpose0, %conv2d0_bias, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, num_outputs='1', num_inputs='5', flatten_data='0', func_name='fuse__contrib_conv2d_winograd_without_weight_transform_broadcast_mul_broadcast_add_relu’) %6 = tvm_op(%5, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_max_pool2d’) %7 = tvm_op(%6, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_flatten’) %10 = tvm_op(%7, %dense0_weight, %dense0_bias, num_outputs='1', num_inputs='3', flatten_data='0', func_name='fuse_dense’) %11 = tvm_op(%10, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_softmax’) ret %11 } net = get_net(filters=64) in_shape = (1, 64, 56, 56) net, params = utils.create_workload(net, 1, in_shape[1:]) with nnvm.compiler.build_config(opt_level=3): graph, lib, params = nnvm.compiler.build(net, target, {"data": in_shape}, params=params)

Winograd アルゴリズム を使う • Winograd F(4x4, 3x3) の例 6 x 6 のフィルター変換がコンパイル時に計算される print("Before precompute prune pass") for (k, v) in params.items(): print("%s:" % k, v.shape) with nnvm.compiler.build_config(opt_level=3): graph, lib, params = nnvm.compiler.build(net, target, {"data": in_shape}, params=params) print("After precompute prune pass") for (k, v) in params.items(): print("%s:" % k, v.shape) Before precompute prune pass conv2d0_weight: (64, 64, 3, 3) conv2d0_bias: (64,) ... After precompute prune pass transpose0: (6, 6, 64, 64) conv2d0_bias: (64,) ... フィルター変換の最後に transpose が起こるため、パラメータ名が変 わっている

NCHWc レイアウトで Convolution • NCHW レイアウトが標準 (N: バッチ数, C, チャネル数, H, W: 縦, 横) • CPU では, NCHWc レイアウトによる Convolution が高速 • NCHW8c, NCHW16c など, ベクトル幅の倍数でチャネルを分割 Understanding Memory Formats, MKL-DNN Documatation

レイアウト変換を有効にする • x86 バックエンドでは, opt_level = 3 の場合に NCHWc レイアウトを使う • レイアウト変換ノードを挿入, NCHWc Convolution に置き換え Graph(%data, %conv2d0_weight_OIHW8i8o, %conv2d0_bias_C8c, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, %dense0_weight, %dense0_bias) { %1 = tvm_op(%data, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse___layout_transform___2’) %6 = tvm_op(%1, %conv2d0_weight_OIHW8i8o, %conv2d0_bias_C8c, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, num_outputs='1', num_inputs='5', flatten_data='0', func_name='fuse__contrib_conv2d_NCHWc_broadcast_mul_broadcast_add_relu’) %7 = tvm_op(%6, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_max_pool2d’) %8 = tvm_op(%7, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse___layout_transform___flatten’) %11 = tvm_op(%8, %dense0_weight, %dense0_bias, num_outputs='1', num_inputs='3', flatten_data='0', func_name='fuse_dense’) %12 = tvm_op(%11, num_outputs='1', num_inputs='1', flatten_data='0', func_name='fuse_softmax’) ret %12 } target = "llvm -mcpu=core-avx2" with nnvm.compiler.build_config(opt_level=3): graph, lib, params = nnvm.compiler.build(net, target, shape_dict, params=params) print(graph.ir())

外部ライブラリの使用 • cuDNN などの, 各ベンダーが提供するライブラリを使うこともできる • 例: 全ての conv2d オペレータを cuDNN のカーネルに置き換える • Fully connected layer を cuBLAS の SGEMM に置き換えることも可能 target = "cuda –libs=cudnn" graph, lib, params = nnvm.compiler.build(net, target, shape_dict, params=params)

Graph Transformation: まとめ • レイアウト変換 • スケーリング演算のパラメータへの埋め込み • コンパイル時実行可能なノードの削除 • オペレータの置き換え (Winograd, NCHWc Convolution, cuDNN) • Operator Fusion

Operator Fusion: 概要 • 複数のオペレータ(カーネル) を一つにまとめる • 余分なメモリへのロード・ストア, GPU カーネル起動オーバーヘッドの削減 Conv2d Bias Add Batch norm ReLU Fused op

Operator Fusion: NNVM の実装 • オペレータは主に4つのグループに分かれている • グループ間でフューズ可能か決められている NN op convolution conv transpose pooling fully connected Injective reshape layout transform transpose concatenate Elementwise math op relu elemwise sum sigmoid Broadcast broadcast_add broadcast_mul … expand_dims NN op には Elementwise か Broadcast op をフューズできる NN op 以外は同じグループの op 同士で フューズできる, など

Operator Fusion: 例 • Convolution + Bias Add + Batch norm + ReLU data = sym.Variable(name="data") data = sym.conv2d(data, kernel_size=(3,3), channels=8, use_bias=True) data = sym.batch_norm(data) data = sym.relu(data) Graph(%data, %conv2d0_weight, %conv2d0_bias,%batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand) { %5 = tvm_op(%data, %conv2d0_weight, %conv2d0_bias, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, flatten_data='0', func_name='fuse_conv2d_broadcast_mul_broadcast_add_relu', num_inputs='5', num_outputs='1’) ret %5 }

Operator Fusion: CUDA カーネル for (int rc_inner = 0; rc_inner < 4; ++rc_inner) { for (int ry_inner = 0; ry_inner < 3; ++ry_inner) { for (int rx_inner = 0; rx_inner < 3; ++rx_inner) { for (int ff = 0; ff < 2; ++ff) { for (int xx = 0; xx < 2; ++xx) { compute[((ff * 2) + xx)] = (compute[((ff * 2) + xx)] + (pad_temp_shared[((((((((int)threadIdx.y) * 8) + (((int)threadIdx.x) * 2)) + (rc_inner * 32)) + (ry_inner * 8)) + rx_inner) + xx)] * input1_shared[(((((((int)threadIdx.z) * 72) + (rc_inner * 9)) + (ry_inner * 3)) + rx_inner) + (ff * 36))])); } } } } } for (int ax1_inner_inner_inner = 0; ax1_inner_inner_inner < 2; ++ax1_inner_inner_inner) { for (int ax3_inner_inner_inner = 0; ax3_inner_inner_inner < 2; ++ax3_inner_inner_inner) { tensor[(((((((((int)blockIdx.y) * 108) + (((int)blockIdx.x) * 6)) + (((int)threadIdx.z) * 5832)) + (((int)threadIdx.y) * 54)) + (((int)threadIdx.x) * 2)) + (ax1_inner_inner_inner * 2916)) + ax3_inner_inner_inner)] = max((((compute[(((ax1_inner_inner_inner * 2) + ax3_inner_inner_inner) - (((int)blockIdx.z) * 4))] + input2[((((int)threadIdx.z) * 2) + ax1_inner_inner_inner)]) * input3[((((int)threadIdx.z) * 2) + ax1_inner_inner_inner)]) + input4[((((int)threadIdx.z) * 2) + ax1_inner_inner_inner)]), 0.000000e+00f); } } • 生成された CUDA カーネルより一部を抜粋

Operator Fusion: 極端な例 def get_net(): def weird_activation(data): act1 = sym.relu(1 - sym.log(data)) act2 = sym.relu(sym.exp(2 * data + 1)) act3 = sym.leaky_relu(data) return -0.5 * act1 + 1.5 * act2 + act3 data = sym.Variable(name="data") data = sym.conv2d(data, kernel_size=(3,3), channels=8) data = sym.batch_norm(data) data = weird_activation(data) return data Graph(%data, %conv2d0_weight, %conv2d0_bias, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand) { %5 = tvm_op(%data, %conv2d0_weight, %conv2d0_bias, %batch_norm0_gamma_mul_div_expand, %batch_norm0_add_beta_expand, flatten_data='0', func_name='fuse_conv2d_broadcast_mul_broadcast_add_log___rsub_scalar___relu___mul_scalar_____mul_scalar_____add_scalar ___exp_relu___mul_scalar___broadcast_add_leaky_relu_broadcast_add', num_inputs='5', num_outputs='1’) ret %5 } Conv Batch norm Elemwise sum

[WIP] Graph level Auto Tuning • AutoTVM によって, オペレータ (ノード) 単位のチューニングが可能になった • だが, 入力のレイアウトは固定 (NCHW, NCHW[8, 16, 32…]c) • 最適なレイアウトは入力によって変わりうる → レイアウト変換のコストを考慮したグラフ全体のチューニングが必要 離散最適化問題を解くことになる 同様の問題を扱った論文 Optimal DNN Primitive Selection with Partitioned Boolean Quadratic Programming, CGO 2018 Node A Node B NCHW8c NCHW16c NCHW32c NCHW8c 0 0.1 0.25 NCHW16c 0.12 0 0.2 NCHW32c 0.3 0.18 0 レイアウト変換コスト NCHW8c 3.2 NCHW16c 2.8 NCHW32c 3.0 NCHW8c 2.6 NCHW16c 2.75 NCHW32c 2.5 ノードコスト ノードコスト

[WIP] Relay – NNVM IR v2 • 詳細は LT にて, または論文を参照 • NNVM IR の言語としての記述力は限られている → 正真正銘のプログラミング言語を組み込んでしまえ! • 純粋関数型言語 + Tensor の各次元数を埋め込んだ型システム • 各グラフ変換パスの, Relay IR 変換パスへの再実装が進行中 • Operator fusion, 自動微分など

Contents • TVM の概要 • オペレータレベルの最適化 • グラフレベルの最適化 • パフォーマンス • デプロイ用のツールとしての側面 • オープンソースプロジェクトとしての側面

NVIDIA GPU 公式ブログより引用 入力サイズ: (1, 3, 224, 224) 単位: ミリ秒 cuDNN には勝ち, TensorRT と互角

AMDGPU AMD 版の cuDNN, MIOpen と比較

ARM GPU

Intel Xeon 18 コア + AVX 512 Optimizing CNN Model Inference on CPUs より引用 MXNet + MKLDNN を 1 として、スピードアップを比較 Graph level Auto Tuning を適用した結果

自分で測ってみた: GPU VGG16 Resnet50 Densenet121 AutoTVM 6.68 4.03 4.64 TVM + cuDNN 7.10 5.14 7.14 MXNet + cuDNN 8.07 5.89 8.39 PyTorch + cuDNN 7.79 6.55 12.3 VGG16 Resnet50 Densenet121 AutoTVM 7.44 6.45 9.2 TVM + MIOpen 7.18 6.13 10.3 Qiita 記事より引用 Radeon R9 Nano GTX 1070 ti 入力サイズ: (1, 3, 224, 224) 単位: ミリ秒

自分で測ってみた: x86 VGG16 Resnet50 Resnet101 Densenet121 AutoTVM 69.26 16.49 30.87 14.53 MXNet + MKLDNN 62.96 21.15 38.85 24.06 Core i7-8700K (6C 6T) + AVX 2 VGG16 Resnet50 Resnet101 AutoTVM 25.35 9.54 16.68 MXNet + MKLDNN 24.88 12.64 22.96 Core i9-7940X (14C 14T) + AVX 512

Contents • TVM の概要 • オペレータレベルの最適化 • グラフレベルの最適化 • パフォーマンス • デプロイ用のツールとしての側面 • オープンソースプロジェクトとしての側面

デプロイ用ツールとして使う • コンパイルしたモデルをデプロイするのも簡単 • 最小限のランタイムライブラリとアプリケーションをリンクするだけ • Python 不要, ONNX もいらない • 場合によっては cuDNN などのライブラリを使うことも可能 • C++ Runtime 以外にも、Java, Rust をサポート • Go Runtime の PR も進行中 libtvm.so libtvm_runtime.so Thread pool Codegen LLVM interface Model loading NDArray Device API HalideIR Lowering pass

ハードウェアバックエンド デスクトップ GPU CUDA, OpenCL, ROCm (AMDGPU) モバイル GPU OpenCL, Metal, Vulkan (SPIR-V) CPU x86, ARM (Android, Rasp) (WIP) FPGA VTA, SDAccel, AOCL VTA については、以前の勉強会で発表あり VTA 試してみた 開発中? • Huawei AI チップ バックエンド (RFC) • Qualcomm Hexagon DSP オフロード (RFC) AMD の公式スライドに TVM のロゴが !!

モデルをファイルにシリアライズ • コンパイルされたコードを含む共有ライブラリ • JSON 形式のグラフ構造の定義 • 各オペレータのパラメータ を出力 graph, lib, params = nnvm.compiler.build(net, target, shape_dict, params=params) output_prefix = “model" lib.export_library("{}_deploy.so".format(output_prefix)) with open("{}_deploy.json".format(output_prefix), "w") as fo: fo.write(graph.json()) with open("{}_deploy.params".format(output_prefix), "wb") as fo: fo.write(nnvm.compiler.save_param_dict(params))

C++ からモデルをロード const std::string json_file(“model_deploy.json"); const std::string param_file(“model_deploy.params"); tvm::runtime::Module mod_syslib = tvm::runtime::Module::LoadFromFile(“model_deploy.so"); std::string json_data; std::string params_data; // Read json and parameter files and fill in json_data, params_data ... const int device_type = kDLGPU; const int device_id = 0; auto runtime_create_func = tvm::runtime::Registry::Get("tvm.graph_runtime.create"); tvm::runtime::Module mod = (*runtime_create_func)(json_data, mod_syslib, device_type, device_id);

推論の実行 tvm::runtime::PackedFunc load_params = mod.GetFunction("load_params"); tvm::runtime::PackedFunc set_input = mod.GetFunction("set_input"); tvm::runtime::PackedFunc run = mod.GetFunction("run"); tvm::runtime::PackedFunc get_output = mod.GetFunction("get_output"); DLTensor* x = nullptr; DLTensor* y = nullptr; // Allocate x and y, fill in x … TVMByteArray params_arr; params_arr.data = params_data.data(); params_arr.size = params_data.length(); load_params(params_arr); set_input("data" x); run(); get_output(0, y);

Contents • TVM の概要 • オペレータレベルの最適化 • グラフレベルの最適化 • パフォーマンス • デプロイ用のツールとしての側面 • オープンソースプロジェクトとしての側面

TVM とコミュニティ • コントリビュータやユーザーを主体としたコミュニティづくりを重視 より詳しくは, TVM Community Structure, [RFC] Community Guideline Improvements しばらくコントリビュートしていると, レビューアになってくれと頼まれる

• 企業からの興味 • Amazon, Huawei: 頻繁に PR が来る. 製品に使っている • Facebook, Alibaba: x86, ARM 向けの最適化をやっているらしい • Microsoft: Github の Issue によく現れる. ONNX がらみ? • Qualcomm: Hexagon DSP バックエンドを開発中 • 国内では, NTT からコミットする人が複数. FPGA? • 他のオープンソースプロジェクトからの利用 • MXNet: NNVM の一部を使用. TVM でコンパイルした関数を呼ぶこともできる • miya: Theano を開発したグループが現在開発中の DL 用プログラミング言語. バックエンドとして NNVM を使用するらしい

まとめ • プログラミング言語やコンパイラ技術の, DL への導入が流行っている • 先進的な最適化・チューニング技術 • 多様なハードウェアバックエンドのサポート • ハイパフォーマンス これらを兼ね備えた TVM に今後も期待. ニューラルネットワーク = プログラム DL フレームワーク = 言語処理系

Links • TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018 • https://arxiv.org/abs/1802.04799 • Learning to Optimize Tensor Programs, NIPS 2019 • AutoTVM 論文 • https://arxiv.org/abs/1805.08166 • Relay: A New IR for Machine Learning Frameworks • https://arxiv.org/abs/1810.00952 • VTA: An Open Hardware-Software Stack for Deep Learning • https://arxiv.org/abs/1807.04188 • Optimizing CNN Model Inference on CPUs • https://arxiv.org/abs/1809.02697 • Discussion forum • https://discuss.tvm.ai/ • 公式ドキュメント. チュートリアルなども充実. • https://docs.tvm.ai/ • Automatic Kernel Optimization for Deep Learning on All Hardware Platforms • https://tvm.ai/2018/10/03/auto-opt-all.html 登録はこちら