MLIRでHello, Tensor

2020年12月16日 21:35

この記事はTensorFlow Advent Calendar 2020、17日目の記事です。この記事ではLLVMコミュニティが中心となって開発しているMLIRという新しいコンパイラ基盤の基本的な使い方を解説します。

MLIRとは

MLIRとはコンパイラ基盤となるオープンソースのソフトウェアでその名はMulti-Level Intermediate Representationの頭文字を取ったものです。もともとはGoogleのTensorFlowチームが開発したソフトウェアでLLVM Foundationに2019年に寄贈されました。

コンパイラ基盤としてはLLVMがよく知られていますが、MLIRはLLVMで得られた知見をより抽象的なレベルで実現し、機械学習アプリケーションに代表されるような複雑な数値計算を様々なハードウェア上に最適な形で実行できるようにします。

例えばコンパイラが真に必要な最適化を行おうと思った場合には各プログラミング言語のセマンティクスであったり構文の知識が必要となります。その結果、下記のように独自のIRをそれぞれが作ってさらにそれをLLVMがコンパイルするといった状況になっていました。LLVMという共通基盤を使っているにも関わらず同じような処理や知見を共有できずにいました。

MLIRのMLはMachine LearningのMLからとったものと思われがちですが応用範囲は機械学習に限りません。Dialectと呼ばれるプラグイン機構を使い非常に柔軟なカスタマイズが可能となっています。例えばMLIRの大きな特徴は以下のようなものです。

- 最適化や機械語生成まで行う包括的なフレームワーク
- AVX命令やGPUなど異なったプログラミングモデルを同じ土俵に乗せ最適化することが可能になる
- Dialectというプラグイン機能を使うことで型やOperationを拡張することが可能
- Polyhedral Compilationというネストされたループやベクトル計算のスケジューリング最適化のテクニックを柔軟に取り込むことができる
- LLVMとシームレスに連携が可能

OperationとDialect

MLIRを利用するにあたって幾つか知っておくべき概念や用語があります。多くはLLVMなどコンパイラの開発に携わった人には馴染み深いものかもしれません。ここでは最低限必要なOperationとDialectと呼ばれるものについて説明します。

MLIRではLLVMでいうInstructionをOperationと呼び、ひとつひとつの計算は実際にはこのOperationという単位で実行されます。TensorFlowのOperationのようにMLIRのOperationはInput, Outputに加えてAttributeという固定値をとることができます。構文は高水準なプログラミング言語の関数に似ています。

このOperationを以下のようにModule, Function, Blockの中に入れ子にされてMLIRのプログラムが構成されます。FunctionやBlockを柔軟に用いて制御構造を実現するわけですが、MLIRの場合にはOperationも自由に拡張することができます。その仕組みがDialectです。

DialectはMLIRに対してOperationや型システム、最適化パスなど様々な機能を柔軟拡張するための仕組みです。実はMLIR本体にはこれらの構造を抽象的にあつかう機能しかなく（それがMLIRの存在意義なのですが）通常のコンパイラに実装されているような最適化や具体的なOperationの実装や型はDialectとして実装されています。Dialectは自分で拡張することもできますが標準として提供されているDialectだけでもこれだけあります。

- 'acc' Dialect
- 'affine' Dialect
- 'arm_neon' Dialect
- 'async' Dialect
- 'avx512' Dialect
- 'gpu' Dialect
- 'linalg' Dialect
- 'llvm' Dialect
- 'llvm_arm_neon' Dialect
- 'llvm_arm_sve' Dialect
- 'llvm_avx512' Dialect
- 'nvvm' Dialect
- 'omp' Dialect
- 'pdl' Dialect
- 'pdl_interp' Dialect
- 'quant' Dialect
- 'rocdl' Dialect
- 'scf' Dialect
- 'shape' Dialect
- 'spv' Dialect
- 'std' Dialect
- 'tensor' Dialect
- 'vector' Dialect

この記事ではMLIRの基本的な使い方を学ぶためTensorの要素を出力するだけのOperaionをもった小さなDialectを実装してみたいと思います。

最終的なコードはLewuathe/mlir-helloにあります。

LLVMのビルド

このプログラムではMLIRを含んだ最新のLLVMをビルドする必要があります。あらかじめCMakeとNinjaをインストールしておいてください。

少し時間はかかりますが、以下のようにビルドすることができます。

