LLVMに新しいPassを追加してみる試行 (CountOp Passを追加する試行)

Writing LLVM Passを読んで、LLVMのPassの追加方法が少しわかったので、今度は別のPassを追加してみよう。 参考にしたのは以下のPDF資料。ここでは、OpCountと呼ばれるPassを新たに作成している。 - Writing an LLVM Pass

https://www.inf.ed.ac.uk/teaching/courses/ct/19-20/slides/llvm-2-writing_pass.pdf

このPassは、LLVM IR中に含まれるOpCodeを数え上げ、それぞれのOpCodeがいくつあるのかを観察するためのものだ。ソースコードを眺めてみる。

• llvm/lib/Transforms/CountOp/CountOp.cpp

#include "llvm/ADT/Statistic.h"
#include "llvm/IR/Function.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;

#define DEBUG_TYPE "opCounter"

STATISTIC(OpCounter, "Count Number of Operations");

namespace {
// CountOp
struct CountOp : public FunctionPass {
  std::map<std::string, int> opCounter;
  static char ID;  // Pass identification, replacement for typeid
  CountOp () : FunctionPass(ID) {}

  virtual bool runOnFunction (Function &F) {
    ++OpCounter;
    errs() << "Function " << F.getName() << '\n';
    for (Function::iterator bb = F.begin(), e = F.end(); bb != e; ++bb) {
      for (BasicBlock::iterator i = bb->begin(), e = bb->end(); i != e; ++i) {
        if (opCounter.find(i->getOpcodeName()) == opCounter.end()) {
          opCounter[i->getOpcodeName()] = 1;
        } else {
          opCounter[i->getOpcodeName()] += 1;
        }
      }
    }
    std::map <std::string, int>::iterator i = opCounter.begin();
    std::map <std::string, int>::iterator e = opCounter.end();
    while (i != e) {
      errs() << "  " << i->first << ": " << i->second << '\n';
      i++;
    }
    errs() << '\n';
    opCounter.clear();
    return false;
  }
};
}

char CountOp::ID = 0;
static RegisterPass<CountOp> X("opCounter", "Counts opcodes per functions");

一つずつソースコードを読み解いていかなければならないだろう。

STATISTIC(OpCounter, "Count Number of Operations");

これは統計情報表示用のカウンタ。OpCounterという変数が定義されて、これをrunOnFunction()内でカウントアップすることでこのPassが何回呼ばれたかの統計を表示することができる。

struct CountOp : public FunctionPass {
  std::map<std::string, int> opCounter;
  static char ID;  // Pass identification, replacement for typeid
  CountOp () : FunctionPass(ID) {}

クラスにはchar IDが格納されているのだけれども、これはPassに対してIDを付けるため。 そしてopCounterはオペコードとその表示回数をカウントするためのmapだ。これがこのPassの本体みたいなものだ。

runOnFunction()の本体に入っていく。引数Function &Fから関数の情報を取り出すことになる。

Functionの中にはイテレータとしてBasicBlockが入っているようで、これを最初のループで繰り返している。 そしてBasicBlockの中には命令が入っているようで、これを次のforループで取り出している。 最内if文は、上記で定義したmapを探索して過去に同じOpcodeを登録しているかどうかをチェックしている。 登録していればカウンタを+1し、そうでなければカウンタを新規登録する。

  virtual bool runOnFunction (Function &F) {
    ++OpCounter;
    errs() << "Function " << F.getName() << '\n';
    for (Function::iterator bb = F.begin(), e = F.end(); bb != e; ++bb) {
      for (BasicBlock::iterator i = bb->begin(), e = bb->end(); i != e; ++i) {
        if (opCounter.find(i->getOpcodeName()) == opCounter.end()) {
          opCounter[i->getOpcodeName()] = 1;
        } else {
          opCounter[i->getOpcodeName()] += 1;
        }
      }
    }
/* ... 以下省略 ... */

最後に出力部分だ。mapの中身をすべて取り出して出力している。これで完成となる。

    std::map <std::string, int>::iterator i = opCounter.begin();
    std::map <std::string, int>::iterator e = opCounter.end();
    while (i != e) {
      errs() << "  " << i->first << ": " << i->second << '\n';
      i++;
    }
    errs() << '\n';
    opCounter.clear();
    return false;
}

クラスが完成したら次はPassの登録となる。以下はおまじないのような感じで覚えておけばよかろう。

char CountOp::ID = 0;
static RegisterPass<CountOp> X("opCounter", "Counts opcodes per functions");

llvm/lib/Transform/CMakeLists.txtに以下を追加して再度LLVMをビルドする。

diff --git a/llvm/lib/Transforms/CMakeLists.txt b/llvm/lib/Transforms/CMakeLists.txt
index dda5f6de11e3..42d46324ced6 100644
--- a/llvm/lib/Transforms/CMakeLists.txt
+++ b/llvm/lib/Transforms/CMakeLists.txt
@@ -9,3 +9,4 @@ add_subdirectory(Hello)
 add_subdirectory(ObjCARC)
 add_subdirectory(Coroutines)
 add_subdirectory(CFGuard)
+add_subdirectory(CountOp)

ビルドが成功すればOK。さっそく実行してみる。

./bin/opt -enable-new-pm=0 -f -load ./lib/LLVMCountOp.so -opCounter < ./sample/hello.bc > /dev/null

Function main
  call: 1
  ret: 1

動いた。正しく動作したようだ。