[PGO] Memory intrinsic calls optimization based on profiled size

[xur]

Summary

This patch optimizes two memory intrinsic operations: memset and memcpy based on the profiled size of the operation. The high level transformation is like:

mem_op(..., size)
==>
switch (size) {
  case s1:
     mem_op(..., s1);
     goto merge_bb;
  case s2:
     mem_op(..., s2);
     goto merge_bb;
  ...
  default:
     mem_op(..., size);
     goto merge_bb;
  }
merge_bb:

There are a few internal options that controls when the optimization takes place.

This patch in an incremental patch based on
https://reviews.llvm.org/D28965

[davidxl]

The optimization pass should be split into two phases as IC promotion. The annotation part should probably be merged with the instrumentation patch. The transformation patch should be done in the same pass as IndirectCallPromotion.

[xur]

In D28966#694902, @davidxl wrote:

The optimization pass should be split into two phases as IC promotion. The annotation part should probably be merged with the instrumentation patch. The transformation patch should be done in the same pass as IndirectCallPromotion.

Is there a good reason for doing the transformation late? Here I don't do the annotation and instead, I do the transformation directly in the same pass. The main reason we have annotations in indirect-call-promotion is we need to call it late (in LTO or ThinLTO). Another reason I'm reluctant to use annotation is that we need to maintain/update it (for inline and unroll, for example).

-Rong

[xur]

Integrated David's suggestion to make the optimization a stand along pass (a function pass).
It currently resides in IndirectCallPromotion.cpp. I will have a follow-up patch to change IndirectCallPromotion.cpp to a more appropriate name.

-Rong

[davidxl]

lib/Transforms/Instrumentation/IndirectCallPromotion.cpp

789 ↗

