[PGO] Attach appropriate funclet operand bundles to value profiling instrumentation calls

[chris.chrulski]

Summary

When generating value profiling instrumentation, ensure the call gets the
correct funclet token, otherwise WinEHPrepare will turn the call (and all
subsequent instructions) into unreachable.

[rnk]

Code makes sense and tests look good, but I have a suggestion for how to make it simpler.

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

875–879

The underlying invariant is that non-intrinsic calls will carry funclet bundles. Intrinsic calls do not. Rather than special casing memop size instrumentation, I think it would be better to:

• always calculate colors
• below, use the funclet bundle directly when instrumenting non-intrinsic calls
• use the color map when instrumenting intrinsic calls (memops)

This logic will be less fragile in the face of future value profiling kinds.

[davidxl]

Can you extract Window's specific code into its own helper function if possible?

llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp

627

Add some comment here about the change.

[chris.chrulski]

Thanks. Good suggestions for making this more robust.

[chris.chrulski]

Moved logic into support routine, and made it less specific to value profiling of memintrinsics calls.

[davidxl]

ok with me if Reid is ok with windows specific logic.

[rnk]

lgtm

llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp

void InstrProfiling::lowerValueProfileInst(InstrProfValueProfileInst *Ind) {
   for (uint32_t Kind = IPVK_First; Kind < ValueKind; ++Kind)
     Index += It->second.NumValueSites[Kind];
 
   IRBuilder<> Builder(Ind);
   bool IsRange = (Ind->getValueKind()->getZExtValue() ==
                   llvm::InstrProfValueKind::IPVK_MemOPSize);
   CallInst *Call = nullptr;
   auto *TLI = &GetTLI(*Ind->getFunction());
+
+  // To support value profiling calls within Windows exception handlers, funclet
+  // information contained within operand bundles needs to be copied over to
+  // the library call. This is required for the IR to be processed by the
+  // WinEHPrepare pass.
+  SmallVector<OperandBundleDef, 1> OpBundles;
+  Ind->getOperandBundlesAsDefs(OpBundles);
   if (!IsRange) {
     Value *Args[3] = {Ind->getTargetValue(),
                       Builder.CreateBitCast(DataVar, Builder.getInt8PtrTy()),
                       Builder.getInt32(Index)};
-    Call = Builder.CreateCall(getOrInsertValueProfilingCall(*M, *TLI), Args);
+    Call = Builder.CreateCall(getOrInsertValueProfilingCall(*M, *TLI), Args,
+                              OpBundles);
   } else {
     Value *Args[6] = {
         Ind->getTargetValue(),
         Builder.CreateBitCast(DataVar, Builder.getInt8PtrTy()),
         Builder.getInt32(Index),
         Builder.getInt64(MemOPSizeRangeStart),
         Builder.getInt64(MemOPSizeRangeLast),
         Builder.getInt64(MemOPSizeLarge == 0 ? INT64_MIN : MemOPSizeLarge)};
-    Call =
-        Builder.CreateCall(getOrInsertValueProfilingCall(*M, *TLI, true), Args);
+    Call = Builder.CreateCall(getOrInsertValueProfilingCall(*M, *TLI, true),
+                              Args, OpBundles);
   }
   if (auto AK = TLI->getExtAttrForI32Param(false))
     Call->addParamAttr(2, AK);
   Ind->replaceAllUsesWith(Call);
   Ind->eraseFromParent();
 }
 
 void InstrProfiling::lowerIncrement(InstrProfIncrementInst *Inc) {

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Instrumentation/PGOInstrumentation.h"
 #include "CFGMST.h"
 #include "ValueProfileCollector.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/EHPersonalities.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/CallSite.h"
 #include "llvm/IR/Comdat.h"
BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   // Second one: Add new edge of InstrBB->DestBB.
   Edge &NewEdge1 = MST.addEdge(InstrBB, DestBB, 0);
   NewEdge1.InMST = true;
   E->Removed = true;
 
   return canInstrument(InstrBB);
 }
 
