Zenn
⌨️

LLVM Pass を書いてみよう

2023/03/01に公開
#
コンパイラ
#
LLVM
#
llvmpass
tech

概要

この記事は LLVM Pass の作り方をさっくりと説明するものです。
例として、2のべき乗での掛け算をシフト演算に変換する pass (パス) の作成を扱います。実装には C++ を使います。なお、 LLVM Machine Pass については扱いません。

今回作成するコードはこちらに置いてあります。

LLVM, LLVM Pass とは

LLVM はコンパイラを作るための基盤で、 clang や rustc などで使われています。 LLVM はプログラムを中間表現 LLVM IR として管理しており、 LLVM Pass は LLVM IR を解析したり最適化したりするのに用いられます。

例えば clang のコンパイルは、次のように進みます。

  1. C で書かれたコードを中間表現 LLVM IR に変換する。
  2. LLVM IR に様々な LLVM Pass を適用し、解析・最適化をおこなう。
  3. LLVM IR をマシン語に変換する。

LLVM には初めから、未使用コードを削除するパスや定数伝播をするパスなど、様々なパスが定義されています (パスの一覧) が、自分で新たなパスを作ることもできます。以降では、簡単なパスの作成を通してその作り方を解説していきます。

LLVM Pass を書く

それではパスを書いていきましょう。

1. LLVM のダウンロード

パスの作成には LLVM の提供するライブラリが必要です。公式がビルド済みのものを配布しているので、それを使います (リンク)。今回はバージョン15.0.6を前提とします。例えば x64 上の Ubuntu を使っている場合は、clang+llvm-15.0.6-x86_64-linux-gnu-ubuntu-18.04.tar.xz をダウンロードします。ダウンロードしたファイルは適当な場所に展開して、環境変数 LLVM15 に展開先のパス (path) を設定しておきましょう。

$ export LLVM15=(展開先のパス)

2. テスト用コードの作成

パスを書く前に、パスを適用する LLVM IR を作成します。
空ディレクトリを作成し、 テストに用いるコードを書きましょう。

$ mkdir -p mul2shift/test
$ cd mul2shift/test
$ touch test.c
test.c
#include <stdio.h>
int x256(int i) {
    return i * 256;
}

int main(void) {
    printf("%d\n", x256(3));
    return 0;
}

C言語のコードを LLVM IR に変換するには clang に -S -emit-llvm を渡します。また、インライン展開や乗算の最適化を避けるために -O0 も付けておきます。

$ $LLVM15/bin/clang -S -emit-llvm -O0 test.c

生成された test.ll を見ると、 関数 x256 の中に 256 を掛ける処理が見つかります。今回の目標は、この乗算をシフト演算に変換することです。

test.ll
; Function Attrs: noinline nounwind optnone uwtable
define dso_local i32 @x256(i32 noundef %0) #0 {
  %2 = alloca i32, align 4
  store i32 %0, ptr %2, align 4
  %3 = load i32, ptr %2, align 4
  %4 = mul nsw i32 %3, 256
  ret i32 %4
}

3. ひな形の作成

パスの作成に最低限必要なコードを書いていきます。

3.1 ファイルの作成

あらかじめ必要なファイルを作成しておきましょう。

$ cd mul2shift
$ touch CMakeLists.txt Mul2Shift.cpp

ディレクトリ構造は次のようになります。

.
├── CMakeLists.txt
├── Mul2Shift.cpp
└── test
   ├── test.c
   └── test.ll

3.2 CMakeLists.txt を書く

CMakeLists.txt を書きます。
HINTS "$ENV{LLVM15}/lib/cmake" の行では LLVM の場所を設定しています。ここでは先ほど設定した環境変数 LLVM15 を参照しています。

CMakeLists.txt
cmake_minimum_required(VERSION 3.13)

project(mul2shift)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
set(CMAKE_CXX_STANDARD 17)

find_package(LLVM
    REQUIRED
    CONFIG
    HINTS "$ENV{LLVM15}/lib/cmake"
    NO_DEFAULT_PATH
)

add_definitions(${LLVM_DEFINITIONS})
include_directories(${LLVM_INCLUDE_DIRS})
link_directories(${LLVM_LIBRARY_DIR})

list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
include(AddLLVM)

add_llvm_pass_plugin(Mul2Shift Mul2Shift.cpp)

CMakeLists.txt が書けたら、ビルドスクリプトを生成しましょう。

$ cmake -B build

3.3 Mul2Shift.cpp を書く

パスを定義します。
ここで定義する Mul2Shift がパスです。 入力された IR コード中にある各関数 F に対して Mul2Shift::run が呼ばれます。 run の返り値は、このパスの影響を受けない解析の一覧です。 run は IR を編集することができるため、一部の解析は run 適用後にやり直す必要があるかもしれません。 run の返り値の役割は、そうした情報を呼び出し元に伝えることです。このひな形は適用対象の関数名を stderr に印字するものであり、 IR を編集しないので、すべての解析結果が保存される (無効にならない) ことを呼び出し元に伝えます。

Mul2Shift.cpp
#include <llvm/IR/PassManager.h>
#include <llvm/Passes/PassPlugin.h>
#include <llvm/Passes/PassBuilder.h>
#include <unordered_set> // あとで使います

using namespace llvm;

struct Mul2Shift: public PassInfoMixin<Mul2Shift> {
    static bool isRequired() { return true; }
    PreservedAnalyses run(Function &F, FunctionAnalysisManager &FAM) {
        errs() << "Function name : " << F.getName() << "\n";
        return PreservedAnalyses::all();
    }
};

同じファイルに、パスの登録に使う処理を書き加えます。