git clone https://github.com/llvm/llvm-project.git
cd llvm-project

mkdir build && cd build
cmake -G Ninja ../llvm \
  -DLLVM_ENABLE_PROJECTS=mlir \
  -DLLVM_BUILD_EXAMPLES=ON \
  -DLLVM_ENABLE_ASSERTIONS=ON \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_RTTI=ON \
  -DLLVM_TARGETS_TO_BUILD="host"
cmake --build . --target check-mlir

テストがパスすれば成功です。

Dialectを宣言する

Dialectを作るためにゴリゴリC++コードを書くことも可能ですが、TableGenというLLVMで使われている宣言的な言語を使うことが推奨されています。これを用いるとDialectやOperationのインタフェースを宣言的に書くことができ、ドキュメントも簡単に生成できるようになります。HelloDialect.tdというファイル名で以下を記述します。これに基づいてTableGenがC++コードを生成してくれます。この仕組みをOperation Definition Specification（ODS）と呼びます。

def Hello_Dialect : Dialect {
   let name = "hello";
   let summary = "A hello out-of-tree MLIR dialect.";
   let description = [{
       This dialect is minimal example to implement hello-world kind of sample code
       for MLIR.
   }];
   // 生成されるC++クラスの名前空間
   let cppNamespace = "::hello";
}

// すべてのOperation共通のインタフェースやTraitを宣言する
class Hello_Op<string mnemonic, list<OpTrait> traits = []> :
       Op<Hello_Dialect, mnemonic, traits>;

HelloOpsDialect.h.incというファイルが生成されるので、必要なクラスから参照させます。

Operationを宣言する　

同じようにOperationもODSを使って宣言することができます。HelloOps.tdというファイル名で以下のように記述します。テンソルを構築するためのConstantOpとそれを表示するためのPrintOpを宣言します。

include "HelloDialect.td"
include "mlir/Interfaces/SideEffectInterfaces.td"

// NoSideEffectはOperationの属性を宣言するTrait
def ConstantOp : Hello_Op<"constant", [NoSideEffect]> {
 let summary = "constant";
 let description = [{
   Constant operation turns a literal into an SSA value. The data is attached
   to the operation as an attribute. For example:

   ```mlir
     %0 = "hello.constant"()
     { value = dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf64> }
     : () -> tensor<2x3xf64>
   ```
 }];

 // C++コードに実装するメソッドの宣言
 let builders = [
   OpBuilderDAG<(ins "mlir::DenseElementsAttr":$value), [{
     build($_builder, $_state, value.getType(), value);
   }]>,
   OpBuilderDAG<(ins "double":$value)>
 ];

 // Operationの入力
 let arguments = (ins F64ElementsAttr:$value);
 // Operationの出力
 let results = (outs F64Tensor);
}


def PrintOp : Hello_Op<"print", [NoSideEffect]> {
   let summary = "print operation";
   let description = [{
       The "print" builtin operation prints a given input tensor, and produces
       no results.
   }];

   let arguments = (ins AnyTypeOf<[F64Tensor, F64MemRef]>:$input);

   let assemblyFormat = "$input attr-dict `:` type($input)";
}

HelloOps.h.incというファイルが生成されます。この生成されたファイルには様々な定義が混じっているので必要なものをプリプロセッサ命令で指定します。例えばOperationのインタフェースをだけ取り出す場合にはヘッダファイルに以下のように書きます。

// GET_OP_CLASSESで囲まれている部分のみ展開される。
#define GET_OP_CLASSES
#include "Hello/HelloOps.h.inc"

このテクニックは他の箇所でも使われていて実装ファイルでHelloDialectの初期化を行う場合も以下のように書くと、自動的にODSで定義されたOperationの初期化も行ってくれます。

void HelloDialect::initialize() {
 addOperations<
#define GET_OP_LIST
#include "Hello/HelloOps.cpp.inc"
     >();
}

Lowering to LLVM

constantやprintの定義はありますが、このままでは特に意味のない空っぽのDialectを実装しただけです。このDialectに命を吹き込むためにはLoweringというプロセスが必要となります。HelloというDialectを他のDialectを経由しながら実装可能な機械語にコンパイルします。このプロセスをLoweringといいます。MLIRの面白いところはこのプロセスを拡張である他のDialectに部分的に変換されながらだんだんとハードウェアに近いレイヤーに落ちていくところです。必要な箇所だけ抽象度の高いレイヤーの情報を保ったままハードウェアの世界に降りていくことが可能です。これをPartial Lowering（Partial Conversion）といいます。

