[licm] clone metadata when hoisting conditional branch

Author: mtrofin

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)

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() &&

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}