【ストレンジコード】Filskaコンパイラ開発計画(6) llvmir以外のdialectを経由したloweringを可能にする【MLIR】

シリーズ一覧

• 【ストレンジコード】Filskaコンパイラ開発計画(1) 言語仕様を整理する【MLIR】
• 【ストレンジコード】Filskaコンパイラ開発計画(2) ソースコードをパースする【MLIR】
• 【ストレンジコード】Filskaコンパイラ開発計画(3) LLVM IRを出力する【MLIR】
• 【ストレンジコード】Filskaコンパイラ開発計画(4) 実行ファイル出力【MLIR】
• 【ストレンジコード】Filskaコンパイラ開発計画(5) サブプログラムを複数作成可能にする【MLIR】
• 【ストレンジコード】Filskaコンパイラ開発計画(6) llvmir以外のdialectを経由したloweringを可能にする【MLIR】(本記事)

はじめに

本シリーズでは、『ストレンジコード』に登場する難解プログラミング言語「Filska」をMLIRへ移植します。本家がPython製インタプリタなので、コンパイラにすることで高速化できるのでは？という淡い期待も寄せています。

（ストレンジコードにはFilska以外の難解プログラミング言語もたくさん登場します！言語好きは必見です （ダイマ））

前回に引き続き、今後の機能追加のために修正を行っていきます。
メモリ（他言語でいう変数に相当）をLLVM IR以外にlowering可能な形式に改修します。

背景: レジスタの表現

Filskaでは、値を格納するためにレジスタを使用します。使用できるレジスタは以下の4種類のみです。

• m: サブプログラムごとに1つ、共有不可
• x, y, z: 全サブプログラムで共有

LLVM IRでは静的単一代入形式(SSA)を採用しているため、再代入を表現することができません。そのため、代わりに alloca により割り当てたメモリ領域に load, store することで読み書きしています¹。

ソースコード

{ main
    " mに10を代入
    set,10
    " mを表示
    prt
    " 終了
    hlt
}

生成されるコード（一部抜粋）

; mainサブプログラムのレジスタmを初期化
%1 = alloca double, i64 0, align 8, !dbg !8
; set命令でmを更新
store double 1.000000e+01, ptr %1, align 8, !dbg !9
; prt命令でmを表示
%2 = load double, ptr %1, align 8, !dbg !10
%3 = call double (ptr, ...) @printf(ptr @frmt_spec, double %2), !dbg !10

課題

現状の実装

現状、レジスタ m をFilskaのdialectで表現する際にサブプログラム名を使用しています。
引数のサブプログラム名をLLVM IRへの変換(lowering)時にallocaの戻り値を参照することでレジスタへの読み書きを実現しています。

filskalang dialect

module {
  "filskalang.subprogram"() <{function_type = () -> (), sym_name = "main"}> ({
    // オペランドのサブプログラム名が `m` の識別に利用される（`m` はサブプログラムごとに異なる点に注意！）
    "filskalang.set"() <{subprogramName = "main", value = 1.000000e+01 : f64}> : () -> ()
    "filskalang.prt"() <{subprogramName = "main"}> : () -> ()
    "filskalang.hlt"() : () -> ()
  }) : () -> ()
}

lib/CodeGen/LowerToLLVM.cpp

// 初期化(Program内部)

// レジスタ `m` を定義するためにalloca命令を挿入
Rewriter.setInsertionPointToStart(EntryBlock);
mlir::Value Cst0 = Rewriter.create<mlir::LLVM::ConstantOp>(
    Loc, Rewriter.getI64Type(), Rewriter.getIndexAttr(0));
auto Alloca = Rewriter.create<mlir::LLVM::AllocaOp>(
    Loc, /*resultType*/ mlir::LLVM::LLVMPointerType::get(Context),
    /*elementType*/ Rewriter.getF64Type(), Cst0);
// HACK: mapへ格納し他命令のlowering時に参照可能にする
SubprogramMemory[Subprogram.getName().str()] = Alloca.getRes();

lib/CodeGen/LowerToLLVM.cpp

// 更新 (set)

// サブプログラム名から対応するalloca命令を取得
auto MemoryPointer = SubprogramMemory.at(SetOp.getSubprogramName().str());
auto Value = llvm::APFloat(SetOp.getValue());
// 指定した命令へstore
Rewriter.create<mlir::LLVM::StoreOp>(Loc, ConstantOp, MemoryPointer);

lib/CodeGen/LowerToLLVM.cpp

// 参照 (prt)

auto MemoryPointer = SubprogramMemory.at(PrtOp.getSubprogramName().str());
// 値を取得
auto Load = Rewriter.create<mlir::LLVM::LoadOp>(
    Loc, mlir::Float64Type::get(Context), MemoryPointer);

arith dialectへloweringできない

上記の実装のままではFilskaのdialectをllvm dialect以外にloweringすることができませんでした。

