はじめに

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

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

前回まででシンプルなソースコードをバイナリまでコンパイルできるようになりましたが、突貫工事のため設計上の課題が残っています。今回は今後の機能追加のために修正を行っていきます。

問題点: サブプログラムを複数作成できない

前回までの実装で、サブプログラム（※他言語の関数に相当）を実装しました。しかし、変数や組み込み命令の初期化をサブプログラムで行っているため、2つ以上のサブプログラムを作成すると名前が衝突してしまいました。

Loweringの実装

class SubprogramOpLowering : public mlir::ConversionPattern {
public:
  mlir::LogicalResult
  matchAndRewrite(mlir::Operation *Op, mlir::ArrayRef<mlir::Value> Operands,
                  mlir::ConversionPatternRewriter &Rewriter) const override {
    // ...
    // main関数を作成
    auto Func = Rewriter.create<mlir::LLVM::LLVMFuncOp>(Loc, "main", DummyType);
    Rewriter.inlineRegionBefore(Subprogram.getBody(), Func.getBody(),
                                Func.end());
    mlir::Block *EntryBlock = &Func.getBody().front();
    // ...
  }
};

{ main
    set,10
    prt
    hlt
}

{ dummy
    set,20
}

loc("example/simple.filska":7:8): error: redefinition of symbol named 'main'

サブプログラムのたびにmain関数ができてしまう

// -----// IR Dump After (anonymous namespace)::FilskalangToLLVMLoweringPass Failed () //----- //
"builtin.module"() ({
  "llvm.mlir.global"() <{addr_space = 0 : i32, constant, global_type = !llvm.array<3 x i8>, linkage = #llvm.linkage<internal>, sym_name = "frmt_spec", value = "%f\00", visibility_ = 0 : i64}> ({
  }) : () -> ()
  "llvm.func"() <{CConv = #llvm.cconv<ccc>, function_type = !llvm.func<f64 (ptr, ...)>, linkage = #llvm.linkage<external>, sym_name = "printf", visibility_ = 0 : i64}> ({
  }) : () -> ()
  "llvm.func"() <{CConv = #llvm.cconv<ccc>, function_type = !llvm.func<void ()>, linkage = #llvm.linkage<external>, sym_name = "main", visibility_ = 0 : i64}> ({
    ...
  }) : () -> ()
  "llvm.func"() <{CConv = #llvm.cconv<ccc>, function_type = !llvm.func<void ()>, linkage = #llvm.linkage<external>, sym_name = "main", visibility_ = 0 : i64}> ({
    ...
  }) : () -> ()
}) : () -> ()

解決策：共通処理を行うための命令を作成

(以下紹介する実装はかなり無理のある汚い部分もあります。記事化しておいてなんですが、参考にならないかもしれません...)

上記問題を解決するため、プログラム全体の共通処理に対応付けるための命令を用意しました。プログラム全体をFilska dialectの命令 ProgramOp に対応付けます。

作戦の概略は以下の通りです。

MLIRのdialectでProgramOp（共通処理置き場命令）を作成

まず、Filskaのdialectを生成する際、各 SubprogramOp の前に ProgramOp を生成します。

filskalang dialectのイメージ

// 共通処理用の命令。プログラム中に一度だけ生成される
"filskalang.program"() // ...

// 以下各サブプログラムとその命令一覧
"filskalang.subprogram"() <{function_type = () -> (), sym_name = "main"}> ({
  // ...
}) : () -> ()

"filskalang.subprogram"() <{function_type = () -> (), sym_name = "dummy"}> ({
  // ...
}) : () -> ()

ProgramOpを実際の初期化処理にlowering

続いて、LLVM IRにloweringする際に、 ProgramOp を初期化処理に変換します。現状ではレジスタ(可変な変数に相当、x, y, z, m の4種しか使えない)の初期化を行っています。

LLVM IRのイメージ

  llvm.func @main() {
    // filskalang.ProgramOp を以下の初期化処理に変換
    %0 = llvm.mlir.constant(0 : index) : i64
    %1 = llvm.alloca %0 x f64 : (i64) -> !llvm.ptr
    %2 = llvm.mlir.constant(0 : index) : i64
    %3 = llvm.alloca %2 x f64 : (i64) -> !llvm.ptr
    // SubProgram mainを実行するためにbr命令でジャンプ
    llvm.br ^bb1

  ^bb1:  // サブプログラム mainを以下の処理に変換
    // ...

  ^bb2:  // サブプログラム dummyを以下の処理に変換
    // ※実際にはどこからも呼ばれないため実行されない)
    // ...
  }

