[PGO] Add a functionality to always instrument the func entry BB

[xur]

Summary

Add an option to always instrument function entry BB (default off)
Add an options to do atomic update the first counter in each instrumented function.

[davidxl]

llvm/lib/Transforms/Instrumentation/CFGMST.h

110

Does this guarantee entry be instrumented?

147

Is this the way to make sure entry is always instrumented?

292

is this needed? Can somthing be done in CompputeMST without relying on weight trick?

[davidxl]

ok. Why atomic update just for first entry? are they particularly more contended than other counters?

[wmi]

Thanks for the patch! I guess atomic update is to reduce the cost of the entry counter for application with many threads. How much difference there is w/wo enabling -atomic-first-counter?

[wmi]

I tried it on an internal benchmark.

base: head of trunk
exp: pgo-instrument-entry=true, atomic-first-counter=false
qps of Instrumented binary run: 344.63 vs 340.56
optimized binary run: qps had no significant change in 6 runs. latency had ~0.1% regression in 3 runs and no regression in other 3 runs.

base: head of trunk
exp: pgo-instrument-entry=true, atomic-first-counter=true
qps of Instrumented binary run: 342.53 vs 265.58
optimized binary run: qps had no significant change in 2 runs, and had ~0.11% improvement in 4 runs. latency had ~0.35% improvement in 6 runs.

The test showed always inserting counter in the entry block had minor impact on instrumented binary and optimized binary runs. Using atomic entry counter could improve profile quality but significantly slowdown instrument binary.

[davidxl]

LGTM. Versioning support is needed when it is flipped on by default.

llvm/lib/Transforms/Instrumentation/CFGMST.h

 #define LLVM_LIB_TRANSFORMS_INSTRUMENTATION_CFGMST_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Support/BranchProbability.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <utility>
 #include <vector>
 
 #define DEBUG_TYPE "cfgmst"
 
