mlir-development

MLIR 開発スキル

このスキルでは、ドメイン固有コンパイラーや高水準の最適化パイプラインを構築するための MLIR（Multi-Level Intermediate Representation）開発を扱います。

MLIR の概要

MLIR とは？

MLIR は、再利用可能で拡張可能なコンパイラーコンポーネントの構築を可能にするコンパイラー基盤です。次の機能を提供します。

• 階層型のマルチレベル IR 表現
• 拡張可能な操作および型システム
• 抽象化レベル間の段階的な低水準化
• 豊富な変換インフラストラクチャー

アーキテクチャー

High-Level DSL
     ↓
Domain-Specific Dialects (e.g., TensorFlow, PyTorch)
     ↓
Mid-Level Dialects (e.g., Linalg, Affine)
     ↓
Low-Level Dialects (e.g., LLVM, GPU)
     ↓
Target Code

コアコンセプト

ダイアレクト

ダイアレクトは、操作、型、属性をまとめたものです。

// Define a custom dialect
class MyDialect : public mlir::Dialect {
public:
    explicit MyDialect(mlir::MLIRContext *context)
        : Dialect("my_dialect", context, 
                  mlir::TypeID::get<MyDialect>()) {
        addOperations<
            MyAddOp,
            MyMulOp,
            MyFuncOp
        >();
        addTypes<MyTensorType>();
    }
    
    static llvm::StringRef getDialectNamespace() { 
        return "my_dialect"; 
    }
};

操作

// Define using ODS (Operation Definition Specification)
// In TableGen file (.td)
def MyAddOp : Op<MyDialect, "add", [Pure]> {
    let summary = "Add two tensors";
    let description = [{
        Performs element-wise addition of two tensors.
    }];
    
    let arguments = (ins 
        AnyTensor:$lhs,
        AnyTensor:$rhs
    );
    
    let results = (outs 
        AnyTensor:$result
    );
    
    let assemblyFormat = [{
        $lhs `,` $rhs attr-dict `:` type($result)
    }];
}

型と属性

// Custom type definition
class MyTensorType : public mlir::Type::TypeBase<
    MyTensorType, mlir::Type, MyTensorTypeStorage> {
public:
    using Base::Base;
    
    static MyTensorType get(mlir::MLIRContext *context,
                            llvm::ArrayRef<int64_t> shape,
                            mlir::Type elementType) {
        return Base::get(context, shape, elementType);
    }
    
    llvm::ArrayRef<int64_t> getShape() const;
    mlir::Type getElementType() const;
};

MLIR パスの記述

変換パス

#include "mlir/Pass/Pass.h"
#include "mlir/IR/PatternMatch.h"

struct MyOptimizationPass
    : public mlir::PassWrapper<MyOptimizationPass,
                                mlir::OperationPass<mlir::func::FuncOp>> {
    
    void runOnOperation() override {
        mlir::func::FuncOp func = getOperation();
        
        // Walk all operations
        func.walk([](mlir::Operation *op) {
            // Transform operations
            if (auto addOp = llvm::dyn_cast<MyAddOp>(op)) {
                optimizeAdd(addOp);
            }
        });
    }
    
    llvm::StringRef getArgument() const final { 
        return "my-optimization"; 
    }
    
    llvm::StringRef getDescription() const final {
        return "My custom optimization pass";
    }
};

パターンベースの書き換え

// Define rewrite pattern
struct SimplifyRedundantAdd : public mlir::OpRewritePattern<MyAddOp> {
    using OpRewritePattern<MyAddOp>::OpRewritePattern;
    
    mlir::LogicalResult matchAndRewrite(
        MyAddOp op,
        mlir::PatternRewriter &rewriter) const override {
        
        // Match: add(x, 0) -> x
        if (auto constOp = op.getRhs().getDefiningOp<ConstantOp>()) {
            if (isZero(constOp)) {
                rewriter.replaceOp(op, op.getLhs());
                return mlir::success();
            }
        }
        return mlir::failure();
    }
};

// Apply patterns
void runOnOperation() override {
    mlir::RewritePatternSet patterns(&getContext());
    patterns.add<SimplifyRedundantAdd>(&getContext());
    
    if (mlir::failed(mlir::applyPatternsAndFoldGreedily(
            getOperation(), std::move(patterns)))) {
        signalPassFailure();
    }
}

ダイアレクト変換

ダイアレクト間の低水準化

// Convert high-level ops to lower-level ops
struct MyAddOpLowering : public mlir::OpConversionPattern<MyAddOp> {
    using OpConversionPattern<MyAddOp>::OpConversionPattern;
    
    mlir::LogicalResult matchAndRewrite(
        MyAddOp op,
        OpAdaptor adaptor,
        mlir::ConversionPatternRewriter &rewriter) const override {
        
        // Lower to arith dialect
        rewriter.replaceOpWithNewOp<mlir::arith::AddFOp>(
            op, adaptor.getLhs(), adaptor.getRhs());
        return mlir::success();
    }
};