サブプログラムを関数ではなくブロックに対応付けているのは、コールスタックのスタックオーバーフローを防ぐためです。サブプログラムは関数と異なり、ジャンプ先からジャンプ元へ戻ることはありません。そのため、サブプログラムのジャンプを関数呼び出しとして実装するとコールスタックが溜まり続けてしまいます。

また、ブロックを使用するため初期化処理の末尾でmain サブプログラム用のブロックへジャンプする必要があります。言い換えると、初期化処理時に main 用ブロック（上記の bb1）が存在している必要があります。そのため、ハックとして main の SubprogramOp ではなくProgramOp のlowering時にこのブロックを初期化しています。

lib/CodeGen/LowerToLLVM.cpp

class ProgramOpLowering : public mlir::ConversionPattern {
public:
  mlir::LogicalResult
  matchAndRewrite(mlir::Operation *Op, mlir::ArrayRef<mlir::Value> Operands,
                  mlir::ConversionPatternRewriter &Rewriter) const override {
    // ...

    // 初期化処理とは別に、サブプログラムmain に対応するブロックも生成
    auto MainBlock = &MainFunc.getBody().front();
    auto InitBlock = Rewriter.createBlock(ExitBlock);

    Rewriter.setInsertionPointToEnd(InitBlock);
    // 初期化処理の最後にmainへのbr命令を追加
    Rewriter.create<mlir::LLVM::BrOp>(Loc, MainBlock);

    // 命令の挿入場所を元に戻す
    Rewriter.setInsertionPointToStart(Program->getBlock());

    // ...
  }
};

サブプログラムの無限ループを実装

上記はサブプログラムの実装として不十分です。Filskaのサブプログラムは終端の命令が評価された後最初の命令の評価に戻る（無限ループする）ため、ブロック最後の命令は自身へのbr命令である必要があります。

  ^bb1:  // サブプログラム mainを以下の処理に変換
    // ...
    llvm.br ^bb1

そこで、filskalang dialectの命令一覧を挿入後末尾にbranch命令を追加しています。

lib/CodeGen/LowerToLLVM.cpp

class SubprogramOpLowering : public mlir::ConversionPattern {
public:
  mlir::LogicalResult
  matchAndRewrite(mlir::Operation *Op, mlir::ArrayRef<mlir::Value> Operands,
                  mlir::ConversionPatternRewriter &Rewriter) const override {
    // ...

    // HACK: main用ブロックはProgramOpのloweringですでに生成しているため新規作成しない
    mlir::Block *Blk;
    if (Subprogram.getName().str() == "main") {
      Blk = MainBlock;
    } else {
      Blk = Rewriter.createBlock(MainBlock);
    }

    // filskalang.SubprogramOpの命令一覧をmain用ブロックに追加
    Rewriter.inlineBlockBefore(&Subprogram.getBody().front(), Blk,
                               Blk->getTerminator()->getIterator());

    Rewriter.setInsertionPointToEnd(Blk);
    // ブロック末尾に自身へのbr命令を追加
    Rewriter.create<mlir::LLVM::BrOp>(Loc, Blk);
  }
};

hlt命令に対応

無限ループに対応したため、Filskaプログラムを終了させる hlt 命令も実装します。
LLVM IRでは早期リターンを行うことができない(return や br 等の命令(terminator) はブロック末尾にただ1つ存在できる)ため、return用のブロックを別途作成しそこへ br 命令でジャンプするようにします。

イメージ

llvm.func @main() {
  // ...

  ^bb1:  // サブプログラム main用ブロック
    // ...
    // hltがある場合自身へジャンプする代わりに終了用ブロックへジャンプ
    llvm.br ^bb2

  ^bb2: // 終了用ブロック
    // 何もせずreturnすることでプログラムを終了
    llvm.return
}

終了用ブロックはプログラムに1つ存在すればよいため ProgramOp のloweringで生成します。
（グローバル変数を使っていて汚いので実装は要修正）

lib/CodeGen/LowerToLLVM.cpp