(On Diff #92382)

Add a comment on the kinds.

NormalGroup--> NonLargeGroup?

794 ↗

(On Diff #92382)

Is it correct to guard the parsing here?

Should the parsing be called at the start of the pass with PreciseStart etc as a member variable?

824 ↗

(On Diff #92382)

DC --> RemainCount?

843 ↗

(On Diff #92382)

Should it use TotalCount here?

[xur]

lib/Transforms/Instrumentation/IndirectCallPromotion.cpp

794 ↗

(On Diff #92382)

moved to class MemOPSizeOpt

843 ↗

(On Diff #92382)

The percent threshold use the ratio of current count and remain count.
I refactor the code a bit (by using a function). The new code also compares the count with the count-threshold.

[xur]

This new patch integrated David's review suggestion to make this optimization a separated pass.
This pass will be called after inlining to avoid the potential interference to in-lining decisions. It will use both the metadata annotation for memop size profiling and the enclosing BB's edge count to find the actual count.

[davidxl]

lib/Transforms/Instrumentation/IndirectCallPromotion.cpp

834 ↗

(On Diff #93408)

Use APInt or SaturatingMultipleAdd to deal with possible overflow.

849 ↗

(On Diff #93408)

Should check the actual count.

883 ↗

(On Diff #93408)

can shrink the MemOpScaleCount check into getScaledCount

[xur]

integrated David's review comments

[davidxl]

lgtm

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

 // Command line option to enable/disable select instruction instrumentation.
 static cl::opt<bool> PGOInstrSelect("pgo-instr-select", cl::init(true),
                                     cl::Hidden);
 
 // Command line option to enable/disable memop intrinsic calls..
 static cl::opt<bool> PGOInstrMemOP("pgo-instr-memop", cl::init(true),
                                     cl::Hidden);
 
+// The minimum call count to optimize memory intrinsic calls.
+static cl::opt<unsigned>
+    MemOPCountThreshold("memop-count-threshold", cl::Hidden, cl::ZeroOrMore,
+                        cl::init(1000),
+                        cl::desc("The minimum count to optimize memory "
+                                 "intrinsic calls"));
+
+// The percent threshold to optimize memory intrinsic calls.
+static cl::opt<unsigned>
+    MemOPPercentThreshold("memop-percent-threshold", cl::init(66), cl::Hidden,
+                          cl::ZeroOrMore,
+                          cl::desc("The percentage threshold for the optimize "
+                                   " memory intrinsic calls"));
+
+// Maximum number of versions for optimizing memory intrinsic call.
+static cl::opt<unsigned>
+    MemOPMaxVersion("memop-max-version", cl::init(3), cl::Hidden,
+                          cl::ZeroOrMore,
+                          cl::desc("The max version for the optimize memory "
+                                   " intrinsic calls"));
+
 namespace {
 
 /// The select instruction visitor plays three roles specified
 /// by the mode. In \c VM_counting mode, it simply counts the number of
 /// select instructions. In \c VM_instrument mode, it inserts code to count
 /// the number times TrueValue of select is taken. In \c VM_annotate mode,
 /// it reads the profile data and annotate the select instruction with metadata.
 enum VisitMode { VM_counting, VM_instrument, VM_annotate, VM_optimize };
struct MemIntrinsicVisitor : public InstVisitor<MemIntrinsicVisitor> {
   Function &F;
   unsigned NMemIs = 0;           // Number of memIntrinsics instrumented.
   VisitMode Mode = VM_counting;  // Visiting mode.
   unsigned CurCtrId = 0;         // Current counter index.
   unsigned TotalNumCtrs = 0;     // Total number of counters
   GlobalVariable *FuncNameVar = nullptr;
   uint64_t FuncHash = 0;
   PGOUseFunc *UseFunc = nullptr;
+  std::vector<MemIntrinsic *> Candidates;
 
   MemIntrinsicVisitor(Function &Func) : F(Func) {}
 
   void countMemIntrinsics(Function &Func) {
     NMemIs = 0;
     Mode = VM_counting;
     visit(Func);
   }
Builder.CreateCall(
         {llvm::ConstantExpr::getBitCast(FuncInfo.FuncNameVar, I8PtrTy),
          Builder.getInt64(FuncInfo.FunctionHash),
          Builder.CreatePtrToInt(Callee, Builder.getInt64Ty()),
          Builder.getInt32(llvm::InstrProfValueKind::IPVK_IndirectCallTarget),
          Builder.getInt32(NumIndirectCallSites++)});
   }
   NumOfPGOICall += NumIndirectCallSites;
 
-  // Now instrument memop instrinsic calls:
+  // Now instrument memop intrinsic calls:
   FuncInfo.MIVisitor.instrumentMemIntrinsics(F, NumCounters,
                                              FuncInfo.FuncNameVar,
                                              FuncInfo.FunctionHash);
 }
 
 // This class represents a CFG edge in profile use compilation.
 struct PGOUseEdge : public PGOEdge {
   bool CountValid;
Builder.CreateCall(
       Intrinsic::getDeclaration(M, Intrinsic::instrprof_value_range_profile),
       {llvm::ConstantExpr::getBitCast(FuncNameVar, I8PtrTy),
        Builder.getInt64(FuncHash), Builder.CreatePtrToInt(Length, Int64Ty),
        Builder.getInt32(llvm::InstrProfValueKind::IPVK_MemOPSize),
        Builder.getInt32(CurCtrId)});
   ++CurCtrId;
 }
 
+void MemIntrinsicVisitor::optimizeMemIntrinsics(Function &Func, PGOUseFunc *UF) {
+  countMemIntrinsics(Func);
+  auto NumMemOPCalls = getNumOfMemIntrinsics();
+  unsigned NumValueSites = UF->getProfileRecord().getNumValueSites(IPVK_MemOPSize);
+  if (NumValueSites != NumMemOPCalls) {
+    std::string Msg =
+        std::string("Inconsistent number of memory intrinsic calls: ") +
+        Func.getName().str();
+    auto &Ctx = Func.getContext();
+    Ctx.diagnose(
+        DiagnosticInfoPGOProfile(Func.getParent()->getName().data(), Msg, DS_Warning));
+    return;
+  }
+
+  // Now optimize memory intrinsic calls.
+  Candidates.clear();
+  Mode = VM_optimize;
+  UseFunc = UF;
+  visit(Func);
+  for (auto &I : Candidates)
+    optimizeOneMemIntrinsic(*I);
+}
+
+static const char *getMIName(const MemIntrinsic &MI) {
+  switch (MI.getIntrinsicID()) {
+    case Intrinsic::memcpy:
+     return "memcpy";
+    case Intrinsic::memmove:
+     return "memmove";
+    case Intrinsic::memset:
+     return "memset";
+    default:
+     return "unknown";
+  }
+}
+void MemIntrinsicVisitor::optimizeOneMemIntrinsic(MemIntrinsic &MI) {
+  const InstrProfRecord &InstrProfRecord = UseFunc->getProfileRecord();
+  unsigned Index = CurCtrId++;
+  uint32_t NV = InstrProfRecord.getNumValueDataForSite(IPVK_MemOPSize, Index);
+  if (!NV)
+    return;
+
+  if (MI.getIntrinsicID() == Intrinsic::memmove)
+    return;
+
+  uint64_t Sum = 0;
+  std::unique_ptr<InstrProfValueData[]> VD =
+      InstrProfRecord.getValueForSite(IPVK_MemOPSize, Index, &Sum);
+ 
+  if (Sum < MemOPCountThreshold)
+    return;
+
+  ArrayRef<InstrProfValueData> VDs(VD.get(), NV);
+
+  uint64_t SumForSmallSizes = 0;
+  uint64_t Threshold = Sum / 100;
+  SmallVector<uint64_t, 16> SizeIds;
+  SmallVector<uint64_t, 16> CaseCounts;
+  uint64_t MaxCount = 0;
+  unsigned Version = 0;
+  // Default case is in the front -- save the slot here.
+  CaseCounts.push_back(0);
+  for (auto &VD : VDs) {
+    int64_t V = VD.Value;
+    uint64_t C = VD.Count;
+    // only count these greater than 1% of the Sum;
+    if (V > 0 && C > Threshold) {
+      SumForSmallSizes += C;
+      SizeIds.push_back(V);
+      CaseCounts.push_back(C);
+      if (C > MaxCount)
+        MaxCount = C;
+      if (++Version > MemOPMaxVersion && MemOPMaxVersion !=0)
+        break;
+    }
+  }
+  // Add count for the default case.
+  uint64_t DC = Sum - SumForSmallSizes;
+  CaseCounts[0] = DC;
+  if (DC > MaxCount)
+    MaxCount = DC;
+
+  DEBUG(dbgs() << "Read one memory intrinsic profile: Index=" << Index
+               << ":  " << SumForSmallSizes << " vs " << Sum << "\n");
+  DEBUG(for (auto &VD : VDs) { dbgs() << "  (" << VD.Value << ","
+                                      << VD.Count << ")\n"; });
+
+  uint64_t CountThreshold = Sum / 100 * MemOPPercentThreshold;
+  if (Sum < 100)
+    CountThreshold = Sum * MemOPPercentThreshold /100;
+  if (SumForSmallSizes <= CountThreshold)
+    return;
+
+  DEBUG(dbgs() << "Optimize one memory intrinsic call\n");
+
+  // mem_op(..., size)
+  // ==>
+  // switch (size) {
+  //   case s1:
+  //      mem_op(..., s1);
+  //      goto merge_bb;
+  //   case s2:
+  //      mem_op(..., s2);
+  //      goto merge_bb;
+  //   ...
+  //   default:
+  //      mem_op(..., size);
+  //      goto merge_bb;
+  // }
+  // merge_bb:
+
+  BasicBlock *BB = MI.getParent();
+  DEBUG(dbgs() << "\n\n== Basic Block Before ==\n");
+  DEBUG(dbgs() << *BB << "\n");
+
+  BasicBlock *DefaultBB = SplitBlock(BB, &MI);
+  BasicBlock::iterator It(MI);
+  ++It;
+  assert(It != DefaultBB->end());
+  BasicBlock *MergeBB = SplitBlock(DefaultBB, &(*It));
+  DefaultBB->setName("MemOP.Default");
+  MergeBB->setName("MemOP.Merge");
+
+  auto &Ctx = F.getContext();
+  IRBuilder<> IRB(BB);
+  BB->getTerminator()->eraseFromParent();
+  Value *SizeVar = MI.getLength();
+  SwitchInst *SI = IRB.CreateSwitch(SizeVar, DefaultBB, SizeIds.size());
+
+  DEBUG(dbgs() << "\n\n== Basic Block After==\n");
+
+  for (uint64_t SizeId: SizeIds) {
+    ConstantInt *CaseSizeId = ConstantInt::get(Type::getInt64Ty(Ctx), SizeId);
+    BasicBlock *CaseBB = BasicBlock::Create(Ctx,
+                                            Twine("MemOP.Case.") + Twine(SizeId),
+                                            &F, DefaultBB);
+    Instruction *NewInst = MI.clone();
+    // Fix the argument.
+    dyn_cast<MemIntrinsic>(NewInst)->setLength(CaseSizeId);
+    CaseBB->getInstList().push_back(NewInst);
+    IRBuilder<> IRBCase(CaseBB);
+    IRBCase.CreateBr(MergeBB);
+    SI->addCase(CaseSizeId, CaseBB);
+    DEBUG(dbgs() << *CaseBB << "\n");
+  }
+  setProfMetadata(F.getParent(), SI, CaseCounts, MaxCount);
+
+  DEBUG(dbgs() << *BB << "\n");
+  DEBUG(dbgs() << *DefaultBB << "\n");
+  DEBUG(dbgs() << *MergeBB << "\n");
+
+  emitOptimizationRemark(
+      F.getContext(), "memop-opt", F, MI.getDebugLoc(),
+      Twine("optimize ") + getMIName(MI) + " with count " +
+      Twine(SumForSmallSizes) + " out of " + Twine(Sum) + " for " +
+      Twine(Version) + " versions");
+
+}
+
 void MemIntrinsicVisitor::visitMemIntrinsic(MemIntrinsic &MI) {
   if (!PGOInstrMemOP)
     return;
   Value *Length = MI.getLength();
   // Not instrument constant length calls.
   if (dyn_cast<ConstantInt>(Length))
     return;
 
   NMemIs++;
   switch (Mode) {
   case VM_counting:
     return;
   case VM_instrument:
     instrumentOneMemIntrinsic(MI);
     return;
+  case VM_optimize:
+    Candidates.push_back(&MI);
+    return;
   default:
     break;
   }
   llvm_unreachable("Unknown visiting mode");
 }
 
 // Traverse all the indirect callsites and annotate the instructions.
 void PGOUseFunc::annotateIndirectCallSites() {
void PGOUseFunc::annotateIndirectCallSites() {
   for (auto &I : IndirectCallSites) {
     DEBUG(dbgs() << "Read one indirect call instrumentation: Index="
                  << IndirectCallSiteIndex << " out of " << NumValueSites
                  << "\n");
     annotateValueSite(*M, *I, ProfileRecord, IPVK_IndirectCallTarget,
                       IndirectCallSiteIndex, MaxNumAnnotations);
     IndirectCallSiteIndex++;
   }
+
+  // Now optimize memory intrinsic calls.
+  FuncInfo.MIVisitor.optimizeMemIntrinsics(F, this);
 }
 } // end anonymous namespace
 
 // Create a COMDAT variable INSTR_PROF_RAW_VERSION_VAR to make the runtime
 // aware this is an ir_level profile so it can set the version flag.
 static void createIRLevelProfileFlagVariable(Module &M) {
   Type *IntTy64 = Type::getInt64Ty(M.getContext());
   uint64_t ProfileVersion = (INSTR_PROF_RAW_VERSION | VARIANT_MASK_IR_PROF);