MLIRでは本来命令を段階的に変換することができ、算術関連の命令(add, sub, neg etc.) を arith dialect に変換することで

filskalang -> arith -> llvm

のような経路でloweringができます。こうすることで、自前実装不要でarith dialectに用意されている最適化の恩恵を受けられます。

しかし、arith dialectの各命令は F64 等の数値のオペランドしか受け取ることができず、サブプログラム名を渡す現状の方法は使えません。

こんなことは無理

%0 = arith.negf (subprogramName = "main") : f64

解決案

そこで、以下の方法でloweringを行うことにしました²。

• filska dialect
  • レジスタを表す命令 filskalang.register を作る
  • 上記命令の戻り値をarith dialectで使用可能な F64 型にする
  • 戻り値を filslalang.prt, filslalang.neg 等のオペランドに渡す
• arith dialect
  • filslalang.neg 等を arithの対応する命令 (例： arith.negf ) に置換する
  • 置換後の命令に同じオペランドを指定する
• llvm dialect
  • filskalang.register を llvm.load に置換することで、arith dialectのオペランドをloadの戻り値にする
  • arith -> llvmのloweringで算術命令が対応するllvm dialectの命令に置換される (例： llvm.negf )

実装

上記の案を実際試してみます。arithへのloweringとして、Filskaの neg 命令を実装します。

• neg: サブプログラムのレジスタ m に格納されている値を符号反転して m に書き戻す

{ main
    set,10
    neg
    prt
    hlt
}

-10

レジスタを表す命令を追加

まずはレジスタを表す命令をFilskaのdialectに追加します。

include/filskalang/CodeGen/Ops.td

def RegisterOp : Filskalang_Op<"register"> {
  let summary = "register operation";
  let description = [{
    The "filskalang.register" builtin operation represents the value in the register `x`, `y`, `z`, or `m`.
  }];

  # 注： `m` はサブプログラムごとに異なる値を指すため、
  let arguments = (ins SymbolNameAttr:$name);
  let results = (outs F64);
}

arith dialectの命令のオペランドに渡すには戻り値が F64 でありさえすればよいので、 let results = (outs F64); で戻り値型を指定しています。

neg 命令のオペランドにはサブプログラム名の代わりに上記の戻り値を指定します。

include/filskalang/CodeGen/Ops.td

def NegOp : Filskalang_Op<"neg"> {
  let summary = "negate operation";
  let description = [{
    The "filskalang.neg" builtin operation represents the "neg" instruction
  }];

  let arguments = (ins F64:$arg);
  let results = (outs F64);
}

この時点で、ソースコードは以下のFilska Dialectへ変換されます。

関連個所のみ抜粋

set, 10
neg

関連個所のみ抜粋

# ...
%0 = "filskalang.register"() <{name = "main"}> : () -> f64
"filskalang.neg"(%0) : (f64) -> f64

arith dialectへのlowering

続いて arith dialectへのloweringを行います。オペランドはすでにarith dialectで使用可能な形になっているため、命令を置換するだけで実装完了です。

lib/CodeGen/LowerToArith.cpp

struct NegOpLowering : public mlir::ConversionPattern {
public:
  mlir::LogicalResult
  matchAndRewrite(mlir::Operation *Op, mlir::ArrayRef<mlir::Value> Operands,
                  mlir::ConversionPatternRewriter &Rewriter) const override {
    auto Loc = Op->getLoc();
    auto NegOp = mlir::cast<mlir::filskalang::NegOp>(Op);

    auto NegFOp = Rewriter.create<mlir::arith::NegFOp>(Loc, NegOp.getArg());

    Rewriter.replaceOp(Op, NegFOp);
    return mlir::success();
  }
};

命令は以下のように変換されます。

arith dialect(関連個所のみ抜粋)

# ...
%0 = "filskalang.register"() <{name = "main"}> : () -> f64
arith.negf %0 : f64

llvm dialectへのlowering時にメモリ処理へ変換

最後に、filskalang.register をLLVM IRのメモリ操作命令に対応付けます。

具体的には

• filskalang.program (プログラム全体) のlowering時に llvm.alloca で m を初期化
• filskalang.register のlowering時に、上記で確保したメモリを llvm.load で読み込み

という流れで変換します³。

最終的に、以下のような命令が出力されます。

llvm dialect(関連個所のみ抜粋)

# ...
%1 = llvm.alloca %0 x f64 : (i64) -> !llvm.ptr
# ...
%3 = llvm.load %1 : !llvm.ptr -> f64
%4 = llvm.fneg %3  : f64

再度計算結果を m へ書き戻す

LLVM IRへ無事loweringできるようになりましたが、まだ想定通りの動作にはなっていません。
Filskaの neg 命令は m を符号反転し その結果を m に格納する なので、格納する部分も実装します。

戻り値を再度 m に格納する命令をdialectに追加します。

include/filskalang/CodeGen/Ops.td

