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

シリーズ一覧

- [【ストレンジコード】Filskaコンパイラ開発計画(1) 言語仕様を整理する【MLIR】](https://qiita.com/Syuparn/items/e8f4046c4a5b551db310)
- [【ストレンジコード】Filskaコンパイラ開発計画(2) ソースコードをパースする【MLIR】](https://qiita.com/Syuparn/items/ca502ec32aded84ebac1)
- [【ストレンジコード】Filskaコンパイラ開発計画(3) LLVM IRを出力する【MLIR】](https://qiita.com/Syuparn/items/1b6fc01bbe7370578756)
- [【ストレンジコード】Filskaコンパイラ開発計画(4) 実行ファイル出力【MLIR】](https://qiita.com/Syuparn/items/24ac10c76b73485ff85d)
- [【ストレンジコード】Filskaコンパイラ開発計画(5) サブプログラムを複数作成可能にする【MLIR】](https://qiita.com/drafts/35ee40e8b77bb8657cee)
- 【ストレンジコード】Filskaコンパイラ開発計画(6) llvmir以外のdialectを経由したloweringを可能にする【MLIR】(本記事)

# はじめに

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

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

https://www.shuwasystem.co.jp/book/9784798069746.html

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

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

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

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

```filska:ソースコード
{ main
    " mに10を代入
    set,10
    " mを表示
    prt
    " 終了
    hlt
}
```

```llvm:生成されるコード（一部抜粋）
; 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"() : () -> ()
  }) : () -> ()
}
```

```cpp: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();
```

```cpp: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);
```

```cpp: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を行うことにしました[^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` に書き戻す

```filska
{ main
    set,10
    neg
    prt
    hlt
}
```

```
-10
```

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

```tablegen: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` 命令のオペランドにはサブプログラム名の代わりに上記の戻り値を指定します。

```tablegen: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へ変換されます。

```filska:関連個所のみ抜粋
set, 10
neg
```

```:関連個所のみ抜粋
# ...
%0 = "filskalang.register"() <{name = "main"}> : () -> f64
"filskalang.neg"(%0) : (f64) -> f64
```

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

```cpp: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` で読み込み

という流れで変換します[^set]。

https://github.com/Syuparn/filskalang/blob/83e670165fa1e24661a49f0288c86b6baf65e0dd/lib/CodeGen/LowerToLLVM.cpp

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

```: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に追加します。

```tablegen: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されて、最終的にはこうなります。

```llvm dialect(一部抜粋)
%3 = llvm.load %1 : !llvm.ptr -> f64
%4 = llvm.fneg %3  : f64
llvm.store %4, %1 : f64, !llvm.ptr
```

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

```filska
{ main
    set,10
    neg
    prt
    hlt
}
```

```bash
$ ./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で何もせず消えます)

```tablegen: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を指定していないのが原因でした。

```diff: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することで確かめる予定です。

[^ssa]: メモリ割り当ては一見非効率ですが、実際には `mem2reg` の最適化によってSSAのレジスタ割り当ての形式に置き換わるためあまり心配いりません。
[^set]: 書き込みの `filskalang.set` は他命令に干渉しないため、サブプログラム名をオペランドに取る実装のままにしています。
[^lowering]: filska dialect生成の時点で `llvm.alloca` 命令を作ってしまい、その戻り値を `filskalang.set`, `filskalang.prt` のオペランドに持たせるというアイデアも考えましたが、filska dialectがllvm dialectと蜜結合してしまうので取りやめました。