Writing an LLVM MachineFunction Pass (Learning LLVM, Part 4)

,

About Machine IR (MIR)

In the previous blogs we got a brief intro to LLVM IR and the associated concepts, such as Module, BasicBlock, Function, Instruction, passes, etc. This post serves as a brief introduction to Machine IR and the related terminology. The process of how Machine IR is actually generated by the compiler will be covered in detail in upcoming posts.

MIR a human-readable serialization format used in the LLVM compiler infrastructure to represent machine-specific intermediate representation, with YAML being used as the data serialization format. The first yaml document can contain an LLVM IR module and the remaining documents contain machine functions.
Please note that unlike LLVM IR, MIR doesn’t have the concept of modules. A MachineModuleInfo class exists but it’s more like a tracker rather than the parent container. A Module in the LLVM IR acts as a container and stores data like the list of functions, globals, named metadata,etc. However, MachineModuleInfo acts as a tracker, storing a pointer to the corresponding LLVM IR module, and a mapping between LLVM IR functions and Machine functions.

1
2
3
4
5
6
7
// llvm/CodeGen/MachineModuleInfo.h
class MachineModuleInfo {
  // ...
  // This is the LLVM Module being worked on.
  const Module *TheModule = nullptr;
  // ...
  DenseMap<const Function*, std::unique_ptr<MachineFunction>> MachineFunctions;

Please note that the LLVM version used for adding references to source code in this blog is LLVM 22.1.8.

Machine Basic Block

Recall the definition of a basic block, “a basic block is a straight-line sequence of instructions with no branches, meaning that execution starts at a single entry point and proceeds sequentially to a single exit point, where it then continues to the next basic block. Basic blocks belong to functions and cannot have jumps into their middle, ensuring that once execution starts, it will proceed through all instructions in the block. A basic block must have exactly one terminator instruction. This instruction tells the possible destinations of this basic block.”

A MachineBasicBlock is similar but:

