Updates branch_weights annotation for call instructions during inlining.

[danielcdh]

Summary

Inliner should update the branch_weights annotation to scale it to proper value.

[danielcdh]

update

[eraman]

include/llvm/IR/Instruction.h

257

Do not use double here and instead specify it as a fraction.

lib/IR/Instruction.cpp

691

What are the assumptions on the MD_prof metadata here? Could you assert V here? If it is indeed possible that ith operand is not constant int, could the next operand be constant int (in which case you shouldn't return here)

lib/Transforms/Utils/InlineFunction.cpp

1427–1454

Document this method. Also, perhaps choose a more meaningful name.

1442

No need to keep track of the cloned blocks and update the calls inside them. VMap also includes the instructions - so just check if they are calls and update the weights.

1453

There is a subtle issue here. Let's say there is a block in the callee that is not executed in this call context and that block has a call instruction. You'll still be reducing the weight of that call even though you haven't cloned it. Similar issue does exist in updating the bfi incrementally, but that's okay because BFI does get recomputed later. In this case, once the weight gets incorrectly updated it remains so forever. What you should be doing is update the weights only when the call has been cloned.

1470

This code that gets either the metadata weight or the block weight should be refactored and moved out of here.

1482

This gets confusing. You're now updating the entry count based on the metadata. When sample profile is used, is it expected that the sum of the calls' profile weights equal the entry count? Do you smooth out the entry count based on call instructions' profile weight when you initially load the profile. In any case, more comments are helpful.

test/Transforms/Inline/prof-update.ll

2

Please expand this test or add a new test case to handle more complex cases where some calls are not cloned into the caller.

[danielcdh]

update

[danielcdh]

lib/IR/Instruction.cpp

691

removed the check.

lib/Transforms/Utils/InlineFunction.cpp

1482

Code updated. Yes, the sample profile collection makes sure the function entry count meet with all call edges.

Let me know if you think there still need to add comments.

[eraman]

lib/IR/Instruction.cpp

693

For counts from instrumentation based profile, I use a 128 bit APInt to do the arithmetic and then educe it to 64 bit. In this case, since the values being multiplied are sample profile weights, I suppose the likelihood of Val * S overflowing 64 bits is very small and the APInt may be unnecessary. Even then, use saturated multiply so that the values are at least sensible even in the unlikely case of overflow.

lib/Transforms/Utils/InlineFunction.cpp

1440

You have to check if Entry.second is a call. It is possible for call instructions to be simplified to a value during the cloning.

1446

Useful to add a comment as to why this is done only when the instruction is in VMap.

test/Transforms/Inline/prof-update.ll

2

Nit: I prefer writing comments to explain why I expect the values in the CHECK statements. This helps anyone trying to modify the test later.

[danielcdh]

update

[danielcdh]

update

[danielcdh]

update

[eraman]

lib/IR/Instruction.cpp

694

Perhaps it is not clear above, but SaturatingMultiply may be sufficient in your case. I'm ok with using APInt as well as in the most common case (multiplication result fits within 64 bits), the division is not expensive. Perhaps add a comment above the udiv stating that this could potentially be expensive, but most likely the product is going to fit within 64 bits.

lib/Transforms/Utils/InlineFunction.cpp

1449

The check above is not needed. After all, if the block is reachable in the call context, all instructions inside are reachable too. But there is one tricky case. Let's say that a call has been simplified into a value (this can happen during pruning), Should you update the call count or not? If you should not update, then the check above should be isa<CallInst>(VMap.Count(CI)). I think updating is the right thing here though.

[zzheng]

If X is Optional, please use X.has_value() to check if it has value and X.getValue() to check for 0 value.

lib/Transforms/Utils/InlineFunction.cpp

1432

Do you mean to use !CalleeEntryCount.has_value()?

1436

I'm confused, do you mean to test if CallSiteCount has no value or it has a value of 0?

1440

Isn't it redundant to check ptr != nullptr?

Is there any chance that Entry.second is NULL and what happen on (*Entry.second) in that case?

[danielcdh]

update

lib/Transforms/Utils/InlineFunction.cpp

1440

It indeed happened once when I compile large code, that's why I added this check. I'm not familiar with cloning code, Easwaran may have some insights on why this could be nullptr? Or it's a bug in function clone?

1449

The prof count only attaches to branch instructions. For the case you mentioned, it's already checked by the "if (CallInst *CI = dyn_cast<CallInst>(&I))", thus I simply removed the second check.

[zzheng]

a couple more nitpicks... please check all use of Optional variables

lib/Transforms/Utils/InlineFunction.cpp

1466

!CalleeCount.hasValue()

1470

!CallCount.hasValue()

[danielcdh]

update

[eraman]

LGTM

lib/Transforms/Utils/InlineFunction.cpp

1440

The value type of the VMap is WeakVH and the value it holds becomes null when it gets destroyed.

include/llvm/IR/Instruction.h

public:
   /// Returns false if no metadata or invalid metadata was found.
   bool extractProfMetadata(uint64_t &TrueVal, uint64_t &FalseVal) const;
 
   /// Retrieve total raw weight values of a branch.
   /// Returns true on success with profile total weights filled in.
   /// Returns false if no metadata was found.
   bool extractProfTotalWeight(uint64_t &TotalVal) const;
 
+  /// Updates branch_weights metadata with the given ratio \p R. If \p Reverse
+  /// is true, update branch_weight with 1 - \p R.
+  void updateProfWeight(double R, bool Reverse);
+
   /// Set the debug location information for this instruction.
   void setDebugLoc(DebugLoc Loc) { DbgLoc = std::move(Loc); }
 
   /// Return the debug location for this node as a DebugLoc.
   const DebugLoc &getDebugLoc() const { return DbgLoc; }
 
   /// Set or clear the nsw flag on this instruction, which must be an operator
   /// which supports this flag. See LangRef.html for the meaning of this flag.

lib/IR/Instruction.cpp

 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 using namespace llvm;
 
 Instruction::Instruction(Type *ty, unsigned it, Use *Ops, unsigned NumOps,
                          Instruction *InsertBefore)
   : User(ty, Value::InstructionVal + it, Ops, NumOps), Parent(nullptr) {
 

 #include "llvm/IR/Instruction.def"
 #undef HANDLE_INST
   }
 
   New->SubclassOptionalData = SubclassOptionalData;
   New->copyMetadata(*this);
   return New;
 }
