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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Mircea Trofin (mtrofin)

<details>
<summary>Changes</summary>

The metadata information for the hoisted conditional branch should be copied from the original branch, not from the current terminator of the block it's hoisted to. Note that debug info was previously copied by `ReplaceInstWithInst` as described in the previous statement. `MD_prof` was dropped outright.

(This was identified through `profcheck` (see Issue #<!-- -->147390), and this PR addresses most of the test failures (when running under profcheck) under `Transforms/LICM`.)

---
Full diff: https://github.com/llvm/llvm-project/pull/152232.diff


3 Files Affected:

- (modified) llvm/lib/Analysis/BranchProbabilityInfo.cpp (+7) 
- (modified) llvm/lib/Transforms/Scalar/LICM.cpp (+14-3) 
- (added) llvm/test/Transforms/LICM/hoist-phi-metadata.ll (+46) 


``````````diff
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}

``````````

</details>


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