+// When generating value profiling calls on Windows routines that make use of
+// handler funclets for exception processing an operand bundle needs to attached
+// to the called function. This routine will set \p OpBundles to contain the
+// funclet information, if any is needed, that should be placed on the generated
+// value profiling call for the value profile candidate call.
+static void
+populateEHOperandBundle(VPCandidateInfo &Cand,
+                        DenseMap<BasicBlock *, ColorVector> &BlockColors,
+                        SmallVectorImpl<OperandBundleDef> &OpBundles) {
+  auto *OrigCall = dyn_cast<CallBase>(Cand.AnnotatedInst);
+  if (OrigCall && !isa<IntrinsicInst>(OrigCall)) {
+    // The instrumentation call should belong to the same funclet as a
+    // non-intrinsic call, so just copy the operand bundle, if any exists.
+    Optional<OperandBundleUse> ParentFunclet =
+        OrigCall->getOperandBundle(LLVMContext::OB_funclet);
+    if (ParentFunclet)
+      OpBundles.emplace_back(OperandBundleDef(*ParentFunclet));
+  } else {
+    // Intrinsics or other instructions do not get funclet information from the
+    // front-end. Need to use the BlockColors that was computed by the routine
+    // colorEHFunclets to determine whether a funclet is needed.
+    if (!BlockColors.empty()) {
+      const ColorVector &CV = BlockColors.find(OrigCall->getParent())->second;
+      assert(CV.size() == 1 && "non-unique color for block!");
+      Instruction *EHPad = CV.front()->getFirstNonPHI();
+      if (EHPad->isEHPad())
+        OpBundles.emplace_back("funclet", EHPad);
+    }
+  }
+}
+
 // Visit all edge and instrument the edges not in MST, and do value profiling.
 // Critical edges will be split.
 static void instrumentOneFunc(
     Function &F, Module *M, BranchProbabilityInfo *BPI, BlockFrequencyInfo *BFI,
     std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
     bool IsCS) {
   // Split indirectbr critical edges here before computing the MST rather than
   // later in getInstrBB() to avoid invalidating it.
FuncInfo.SIVisitor.instrumentSelects(F, &I, NumCounters, FuncInfo.FuncNameVar,
                                        FuncInfo.FunctionHash);
   assert(I == NumCounters);
 
   if (DisableValueProfiling)
     return;
 
   NumOfPGOICall += FuncInfo.ValueSites[IPVK_IndirectCallTarget].size();
 
+  // Intrinsic function calls do not have funclet operand bundles needed for
+  // Windows exception handling attached to them. However, if value profiling is
+  // inserted for one of these calls, then a funclet value will need to be set
+  // on the instrumentation call based on the funclet coloring.
+  DenseMap<BasicBlock *, ColorVector> BlockColors;
+  if (F.hasPersonalityFn() &&
+      isFuncletEHPersonality(classifyEHPersonality(F.getPersonalityFn())))
+    BlockColors = colorEHFunclets(F);
+
   // For each VP Kind, walk the VP candidates and instrument each one.
   for (uint32_t Kind = IPVK_First; Kind <= IPVK_Last; ++Kind) {
     unsigned SiteIndex = 0;
     if (Kind == IPVK_MemOPSize && !PGOInstrMemOP)
       continue;
 
     for (VPCandidateInfo Cand : FuncInfo.ValueSites[Kind]) {
       LLVM_DEBUG(dbgs() << "Instrument one VP " << ValueProfKindDescr[Kind]
                         << " site: CallSite Index = " << SiteIndex << "\n");
 
       IRBuilder<> Builder(Cand.InsertPt);
       assert(Builder.GetInsertPoint() != Cand.InsertPt->getParent()->end() &&
              "Cannot get the Instrumentation point");
 
       Value *ToProfile = nullptr;
       if (Cand.V->getType()->isIntegerTy())
         ToProfile = Builder.CreateZExtOrTrunc(Cand.V, Builder.getInt64Ty());
       else if (Cand.V->getType()->isPointerTy())
         ToProfile = Builder.CreatePtrToInt(Cand.V, Builder.getInt64Ty());
       assert(ToProfile && "value profiling Value is of unexpected type");
 
+      SmallVector<OperandBundleDef, 1> OpBundles;
+      populateEHOperandBundle(Cand, BlockColors, OpBundles);
       Builder.CreateCall(
           Intrinsic::getDeclaration(M, Intrinsic::instrprof_value_profile),
           {ConstantExpr::getBitCast(FuncInfo.FuncNameVar, I8PtrTy),
            Builder.getInt64(FuncInfo.FunctionHash), ToProfile,
-           Builder.getInt32(Kind), Builder.getInt32(SiteIndex++)});
+           Builder.getInt32(Kind), Builder.getInt32(SiteIndex++)},
+          OpBundles);
     }
   } // IPVK_First <= Kind <= IPVK_Last
 }
 
 namespace {
 
 // This class represents a CFG edge in profile use compilation.
 struct PGOUseEdge : public PGOEdge {

llvm/test/Transforms/PGOProfile/indirect_call_profile_funclet.ll

+; RUN: opt < %s -pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
+; RUN: opt < %s -pgo-instr-gen -instrprof -S | FileCheck %s --check-prefix=LOWER
+
+; RUN: opt < %s -passes=pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
+; RUN: opt < %s -passes=pgo-instr-gen,instrprof -S | FileCheck %s --check-prefix=LOWER
+
+; This test is to verify that PGO runtime library calls get created with the
+; appropriate operand bundle funclet information when an indirect call
+; was marked with as belonging to a particular funclet.
+
+; Test case based on this source:
+;  extern void may_throw(int);
+;
+;  class base {
+;  public:
+;    base() : x(0) {};
+;    int get_x() const { return x; }
+;    virtual void update() { x++; }
+;    int x;
+;  };
+;
+;  class derived : public base {
+;  public:
+;    derived() {}
+;    virtual void update() { x--; }
+;  };
+;
+;  void run(base* b, int count) {
+;    try {
+;      may_throw(count);
+;    }
+;    catch (...) {
+;      // Virtual function call in exception handler for value profiling.
+;      b->update();
+;    }
+;  }
+
+%class.base = type { i32 (...)**, i32 }
+define dso_local void @"?run@@YAXPEAVbase@@H@Z"(%class.base* %b, i32 %count) personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+entry:
+  invoke void @"?may_throw@@YAXH@Z"(i32 %count)
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %entry
+  %tmp = catchswitch within none [label %catch] unwind to caller
+
+catch:                                            ; preds = %catch.dispatch
+  %tmp1 = catchpad within %tmp [i8* null, i32 64, i8* null]
+  %tmp2 = bitcast %class.base* %b to void (%class.base*)***
+  %vtable = load void (%class.base*)**, void (%class.base*)*** %tmp2, align 8
+  %tmp3 = load void (%class.base*)*, void (%class.base*)** %vtable, align 8
+  call void %tmp3(%class.base* %b) [ "funclet"(token %tmp1) ]
+  catchret from %tmp1 to label %try.cont
+
+try.cont:                                         ; preds = %catch, %entry
+  ret void
+}
+
+; GEN: catch:
+; GEN: call void @llvm.instrprof.value.profile(
+; GEN-SAME: [ "funclet"(token %tmp1) ]
+
+; LOWER: catch:
+; LOWER: call void @__llvm_profile_instrument_target(
+; LOWER-SAME: [ "funclet"(token %tmp1) ]
+
+declare dso_local void @"?may_throw@@YAXH@Z"(i32)
+declare dso_local i32 @__CxxFrameHandler3(...)

llvm/test/Transforms/PGOProfile/memop_profile_funclet.ll

+; RUN: opt < %s -pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
+; RUN: opt < %s -pgo-instr-gen -instrprof -S | FileCheck %s --check-prefix=LOWER
+
+; RUN: opt < %s -passes=pgo-instr-gen -S | FileCheck %s --check-prefix=GEN
+; RUN: opt < %s -passes=pgo-instr-gen,instrprof -S | FileCheck %s --check-prefix=LOWER
+
+; This test is to verify that PGO runtime library calls get created with the
+; appropriate operand bundle funclet information when a memory intrinsic
+; being value profiled is called within an exception handler.
+
+; Test case based on this source:
+;  #include <memory.h>
+;
+;  extern void may_throw(int);
+;
+;  #define MSG "0123456789012345\0"
+;  unsigned len = 16;
+;  char msg[200];
+;
+;  void run(int count) {
+;    try {
+;      may_throw(count);
+;    }
+;    catch (...) {
+;      memcpy(msg, MSG, len);
+;      throw;
+;    }
+;  }
+
+%eh.ThrowInfo = type { i32, i32, i32, i32 }
+
+$"??_C@_0BC@CABPINND@Exception?5caught?$AA?$AA@" = comdat any
+
+@"?len@@3IA" = dso_local global i32 16, align 4
+@"?msg@@3PADA" = dso_local global [200 x i8] zeroinitializer, align 16
+@"??_C@_0BC@CABPINND@Exception?5caught?$AA?$AA@" = linkonce_odr dso_local unnamed_addr constant [18 x i8] c"0123456789012345\00\00", comdat, align 1
+
+define dso_local void @"?run@@YAXH@Z"(i32 %count) personality i8* bitcast (i32 (...)* @__CxxFrameHandler3 to i8*) {
+entry:
+  invoke void @"?may_throw@@YAXH@Z"(i32 %count)
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %entry
+  %tmp = catchswitch within none [label %catch] unwind to caller
+
+catch:                                            ; preds = %catch.dispatch
+  %tmp1 = catchpad within %tmp [i8* null, i32 64, i8* null]
+  %tmp2 = load i32, i32* @"?len@@3IA", align 4
+  %conv = zext i32 %tmp2 to i64
+  call void @llvm.memcpy.p0i8.p0i8.i64(
+    i8* getelementptr inbounds ([200 x i8], [200 x i8]* @"?msg@@3PADA", i64 0, i64 0),
+    i8* getelementptr inbounds ([18 x i8], [18 x i8]* @"??_C@_0BC@CABPINND@Exception?5caught?$AA?$AA@", i64 0, i64 0),
+    i64 %conv, i1 false)
+  call void @_CxxThrowException(i8* null, %eh.ThrowInfo* null) #3 [ "funclet"(token %tmp1) ]
+  unreachable
+
+try.cont:                                         ; preds = %entry
+  ret void
+}
+
+; GEN: catch:
+; GEN: call void @llvm.instrprof.value.profile(
+; GEN-SAME: [ "funclet"(token %tmp1) ]
+
+; LOWER: catch:
+; LOWER: call void @__llvm_profile_instrument_range(
+; LOWER-SAME: [ "funclet"(token %tmp1) ]
+
+declare dso_local void @"?may_throw@@YAXH@Z"(i32)
+declare dso_local i32 @__CxxFrameHandler3(...)
+
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8*, i8*, i64, i1)
+declare dso_local void @_CxxThrowException(i8*, %eh.ThrowInfo*)