今回はLLVM IRまで落ちてしまえば、LLVMのランタイムで実行できるためLLVM IRまで変換を行います。Hello DialectはMLIRのプリミティブタイプであるTensorを使っていますが、これを直ちにLLVM IRまで変換するのはちょっと大変なので途中Tensorの計算をサポートしているAffinetまたはStandardというDialectを経由することにします。

Hello -> Affine, Standard -> LLVM IR

という順番で変換を行います。以下のHelloToAffineLowerPassというクラスを定義します。runOnFunctionメソッドを実装することでFunctionを変換することができます。この変換プロセスをPassと言います。

class HelloToAffineLowerPass : public mlir::PassWrapper<HelloToAffineLowerPass, mlir::FunctionPass> {
 void getDependentDialects(mlir::DialectRegistry &registry) const override {
   registry.insert<mlir::AffineDialect, mlir::StandardOpsDialect>();
 }

 // Functionの変換を行います。
 void runOnFunction() final;
};
}

Partial LoweringではすべてのOperationを変換するわけではないので、そのまま残すべきものと変換すべきものを明らかにする必要があります。これをConversion Targetといいます。Conversion Targetに含まれるものをLegal、そうでないものをIllegalとして書く必要があります。今回はHello DialectからAffineにすべて変換しますが、PrintだけはLLVM IRに直接変換できるのでここでは残しておきます。

  mlir::ConversionTarget target(getContext());

 // Helloは残してはいけない
 target.addIllegalDialect<hello::HelloDialect>();
 
 // AffineかStandard Dialectに変換しなければならない
 target.addLegalDialect<mlir::AffineDialect, mlir::StandardOpsDialect>();
 
 // PrintOpは例外的に残してもよい
 target.addLegalOp<hello::PrintOp>();

このPassを経るとPrintOp以外はAffineあるいはStandardで記述されていることになります。

まったく同じようにAffineからLLVMに変換するようなPassも実装します。このPassはPrintOpも変換します。

class HelloToLLVMLoweringPass : public mlir::PassWrapper<HelloToLLVMLoweringPass, mlir::OperationPass<mlir::ModuleOp>> {
public:
 void getDependentDialects(mlir::DialectRegistry &registry) const override {
   registry.insert<mlir::LLVM::LLVMDialect, mlir::scf::SCFDialect>();
 }

 // Operationの変換を行う
 void runOnOperation() final;
};


void HelloToLLVMLoweringPass::runOnOperation() {
 mlir::LLVMConversionTarget target(getContext());
 target.addLegalOp<mlir::ModuleOp, mlir::ModuleTerminatorOp>();

 mlir::LLVMTypeConverter typeConverter(&getContext());

 // Affine, StandardからLLVMへの変換Passはすでにあるものを利用する
 mlir::OwningRewritePatternList patterns;
 populateAffineToStdConversionPatterns(patterns, &getContext());
 populateLoopToStdConversionPatterns(patterns, &getContext());
 populateStdToLLVMConversionPatterns(typeConverter, patterns);

 // PrintOpの変換は自前で実装する
 patterns.insert<hello::PrintOpLowering>(&getContext());
 auto module = getOperation();
 if (failed(applyFullConversion(module, target, std::move(patterns)))) {
   signalPassFailure();
 }
}

これでHelloからLLVM IRまでの変換パスができたことになります。

長くなりましたが、最後にこれらのPassをPassManagerを使って登録するとLoweringの準備が完了します。

int loadAndProcessMLIR(mlir::MLIRContext &context, mlir::OwningModuleRef &module) {
 // MLIRファイルを読み込む
 if (int error = loadMLIR(context, module)) {
   return error;
 }

 // 作ったPassを登録する
 mlir::PassManager passManager(&context);
 mlir::OpPassManager &optPm = passManager.nest<mlir::FuncOp>();
 optPm.addPass(hello::createLowerToAffinePass());
 passManager.addPass(hello::createLowerToLLVMPass());

 if (mlir::failed(passManager.run(*module))) {
   return 4;
 }

 return 0;
}

Moduleには変換後のIRが入っています。実行してみましょう。

hello-optでLLVM IRを生成する

