How to Write A CORRECT LLVM Pass As A Plugin

06 Aug, 2023

Recently I started to learn about LLVM. I have read many materials about how to write LLVM Passes in the middle end. However, because of the complexity of the mechanism, many materials, including the official tutorial, are not very clear, fully updated, and consistent with the latest version of llvm toolchain. If you have followed those materials, as a beginner, you probably have noticed that it is not as easy to construct a proper simple LLVM Pass as you expected. By "expected", I mean I would like to have custom passes as loadable plugins located in anywhere of local system with no need to rebuild from the LLVM source code each time. There is a new pass manager, not that new, which is replacing the legacy pass manager. Maybe you have tried add_llvm_pass_plugin(MyPassName source.cpp) in CMakeLists.txt, but the command cannot be found by the system. Actually it is fully usable but only with the latest version of LLVM(18.0). For people who are using 16.0 version of LLVM, this could be a bit confusing. To learn how to write a properly working Pass as Plugins, I would suggest you firstly look at the following two examples:

HelloWorld

#include "llvm/Transforms/Utils/HelloWorld.h"

using namespace llvm;

PreservedAnalyses HelloWorldPass::run(Function &F,
                                      FunctionAnalysisManager &AM) {
  errs() << F.getName() << "\n";
  return PreservedAnalyses::all();
}

#include "llvm/IR/Function.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Pass.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

static cl::opt<bool> Wave("wave-goodbye", cl::init(false),
                          cl::desc("wave good bye"));

namespace {

bool runBye(Function &F) {
  if (Wave) {
    errs() << "Bye: ";
    errs().write_escaped(F.getName()) << '\n';
  }
  return false;
}

struct LegacyBye : public FunctionPass {
  static char ID;
  LegacyBye() : FunctionPass(ID) {}
  bool runOnFunction(Function &F) override { return runBye(F); }
};

struct Bye : PassInfoMixin<Bye> {
  PreservedAnalyses run(Function &F, FunctionAnalysisManager &) {
    if (!runBye(F))
      return PreservedAnalyses::all();
    return PreservedAnalyses::none();
  }
};

} // namespace

char LegacyBye::ID = 0;

static RegisterPass<LegacyBye> X("goodbye", "Good Bye World Pass",
                                 false /* Only looks at CFG */,
                                 false /* Analysis Pass */);

/* New PM Registration */
llvm::PassPluginLibraryInfo getByePluginInfo() {
  return {LLVM_PLUGIN_API_VERSION, "Bye", LLVM_VERSION_STRING,
          [](PassBuilder &PB) {
            PB.registerVectorizerStartEPCallback(
                [](llvm::FunctionPassManager &PM, OptimizationLevel Level) {
                  PM.addPass(Bye());
                });
            PB.registerPipelineParsingCallback(
                [](StringRef Name, llvm::FunctionPassManager &PM,
                   ArrayRef<llvm::PassBuilder::PipelineElement>) {
                  if (Name == "goodbye") {
                    PM.addPass(Bye());
                    return true;
                  }
                  return false;
                });
          }};
}

#ifndef LLVM_BYE_LINK_INTO_TOOLS
extern "C" LLVM_ATTRIBUTE_WEAK ::llvm::PassPluginLibraryInfo
llvmGetPassPluginInfo() {
  return getByePluginInfo();
}
#endif

By these examples, I have successfully gotten my own loadable Plugin MyHelloPass:

#include "llvm/Analysis/PostDominators.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstVisitor.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Passes/PassPlugin.h"

#define DEBUG_TYPE "test-hello-plugin"

using namespace llvm;

namespace {
    struct MyHelloPluginPass : PassInfoMixin<MyHelloPluginPass> {
        PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM) {
            errs() << "Running MyHelloPluginPass on function " << F.getName() << "\n";

            return PreservedAnalyses::none();
        }
    };
} // namespace

/// Registration
PassPluginLibraryInfo getPassPluginInfo() {
    const auto callback = [](PassBuilder &PB) {
        PB.registerPipelineParsingCallback(
                [](StringRef Name, FunctionPassManager &FPM, auto) {
                    if (Name == "my-hello-plugin") {
                        FPM.addPass(MyHelloPluginPass());
                        return true;
                    }
                    return false;
                });
    };
    return {LLVM_PLUGIN_API_VERSION, "MyHelloPluginPass",
            LLVM_VERSION_STRING, callback};
};

extern "C" LLVM_ATTRIBUTE_WEAK PassPluginLibraryInfo llvmGetPassPluginInfo() {
    return getPassPluginInfo();
}

#undef DEBUG_TYPE

And in CMakeLists.txt, we don't need add_llvm_pass_plugin, what we need is only add_library(MyHelloPlugin SHARED MyHelloPlugin.cpp).

I hope this blog could help you, you could find more info on my github repo.

#compiler #llvm