[llvm] [licm] clone metadata when hoisting conditional branch (PR #152232)

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin created https://github.com/llvm/llvm-project/pull/152232

None

>From 667dcb07cab616e52e9ecda5b048fc15dabf556d Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 5 Aug 2025 16:24:12 -0700
Subject: [PATCH] [licm] clone metadata when hoisting conditional branch

---
 llvm/lib/Analysis/BranchProbabilityInfo.cpp   |  7 +++
 llvm/lib/Transforms/Scalar/LICM.cpp           | 17 +++++--
 .../Transforms/LICM/hoist-phi-metadata.ll     | 46 +++++++++++++++++++
 3 files changed, 67 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Transforms/LICM/hoist-phi-metadata.ll

diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index 671e023415a6e..1467bcd61ec8b 100644
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -59,6 +59,13 @@ static cl::opt<std::string> PrintBranchProbFuncName(
     cl::desc("The option to specify the name of the function "
              "whose branch probability info is printed."));
 
+// Flag used for an ablation performance test, Issue #147390. Placing it here
+// because referencing Analysis should be feasible from anywhere. Will be
+// removed after the ablation test.
+cl::opt<bool> DisableProfilingInfoCorrectPropagation(
+    "profcheck-disable-fixes", cl::Hidden, cl::init(false),
+    cl::desc("Temporary flag, will be used for an ablation test"));
+
 INITIALIZE_PASS_BEGIN(BranchProbabilityInfoWrapperPass, "branch-prob",
                       "Branch Probability Analysis", false, true)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index c3f80f901a120..3a17e34573cc4 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -218,6 +218,8 @@ using PointersAndHasReadsOutsideSet =
 static SmallVector<PointersAndHasReadsOutsideSet, 0>
 collectPromotionCandidates(MemorySSA *MSSA, AliasAnalysis *AA, Loop *L);
 
+extern cl::opt<bool> DisableProfilingInfoCorrectPropagation;
+
 namespace {
 struct LoopInvariantCodeMotion {
   bool runOnLoop(Loop *L, AAResults *AA, LoopInfo *LI, DominatorTree *DT,
@@ -857,9 +859,18 @@ class ControlFlowHoister {
     }
 
     // Now finally clone BI.
-    ReplaceInstWithInst(
-        HoistTarget->getTerminator(),
-        BranchInst::Create(HoistTrueDest, HoistFalseDest, BI->getCondition()));
+    auto *NewBI =
+        BranchInst::Create(HoistTrueDest, HoistFalseDest, BI->getCondition());
+    ReplaceInstWithInst(HoistTarget->getTerminator(), NewBI);
+    // Copy all the metadata. In particular:
+    // - debug info (critical to Sample-based profiling) should be the same as
+    // the original branch, not that of HoistTarget->getTerminator(), which is
+    // what ReplaceInstWithInst would use.
+    // - md_prof should also come from the original branch - since the condition
+    // was hoisted, the branch probabilities shouldn't change.
+    if (!DisableProfilingInfoCorrectPropagation)
+      NewBI->copyMetadata(*BI);
+
     ++NumClonedBranches;
 
     assert(CurLoop->getLoopPreheader() &&
diff --git a/llvm/test/Transforms/LICM/hoist-phi-metadata.ll b/llvm/test/Transforms/LICM/hoist-phi-metadata.ll
new file mode 100644
index 0000000000000..031a00d11bbf3
--- /dev/null
+++ b/llvm/test/Transforms/LICM/hoist-phi-metadata.ll
@@ -0,0 +1,46 @@
+; Test that hoising conditional branches copies over profiling metadata
+; RUN: opt -S -passes=licm -licm-control-flow-hoisting=1 %s -o - | FileCheck %s
+
+; CHECK-LABEL: @triangle_phi
+define void @triangle_phi(i32 %x, ptr %p) {
+; CHECK-LABEL: entry:
+; CHECK: %cmp1 = icmp sgt i32 %x, 0
+; CHECK: br i1 %cmp1, label %[[IF_LICM:.*]], label %[[THEN_LICM:.*]], !prof !0
+entry:
+  br label %loop
+
+; CHECK: [[IF_LICM]]:
+; CHECK: %add = add i32 %x, 1
+; CHECK: br label %[[THEN_LICM]]
+
+; CHECK: [[THEN_LICM]]:
+; CHECK: phi i32 [ %add, %[[IF_LICM]] ], [ %x, %entry ]
+; CHECK: store i32 %phi, ptr %p
+; CHECK: %cmp2 = icmp ne i32 %phi, 0
+; CHECK: br label %loop
+
+loop:
+  %cmp1 = icmp sgt i32 %x, 0
+  br i1 %cmp1, label %if, label %then, !prof !0
+
+if:
+  %add = add i32 %x, 1
+  br label %then
+
+; CHECK-LABEL: then:
+; CHECK: br i1 %cmp2, label %loop, label %end, !prof !1
+then:
+  %phi = phi i32 [ %add, %if ], [ %x, %loop ]
+  store i32 %phi, ptr %p
+  %cmp2 = icmp ne i32 %phi, 0
+  br i1 %cmp2, label %loop, label %end, !prof !1
+
+; CHECK-LABEL: end:
+end:
+  ret void
+}
+
+; CHECK: !0 = !{!"branch_weights", i32 5, i32 7}
+; CHECK: !1 = !{!"branch_weights", i32 13, i32 11}
+!0 = !{!"branch_weights", i32 5, i32 7}
+!1 = !{!"branch_weights", i32 13, i32 11}