リポジトリではhello-optというCLIツールがあるのでこれでLLVM IRを生成させてみます。以下のようなHello Dialectで書かれたMLIRを用意します。

func @main() {
   %0 = "hello.constant"() {value = dense<1.0> : tensor<2x3xf64>} : () -> tensor<2x3xf64>
   "hello.print"(%0) : (tensor<2x3xf64>) -> ()
   return
}

2x3の大きさのテンソルを標準出力に表示するだけのプログラムです。constantとprintは自分で作ったOperation。returnはもともとStandard Dialectで定義されたOperationです。この時点で異なるDialectが混在していることがわかります。

mlir-hellloリポジトリを以下のようにビルドします。

mkdir build && cd build
LLVM_DIR=$HOME/dev/llvm-project/build/lib/cmake/llvm \
  MLIR_DIR=$HOME/dev/llvm-project/build/lib/cmake/mlir \
  cmake -G Ninja ..
  
cmake --build . --target hello-opt

hello-optでprint.mlirを変換します。以下のようなLLVM IRが出力されます。

./bin/hello-opt ../test/Hello/print.mlir

; ModuleID = 'LLVMDialectModule'
source_filename = "LLVMDialectModule"
target datalayout = "e-m:o-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-apple-darwin19.6.0"

@frmt_spec = internal constant [4 x i8] c"%f \00"

; Function Attrs: nofree nounwind
declare noundef i32 @printf(i8* nocapture noundef readonly, ...) local_unnamed_addr #0

; Function Attrs: nounwind
define void @main() local_unnamed_addr #1 !dbg !3 {
.preheader:
 %0 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00), !dbg !7
 %1 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00), !dbg !7
 %2 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00), !dbg !7
 %putchar = tail call i32 @putchar(i32 10), !dbg !7
 %3 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00), !dbg !7
 %4 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00), !dbg !7
 %5 = tail call i32 (i8*, ...) @printf(i8* nonnull dereferenceable(1) getelementptr inbounds ([4 x i8], [4 x i8]* @frmt_spec, i64 0, i64 0), double 1.000000e+00), !dbg !7
 %putchar.1 = tail call i32 @putchar(i32 10), !dbg !7
 ret void, !dbg !9
}

; Function Attrs: nofree nounwind
declare noundef i32 @putchar(i32 noundef) local_unnamed_addr #0

attributes #0 = { nofree nounwind }
attributes #1 = { nounwind }

!llvm.dbg.cu = !{!0}
!llvm.module.flags = !{!2}

!0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1, producer: "mlir", isOptimized: true, runtimeVersion: 0, emissionKind: FullDebug)
!1 = !DIFile(filename: "LLVMDialectModule", directory: "/")
!2 = !{i32 2, !"Debug Info Version", i32 3}
!3 = distinct !DISubprogram(name: "main", linkageName: "main", scope: null, file: !4, line: 3, type: !5, scopeLine: 3, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !6)
!4 = !DIFile(filename: "../test/Hello/print.mlir", directory: "/Users/sasaki/dev/mlir-hello/build")
!5 = !DISubroutineType(types: !6)
!6 = !{}
!7 = !DILocation(line: 5, column: 5, scope: !8)
!8 = !DILexicalBlockFile(scope: !3, file: !4, discriminator: 0)
!9 = !DILocation(line: 6, column: 5, scope: !8)

これを実行するにはJIT環境でLLVM bitcodeを実行するlliというツールを使います。出力をlliに食わせてみましょう。

./bin/hello-opt ../test/Hello/print.mlir | lli
1.000000 1.000000 1.000000
1.000000 1.000000 1.000000

無事動きました！

ちなみにODSフレームワークで記述しておくと以下のようにドキュメントを生成することも簡単にできます。

cmake --build . --target mlir-doc

まとめ

長くなりましたがMLIRを使ってTensorを出力するDialectを作ってみました。MLIRのチュートリアル的なものを目指したため今回のプログラムではMLIRのほんの一部の機能しか使っていません。他にもInterfaceやVerifierなどカバーしていないトピックがたくさんあります。興味ありましたら、公式ドキュメントをご参照ください。

参照

- 2020-02-26 - CGO 2020 Talk
- MLIR Tutorial
- Lewuathe/mlir-hello

お読みいただきありがとうございます。もしよければスキやシェアしていただけると大変うれしいです！