+
+void Instruction::updateProfWeight(double R, bool Reverse) {
+  auto *ProfileData = getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData)
+    return;
+
+  auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
+  if (!ProfDataName || !ProfDataName->getString().equals("branch_weights"))
+    return;
+
+  SmallVector<uint32_t, 4> Weights;
+  for (unsigned i = 1; i < ProfileData->getNumOperands(); i++) {
+    auto *V = mdconst::dyn_extract<ConstantInt>(ProfileData->getOperand(i));
+    if (!V)
+      return;
+    uint64_t Val = V->getValue().getZExtValue();
+    Weights.push_back(Reverse ? Val - Val * R : Val * R);
+  }
+  MDBuilder MDB(getContext());
+  setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
+}

lib/Transforms/Utils/InlineFunction.cpp

for (auto const &Entry : VMap) {
     }
     CallerBFI->setBlockFreq(ClonedBB, Freq);
   }
   BasicBlock *EntryClone = cast<BasicBlock>(VMap.lookup(&CalleeEntryBlock));
   CallerBFI->setBlockFreqAndScale(
       EntryClone, CallerBFI->getBlockFreq(CallSiteBlock).getFrequency(),
       ClonedBBs);
 }
 
+static void updateCallProfile(Function *Callee,
+                              const ValueToValueMapTy &VMap,
+                              const Optional<uint64_t> &CalleeEntryCount,
+                              const Instruction *TheCall) {
+  uint64_t CallCount = 0;
+  if (!CalleeEntryCount || CalleeEntryCount.getValue() < 1 ||
+      !TheCall->extractProfTotalWeight(CallCount) || CallCount < 1)
+    return;
+  double Ratio = std::min((double) CallCount / CalleeEntryCount.getValue(),
+                          1.0);
+
+  SmallPtrSet<BasicBlock *, 16> ClonedBBs;
+  for (auto const &Entry : VMap) {
+    if (!isa<BasicBlock>(Entry.first) || !Entry.second)
+      continue;
+    BasicBlock *ClonedBB = cast<BasicBlock>(Entry.second);
+    if (ClonedBBs.insert(ClonedBB).second)
+      for (Instruction &I : *ClonedBB)
+        if (CallInst *CI = dyn_cast<CallInst>(&I))
+          CI->updateProfWeight(Ratio, false);
+  }
+  for (BasicBlock &BB : *Callee)
+    for (Instruction &I : BB)
+      if (CallInst *CI = dyn_cast<CallInst>(&I))
+        CI->updateProfWeight(Ratio, true);
+}
+
 /// Update the entry count of callee after inlining.
 ///
 /// The callsite's block count is subtracted from the callee's function entry
 /// count.
