mlir

MLIR

目的

エージェントが MLIR（Multi-Level IR）を扱えるようにガイドします。ops、regions、blocks、values、組み込み dialect（arith、func、memref、affine、linalg）、ODS を使ったカスタム dialect の記述、ConversionPattern を使った lowering pass、mlir-opt CLI、ML コンパイラーのユースケース（Torch-MLIR、IREE）を扱います。

使用する場面

• ドメイン固有のコンパイラー IR（グラフィックス、ML、ハードウェア DSL）の構築
• 高水準 ops を LLVM または GPU dialect に lowering する
• 段階的な lowering パイプライン（linalg → loops → LLVM）の記述
• ML のデプロイに向けた IREE または Torch-MLIR との統合
• dialect 間で再利用できる変換 pass の作成
• 適切な抽象化レベルでのコンパイラー最適化の試作

ワークフロー

1. MLIR の構造

Module
└── func.func @main()
    └── region
        └── block ^bb0:
            └── operations (ops) producing SSA values

主な概念：

• Operation — 命令のようなノード（arith.addi、memref.load）
• Region — ブロックのコンテナー（関数、制御フロー）
• Block — 順序付けられた ops を持つ CFG ノード
• Value — op の SSA 結果またはブロック引数

2. 組み込み dialect

Dialect	目的
arith	整数／浮動小数点の算術演算
func	関数の定義と呼び出し
memref	形状とストライドを持つバッファーの抽象化
affine	アフィンループネスト、マップ／セット制約
linalg	構造化線形代数 ops
scf	構造化制御フロー（for、if）
llvm	最終的な lowering 用の LLVM IR dialect
gpu	GPU カーネルの起動

// example.mlir
func.func @add(%a: memref<4xf32>, %b: memref<4xf32>, %c: memref<4xf32>) {
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : index
  scf.for %i = %c0 to %c4 step %c1 {
    %av = memref.load %a[%i] : memref<4xf32>
    %bv = memref.load %b[%i] : memref<4xf32>
    %sum = arith.addf %av, %bv : f32
    memref.store %sum, %c[%i] : memref<4xf32>
  }
  return
}

3. mlir-opt CLI

# Parse and print
mlir-opt example.mlir

# Run canonicalization
mlir-opt example.mlir -canonicalize

# Lower affine to scf
mlir-opt affine.mlir -lower-affine

# Full pipeline toward LLVM
mlir-opt input.mlir \
  --linalg-bufferize \
  --convert-linalg-to-loops \
  --convert-scf-to-cf \
  --convert-arith-to-llvm \
  --convert-memref-to-llvm \
  --convert-func-to-llvm \
  -o llvm.mlir

4. ODS — Operation Definition Specification

// MyOps.td
include "mlir/IR/OpBase.td"

def My_Dialect : Dialect {
    let name = "my";
    let summary = "My custom dialect";
}

class My_Op<string mnemonic, list<Trait> traits = []> :
    Op<My_Dialect, mnemonic, traits>;

def AddOp : My_Op<"add", [Pure]> {
    let summary = "Add two values";
    let arguments = (ins AnyType:$lhs, AnyType:$rhs);
    let results = (outs AnyType:$result);
    let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($result)";
}

# Generate C++ from TableGen
mlir-tblgen -gen-op-defs MyOps.td -I include/ -o MyOps.cpp.inc

5. カスタム dialect の C++ 実装

#include "mlir/IR/DialectImplementation.h"
#include "MyDialect.h"

#include "MyOps.cpp.inc"

void MyDialect::initialize() {
    addOperations<
#define GET_OP_LIST
#include "MyOps.cpp.inc"
    >();
}

#define GET_OP_CLASSES
#include "MyOps.cpp.inc"

6. Lowering pass

#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Transforms/DialectConversion.h"

struct AddOpLowering : OpConversionPattern<my::AddOp> {
    using OpConversionPattern::OpConversionPattern;

    LogicalResult matchAndRewrite(my::AddOp op, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const override {
        rewriter.replaceOpWithNewOp<arith::AddIOp>(op, adaptor.getLhs(), adaptor.getRhs());
        return success();
    }
};

void populateLoweringPatterns(RewritePatternSet &patterns) {
    patterns.add<AddOpLowering>(patterns.getContext());
}

// In pass:
mlir::ConversionTarget target(*context);
target.addIllegalDialect<my::MyDialect>();
target.addLegalDialect<arith::ArithDialect>();

if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();

7. ML コンパイラー向けの linalg

%0 = linalg.matmul ins(%A, %B : tensor<128x256xf32>, tensor<256x64xf32>)
                   outs(%C : tensor<128x64xf32>) -> tensor<128x64xf32>

Lowering の経路：linalg → scf loops → affine → llvm

8. Torch-MLIR と IREE

# Torch-MLIR: PyTorch → MLIR
python -m torch_mlir.tools.import-onnx --onnx-model model.onnx -o model.mlir

# IREE: MLIR → GPU/CPU executable
iree-compile --iree-hal-target-backends=llvm-cpu model.mlir -o model.vmfb
iree-run-module --module=model.vmfb --function=main

よくある問題

症状	原因	解決策
Dialect が登録されていない	registerDialect が不足している	ツール／pass の初期化時に登録する
ODS のビルドに失敗する	TableGen のインクルードパス	mlir/IR/OpBase.td に対する -I を確認する
Lowering が完了しない	不正な ops が残っている	--mlir-print-ir-after-failure でデバッグする
パターン内の型が一致しない	adaptor の型が正しくない	型付きの OpAdaptor アクセサーを使う
mlir-opt がクラッシュする	IR が無効	-verify-each を実行する
Lowering 後に関数が空になる	すべての ops が不正とされ、変換されていない	不足しているパターンを追加する

関連スキル

• skills/compiler-internals/llvm-passes — LLVM pass に相当するもの
• skills/compiler-internals/compiler-frontend — AST から MLIR へのインポート
• skills/compiler-internals/jit-compilation — MLIR→LLVM の JIT コンパイル
• skills/compilers/llvm — LLVM IR の出力先
• skills/gpu/cuda — GPU dialect の lowering 先
• skills/gpu/triton-lang — 代替となる GPU カーネル IR