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

[Mircea Trofin via llvm-commits]

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

>From 96a8b815664b729501305c273af35f9c42864570 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     | 61 +++++++++++++++++++
 3 files changed, 82 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..e56dd42d2970a 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."));
 
+// FIXME: 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..97704b97ed24d
--- /dev/null
+++ b/llvm/test/Transforms/LICM/hoist-phi-metadata.ll
@@ -0,0 +1,61 @@
+; 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:.*]], !dbg [[ORIG_DBG:![0-9]+]], !prof [[ORIG_PROF:![0-9]+]]
+entry:
+  br label %loop, !dbg !5
+
+; 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
+
+; CHECK-LABEL: loop:
+loop:
+  %cmp1 = icmp sgt i32 %x, 0
+; CHECK: br i1 %cmp1, label %if, label %then, !dbg [[ORIG_DBG]], !prof [[ORIG_PROF]]
+  br i1 %cmp1, label %if, label %then, !dbg !6, !prof !8
+if:
+  %add = add i32 %x, 1
+  br label %then
+
+; CHECK-LABEL: then:
+then:
+  %phi = phi i32 [ %add, %if ], [ %x, %loop ]
+  store i32 %phi, ptr %p
+  %cmp2 = icmp ne i32 %phi, 0
+; CHECK-NEXT: br i1 %cmp2, label %loop, label %end, !dbg [[OTHER_DBG:![0-9]+]], !prof [[OTHER_PROF:![0-9]+]]
+  br i1 %cmp2, label %loop, label %end, !dbg !7, !prof !9
+
+; CHECK-LABEL: end:
+end:
+  ret void
+}
+
+; CHECK-DAG: [[ORIG_DBG]] = !DILocation(line: 2, column: 22
+; CHECK-DAG: [[ORIG_PROF]] = !{!"branch_weights", i32 5, i32 7}
+; CHECK-DAG: [[OTHER_DBG]] = !DILocation(line: 3, column: 22
+; CHECK-DAG: [[OTHER_PROF]] = !{!"branch_weights", i32 13, i32 11}
+
+
+!llvm.module.flags = !{!2, !3}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1)
+!1 = !DIFile(filename: "t", directory: "/")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = distinct !DISubprogram(name: "triangle_phi", linkageName: "triangle_phi", scope: !1, file: !1, line: 1, unit: !0)
+!5 = !DILocation(line: 1, column: 22, scope: !4)
+!6 = !DILocation(line: 2, column: 22, scope: !4)
+!7 = !DILocation(line: 3, column: 22, scope: !4)
+!8 = !{!"branch_weights", i32 5, i32 7}
+!9 = !{!"branch_weights", i32 13, i32 11}