-static void updateCalleeCount(BlockFrequencyInfo &CallerBFI, BasicBlock *CallBB,
-                              Function *Callee) {
+static void updateCalleeCount(BlockFrequencyInfo *CallerBFI, BasicBlock *CallBB,
+                              Instruction *CallInst, Function *Callee) {
   // If the callee has a original count of N, and the estimated count of
   // callsite is M, the new callee count is set to N - M. M is estimated from
   // the caller's entry count, its entry block frequency and the block frequency
   // of the callsite.
   Optional<uint64_t> CalleeCount = Callee->getEntryCount();
   if (!CalleeCount)
     return;
-  Optional<uint64_t> CallSiteCount = CallerBFI.getBlockProfileCount(CallBB);
-  if (!CallSiteCount)
+  uint64_t CallCount = 0;
+  if (!CallInst->extractProfTotalWeight(CallCount))
+    if (CallerBFI) {
+      Optional<uint64_t> CallSiteCount = CallerBFI->getBlockProfileCount(CallBB);
+      if (CallSiteCount)
+        CallCount = CallSiteCount.getValue();
+    }
+  if (CallCount == 0)
     return;
   // Since CallSiteCount is an estimate, it could exceed the original callee
   // count and has to be set to 0.
-  if (CallSiteCount.getValue() > CalleeCount.getValue())
+  if (CallCount > CalleeCount.getValue())
     Callee->setEntryCount(0);
   else
-    Callee->setEntryCount(CalleeCount.getValue() - CallSiteCount.getValue());
+    Callee->setEntryCount(CalleeCount.getValue() - CallCount);
 }
 
 /// This function inlines the called function into the basic block of the
 /// caller. This returns false if it is not possible to inline this call.
 /// The program is still in a well defined state if this occurs though.
 ///
 /// Note that this only does one level of inlining.  For example, if the
 /// instruction 'call B' is inlined, and 'B' calls 'C', then the call to 'C' now
Function::iterator FirstNewBlock;
     // (which can happen, e.g., because an argument was constant), but we'll be
     // happy with whatever the cloner can do.
     CloneAndPruneFunctionInto(Caller, CalledFunc, VMap,
                               /*ModuleLevelChanges=*/false, Returns, ".i",
                               &InlinedFunctionInfo, TheCall);
     // Remember the first block that is newly cloned over.
     FirstNewBlock = LastBlock; ++FirstNewBlock;
 
-    if (IFI.CallerBFI != nullptr && IFI.CalleeBFI != nullptr) {
+    if (IFI.CallerBFI != nullptr && IFI.CalleeBFI != nullptr)
       // Update the BFI of blocks cloned into the caller.
       updateCallerBFI(OrigBB, VMap, IFI.CallerBFI, IFI.CalleeBFI,
                       CalledFunc->front());
-      // Update the profile count of callee.
-      updateCalleeCount(*IFI.CallerBFI, OrigBB, CalledFunc);
-    }
+
+    updateCallProfile(CalledFunc, VMap, CalledFunc->getEntryCount(), TheCall);
+    updateCalleeCount(IFI.CallerBFI, OrigBB, TheCall, CalledFunc);
 
     // Inject byval arguments initialization.
     for (std::pair<Value*, Value*> &Init : ByValInit)
       HandleByValArgumentInit(Init.first, Init.second, Caller->getParent(),
                               &*FirstNewBlock, IFI);
 
     Optional<OperandBundleUse> ParentDeopt =
         CS.getOperandBundle(LLVMContext::OB_deopt);

test/Transforms/Inline/prof-update.ll

+; RUN: opt < %s -inline -S | FileCheck %s
+; Checks if inliner updates branch_weights annotation for call instructions.
+
+declare void @ext();
+
+; CHECK: define void @callee() !prof ![[ENTRY_COUNT:[0-9]*]]
+define void  @callee() !prof !1 {
+; CHECK: call void @ext(), !prof ![[COUNT_CALLEE:[0-9]*]]
+  call void @ext(), !prof !2
+  ret void
+}
+
+; CHECK: define void @caller()
+define void @caller() {
+; CHECK: call void @ext(), !prof ![[COUNT_CALLER:[0-9]*]]
+  call void @callee(), !prof !3
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"MaxFunctionCount", i32 2000}
+!1 = !{!"function_entry_count", i64 1000}
+!2 = !{!"branch_weights", i64 2000}
+!3 = !{!"branch_weights", i64 400}
+attributes #0 = { alwaysinline }
+; CHECK: ![[ENTRY_COUNT]] = !{!"function_entry_count", i64 600}
+; CHECK: ![[COUNT_CALLEE]] = !{!"branch_weights", i32 1200}
+; CHECK: ![[COUNT_CALLER]] = !{!"branch_weights", i32 800}