+using namespace llvm;
+static cl::opt<bool> PGOInstrumentEntry(
+    "pgo-instrument-entry", cl::init(false), cl::Hidden,
+    cl::desc("Force to instrument function entry basicblock."));
+
 namespace llvm {
 
 /// An union-find based Minimum Spanning Tree for CFG
 ///
 /// Implements a Union-find algorithm to compute Minimum Spanning Tree
 /// for a given CFG.
 template <class Edge, class BBInfo> class CFGMST {
 public:
public:
   // Traverse the CFG using a stack. Find all the edges and assign the weight.
   // Edges with large weight will be put into MST first so they are less likely
   // to be instrumented.
   void buildEdges() {
     LLVM_DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n");
 
     const BasicBlock *Entry = &(F.getEntryBlock());
     uint64_t EntryWeight = (BFI != nullptr ? BFI->getEntryFreq() : 2);
+    // If we want to instrument the entry count, lower the weight to 0.
+    if (PGOInstrumentEntry)
+      EntryWeight = 0;
     Edge *EntryIncoming = nullptr, *EntryOutgoing = nullptr,
          *ExitOutgoing = nullptr, *ExitIncoming = nullptr;
     uint64_t MaxEntryOutWeight = 0, MaxExitOutWeight = 0, MaxExitInWeight = 0;
 
     // Add a fake edge to the entry.
     EntryIncoming = &addEdge(nullptr, Entry, EntryWeight);
     LLVM_DEBUG(dbgs() << "  Edge: from fake node to " << Entry->getName()
                       << " w = " << EntryWeight << "\n");
 
     // Special handling for single BB functions.
for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
           if (Critical) {
             if (scaleFactor < UINT64_MAX / CriticalEdgeMultiplier)
               scaleFactor *= CriticalEdgeMultiplier;
             else
               scaleFactor = UINT64_MAX;
           }
           if (BPI != nullptr)
             Weight = BPI->getEdgeProbability(&*BB, TargetBB).scale(scaleFactor);
+          if (Weight == 0)
+            Weight++;
           auto *E = &addEdge(&*BB, TargetBB, Weight);
           E->IsCritical = Critical;
           LLVM_DEBUG(dbgs() << "  Edge: from " << BB->getName() << " to "
                             << TargetBB->getName() << "  w=" << Weight << "\n");
 
           // Keep track of entry/exit edges:
           if (&*BB == Entry) {
             if (Weight > MaxEntryOutWeight) {

 
 public:
   CFGMST(Function &Func, BranchProbabilityInfo *BPI_ = nullptr,
          BlockFrequencyInfo *BFI_ = nullptr)
       : F(Func), BPI(BPI_), BFI(BFI_) {
     buildEdges();
     sortEdgesByWeight();
     computeMinimumSpanningTree();
+    if (PGOInstrumentEntry && (AllEdges.size() > 1))
+      std::iter_swap(std::move(AllEdges.begin()),
+                     std::move(AllEdges.begin() + AllEdges.size() - 1));
   }
 };
 
 } // end namespace llvm
 
 #undef DEBUG_TYPE // "cfgmst"
 
 #endif // LLVM_LIB_TRANSFORMS_INSTRUMENTATION_CFGMST_H

llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp

cl::opt<bool> AtomicCounterUpdateAll(
     cl::init(false));
 
 cl::opt<bool> AtomicCounterUpdatePromoted(
     "atomic-counter-update-promoted", cl::ZeroOrMore,
     cl::desc("Do counter update using atomic fetch add "
              " for promoted counters only"),
     cl::init(false));
 
+cl::opt<bool> AtomicFirstCounter(
+    "atomic-first-counter", cl::ZeroOrMore,
+    cl::desc("Use atomic fetch add for first counter in a function (usually "
+             "the entry counter)"),
+    cl::init(false));
+
 // If the option is not specified, the default behavior about whether
 // counter promotion is done depends on how instrumentaiton lowering
 // pipeline is setup, i.e., the default value of true of this option
 // does not mean the promotion will be done by default. Explicitly
 // setting this option can override the default behavior.
 cl::opt<bool> DoCounterPromotion("do-counter-promotion", cl::ZeroOrMore,
                                  cl::desc("Do counter register promotion"),
                                  cl::init(false));
if (!LI) {
         Bias->setVisibility(GlobalVariable::HiddenVisibility);
       }
       LI = Builder.CreateLoad(Int64Ty, Bias);
     }
     auto *Add = Builder.CreateAdd(Builder.CreatePtrToInt(Addr, Int64Ty), LI);
     Addr = Builder.CreateIntToPtr(Add, Int64PtrTy);
   }
 
-  if (Options.Atomic || AtomicCounterUpdateAll) {
+  if (Options.Atomic || AtomicCounterUpdateAll ||
+      (Index == 0 && AtomicFirstCounter)) {
     Builder.CreateAtomicRMW(AtomicRMWInst::Add, Addr, Inc->getStep(),
                             AtomicOrdering::Monotonic);
   } else {
     Value *IncStep = Inc->getStep();
     Value *Load = Builder.CreateLoad(IncStep->getType(), Addr, "pgocount");
     auto *Count = Builder.CreateAdd(Load, Inc->getStep());
     auto *Store = Builder.CreateStore(Count, Addr);
     if (isCounterPromotionEnabled())

llvm/test/Transforms/PGOProfile/instr_entry_bb.ll

+; RUN: opt < %s -pgo-instr-gen -pgo-instrument-entry -S | FileCheck %s --check-prefix=GEN
+; RUN: opt < %s -passes=pgo-instr-gen -pgo-instrument-entry -S | FileCheck %s --check-prefix=GEN
+; RUN: opt < %s -pgo-instr-gen -pgo-instrument-entry -instrprof -atomic-first-counter -S | FileCheck %s --check-prefix=GENA
+; RUN: opt < %s -passes=pgo-instr-gen,instrprof -pgo-instrument-entry -atomic-first-counter -S | FileCheck %s --check-prefix=GENA
+
+; RUN: llvm-profdata merge %S/Inputs/branch2.proftext -o %t.profdata
+; RUN: opt < %s -pgo-instr-use -pgo-test-profile-file=%t.profdata -pgo-instrument-entry -S | FileCheck %s --check-prefix=USE
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -pgo-instrument-entry -S | FileCheck %s --check-prefix=USE
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; GEN: $__llvm_profile_raw_version = comdat any
+; GEN: @__llvm_profile_raw_version = constant i64 {{[0-9]+}}, comdat
+; GEN: @__profn_test_br_2 = private constant [9 x i8] c"test_br_2"
+
+define i32 @test_br_2(i32 %i) {
+entry:
+; GEN: entry:
+; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([9 x i8], [9 x i8]* @__profn_test_br_2, i32 0, i32 0), i64 29667547796, i32 2, i32 0)
+; GENA: entry:
+; GENA: %{{[0-9+]}} = atomicrmw add i64* getelementptr inbounds ([2 x i64], [2 x i64]* @__profc_test_br_2, i64 0, i64 0), i64 1 monotonic
+; USE: br i1 %cmp, label %if.then, label %if.else
+; USE-SAME: !prof ![[BW_ENTRY:[0-9]+]]
+; USE: ![[BW_ENTRY]] = !{!"branch_weights", i32 0, i32 1}
+  %cmp = icmp sgt i32 %i, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+; GEN: if.then:
+; GEN-NOT: llvm.instrprof.increment
+  %add = add nsw i32 %i, 2
+  br label %if.end
+
+if.else:
+; GEN: if.else:
+; GEN: call void @llvm.instrprof.increment(i8* getelementptr inbounds ([9 x i8], [9 x i8]* @__profn_test_br_2, i32 0, i32 0), i64 29667547796, i32 2, i32 1)
+; GENA: if.else:
+; GENA:  %pgocount = load i64, i64* getelementptr inbounds ([2 x i64], [2 x i64]* @__profc_test_br_2, i64 0, i64 1), align 8
+; GENA:  [[V:%[0-9]*]] = add i64 %pgocount, 1
+; GENA:  store i64 [[V]], i64* getelementptr inbounds ([2 x i64], [2 x i64]* @__profc_test_br_2, i64 0, i64 1), align 8
+  %sub = sub nsw i32 %i, 2
+  br label %if.end
+
+if.end:
+; GEN: if.end:
+; GEN-NOT: llvm.instrprof.increment
+  %retv = phi i32 [ %add, %if.then ], [ %sub, %if.else ]
+  ret i32 %retv
+; GEN: ret
+}