MLIR のメモ(主に Vulkan で推論実行をターゲット)

背景

• Vulkan グラフィックスアプリに, 機械学習を取り込みたい
• Android や, エッジデバイスで Production で機械学習推論をしたい(主に音声認識, image processing, 3D reconstruction)

現状では MLIR 自体はリリースに近いようですが, Vulkan backend は開発中です.

なぜ MLIR か

• LLVM のプロジェクトに取り込まれ, 活発な開発とメンテがされると想定できる
• IR なので, 高級開発言語(C/C++, Python など)のライブラリ依存性を下げることができる
  • Vulkan(SPIR-V) or CPU 実行の場合, 実行環境をミニマムにできる
    • ビルド時も, OpenCL/CUDA SDK をインストールして... などの問題が減る
  • 高級言語では指定しづらいものでも, IR レベルでは指定しやすいものがあったりする
    • quantize とか

SPIR-V ではダメなのか?

SPIR-V の中間言語はいじるのめんどい(SSA ではあるが, スタックマシン的な記述)
テンソルなどの概念が無い.

懸念点, 制約事項

• 上位のインターフェースは自前なりで作る必要がある
• モデルデータのロードなども, 別途書く必要がある
  • IREE などは, この辺りも含んでの対応のようですね.

とはいえ, だいたい production で使いたいときは, 利用したい機能(ドメイン)が確定していると思いますので, 一般化された onnx, torchscript などに振り回されず開発専念 & デバッグできて楽かもしれません.

MLIR 関連

• IREE : SPIR-V -> Vulkan 実行環境 https://github.com/google/iree
  • まだ出始め. プロジェクトキャンセルのリスクあり
• PlaidML : https://github.com/plaidml/plaidml
  • 元々は OpenCL だったと思うが, Stripe -> MLIR -> GPU 実行に移行中?
• TVM
  • 同様の取組として Relay IR を考案している. https://docs.tvm.ai/dev/relay_intro.html

情報

まずはこのあたりがよいでしょうか.

MLIR 開発

バイナリや C API あたり叩いてだと面倒なので, まずは ASCII 形式でいろいろためしてみるのがよいでしょう.

Python で MLIR ascii を読み書きできるライブラリ, pymlir があります.

自前でパーサを用意しているので, LLVM ライブラリに依存しないのがいいですね!