[llvm] [LoopInfo] Preserve profile information in makeLoopInvariant (PR #174171)

[Aiden Grossman via llvm-commits]

https://github.com/boomanaiden154 updated https://github.com/llvm/llvm-project/pull/174171

>From 2cee2e3482f272910b0c1e1643a21f7603588321 Mon Sep 17 00:00:00 2001
From: Aiden Grossman <aidengrossman at google.com>
Date: Fri, 2 Jan 2026 01:15:29 +0000
Subject: [PATCH] [LoopInfo] Preserve profile information in makeLoopInvariant

When hoisting loop invariant instructions, we can preserve profile
metadata because it depends solely on the condition (which is loop
invariant) rather than where we are in the control flow graph.
---
 llvm/lib/Analysis/LoopInfo.cpp                    | 13 +++++++++++--
 .../invalidate-scev-after-hoisting.ll             | 15 +++++++++++----
 llvm/utils/profcheck-xfail.txt                    |  1 -
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Analysis/LoopInfo.cpp b/llvm/lib/Analysis/LoopInfo.cpp
index d84721b7f8f4b..a780f3231a546 100644
--- a/llvm/lib/Analysis/LoopInfo.cpp
+++ b/llvm/lib/Analysis/LoopInfo.cpp
@@ -54,6 +54,10 @@ static cl::opt<bool, true>
     VerifyLoopInfoX("verify-loop-info", cl::location(VerifyLoopInfo),
                     cl::Hidden, cl::desc("Verify loop info (time consuming)"));
 
+namespace llvm {
+extern cl::opt<bool> ProfcheckDisableMetadataFixes;
+} // end namespace llvm
+
 //===----------------------------------------------------------------------===//
 // Loop implementation
 //
@@ -112,8 +116,13 @@ bool Loop::makeLoopInvariant(Instruction *I, bool &Changed,
   // There is possibility of hoisting this instruction above some arbitrary
   // condition. Any metadata defined on it can be control dependent on this
   // condition. Conservatively strip it here so that we don't give any wrong
-  // information to the optimizer.
-  I->dropUnknownNonDebugMetadata();
+  // information to the optimizer. We preserve profile metadata as instructions
+  // that can take profile metadata (select and br) with loop invariant
+  // conditions will maintain their weights.
+  if (ProfcheckDisableMetadataFixes)
+    I->dropUnknownNonDebugMetadata();
+  else
+    I->dropUnknownNonDebugMetadata({LLVMContext::MD_prof});
 
   if (SE)
     SE->forgetBlockAndLoopDispositions(I);
diff --git a/llvm/test/Transforms/LoopDeletion/invalidate-scev-after-hoisting.ll b/llvm/test/Transforms/LoopDeletion/invalidate-scev-after-hoisting.ll
index bdd51c2b6bc53..79c3773ebb686 100644
--- a/llvm/test/Transforms/LoopDeletion/invalidate-scev-after-hoisting.ll
+++ b/llvm/test/Transforms/LoopDeletion/invalidate-scev-after-hoisting.ll
@@ -1,4 +1,4 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
 ; RUN: opt -passes='loop(indvars,loop-deletion),verify<scalar-evolution>,print<scalar-evolution>' -S %s 2>&1| FileCheck %s
 
 ; Make sure the SCEV for %invar is invalidated properly when the instruction is
@@ -123,13 +123,13 @@ outer.latch:
   br label %outer.header
 }
 
-define void @test_pr58314() {
+define void @test_pr58314() !prof !0 {
 ; CHECK-LABEL: @test_pr58314(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[OUTER_HEADER:%.*]]
 ; CHECK:       outer.header:
 ; CHECK-NEXT:    [[C:%.*]] = icmp ne i16 0, 0
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[C]], i1 false, i1 true
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[C]], i1 false, i1 true, !prof [[PROF1:![0-9]+]]
 ; CHECK-NEXT:    br label [[INNER:%.*]]
 ; CHECK:       inner:
 ; CHECK-NEXT:    br i1 true, label [[INNER]], label [[OUTER_LATCH:%.*]]
@@ -147,7 +147,7 @@ outer.header:
 
 inner:
   %c = icmp ne i16 0, 0
-  %sel = select i1 %c, i1 false, i1 true
+  %sel = select i1 %c, i1 false, i1 true, !prof !1
   br i1 true, label %inner, label %outer.latch
 
 outer.latch:
@@ -156,3 +156,10 @@ outer.latch:
 exit:
   ret void
 }
+
+!0 = !{!"function_entry_count", i64 1000}
+!1 = !{!"branch_weights", i32 4, i32 1}
+;.
+; CHECK: [[META0:![0-9]+]] = !{!"function_entry_count", i64 1000}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 4, i32 1}
+;.
diff --git a/llvm/utils/profcheck-xfail.txt b/llvm/utils/profcheck-xfail.txt
index beffde14d5f8b..56612973f9574 100644
--- a/llvm/utils/profcheck-xfail.txt
+++ b/llvm/utils/profcheck-xfail.txt
@@ -454,7 +454,6 @@ Transforms/IROutliner/region-inputs-in-phi-nodes.ll
 Transforms/LoopBoundSplit/bug51866.ll
 Transforms/LoopBoundSplit/bug-loop-bound-split-phi-in-exit-block.ll
 Transforms/LoopBoundSplit/loop-bound-split.ll
-Transforms/LoopDeletion/invalidate-scev-after-hoisting.ll
 Transforms/LoopIdiom/AArch64/byte-compare-index.ll
 Transforms/LoopIdiom/AArch64/find-first-byte.ll
 Transforms/LoopIdiom/RISCV/byte-compare-index.ll