Set the prof weight correctly for call instructions in DeadArgumentElimination.

Summary: In DeadArgumentElimination, the call instructions will be replaced. We also need to set the prof weights so that function inlining can find the correct profile.

Reviewers: eraman

Reviewed By: eraman

Subscribers: llvm-commits

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

llvm-svn: 298660

4 files changed, 40 insertions, 0 deletions

diff --git a/llvm/include/llvm/IR/Instruction.h b/llvm/include/llvm/IR/Instruction.h
index 9dbe2ec0e90..c1be9da498d 100644
--- a/llvm/include/llvm/IR/Instruction.h
+++ b/llvm/include/llvm/IR/Instruction.h
@@ -255,6 +255,9 @@ public:
   /// Updates branch_weights metadata by scaling it by \p S / \p T.
   void updateProfWeight(uint64_t S, uint64_t T);
 
+  /// Sets the branch_weights metadata to \p W for CallInst.
+  void setProfWeight(uint64_t W);
+
   /// Set the debug location information for this instruction.
   void setDebugLoc(DebugLoc Loc) { DbgLoc = std::move(Loc); }
 
diff --git a/llvm/lib/IR/Instruction.cpp b/llvm/lib/IR/Instruction.cpp
index fc453d4f8d7..fc8a0566ab5 100644
--- a/llvm/lib/IR/Instruction.cpp
+++ b/llvm/lib/IR/Instruction.cpp
@@ -652,3 +652,12 @@ void Instruction::updateProfWeight(uint64_t S, uint64_t T) {
   MDBuilder MDB(getContext());
   setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
 }
+
+void Instruction::setProfWeight(uint64_t W) {
+  assert((isa<CallInst>(this) || isa<InvokeInst>(this)) &&
+         "Can only set weights for call and invoke instrucitons");
+  SmallVector<uint32_t, 1> Weights;
+  Weights.push_back(W);
+  MDBuilder MDB(getContext());
+  setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
+}
diff --git a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
index 1f2216d37a1..fe79efce290 100644
--- a/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
+++ b/llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp
@@ -194,6 +194,9 @@ bool DeadArgumentEliminationPass::DeleteDeadVarargs(Function &Fn) {
           cast<CallInst>(Call)->getTailCallKind());
     }
     New->setDebugLoc(Call->getDebugLoc());
+    uint64_t W;
+    if (Call->extractProfTotalWeight(W))
+      New->setProfWeight(W);
 
     Args.clear();
 
@@ -901,6 +904,9 @@ bool DeadArgumentEliminationPass::RemoveDeadStuffFromFunction(Function *F) {
           cast<CallInst>(Call)->getTailCallKind());
     }
     New->setDebugLoc(Call->getDebugLoc());
+    uint64_t W;
+    if (Call->extractProfTotalWeight(W))
+      New->setProfWeight(W);
 
     Args.clear();
 
diff --git a/llvm/test/Transforms/DeadArgElim/call_profile.ll b/llvm/test/Transforms/DeadArgElim/call_profile.ll
new file mode 100644
index 00000000000..6acb6f000db
--- /dev/null
+++ b/llvm/test/Transforms/DeadArgElim/call_profile.ll
@@ -0,0 +1,22 @@
+; RUN: opt -deadargelim -S < %s | FileCheck %s
+
+; Checks if !prof metadata is corret in deadargelim.
+
+define void @caller() #0 {
+; CHECK: call void @test_vararg(), !prof ![[PROF:[0-9]]]
+; CHECK: call void @test(), !prof ![[PROF]]
+  call void (i32, ...) @test_vararg(i32 1), !prof !0
+  call void @test(i32 1), !prof !0
+  ret void
+}
+
+define internal void @test_vararg(i32, ...) #1 {
+  ret void
+}
+
+define internal void @test(i32 %a) #1 {
+  ret void
+}
+
+; CHECK:![[PROF]] = !{!"branch_weights", i32 30}
+!0 = !{!"branch_weights", i32 30}