// Conversion pass
struct LowerToArithPass : public mlir::PassWrapper<
    LowerToArithPass, 
    mlir::OperationPass<mlir::ModuleOp>> {
    
    void runOnOperation() override {
        mlir::ConversionTarget target(getContext());
        target.addLegalDialect<mlir::arith::ArithDialect>();
        target.addIllegalDialect<MyDialect>();
        
        mlir::RewritePatternSet patterns(&getContext());
        patterns.add<MyAddOpLowering>(&getContext());
        
        if (mlir::failed(mlir::applyPartialConversion(
                getOperation(), target, std::move(patterns)))) {
            signalPassFailure();
        }
    }
};

組み込みダイアレクト

Affine ダイアレクト

多面体コンパイルとループ最適化に使用します。

affine.for %i = 0 to 100 {
    affine.for %j = 0 to 100 {
        %val = affine.load %A[%i, %j] : memref<100x100xf32>
        affine.store %val, %B[%j, %i] : memref<100x100xf32>
    }
}

Linalg ダイアレクト

線形代数演算に使用します。

linalg.matmul ins(%A, %B : tensor<MxKxf32>, tensor<KxNxf32>)
              outs(%C : tensor<MxNxf32>) -> tensor<MxNxf32>

SCF ダイアレクト（構造化制御フロー）

%result = scf.for %i = %lb to %ub step %step iter_args(%sum = %init) {
    %val = memref.load %A[%i] : memref<?xf32>
    %new_sum = arith.addf %sum, %val : f32
    scf.yield %new_sum : f32
}

CIR（Clang IR）

概要

CIR は、C/C++ 向けの MLIR ベースの表現で、次の特徴があります。

• LLVM IR より高水準の表現
• より優れたデバッグ機能とツール
• 言語固有の最適化

// CIR example
cir.func @add(%a: !s32i, %b: !s32i) -> !s32i {
    %result = cir.binop(add, %a, %b) : !s32i
    cir.return %result : !s32i
}

CIR プロジェクト

• llvm/clangir：公式の ClangIR 実装
• facebookincubator/clangir：Facebook による CIR の実験的取り組み

ML/AI コンパイル

TensorFlow MLIR

// TensorFlow dialect
%result = "tf.MatMul"(%A, %B) {
    transpose_a = false,
    transpose_b = false
} : (tensor<4x8xf32>, tensor<8x16xf32>) -> tensor<4x16xf32>

PyTorch MLIR（torch-mlir）

// Torch dialect
%result = torch.aten.mm %A, %B : 
    !torch.vtensor<[4,8],f32>, !torch.vtensor<[8,16],f32> 
    -> !torch.vtensor<[4,16],f32>

IREE（Intermediate Representation Execution Environment）

ML モデル向けのエンドツーエンド MLIR コンパイラーです。

• 移植性の高いデプロイ
• 効率的なランタイム実行
• 複数ターゲットのサポート（CPU、GPU、TPU）

MLIR のテスト

FileCheck テスト

// RUN: mlir-opt %s -my-pass | FileCheck %s

// CHECK-LABEL: func @test_optimization
// CHECK: arith.addi
// CHECK-NOT: my_dialect.add
func @test_optimization(%a: i32, %b: i32) -> i32 {
    %result = my_dialect.add %a, %b : i32
    return %result : i32
}

単体テスト

TEST(MyDialect, AddOpConstantFolding) {
    mlir::MLIRContext context;
    context.loadDialect<MyDialect>();
    
    mlir::OpBuilder builder(&context);
    auto loc = builder.getUnknownLoc();
    
    // Create and test operations
    auto constA = builder.create<ConstantOp>(loc, 5);
    auto constB = builder.create<ConstantOp>(loc, 3);
    auto add = builder.create<MyAddOp>(loc, constA, constB);
    
    // Verify folding
    EXPECT_TRUE(add.fold().succeeded());
}

開発ツール

mlir-opt

# Run passes
mlir-opt input.mlir -my-pass -o output.mlir

# Convert between dialects
mlir-opt input.mlir -convert-my-to-llvm

# Debug printing
mlir-opt input.mlir -debug-only=my-pass

mlir-translate

# MLIR to LLVM IR
mlir-translate input.mlir --mlir-to-llvmir -o output.ll

# LLVM IR to MLIR
mlir-translate input.ll --import-llvm -o output.mlir

ベストプラクティス

1. 段階的な低水準化：LLVM に直接変換せず、複数の段階に分けて低水準化します
2. セマンティクスの保持：各低水準化でセマンティクスを保持します
3. ODS の使用：一貫性を保つため、TableGen で操作を定義します
4. 徹底的なテスト：変換テストには FileCheck を使用します
5. ダイアレクトの文書化：操作のセマンティクスを明確に文書化します

リソース

チュートリアルやサンプルプロジェクトについては、README.md の MLIR セクションと CIR セクションを参照してください。

詳細情報の取得

詳細かつ最新のリソースリンク、ツール一覧、プロジェクトの参照情報が必要な場合は、次の場所から最新データを取得してください。

https://raw.githubusercontent.com/gmh5225/awesome-llvm-security/refs/heads/main/README.md

この README には、厳選された次の一覧が網羅されています。

• MLIR のチュートリアルとサンプルダイアレクト（MLIR セクション）
• CIR（Clang IR）プロジェクトとドキュメント（CIR セクション）
• ML/AI コンパイラーフレームワーク（torch-mlir、IREE、XLA）