def MetaSetOp : Filskalang_Op<"metaset"> {
  let summary = "meta set operation";
  let description = [{
    The "filskalang.metaset" operation is used to push back the caluculation result to register `m`.
  }];

  let arguments = (ins F64:$value, SymbolNameAttr:$subprogramName);
}

上記 MetaSetOp はオペランドに neg 等の戻り値を取ります。生成されるMLIRは以下の通りです。

filskalang dialect(一部抜粋)

%1 = "filskalang.neg"(%0) : (f64) -> f64
"filskalang.metaset"(%1) <{subprogramName = "main"}> : (f64) -> ()
%2 = "filskalang.register"() <{name = "main"}> : () -> f64

loweringされて、最終的にはこうなります。

%3 = llvm.load %1 : !llvm.ptr -> f64
%4 = llvm.fneg %3  : f64
llvm.store %4, %1 : f64, !llvm.ptr

これで想定通りの結果が得られました。

{ main
    set,10
    neg
    prt
    hlt
}

$ ./filskac tests/src/neg.filska 
$ ./tests/src/neg
-10.000000

はまったところ

terminatorが存在しない

arith dialectにloweringする際に以下エラーが発生しました。

loc("tests/src/subprogram.filska":4:8): error: block with no terminator, has "filskalang.hlt"() : () -> ()

MLIRではブロックの最後の命令がterminator (llvm.return 等のブロック末尾で使用することが想定される命令）である必要があるため、そのバリデーションチェックが働いたと考えられます。

一方、Filskaのサブプログラムは終端命令の評価が終わると先頭にジャンプするため、最後の命令という概念が存在しません。

そこで、terminatorの制約を満たすためだけの命令を末尾に挿入することにしました。
(役割は無いためloweringで何もせず消えます)

include/filskalang/CodeGen/Ops.td

# NOTE: Terminatorとして使えるよう指定
def DummyTerminatorOp : Filskalang_Op<"dummyterminator", [Terminator]> {
  let summary = "dummy terminator operation";
  let description = [{
    The "filskalang.dummyterminator" operation is a placeholder operator at the end of each subprogram.
    It is only used to satisfy the terminator constraint in an MLIR block;
    the block must have one and only one terminator operator at the end.
  }];
}

生成される命令

module {
  "filskalang.program"() <{function_type = () -> (), sym_name = "program"}> ({
    "filskalang.dummyterminator"() : () -> ()
  }) : () -> ()
  "filskalang.subprogram"() <{function_type = () -> (), sym_name = "main"}> ({
    "filskalang.set"() <{subprogramName = "main", value = 1.000000e+01 : f64}> : () -> ()
    %0 = "filskalang.register"() <{name = "main"}> : () -> f64
    %1 = "filskalang.neg"(%0) : (f64) -> f64
    "filskalang.metaset"(%1) <{subprogramName = "main"}> : (f64) -> ()
    %2 = "filskalang.register"() <{name = "main"}> : () -> f64
    "filskalang.prt"(%2) : (f64) -> ()
    "filskalang.hlt"() : () -> ()
    // terminator制約を満たすためだけに挿入
    "filskalang.dummyterminator"() : () -> ()
  }) : () -> ()
}

エラー it isn't known in this MLIRContext

lowering中に命令が認識できないというエラーが発生しました。

LLVM ERROR: Building op `arith.negf` but it isn't known in this MLIRContext: the dialect may not be loaded or this operation hasn't been added by the dialect. See also https://mlir.llvm.org/getting_started/Faq/#registered-loaded-dependent-whats-up-with-dialects-management
Aborted

loweringPassの getDependentDialects にdialectを指定していないのが原因でした。

lib/CodeGen/LowerToArith.cpp

struct FilskalangToArithLoweringPass
    : public mlir::PassWrapper<FilskalangToArithLoweringPass,
                               mlir::OperationPass<mlir::ModuleOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FilskalangToArithLoweringPass)

  void getDependentDialects(mlir::DialectRegistry &registry) const override {
 +   registry.insert<mlir::arith::ArithDialect>();
  }
  void runOnOperation() final;
};

おわりに

以上、FilskaのMLIR移植の進捗でした。今回も前回に引き続き修正が中心でした。勉強しながらの実装なので、徐々にダーティハックに染まってしまっています...
メモリをdialectの命令で表現する方法が本当にうまくいくか、次回以降 add 命令を filskalang dialectからarith dialectへloweringすることで確かめる予定です。

1. メモリ割り当ては一見非効率ですが、実際には mem2reg の最適化によってSSAのレジスタ割り当ての形式に置き換わるためあまり心配いりません。 ↩

2. filska dialect生成の時点で llvm.alloca 命令を作ってしまい、その戻り値を filskalang.set, filskalang.prt のオペランドに持たせるというアイデアも考えましたが、filska dialectがllvm dialectと蜜結合してしまうので取りやめました。 ↩

3. 書き込みの filskalang.set は他命令に干渉しないため、サブプログラム名をオペランドに取る実装のままにしています。 ↩