DEV Community

Cover image for Implementing Loop Invariant code Motion with LLVM
cppchedy
cppchedy

Posted on

Implementing Loop Invariant code Motion with LLVM

Introduction

In the previous articles, we learned how to write analysis and transformation passes with LLVM. We started with a simple pass that counts, in a given module, how many add instructions there are. Next, we saw a simple implementation of Dead code elimination optimization and constant propagation. Our focus was on practicality: we simplify by making a few assumptions and start showing the code while keeping the theory to a minimum.

We are keeping the same focus here: being practical with minimum theory. Additionally, we solve an optimization problem targeting a foundational element in compiler theory, namely loops. More precisely, We explore LICM, a classic optimization present in almost every serious compiler implementation.

In the following section, we introduce LICM briefly and set some assumptions to simplify the implementation. Keep in mind, the goal is to show LLVM facilities, not to fully cover the topic. You can always find a production-ready pass in LLVM if you need one.

Loop-invariant Code Motion

At a high level, Loop-Invariant Code Motion (LICM) is an optimization that moves computations that do not change across loop iterations outside of the loop. This avoids recomputing the same value repeatedly.

Let’s look at this example in LLVM IR.

define i32 @foo(i32 %n, i32 %x) {
entry:
  br label %loop

loop:
  %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
  %result = phi i32 [ 0, %entry ], [ %result.next, %loop ]

  %invariant = add i32 %x, 10
  %mul = mul i32 %i, %invariant
  %result.next = add i32 %result, %mul

  %inc = add i32 %i, 1
  %cmp = icmp slt i32 %inc, %n
  br i1 %cmp, label %loop, label %exit

exit:
  ret i32 %result.next
}

Enter fullscreen mode Exit fullscreen mode

And here a possible C++ source-level counterpart:

int foo(int n, int x) {
    int result = 0;

    for (int i = 0; i < n; ++i) {
        result += i * (x + 10);
    }

    return result;
}
Enter fullscreen mode Exit fullscreen mode

This is a simple for loop. The computation inside the loop uses addition and multiplication, and one of those computations is loop-invariant. LICM will detect the one and then hoist it outside.

The previous LLVM IR is transformed to:

define i32 @foo(i32 %n, i32 %x) {
entry:
  %invariant = add i32 %x, 10
  br label %loop

loop:
  %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
  %result = phi i32 [ 0, %entry ], [ %result.next, %loop ]

  %mul = mul i32 %i, %invariant
  %result.next = add i32 %result, %mul

  %inc = add i32 %i, 1
  %cmp = icmp slt i32 %inc, %n
  br i1 %cmp, label %loop, label %exit

exit:
  ret i32 %result.next
}
Enter fullscreen mode Exit fullscreen mode

With that we can move to our pass's assumptions on the input we accept.

Setting expectations

We adopt a couple of assumptions for simplification and more importantly to keep the theory minimum. We are not looking to implement a production-grade LICM. It already exists in LLVM. We are trying to expose LLVM capabilities while building something interesting.

Without further ado, We assume that our pass:

  • Only deal with binary operations that are safe to speculate. For example, we exclude division because it can trap for certain operands.
  • Has no memory, No pointers or aliasing problems
  • Has no nested loops or control-flow inside the loop's body
  • Has no calls

Practically, It means that we only accept computation loops with only addition, multiplication and/or subtraction. we don't have to deal with aliasing, we don't need to reason about dominance and we restrict our invariant propagation to dependency chains consisting of binary operations. Also, with these assumptions, we get a single basic block as the loop body. This gives us a simple structure where we can process instructions in program order and rely on our single traversal to propagate invariance through dependency chains.

With these assumptions, we can get something concrete fast on which you can build on while exploring the other concepts we skipped over on your own.

Note: These are assumptions about the input rather than properties that the pass verifies.

Implementing the Pass

