AI Engine に適したCNN アーキテクチャの検討と実装

Transcript

1. AI Engineに適した CNNアーキテクチャの検討と実装 2023/6/8 ザイリンクス株式会社 安藤 潤
2. 2 | [Public] 本発表について • 目的 • AI Engineのアーキテクチャとプログラミングモデルを紹介します • CNN実装の経験からAI Engineプログラミングの難しさや面白さを共有します • ご注意 • CNN実装はプログラミング習得のために個人の趣味で行いました
3. 3 | [Public] 発表の流れ 1. AI Engine紹介 • アーキテクチャ • プログラミングモデル 2. CNNアーキテクチャの検討と実装 • 動機 • アーキテクチャ検討 • 実装の試行錯誤 • 性能見積もり 3. まとめ
4. AI Engine紹介
5. 5 | [Public] AI Engine MEMORY MEMORY AI Engine AI Engine MEMORY MEMORY AI Engine Memory AI Engine Memory AI Engine Memory AI Engine Memory AI Engine Xilinx AI Engine • 概要 • 7nm世代Versalファミリーに搭載 • インターコネクト、メモリで相互接続された数十～数百の1GHz+ VLIW/SIMDプロセッサ • 信号処理、機械学習に特化 • 第二世代（AIE-ML）まで発表済み Versal® ACAP AI Engine
6. 6 | [Public] AMDプロセッサとの統合 • Ryzen AI • Ryzen 7040シリーズ・モバイル・プロセッサにAI Engineを搭載 • Windowsデモ : https://github.com/amd/RyzenAI-cloud-to-client-demo https://www.amd.com/en/products/ryzen-ai
7. 7 | [Public] AI Engineアーキテクチャ • データフローアーキテクチャ • プロセッサが二次元アレイ状に配置 • コア間が直接データをやり取りしながら計算 Memory AI Engine Memory AI Engine Memory AI Engine Memory AI Engine Memory AI Engine Memory AI Engine Memory AI Engine AI Engine Memory AI Engine Memory L1 core L0 Block 0 Block 1 L1 L2 DRAM L0 L0 L0 L0 L0 core core core core core CPU / GPU AI Engine
8. 8 | [Public] AI Engineコア • VLIW/SIMDプロセッサ • ベクトルユニットを効率的に動かすことが肝心 • 128 MAC / cycle (INT8) • 専用API、Intrinsicsを使ってC++でコーディング 固定小数点演算ベクトルユニット (AM009) xbuff zbuff 並び替え 乗算 加算ツリー 累積 MAC (multiply–accumulate) ベクトルレジスタ AI Engine MEMORY MEMORY AI Engine AI Engine MEMORY MEMORY AI Engine Memory AI Engine Memory AI Engine Memory AI Engine Memory AI Engine Store Unit Scalar Unit Scalar Register File Scalar ALU Non-linear Functions Instruction Fetch & Decode Unit AGU Vector Unit Vector Register File Fixed-Point Vector Unit Floating-Point Vector Unit Load Unit B AGU Load Unit A AGU Memory Interface Stream Interface AI Engine Core
9. 9 | [Public] AI Engineプログラミングモデル • アルゴリズムをカーネルに分割しデータフローとして定義 • データ駆動：各カーネルはデータを受け取って処理して次に渡す、を繰り返す • カーネル間の接続はウィンドウ（PIPO）またはストリーム（FIFO） • タスク並列：入力データに対して各カーネルが同時に処理を実行 • FPGA設計と同じ考え方 • 開発の流れ：アーキテクチャ設計 → カーネル実装 → グラフ実装 Kernel B Kernel C Kernel D PIPO PIPO Kernel A PIPO PIPO 洗う担当 切る担当 炒める担当 味付け担当 盛り付け担当 流れ作業で食品を調理をするイメージ （食材駆動、作業並列） データフローには分岐や合流があってもよい
10. 10 | [Public] AI Engineプログラミングモデル • AI Engineアレイ • AIEタイル間はインターコネクトで接続 • AIEコアは隣接する上下左右のメモリにアクセスできる AIE core Data Memory AIE core Data Memory IC IC AIE core Data Memory IC AIE core Data Memory IC AIE core Data Memory IC AIE core Data Memory IC AIE Tile AIE Tile AIE Tile AIE Tile AIE Tile AIE Tile
11. 11 | [Public] AI Engineプログラミングモデル • データフローを配置配線 • グラフ記述時に制約を与える • 各カーネルの実行時間、必要なメモリサイズなどを考慮して決める Kernel A AIE core Data Memory AIE core Data Memory IC IC Kernel B AIE core Data Memory IC Kernel C AIE core Data Memory IC Kernel D PIPO PIPO PIPO PIPO AIE core Data Memory IC AIE core Data Memory IC PIPO Kernel B Kernel C Kernel D PIPO PIPO Kernel A PIPO PIPO 配置配線 To PL From PL AIE Tile AIE Tile AIE Tile AIE Tile AIE Tile AIE Tile
12. AI Engineに適した CNNアーキテクチャの検討と実装
13. 13 | [Public] CNN実装の動機 • 技術習得 • AI Engineはキーとなるテクノロジー、手を動かして理解を深めたい • コンピュータ将棋 • DLを利用した将棋AIアルゴリズムが流行 • 9x9の盤面をCNNに入力し手番側の勝率と次の一手を予測、 モンテカルロ木探索で最良の手を求める • 第32回/第33回世界コンピュータ将棋選手権 • EPYC 7742 + A100 x8 サーバー、9台のクラスタ構成のチームが優勝 • モデルの学習、推論に高価な機材が必要 • 興味 • AI Engineで強い将棋AIを実装してみたい！ 山岡忠夫・加納邦彦（2021）． 強い将棋ソフトの創りかた マイナビ出版
14. 14 | [Public] アプローチ、目標 • 既存のアプローチ • 汎用エンジン : Vitis AI • あらかじめ設計済みのConvユニット＋コントローラ • モデルから命令列にコンパイル • コード生成 : FINN • レイヤーごとにテンプレートを用意 • モデルからコードを生成 • 本実装のアプローチ • AI Engineに適した将棋AI用CNNアーキテクチャを実装する • アーキテクチャに合わせてモデルを学習する • 目標 • 性能に妥協しない：バッチなしで低レイテンシ、高スループットを達成すること • ハードウェアを活かす：INT8 MACをできるだけ効率よく回す • ターゲットデバイス • VCK5000 Versal開発カード VCK5000
15. 15 | [Public] アーキテクチャ検討 • ResNet Plainアーキテクチャの問題点 • Conv2D重みが大きすぎる • チャネル数が256のとき、3x3畳み込みで3x3x256x256=576KB • SRAMでは容量不足、DRAMでは帯域不足 • DRAM帯域ネックを隠蔽するにはバッチ処理が必要 • モデルの軽量化 • Separable Convolution（MobileNet、ConvNeXt） • Pixelwise ConvolutionとDepthwise Convolutionに分離 • AIEのINT8 MACはデータ並び替えが32bit単位に制限 ResNet Plainアーキテクチャ xbuff zbuff 並び替え 乗算 加算ツリー 累積 ベクトルレジスタ
16. 16 | [Public] アーキテクチャ検討 • モデルの軽量化 • ResNeXt • 4チャネルごとに3x3畳み込み（Grouped Convolution） • INT8 4つ＝32bitとなりデータ並び替えに都合が良い • 重みが十分小さい • 採用 ResNet Bottleneck ResNeXt ResNet Plainアーキテクチャ
17. 17 | [Public] ResNeXt実装（試行錯誤）の流れ • カーネル実装 • Conv1x1（Pixelwise Convolution） • 行列積APIを利用 • Conv3x3（Grouped Convolution） • 行列積に変換して計算する方式 • 直接計算する方式 • グラフ実装 • ResNeXtブロック • ResNeXt全体 ResNeXt block ResNeXt block ResNeXt block …
18. 18 | [Public] 行列積によるConv3x3実装 • 行列積APIを活用 • シンプルな行列積はMAC稼働率80% • データ並び替えが必要 • 畳み込み演算を行列積で行うためのim2col変換 • 行列積APIのためのタイル化 • 並び替え実装の検討 • for文による素朴な実装は非効率 • ベクトルレジスタにデータをロードしてselect関数で取り出してメモリにストア • im2colとタイル化を同時に考慮したコードを書きたいがとても大変 • コード自動生成 • 各タイルを出力するためには複数の入力データが必要、レジスタにロード済みのデータは再利用したい • どのような順序でタイルを出力すれば入力データのロード回数を少なくできるか • 巡回セールスマン問題の解法を利用して出力順を決定 • 出力タイルがノード、ノード間の移動コストは新たにロードしなくてはならないデータの数 • データがうまく再利用されるようにレジスタへのロードをスケジューリング • 十分に少ないロード回数で出力タイルを加工するコードが生成できた
19. 19 | [Public] 行列積によるConv3x3実装 • 結果 • 想定よりも数倍遅い • 原因 • 多量の複雑な並び替え • 並び替えを指示する引数はスカラレ ジスタを介して渡される • スカラレジスタに32bitの値を設定す るには2サイクル必要 • 並び替えのためだけにサイクル 数を消費するのは無駄 • Intrinsicsを使って直接計算する 自動生成したコード
20. 20 | [Public] 効率的な畳み込み演算の検討 • 行列積API（INT8）のアセンブリを調査 • どのようにデータ並び替えているか • xbuffはレーン内にデータを順に並べることができないが zbuffはできることを利用していた • 行列乗算の左に入力データ、右に重みデータ • 入力データはタイル化されていればよい • 重みデータは事前にタイル化しておく lanes=16 xstart=0, xoffs=[0,0,0,0,0,0,0,0], xstep=8, xsquare=[0,2,1,3] zstart=0, zoffs=[0,0,4,4,8,8,12,12], zstep=2 [xbuff] 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 [zbuff] 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 16 17 18 19 20 21 22 23 16 17 18 19 20 21 22 23 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 24 25 26 27 28 29 30 31 24 25 26 27 28 29 30 31 24 25 26 27 28 29 30 31 mmul （4x8x8）のmul/mac configuration 重みデータ 入力データ 重みデータ 入力データ Lane 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 Lane 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
21. 21 | [Public] 直接計算によるConv3x3実装 • 計算方法 • パディングを入れて、4画素4ch単位（16レーン）でMAC • 重みは事前に並び替えておく • MAC効率 • 重みゼロ 3/4 × パディング 9/12 ＝ 56.25% xstart=0 xoffs=[0]*8 xstep=8 zstart=12 zoffs=[8,8,10,10,12,12,14,14] zstep=2 [xbuff] 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 0 4 8 12 16 20 24 28 1 5 9 13 17 21 25 29 2 6 10 14 18 22 26 30 3 7 11 15 19 23 27 31 [zbuff] 28 29 30 31 0 1 2 3 28 29 30 31 0 1 2 3 28 29 30 31 0 1 2 3 28 29 30 31 0 1 2 3 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7 4 5 6 7 8 9 10 11 4 5 6 7 8 9 10 11 4 5 6 7 8 9 10 11 4 5 6 7 8 9 10 11 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 xstart=32 xoffs=[0]*8 xstep=8 zstart=4 zoffs=[0,0,2,2,4,4,6,6] zstep=2 [xbuff] 32 36 40 44 48 52 56 60 33 37 41 45 49 53 57 61 34 38 42 46 50 54 58 62 35 39 43 47 51 55 59 63 32 36 40 44 48 52 56 60 33 37 41 45 49 53 57 61 34 38 42 46 50 54 58 62 35 39 43 47 51 55 59 63 32 36 40 44 48 52 56 60 33 37 41 45 49 53 57 61 34 38 42 46 50 54 58 62 35 39 43 47 51 55 59 63 32 36 40 44 48 52 56 60 33 37 41 45 49 53 57 61 34 38 42 46 50 54 58 62 35 39 43 47 51 55 59 63 [zbuff] 4 5 6 7 8 9 10 11 4 5 6 7 8 9 10 11 4 5 6 7 8 9 10 11 4 5 6 7 8 9 10 11 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 8 9 10 11 12 13 14 15 12 13 14 15 16 17 18 19 12 13 14 15 16 17 18 19 12 13 14 15 16 17 18 19 12 13 14 15 16 17 18 19 16 17 18 19 20 21 22 23 16 17 18 19 20 21 22 23 16 17 18 19 20 21 22 23 16 17 18 19 20 21 22 23 重みゼロ Input Output 9x9将棋盤の1行
22. 22 | [Public] Batch Normalization / ReLU • 畳み込みの他に必要な処理を実装 • Bias • BNで発生したBiasをアキュムレータの初期値に設定 • Scale • BNと量子化に必要なScaleを乗算とシフトで実装 • アキュムレートレジスタの内容を一度ベクトルレジスタに移動す る必要がある • ReLU • ubsrs関数（Shift-Round-Saturate）でacc48からuint8（0~255）に 切り詰めることができる • パディング • select16関数でパディング部分をゼロ埋めする
23. 23 | [Public] メモリレイアウトの問題 • 行列積APIのメモリレイアウト • 4x8x4モード : 入力と出力のメモリレイアウトが異なる → Conv1x1の前後に並び替えが必要 • 2x8x8モード : 入力と出力が8chで一致 → Conv3x3（grouped conv）も8chに統一 C HW C HW
24. 24 | [Public] Conv3x3の8ch化 • ResNeXtアーキテクチャを変更 • Conv3x3 Grouped Convを8chごとに変更 → 重みサイズ、計算量が2倍 • モデル精度に大きな変化はなさそう • 計算方法 • 2画素8ch単位でMAC → パディングは1画素分 • 重みゼロ部分が不要 • 1行（80要素）分をアキュムレートレジスタにのせてまとめて計算 • MAC効率 • パディング 9/10 ＝ 90% W=10
25. 25 | [Public] Conv3x3の8ch化 • MAC稼働率向上の工夫 • Batch Normalizationを削除 • floatモデル上でBNをConv2Dに吸収 • 乗算を削除 • 余計な乗算をBiasに吸収、シフトのみで量子化 • パディング挿入を削除 • 端の計算時に結果がゼロになるように細工 • コーディング上の最適化 • 変数の寿命を明示的に短くする、など • 結果 • 出力32chあたり 4200cycle → 4000cycle • 計算量は倍になったがMAC効率が改善 • 畳み込み以外の演算が不要 • MAC稼働率 • 28% → 37% 改善後 改善前
26. 26 | [Public] ResNeXtブロック実装 • Conv1x1 • パラメータ（32KB）がデータメモリに収まるようにカーネルを4つに分割し並列化 • 出力チャネル方向で処理を分担 • 1行ごとの処理にして入出力バッファ容量を削減
27. 27 | [Public] ResNeXtブロック実装 • Conv3x3 • Conv1x1と同じく4並列 • Conv1x1が9回実行→ Conv3x3が1回実行 • マルチレート解析により自動検出 • 後段Conv1x1のメモリレイアウトに合わせて出力 • ショートカット接続 • デッドロック防止のためバッファが必要
28. 28 | [Public] ResNeXtブロック実装 • 配置配線 • ツール任せでは効率的な配置配線にならない • 隣接メモリを介したピンポンにならない • 手動でカーネル、バッファの配置を検討 • 2x2タイルにConv 3つを配置 • この他にスタック・ヒープ領域も必要 • 同一タイルに異なるレートのカーネルは 配置不可 • 問題 • カーネルの入出力が多いと DMAリソースが不足してエラー • S2MM、MM2S、マルチレート
29. 29 | [Public] ResNeXtブロック実装 • DMAリソース不足対策 • Mergeカーネルを追加して カーネルの入力数を削減
30. 30 | [Public] ResNeXtブロック実装 • 配置配線の再検討 • Mergeカーネル • メモリから直接入力してDMA節約 • 1並列分が2x2内に収まるように配置 • ブロックの配置配線に成功
31. 31 | [Public] ResNeXt全体実装 • グラフ記述DSLを作成 • 様々な配置配線を試行しやすくするため • Pythonでグラフを記述しC++コードを生成 • 500行のコードから11,500行のコードを出力
32. 32 | [Public] ResNeXt全体実装 • 紹介しきれない数々の困難を乗り越え実装に成功 • 入力層 → ResNeXt 23ブロック → 出力層（Policy/Value） Policy / Value Input layer ResNeXt block ResNeXt block
33. 33 | [Public] 性能見積もり • AIEシミュレータ • サイクルアキュレート • 実機とおおよそ一致 • トレースを取得できる • シミュレーション結果 • レイテンシ • 入力層 : 10us • ResNeXtブロック : 15us • 出力層 : 10us • ネットワーク全体 : 10 + 15 * 23 + 10 = 365us block0 block1 conv0 conv1 conv2
34. 34 | [Public] 性能見積もり • シミュレーション結果 • データ間隔 : 14.3us • 各レイヤーのレイテンシは最大9us程度 • Conv1x1→3x3にマルチレート起因の待ちが発生 • スループット : 70k/s • ResNeXt, 256チャネル, 23ブロック • INT8, Batch=1 • 126,385,920 MAC • (256x128+3x3x8x8x4+128x256)x9x9x23 • 参考：A100*1 • スループット : 70k/s • ResNet Plain, 192チャネル, 10ブロック • TensorRT(INT8), Batch=128 • 537,477,120 MAC • 3x3x192x192x9x9x2x10 *1 https://tadaoyamaoka.hatenablog.com/entry/2020/11/07/154352
35. まとめ
36. 36 | [Public] まとめ • 将棋AIのCNNをAI Engineに実装した • 様々な実装を試行錯誤 • MAC稼働率を上げるのは難しかった • Conv1x1: 53% • Conv3x3: 37% • 400コアのほとんどを使用 • レイテンシ、スループットは満足 • INT8モデルを学習 • CNNアーキテクチャは改善の余地ありそう • AI Engineプログラム最適化の知見 • アルゴリズム • オンチップメモリを活用 • 計算を簡単化 • カーネル • MACを効率的に動かす • データフロー • カーネル分割（データ並列、タスク並列） • メモリレイアウトを統一 • 配置配線 • 効率的なカーネル間接続 • リソース制約を考慮した配置
38. 38 | [Public] 今後の課題 • カーネル配置 • 4コア中3コアしかMACが動いていない • ショートカットのないモデルを試してみる？ • RepVGG、MobileOne • Conv1x1と3x3でレイテンシが異なる • 稼働率の低いコアがいる • CNNモデル • ResNeXtが将棋AIに向いていない疑惑 • ResNet Plainと比べてロスの収束が遅い、精度低い • 推論時と同じネットワークで学習しているせいかも！ • 実機 • ホストとはQDMAでストリーミング通信で きることが望ましい • 現状プラットフォームが非対応 • V70 • 第2世代AIE、AIに特化 • INT8 MAC倍 • INT4+SparsityでINT8の4倍 • BF16はINT8の半分 • DM倍 • レジスタ3倍 • メモリタイル追加 • 他のモデル • 姿勢推定、物体検出、超解像 • LLM（RWKV）, Whisper
39. 39 | [Public] 量子化CNNモデルの作成方法 • Brevitas • 低ビット量子化モデルをQuantization-aware training (QAT) するPyTorchライブラリ • 手順 • Brevitas上でINT8モデルを作成 • 同等のFloatモデルを学習 • BNをConv2Dに吸収 • Floatモデルの重みをINT8モデルに読み込んでキャリブレーション • 劣化した精度をQATで回復 • 標準ONNXモデルとして出力 • この時点では演算やパラメータはfloatのまま • ONNXモデルを変換 • 冗長なレイヤーを統合・削除 • 固定値を伝搬して量子化 • 整数演算だけで実行可能なモデルができる • モデルのパラメータをC++ヘッダーファイルとして出力
40. 40 | [Public] 紹介しきれない数々の困難 • aiecompiler内部の処理で引数文字列長がMAX_ARG_STRLEN（128KB）超え • Linuxカーネルを修正 • InsertFIFOが終わらない • Re-convergenceパスが組み合わせ爆発して数え上げが終わらない • 配線をPLに出すことでグラフを分割 • 内部エラー • 最適化オプションを無効にして回避