[InstCombine] Fix !prof metadata preservation for invokes

[rnk]

Summary

Bug noticed by inspection.

Extend the test to handle invokes as well as calls, and rewrite it to
not depend on the inliner and other passes.

Also simplify the call site replacement code with CallSite, similar to
what I did to dead arg elimination and arg promotion (rL300235 and
rL300229).

[davidxl]

good cleanup, lgtm

[yrouban]

llvm/trunk/lib/Transforms/InstCombine/InstCombineCalls.cpp

4142 ↗

(On Diff #95195)

Reid, any reason to make "branch_weights" here from "VP"?
Do you mind if I preserve either "branch_weights" or "VP" by exact copying?

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

assert((ArgAttrs.size() == FT->getNumParams() || FT->isVarArg()) &&
          "missing argument attributes");
   LLVMContext &Ctx = Callee->getContext();
   AttributeList NewCallerPAL = AttributeList::get(
       Ctx, FnAttrs, AttributeSet::get(Ctx, RAttrs), ArgAttrs);
 
   SmallVector<OperandBundleDef, 1> OpBundles;
   CS.getOperandBundlesAsDefs(OpBundles);
 
-  Instruction *NC;
+  CallSite NewCS;
   if (InvokeInst *II = dyn_cast<InvokeInst>(Caller)) {
-    NC = Builder->CreateInvoke(Callee, II->getNormalDest(), II->getUnwindDest(),
-                               Args, OpBundles);
-    NC->takeName(II);
-    cast<InvokeInst>(NC)->setCallingConv(II->getCallingConv());
-    cast<InvokeInst>(NC)->setAttributes(NewCallerPAL);
+    NewCS = Builder->CreateInvoke(Callee, II->getNormalDest(),
+                                  II->getUnwindDest(), Args, OpBundles);
   } else {
-    CallInst *CI = cast<CallInst>(Caller);
-    NC = Builder->CreateCall(Callee, Args, OpBundles);
-    NC->takeName(CI);
-    // Preserve the weight metadata for the new call instruction. The metadata
-    // is used by SamplePGO to check callsite's hotness.
-    uint64_t W;
-    if (CI->extractProfTotalWeight(W))
-      NC->setProfWeight(W);
-
-    cast<CallInst>(NC)->setTailCallKind(CI->getTailCallKind());
-    cast<CallInst>(NC)->setCallingConv(CI->getCallingConv());
-    cast<CallInst>(NC)->setAttributes(NewCallerPAL);
-  }
+    NewCS = Builder->CreateCall(Callee, Args, OpBundles);
+    cast<CallInst>(NewCS.getInstruction())
+        ->setTailCallKind(cast<CallInst>(Caller)->getTailCallKind());
+  }
+  NewCS->takeName(Caller);
+  NewCS.setCallingConv(CS.getCallingConv());
+  NewCS.setAttributes(NewCallerPAL);
+
+  if (Caller->extractProfTotalWeight(W))
+    NewCS->setProfWeight(W);
 
   // Insert a cast of the return type as necessary.
+  Instruction *NC = NewCS.getInstruction();
   Value *NV = NC;
   if (OldRetTy != NV->getType() && !Caller->use_empty()) {
     if (!NV->getType()->isVoidTy()) {
       NV = NC = CastInst::CreateBitOrPointerCast(NC, OldRetTy);
       NC->setDebugLoc(Caller->getDebugLoc());
 
       // If this is an invoke instruction, we should insert it after the first
       // non-phi, instruction in the normal successor block.

llvm/test/Transforms/InstCombine/cast-call-combine-prof.ll

-; RUN: opt -instcombine -inline -S -inline-threshold=0 -hot-callsite-threshold=100 < %s | FileCheck %s
-; Checks if VP profile is used for hotness checks in inlining after instcombine
-; converted the call to a direct call.
-
-declare void @bar(i16 *)
-
-define void @foo(i16* %a) {
-  call void @bar(i16* %a)
-  call void @bar(i16* %a)
+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+; Check that instcombine preserves !prof metadata when removing function
+; prototype casts.
+
+declare i32 @__gxx_personality_v0(...)
+declare void @__cxa_call_unexpected(i8*)
+declare void @foo(i16* %a)
+
+; CHECK-LABEL: @test_call()
+; CHECK: call void @foo(i16* null), !prof ![[PROF:[0-9]+]]
+define void @test_call() {
+  call void bitcast (void (i16*)* @foo to void (i8*)*) (i8* null), !prof !0
   ret void
 }
 
-; CHECK-LABEL: @test()
-; CHECK-NEXT: call void @bar
-; CHECK-NEXT: call void @bar
-define void @test() {
-  call void bitcast (void (i16*)* @foo to void (i8*)*) (i8* null), !prof !0
+; CHECK-LABEL: @test_invoke()
+; CHECK: invoke void @foo(i16* null)
+; CHECK-NEXT: to label %done unwind label %lpad, !prof ![[PROF]]
+define void @test_invoke() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+  invoke void bitcast (void (i16*)* @foo to void (i8*)*) (i8* null)
+          to label %done unwind label %lpad, !prof !0
+
+done:
   ret void
+
+lpad:
+  %lp = landingpad { i8*, i32 }
+          filter [0 x i8*] zeroinitializer
+  %ehptr = extractvalue { i8*, i32 } %lp, 0
+  tail call void @__cxa_call_unexpected(i8* %ehptr) noreturn nounwind
+  unreachable
 }
 
+; CHECK: ![[PROF]] = !{!"branch_weights", i32 2000}
 !0 = !{!"VP", i32 0, i64 2000, i64 -3913987384944532146, i64 2000}
 
 !llvm.module.flags = !{!1}
 
 !1 = !{i32 1, !"ProfileSummary", !2}
 !2 = !{!3, !4, !5, !6, !7, !8, !9, !10}
 !3 = !{!"ProfileFormat", !"InstrProf"}
 !4 = !{!"TotalCount", i64 10000}