Mul2Shift.cpp
extern "C" LLVM_ATTRIBUTE_WEAK
PassPluginLibraryInfo llvmGetPassPluginInfo() {
    return {
        LLVM_PLUGIN_API_VERSION, "Mul2Shift", "v0.1",
        [](PassBuilder &PB) {
            PB.registerPipelineParsingCallback(
                [](
                    StringRef Name,
                    FunctionPassManager &FPM,
                    ArrayRef<PassBuilder::PipelineElement>
                ) {
                    if (Name == "mul2shift") {
                        FPM.addPass(Mul2Shift());
                        return true;
                    }
                    return false;
                }
            );
        }
    };
}

変数名がアッパーキャメルケースであることに抵抗感があるかもしれませんが、これはLLVM のコーディング規約に従った結果です。

3.4 ひな形の動作確認

ひな形の動作を確認しましょう。
パスの適用には opt コマンドを使います。

# ビルド
$ cmake --build ./build
# パスの適用
$ $LLVM15/bin/opt --load-pass-plugin=build/Mul2Shift.so --passes="function(mul2shift)" -S -o /dev/null
 test/test.ll

opt の結果から、 Mul2Shift::run が関数 x256 と関数 main に対して実行されていることがわかります。

opt の結果
Function name : x256
Function name : main

4. シフト演算への変換を実装

Mul2Shift::run の実装を、2のべき乗での掛け算をシフト演算に変換する処理に書き換えましょう。

まずは、 WorkList に最適化が可能な命令を集めます。
二重の for 文は関数 F 内の命令を順に処理するためのものです。for 文の中では、命令 I のうち、次の条件をすべて満たすものを WorkList に追加しています。

  • オペコードが乗算 (Instruction::Mul) である。
  • 第1オペランド (乗数) が整数定数 (ConstantInt)である。
  • 第1オペランド (乗数) が2のべき乗である。
run
std::unordered_set<Instruction*> WorkList;

for (BasicBlock &BB: F) {
    for (Instruction &I: BB) {
        if (I.getOpcode() != Instruction::Mul) continue;
        auto const *CI = dyn_cast<ConstantInt>(I.getOperand(1));
        if (!CI) continue;
        auto const V = CI->getValue();
        if (V.exactLogBase2() == -1) continue;
        WorkList.insert(&I);
    }
}

続いて、集めた乗算命令をシフト演算に変換します。
Log2Multiplier は乗数の対数をとったものです。この値を IRBuilder::CreateShl に渡すことでシフト命令 Shl を作成・挿入します。挿入位置は IRBuilder のコンストラクタで指定した命令 Mul の直前です。続いて Mul を削除します。まずは replaceAllUsesWith を用いて乗算の結果を使う処理をシフト命令の結果を使う処理に置き換えます。これにより Mul に依存する命令がなくなるため、 Mul を安全に削除・解放 (eraseFromParent) できるようになります。

run
for (auto *Mul: WorkList) {
    auto Log2Multiplier = cast<ConstantInt>(Mul->getOperand(1))->getValue().exactLogBase2();
    IRBuilder<> Builder(Mul);
    auto *Shl = Builder.CreateShl(Mul->getOperand(0), Log2Multiplier);
    Mul->replaceAllUsesWith(Shl);
    Mul->eraseFromParent();
}

今回は IR を変更したため、すべての解析結果が保存されないことを呼び出し元に伝えます。実際には保存される解析もあるかと思いますが、それを調べ上げるのは大変なので、これで妥協します (なにか良い方法があれば教えてください)。

run
return PreservedAnalyses::none();

4. 動作確認

作成したパスの動作を確認しましょう。

# ビルド
$ cmake --build ./build
# パスの適用
$ $LLVM15/bin/opt --load-pass-plugin=build/Mul2Shift.so --passes="function(mul2shift)" -S -o test/modified.ll
 test/test.ll

ひな形の動作確認をした際は opt の出力を捨てていましたが、ここでは test/modified.ll に書き込んでいます。

modified.ll を見ると、 mul 命令が shl 命令に書き換わっていることがわかります。

modified.ll
; Function Attrs: noinline nounwind optnone uwtable
define dso_local i32 @x256(i32 noundef %0) #0 {
  %2 = alloca i32, align 4
  store i32 %0, ptr %2, align 4
  %3 = load i32, ptr %2, align 4
  %4 = shl i32 %3, 8
  ret i32 %4
}

最後にこれをコンパイルして、正しく動くことを確かめましょう。

$ cd test
$ $LLVM15/bin/clang modified.ll
$ ./a.out

768 が表示されたら成功です。お疲れさまでした。

まとめ

乗算をシフト演算に変換するパスの作成を通して LLVM Pass の作り方を説明してきました。 LLVM はプロジェクトが巨大なうえ日本語の情報も少なく、初学者がパスを書くのはなかなかに大変ですが、この記事が何かのお役に立てれば幸いです。

補足

今回作成したパスは IR 上の関数を処理するものでしたが、例えば次のように書くと翻訳単位 (Module) 全体に対して処理を行うこともできます。

struct Mul2Shift: public PassInfoMixin<Mul2Shift> {
    static bool isRequired() { return true; }
    PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM) {
        errs() << M.getName() << "\n";
        return PreservedAnalyses::all();
    }
};

extern "C" LLVM_ATTRIBUTE_WEAK
PassPluginLibraryInfo llvmGetPassPluginInfo() {
    return {
        LLVM_PLUGIN_API_VERSION, "Mul2Shift", "v0.1",
        [](PassBuilder &PB) {
            PB.registerPipelineParsingCallback(
                [](
                    StringRef Name,
                    ModulePassManager &MPM,
                    ArrayRef<PassBuilder::PipelineElement>
                ) {
                    if (Name == "mul2shift") {
                        MPM.addPass(Mul2Shift());
                        return true;
                    }
                    return false;
                }
            );
        }
    };
}

Discussion