The algorithm I am using is quite simple. For each instruction inside a loop, we check its operands: as long as it's a constant or is a value defined by an instruction outside the loop then it's an invariant. By navigating the instructions in the program order, we guarantee that any Invariant instruction is hoisted and its subsequent dependent instructions are correctly evaluated. In other words, processing instructions in program order guarantees that the current instruction's operands (e.g. values defined by an instruction) are already tested for invariance.

Let's move to something concrete. We show the code before breaking it down bit by bit.

struct SimpleLICMPass : public PassInfoMixin<SimpleLICMPass> {

  bool isLoopInvariant(Instruction *I, Loop &L) {

    for (auto &Oprd : I->operands()) {

      if (Instruction *in = dyn_cast<Instruction>(Oprd)) {
        if (L.contains(in)) {
          return false;
        }
      }
    }

    return true;

  }

  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM) {
    LoopInfo &LI = AM.getResult<LoopAnalysis>(F);

    for (auto *L : LI) {

      auto *preheader = L->getLoopPreheader();

      if (!preheader)
        continue;

      for (auto *BB : L->blocks()) {
        for (auto I = BB->begin(), E = BB->end(); I != E; ) {
          auto it = &*I++;

          if (it->isBinaryOp() && isLoopInvariant(it, *L)) {
            it->moveBefore(preheader->getTerminator());
            errs() << "moved " << *it << '\n';
          }
        }
      }
    }

    for (auto &BB : F) {
      errs() << "  BasicBlock:\n";
      
      for (auto &I : BB) {
        errs() << "    Instruction: " << I << "\n";
      }
    }    
    errs() << "\n";

    return PreservedAnalyses::none();
  }
};
Enter fullscreen mode Exit fullscreen mode

We have two methods to explain:

  • run, the entry point of our pass
  • isLoopInvariant, a predicate used to decide if an instruction is invariant

Let's start with run first. We start by using the Analysis manager to get a reference to the result of running the Loop analysis pass. This gives us access to all loops inside the current function.

Next, we iterate over them using a for-range loop. For each loop, we get the preheader, where we will hoist invariant instructions, and test if it exists, if not, we skip to the next.

Now, we iterate over the blocks of the loop and for each block we iterate over its instructions.
For each instruction, we check if it's a binary operation and invariant. We move it to the preheader before its terminator if it's. We do that using moveBefore.

Remember our little bug from the previous article? don't forget auto it = &*I++; at the beginning to ensure we advance before transforming things.

The rest is just printing the transformed function.

Let's move to isLoopInvariant. This method accepts a pointer to the instruction in question and a reference to Loop. We Iterate over the operands of the instruction. if the operand is defined by an instruction, we further test its location (inside or outside the loop). We return false if it's inside the loop. Voila! Look simple enough, right? well, don't forget the assumptions we made on the input :).

Next, Compile the pass like shown in previous articles and run:

$ opt -load-pass-plugin=./LoopInvariantCodeMotion.so -passes=simple-licm -disable-output < licm.ll
Enter fullscreen mode Exit fullscreen mode

and the result is:

moved   %invariant = add i32 %x, 10
  BasicBlock:
    Instruction:   %invariant = add i32 %x, 10
    Instruction:   br label %loop
  BasicBlock:
    Instruction:   %i = phi i32 [ 0, %entry ], [ %inc, %loop ]
    Instruction:   %result = phi i32 [ 0, %entry ], [ %result.next, %loop ]
    Instruction:   %mul = mul i32 %i, %invariant
    Instruction:   %result.next = add i32 %result, %mul
    Instruction:   %inc = add i32 %i, 1
    Instruction:   %cmp = icmp slt i32 %inc, %n
    Instruction:   br i1 %cmp, label %loop, label %exit
  BasicBlock:
    Instruction:   ret i32 %result.next
Enter fullscreen mode Exit fullscreen mode

As shown above the invariant instruction was detected and moved correctly outside the loop.

Conclusion

In this article, we were exposed to some of LLVM loop facilities. We saw how can we use them in implementing LICM pass with prior assumption on input. To understand more, I invite you to loosen these assumptions at your own pace and expand on what the pass can handle.

Top comments (0)