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

Patch by Chris Chrulski

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.

Differential Revision: https://reviews.llvm.org/D73221

lib/Transforms/Instrumentation/InstrProfiling.cpp

@@ -638,11 +638,19 @@ void InstrProfiling::lowerValueProfileInst(InstrProfValueProfileInst *Ind) {
                   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(),
@@ -651,8 +659,8 @@ void InstrProfiling::lowerValueProfileInst(InstrProfValueProfileInst *Ind) {
         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);

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -52,6 +52,7 @@
 #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"
@@ -63,6 +64,7 @@
 #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"
@@ -799,6 +801,37 @@ BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   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(
@@ -839,6 +872,15 @@ static void instrumentOneFunc(
 
   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;
@@ -860,11 +902,14 @@ static void instrumentOneFunc(
         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
 }

test/Transforms/PGOProfile/indirect_call_profile_funclet.ll

@@ -0,0 +1,68 @@
+; 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(...)

test/Transforms/PGOProfile/memop_profile_funclet.ll

@@ -0,0 +1,73 @@
+; 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*)