Sub-optimal placement of loop exit blocks

[rajkrthakur]

Summary

This draft patch assigns probabilistic metadata to the branch for the required block placement

This resolves https://github.com/llvm/llvm-project/issues/50529

TODO - Add unit tests.

[efriedma]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

394

Please avoid using the LLVM C API from C++ code. Please just use the C++ API (for example, MDString::get).

You can get the LLVMContext with something like LatchBI->getContext().

398

These weights seem sort of random; do we have any reason to think they're accurate? The default heuristics for branch weights should assume that loop backedges are likely.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopIterator.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
+#include "llvm-c/Core.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   // (2b) latch is conditional and is an exiting block
   // FIXME: The implementation can be extended to work with more complicated
   // cases, e.g. loops with multiple latches.
   BranchInst *LatchBI = dyn_cast<BranchInst>(LatchBlock->getTerminator());
 
   // A conditional branch which exits the loop, which can be optimized to an
   // unconditional branch in the unrolled loop in some cases.
   bool LatchIsExiting = L->isLoopExiting(LatchBlock);
+
+  if (LatchBI && (LatchBI->isConditional() && LatchIsExiting)) {
+    // Insert block frequency for exit latch. This rearranges the basic blocks
+    // neatly.
+    LLVMMetadataRef branch_weights =
+        LLVMMDStringInContext2(LLVMGetGlobalContext(), "branch_weights", 14);
+    LLVMMetadataRef weight1 =
+        LLVMValueAsMetadata(LLVMConstInt(LLVMIntType(32), 20, 0));
+    LLVMMetadataRef weight2 =
+        LLVMValueAsMetadata(LLVMConstInt(LLVMIntType(32), 80, 0));
+    LLVMMetadataRef mds[] = {branch_weights, weight1, weight2};
+    LLVMMetadataRef metadata =
+        LLVMMDNodeInContext2(LLVMGetGlobalContext(), mds, 3);
+    LLVMValueRef metadata_value =
+        LLVMMetadataAsValue(LLVMGetGlobalContext(), metadata);
+
+    LLVMSetMetadata(wrap(LatchBI), LLVMGetMDKindID("prof", 4), metadata_value);
+  }
+
   if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) {
     LLVM_DEBUG(
         dbgs() << "Can't unroll; a conditional latch must exit the loop");
     return LoopUnrollResult::Unmodified;
   }
 
   // Loops containing convergent instructions cannot use runtime unrolling,
   // as the prologue/epilogue may add additional control-dependencies to