LLVM — Writing Pass Instrumentations for the New PassManager

The more general and flexible version of “-print-before" / "-print-after"

Switching to the new PassManager by default in LLVM (and Clang) is around the corner now. So unsurprisingly most of the features in the old PassManager have been transparently ported to the new one. The classic -print-after-all (and its siblings like -print-before) command line option in opt / llc, alone with the painful memory of 10000+ lines IR dump file, is not an exception. What changed, however, is how this particular feature is implemented in the newer version. That is, the Pass Instrumentation Framework .

In this short tutorial, I will briefly introduce how -print-after / -print-before were implemented in the new PassManager first, then show you how to write your own Pass instrumentations.

Believe or not, the core logic to run each Pass in the new PassManager is pretty easy to understand (Adapted from include/llvm/IR/PassManager.h):

Adapted from include/llvm/IR/PassManager.h (All AnalysisManager related code inside the loop are trimmed)

Line 9 at the snippet above is the main “Pass runner”. In addition to that, line 4 and line 12 suggested that the PassManager will do something before and after running each Pass. And actually, that is exactly where the magic of Pass Instrumentation framework happens.

As the name suggested, the -print-before / -print-after options will print IR code before and after the designated Pass, respectively. The new PassManager code mentioned above further generalizes this concept by providing two extension points for developers to customize the actions they want to perform before and after running a Pass, rather than just printing IR. These extension (or customization) points are so called Pass Instrumentation in the new PM.

To get started on making our own instrumentations, fist let’s have a quick recap on how to write a Pass plugin in the new PM. My blog post Writing LLVM Pass in 2018 already had some pretty decent and detailed descriptions on how to write a Pass for the new PassManager. Feel free to take a look if you haven’t. The part we’re gonna focusing here is the “plugin entry point”, or the registration callback in the Pass plugin:

The lambda function from line 5 to 7 gives you a PassBuilder instance which you can use to register your Pass at a desired position in the Pass pipeline. For example: (stealing from an example from one of my earlier posts :P)

…Or you can register a new Pass instrumentation using the object returning from PassBuilder::getPassInstrumentationCallbacks(). To have a better understanding, let’s create a simple demo mimicking the functionality of -stop-after command line option, which effectively stalls the entire (codegen) pass pipeline after finishing the designated pass, in llc. Here is the first step:

Simple “stop-after” instrument demo— Step 1

The above code showed that after getting the PassInstrumentationCallbacks object from PassBuilder, we passed it to StopAfterInstrument::registerCallbacks, which will register the beforePass and afterPass function in the same class as the callbacks invoked before and after each Pass, respectively.

Though function prototypes of these two callback have same argument types, the return types are different: The first argument represents the textual Pass Id (Note that this string is not the abbrev-and-lowercase one used in the legacy PM, instead it’s longer and more meaningful, like “LowerExpectIntrinsicPass”); The second argument is the input IR, represented in llvm::Any type. You can use the any_isa<Ty>(obj) and any_cast<Ty>(obj) utilities to convert it into normal IR types like llvm::Function; Finally, for callbacks running before a Pass (i.e. the callback for registerBeforePassCallback), the returned boolean indicating whether PassManager should run this Pass or not. This is one of the most important features in new PassManager IMO. It enables us to stop after a certain Pass as this demo is going to achieve:

Simple “stop-after” instrument demo — Step 2

As shown above, we maintain a boolean flag, Stopped, that will be set in afterPass function based on the textual Pass Id and number of triggering in the past, and be consulted in beforePass function to halt the entire Pass pipeline.

Note that we’re not going to skip printer Passes because we’re going to use Clang to test our instrumentation plugin using the following commands and it will be impossible to output the resulting LLVM IR without those Passes.

clang -O3 -fexperimental-new-pass-manager \
      -fpass-plugin=InstrumentPlugin.so \
      -emit-llvm -S ./input.c

In case you’re wondering how to use any_isa<Ty>(...) and any_cast<Ty>(...) in this context:

Which will give you stderr output similar to this:

Skipping FunctionPass EarlyCSEPass @ function main
Skipping FunctionPass LowerExpectIntrinsicPass @ function main
Skipping FunctionPass CallSiteSplittingPass @ function main
Skipping FunctionPass PassManager<llvm::Function> @ function foo
Skipping ModulePass IPSCCPPass
Skipping ModulePass CalledValuePropagationPass
Skipping ModulePass GlobalOptPass
Skipping ModulePass ModuleToFunctionPassAdaptor<llvm::PromotePass>
Skipping ModulePass DeadArgumentEliminationPass
Skipping ModulePass ModuleToFunctionPassAdaptor<llvm::PassManager<llvm::Function> >
Skipping ModulePass ModuleInlinerWrapperPass
Skipping ModulePass PassManager<llvm::Module>

Full demo code can be found here: https://github.com/mshockwave/LLVM-NewPM-PassInstrumentation-Demo