class SubprogramOpLowering : public mlir::ConversionPattern {
public:
  mlir::LogicalResult
  matchAndRewrite(mlir::Operation *Op, mlir::ArrayRef<mlir::Value> Operands,
                  mlir::ConversionPatternRewriter &Rewriter) const override {
    // ...
    auto MainBlock = &MainFunc.getBody().front();
    // 終了用ブロックを追加
    ExitBlock = Rewriter.createBlock(MainBlock);
    auto InitBlock = Rewriter.createBlock(ExitBlock);

    // returnのみ挿入
    Rewriter.setInsertionPointToEnd(ExitBlock);
    Rewriter.create<mlir::LLVM::ReturnOp>(Loc, mlir::ArrayRef<mlir::Value>());
    // ...
  }
};

このブロックへのジャンプは HltOp のlowering時に追加されます。前述の通り早期リターンはできないため、ブロック内の hlt 以降の命令をすべて消すことでエラーを回避しています（ゴリ押し）。

lib/CodeGen/LowerToLLVM.cpp

class HltOpLowering : public mlir::ConversionPattern {
public:
  mlir::LogicalResult
  matchAndRewrite(mlir::Operation *Op, mlir::ArrayRef<mlir::Value> Operands,
                  mlir::ConversionPatternRewriter &Rewriter) const override {
    // ...
    // 終了用ブロックへのbr命令を追加
    auto BrOp = Rewriter.create<mlir::LLVM::BrOp>(Loc, ExitBlock);

    // HACK: br命令の後に別の命令が存在してはいけないため、以降の命令を削除
    // Op.isAfterInBlock が無いため回りくどい方法で比較...
    bool IsAfter = false;
    for (auto &Op : Rewriter.getBlock()->getOperations()) {
      if (!IsAfter && !Op.isBeforeInBlock(BrOp)) {
        // BrOp自身の場合のみこのケースに該当するため、以降は削除してよい
        IsAfter = true;
        continue;
      }
      if (IsAfter) {
        Rewriter.eraseOp(&Op);
      }
    }

    Rewriter.create<mlir::LLVM::BrOp>(Loc, ExitBlock);
    // ...
  }
};

実際に実装したコードが以下になります（長いので省略）。

これでサブプログラムを複数書いても想定通り実行できるようになりました。

{ main
    set,10
    prt
    hlt
}

{ dummy
    set,20
}

module {
  llvm.mlir.global internal constant @frmt_spec("%f\00") {addr_space = 0 : i32}
  llvm.func @printf(!llvm.ptr, ...) -> f64
  llvm.func @main() {
    // 初期化処理
    %0 = llvm.mlir.constant(0 : index) : i64
    %1 = llvm.alloca %0 x f64 : (i64) -> !llvm.ptr
    %2 = llvm.mlir.constant(0 : index) : i64
    %3 = llvm.alloca %2 x f64 : (i64) -> !llvm.ptr
    llvm.br ^bb3
  // 終了用ブロック
  ^bb1:  // pred: ^bb3
    llvm.return
  // dummyに対応するブロック
  ^bb2:  // pred: ^bb2
    %4 = llvm.mlir.constant(2.000000e+01 : f64) : f64
    llvm.store %4, %1 : f64, !llvm.ptr
    llvm.br ^bb2
  // mainに対応するブロック
  ^bb3:  // pred: ^bb0
    %5 = llvm.mlir.constant(1.000000e+01 : f64) : f64
    llvm.store %5, %3 : f64, !llvm.ptr
    %6 = llvm.load %3 : !llvm.ptr -> f64
    %7 = llvm.mlir.addressof @frmt_spec : !llvm.ptr
    %8 = llvm.mlir.constant(0 : index) : i64
    %9 = llvm.getelementptr %7[%8, %8] : (!llvm.ptr, i64, i64) -> !llvm.ptr, !llvm.array<3 x i8>
    %10 = llvm.call @printf(%9, %6) vararg(!llvm.func<f64 (ptr, ...)>) : (!llvm.ptr, f64) -> f64
    llvm.br ^bb1
  }
}

おわりに

今回は設計の修正が中心でした。ブロックの取り回しに苦戦しかなり無理のある実装となってしまっています。Filskaの言語特性を考えるとMLIRは不要で直接LLVM IRを生成した方が綺麗に書けたのかもしれません。とはいえ、実装し始めてしまったものは仕方がありません（そもそもこの実装はMLIRの勉強のためにはじめたもの）。本家Filskaを完全移植できるまで、次回以降も開発を進めていきたいと思います。