  • Unlike a basic block, a Machine Basic Block can have multiple terminator instructions. Let’s see some code-snippets from the LLVM source code to prove that.
1
2
3
4
5
6
// llvm/IR/BasicBlock.h
const Instruction *getTerminator() const LLVM_READONLY {
    if (InstList.empty() || !InstList.back().isTerminator())
        return nullptr;
    return &InstList.back();
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
// llvm/CodeGen/MachineBasicBlock.h
inline iterator_range<iterator> terminators() {
    return make_range(getFirstTerminator(), end());
}

// lib/CodeGen/MachineBasicBlock.cpp
MachineBasicBlock::iterator MachineBasicBlock::getFirstTerminator() {
  iterator B = begin(), E = end(), I = E;
  while (I != B && ((--I)->isTerminator() || I->isDebugInstr()))
    ; /*noop */
  while (I != E && !I->isTerminator())
    ++I;
  return I;
}

This shows that a MachineBasicBlock can have multiple terminator instructions.

The MIR Parser

Now let’s try to load an MIR file. Since there’s nothing like MachineModule, parsing an MIR file gives an llvm::Module object.

1
2
3
4
5
// lib/CodeGen/MIRParser/MIRParser.cpp
std::unique_ptr<Module>
MIRParser::parseIRModule(DataLayoutCallbackTy DataLayoutCallback) {
  return Impl->parseIRModule(DataLayoutCallback);
}

This function creates a new llvm::Module object. In case the yaml file didn’t have any LLVM IR at the first document, an empty module is created.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
// CodeGen/MIRParser/MIRParser.cpp
std::unique_ptr<Module>
MIRParserImpl::parseIRModule(DataLayoutCallbackTy DataLayoutCallback) {
  // ...
  // ...
  if (const auto *BSN =
          dyn_cast_or_null<yaml::BlockScalarNode>(In.getCurrentNode())) {
    SMDiagnostic Error;
    M = parseAssembly(MemoryBufferRef(BSN->getValue(), Filename), Error,
                      Context, &IRSlots, DataLayoutCallback);
    if (!M) {
      reportDiagnostic(diagFromBlockStringDiag(Error, BSN->getSourceRange()));
      return nullptr;
    }
    In.nextDocument();
    if (!In.setCurrentDocument())
      NoMIRDocuments = true;
  } else {
    // Create an new, empty module.
    M = std::make_unique<Module>(Filename, Context);
    if (auto LayoutOverride = DataLayoutCallback(M->getTargetTriple().str(),
                                                 M->getDataLayoutStr()))
      M->setDataLayout(*LayoutOverride);
    NoLLVMIR = true;
  }
  return M;
}

Loading an MIR file

We can use the tool LLVM-Lens to directly retrieve the MIR at a particular stage for a given C/C++ file. Let’s create a file test.c

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
// test.c
int square(int x) {
  return x * x;
}

int add(int a, int b) {
  return a + b;
}

int sum_of_squares(int a, int b) {
  return add(square(a), square(b));
}

int factorial(int n) {
  if (n <= 1)
    return 1;

  return n * factorial(n - 1);
}

Now let’s use llvm-lens to get the MIR

1
2
3
4
5
6
7
# test.py
import llvm_lens

mir = llvm_lens.machine_ir("test.c",stop_after="x86-isel",debug_info=False)

with open("test.mir","w") as f:
    f.write(mir)


Here we can see that the generated MIR contains an LLVM IR module in the first document followed by MachineFunctions in the following documents. Let’s write a C++ program using the LLVM APIs that loads this file and creates an llvm::Module.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
#include <llvm/CodeGen/MIRParser/MIRParser.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/Support/SourceMgr.h>

int main(int argc, char **argv) {
  const char *MIRPath = argc > 1 ? argv[1] : "test.mir";
  llvm::LLVMContext ctx;
  llvm::SMDiagnostic error;

  auto parser = llvm::createMIRParserFromFile(MIRPath, error, ctx);
  if (!parser) {
    error.print("learning_llvm", llvm::errs());
    return 1;
  }

  auto mod = parser->parseIRModule();
  if (!mod) {
    error.print("learning_llvm", llvm::errs());
    return 1;
  }
  return 0;
}

Target, TargetRegistry and TargetMachine

Now that we have the Module, we need to find a way to run a Machine Function pass. Note that a Module doesn’t store any MachineFunction objects so we can’t directly use it to run a Machine Function pass. What we need to do is to first create a MachineModuleInfo object and populate the required information into it. Earlier we saw that this class stores a mapping between Functions and MachineFunctions. Before creating a MachineModuleInfo object, we need to create a TargetMachine, because creating a MachineModuleInfo object without supplying a valid pointer to TargetMachine will crash the program because the constructor dereferences the ptr to TargetMachine, and it’s default initialized to nullptr.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
// llvm/CodeGen/MachineModuleInfo.h
LLVM_ABI explicit MachineModuleInfo(const TargetMachine *TM = nullptr);

// lib/CodeGen/MachineModuleInfo.cpp
MachineModuleInfo::MachineModuleInfo(const TargetMachine *TM)
    : TM(*TM), Context(TM->getTargetTriple(), TM->getMCAsmInfo(),
                       TM->getMCRegisterInfo(), TM->getMCSubtargetInfo(),
                       nullptr, &TM->Options.MCOptions, false) {
  Context.setObjectFileInfo(TM->getObjFileLowering());
  initialize();
}

MachineModuleInfo::MachineModuleInfo(const TargetMachine *TM,
                                     MCContext *ExtContext)
    : TM(*TM), Context(TM->getTargetTriple(), TM->getMCAsmInfo(),
                       TM->getMCRegisterInfo(), TM->getMCSubtargetInfo(),
                       nullptr, &TM->Options.MCOptions, false),
      ExternalContext(ExtContext) {
  Context.setObjectFileInfo(TM->getObjFileLowering());
  initialize();
}

The Target class provides a helper function that can be used to create a TargetMachine.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
// llvm/MC/TargetRegistry.h
class Target{
  // ...
  // ...
  TargetMachine *createTargetMachine(
      const Triple &TT, StringRef CPU, StringRef Features,
      const TargetOptions &Options, std::optional<Reloc::Model> RM,
      std::optional<CodeModel::Model> CM = std::nullopt,
      CodeGenOptLevel OL = CodeGenOptLevel::Default, bool JIT = false) const {
    if (!TargetMachineCtorFn)
      return nullptr;
    return TargetMachineCtorFn(*this, TT, CPU, Features, Options, RM, CM, OL,
                               JIT);
  }
  // ...
  // ...
  }

We can get a Target (defined in llvm/MC/TargetRegistry.h) object using the TargetRegistry. A TargetRegistry is a generic interface for target specific features. It provides a method lookupTarget that gives a Target object based on a Triple. A triple is nothing but a string that stores info about the environment, such as the architecture, the vendor details, operating system, etc. For instance, x86_64-pc-linux-gnu is a triple.

To get a target triple, we can do one of the following two things.

  1. Retrieve the triple from the module loaded for the MIR file, if it’s present there.
  2. Get the default triple the compiler has been configured to produce code for.

Here, the MIR file already contains the target triple. The default triple can be retrieved by calling llvm::sys::getDefaultTargetTriple present in the header llvm/TargetParser/Host.h. We also need to initialize the targets by calling llvm::InitializeAllTargets, which will make all target machines that LLVM is configured to support available through the TargetRegistry.
Now that we know the terminology, let’s create a TargetMachine and create a MachineModuleInfo object.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
llvm::Triple triple(mod->getTargetTriple());
std::string target_error;

const llvm::Target *target =
    llvm::TargetRegistry::lookupTarget(triple, target_error);
if (!target) {
    llvm::errs() << "learning_llvm: " << target_error << "\n";
    return 1;
}

llvm::TargetMachine *TM = target->createTargetMachine(
    triple, "", "", llvm::TargetOptions(), std::nullopt);
if (!TM) {
    llvm::errs() << "learning_llvm: could not create a target machine for "
                    << triple.getTriple() << "\n";
    return 1;
}

Providing required information to MachineModuleInfo

Now that we have the target machine, we’re ready to initialize MachineModuleInfo. Doing this is really simple:

1
llvm::MachineModuleInfo MMI(TM);

Recall that MachineModuleInfo contains a MachineFunctions field that is a mapping between a Function and a MachineFunction.

1
2
3
4
5
6
7
// llvm/CodeGen/MachineModuleInfo.h
class MachineModuleInfo {
  // ...
  DenseMap<const Function*, std::unique_ptr<MachineFunction>> MachineFunctions;
  // ...
}

This mapping is populated via a call to the getOrCreateMachineFunction function.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
// lib/CodeGen/MachineModuleInfo.cpp
MachineFunction &MachineModuleInfo::getOrCreateMachineFunction(Function &F) {
  // Shortcut for the common case where a sequence of MachineFunctionPasses
  // all query for the same Function.
  if (LastRequest == &F)
    return *LastResult;

  auto I = MachineFunctions.insert(
      std::make_pair(&F, std::unique_ptr<MachineFunction>()));
  MachineFunction *MF;
  if (I.second) {
    // No pre-existing machine function, create a new one.
    const TargetSubtargetInfo &STI = *TM.getSubtargetImpl(F);
    MF = new MachineFunction(F, TM, STI, getContext(), NextFnNum++);
    MF->initTargetMachineFunctionInfo(STI);

    // MRI callback for target specific initializations.
    TM.registerMachineRegisterInfoCallback(*MF);

    // Update the set entry.
    I.first->second.reset(MF);
  } else {
    MF = I.first->second.get();
  }

  LastRequest = &F;
  LastResult = MF;
  return *MF;
}

We can indirectly invoke it by calling parser->parseMachineFunctions

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
// llvm/CodeGen/MIRParser/MIRParser.h
class MIRParser {
// ...
  LLVM_ABI bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI);
  // ...
  LLVM_ABI bool parseMachineFunctions(Module &M, ModuleAnalysisManager &MAM);
};

// lib/CodeGen/MIRParser/MIRParser.cpp
bool MIRParser::parseMachineFunctions(Module &M, MachineModuleInfo &MMI) {
  return Impl->parseMachineFunctions(M, MMI);
}

bool MIRParser::parseMachineFunctions(Module &M, ModuleAnalysisManager &MAM) {
  auto &MMI = MAM.getResult<MachineModuleAnalysis>(M).getMMI();
  return Impl->parseMachineFunctions(M, MMI, &MAM);
}

Both of these functions call MIRParserImpl::parseMachineFunctions. Let’s have a look at it.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
bool parseMachineFunctions(Module &M, MachineModuleInfo &MMI,
                            ModuleAnalysisManager *FAM = nullptr);
bool parseMachineFunction(Module &M, MachineModuleInfo &MMI,
                          ModuleAnalysisManager *FAM);

bool MIRParserImpl::parseMachineFunctions(Module &M, MachineModuleInfo &MMI,
                                          ModuleAnalysisManager *MAM) {
  if (NoMIRDocuments)
    return false;

  // Parse the machine functions.
  do {
    if (parseMachineFunction(M, MMI, MAM))
      return true;
    In.nextDocument();
  } while (In.setCurrentDocument());

  return false;
}

bool MIRParserImpl::parseMachineFunction(Module &M, MachineModuleInfo &MMI,
                                         ModuleAnalysisManager *MAM) {
  // ...
  if (!MAM) {
    if (MMI.getMachineFunction(*F) != nullptr)
      return error(Twine("redefinition of machine function '") + FunctionName +
                   "'");

    // Create the MachineFunction.
    MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
  // ...
}

Now let’s call the function and trigger the logic to populate the MachineFunctions dense map.

1
2
3
if (parser->parseMachineFunctions(*mod, MMI)) {
  return 1;
}

Creating and running a Machine Function Pass

Now that we have all the information about MachineFunctions within the MMI, we can create a MachineFunctionPassManager and run the pass. First of all, let’s create a pass that just prints the function names.

1
2
3
4
5
6
7
8
9
10
class DummyPass : public llvm::PassInfoMixin<DummyPass> {
public:
  static llvm::StringRef name() { return "dummy-pass"; }

  llvm::PreservedAnalyses run(llvm::MachineFunction &MF,
                              llvm::MachineFunctionAnalysisManager &MFAM) {
    llvm::outs() << "Running on: " << MF.getName() << "\n";
    return llvm::PreservedAnalyses::all();
  }
};

Now let’s extract the MachineFunctions from the MMI, register and run the pass. There’s an alternate and better way to this as well, which will be covered in the upcoming posts.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
llvm::SmallVector<llvm::MachineFunction *> machine_functions;

for (llvm::Function &F : *mod) {
  machine_functions.push_back(MMI.getMachineFunction(F));
}

llvm::MachineFunctionAnalysisManager MFAM;
llvm::MachineFunctionPassManager MFPM;
MFPM.addPass(DummyPass()); 
MFAM.registerPass([&]{return llvm::PassInstrumentationAnalysis();});

for(llvm::MachineFunction *MF:machine_functions){
  MFPM.run(*MF,MFAM);
}

Notice the line MFAM.registerPass([&]{return llvm::PassInstrumentationAnalysis();});. We cannot just ignore this otherwise the program will crash.
MachineFunctionPassManager is just an alias for PassManager<MachineFunction>, and run pulls PassInstrumentationAnalysis out of the analysis manager before it runs anything.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
// lib/CodeGen/MachinePassManager.cpp
template <>
PreservedAnalyses
PassManager<MachineFunction>::run(MachineFunction &MF,
                                  AnalysisManager<MachineFunction> &MFAM) {
  PassInstrumentation PI = MFAM.getResult<PassInstrumentationAnalysis>(MF);
  PreservedAnalyses PA = PreservedAnalyses::all();
  for (auto &Pass : Passes) {
    if (!PI.runBeforePass<MachineFunction>(*Pass, MF))
      continue;

    PreservedAnalyses PassPA = Pass->run(MF, MFAM);
    MFAM.invalidate(MF, PassPA);
    PI.runAfterPass(*Pass, MF, PassPA);
    PA.intersect(std::move(PassPA));
  }
  PA.preserveSet<AllAnalysesOn<MachineFunction>>();
  return PA;
}

PassInstrumentationAnalysis is an analysis (defined in llvm/IR/PassInstrumentation.h) that has to be registered like any other. The getResult function called above calls getResultImpl which further calls the lookupPass function.

1
2
3
4
5
6
7
// llvm/IR/PassManager.h
PassConceptT &lookUpPass(AnalysisKey *ID) {
  typename AnalysisPassMapT::iterator PI = AnalysisPasses.find(ID);
  assert(PI != AnalysisPasses.end() &&
         "Analysis passes must be registered prior to being queried!");
  return *PI->second;
}

In case when assertions are disabled, such as in release builds the assertion will not be triggered, causing a segmentation fault.

1
2
3
4
5
pwndbg> bt
#0  0x00007ffff4147607 in llvm::AnalysisManager<llvm::MachineFunction>::getResultImpl(llvm::AnalysisKey*, llvm::MachineFunction&) ()
   from /usr/lib/x86_64-linux-gnu/libLLVM.so.22.1
#1  0x00007ffff4149675 in llvm::PassManager<llvm::MachineFunction, llvm::AnalysisManager<llvm::MachineFunction>>::run(llvm::MachineFunction&, llvm::AnalysisManager<llvm::MachineFunction>&) () from /usr/lib/x86_64-linux-gnu/libLLVM.so.22.1
#2  0x0000555555557a09 in main (argc=1, argv=0x7fffffffd9f8

So we need to have the PassInstrumentationAnalysis here. Doing this is simple.

1
MFAM.registerPass([&]{return llvm::PassInstrumentationAnalysis();});

Now that we are familiar with all the concepts involved, let’s compile and run the program.

We can see that the pass printed all the MachineFunctions present in the original MIR file. Complete code for this:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
#include <llvm/CodeGen/MIRParser/MIRParser.h>
#include <llvm/CodeGen/MachineModuleInfo.h>
#include <llvm/CodeGen/MachinePassManager.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/MC/TargetRegistry.h>
#include <llvm/Support/SourceMgr.h>
#include <llvm/Support/TargetSelect.h>
#include <llvm/Target/TargetMachine.h>
#include <llvm/Target/TargetOptions.h>
#include <llvm/TargetParser/Host.h>
#include <llvm/Passes/PassBuilder.h>

class DummyPass : public llvm::PassInfoMixin<DummyPass> {
public:
  static llvm::StringRef name() { return "dummy-pass"; }

  llvm::PreservedAnalyses run(llvm::MachineFunction &MF,
                              llvm::MachineFunctionAnalysisManager &MFAM) {
    llvm::outs() << "Running on: " << MF.getName() << "\n";
    return llvm::PreservedAnalyses::all();
  }
};

int main(int argc, char **argv) {
  const char *MIRPath = argc > 1 ? argv[1] : "test.mir";
  llvm::LLVMContext ctx;
  llvm::SMDiagnostic error;

  auto parser = llvm::createMIRParserFromFile(MIRPath, error, ctx);
  if (!parser) {
    error.print("learning_llvm", llvm::errs());
    return 1;
  }

  auto mod = parser->parseIRModule();
  if (!mod) {
    error.print("learning_llvm", llvm::errs());
    return 1;
  }

  llvm::InitializeNativeTarget();
  llvm::Triple triple(mod->getTargetTriple());
  std::string target_error;

  const llvm::Target *target =
      llvm::TargetRegistry::lookupTarget(triple, target_error);
  if (!target) {
    llvm::errs() << "learning_llvm: " << target_error << "\n";
    return 1;
  }

  llvm::TargetMachine *TM = target->createTargetMachine(
      triple, "", "", llvm::TargetOptions(), std::nullopt);
  if (!TM) {
    llvm::errs() << "learning_llvm: could not create a target machine for "
                 << triple.getTriple() << "\n";
    return 1;
  }

  llvm::MachineModuleInfo MMI(TM);

  if (parser->parseMachineFunctions(*mod, MMI)) {
    return 1;
  }

  llvm::SmallVector<llvm::MachineFunction *> machine_functions;

  for (llvm::Function &F : *mod) {
    machine_functions.push_back(MMI.getMachineFunction(F));
  }

  llvm::MachineFunctionAnalysisManager MFAM;
  llvm::MachineFunctionPassManager MFPM;
  MFPM.addPass(DummyPass()); 
  MFAM.registerPass([&]{return llvm::PassInstrumentationAnalysis();});

  for(llvm::MachineFunction *MF:machine_functions){
    MFPM.run(*MF,MFAM);
  }
  return 0;
}

References

https://llvm.org/docs/MIRLangRef.html
https://github.com/llvm/llvm-project