[LoopInfo] Preserve profile information in makeLoopInvariant

[boomanaiden154]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Aiden Grossman (boomanaiden154)

Changes

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.

---

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

2 Files Affected:

• (modified) llvm/lib/Analysis/LoopInfo.cpp (+4-2)
• (modified) llvm/test/Transforms/LoopDeletion/invalidate-scev-after-hoisting.ll (+11-4)

diff --git a/llvm/lib/Analysis/LoopInfo.cpp b/llvm/lib/Analysis/LoopInfo.cpp
index d84721b7f8f4b..ef12fe50a8b43 100644
--- a/llvm/lib/Analysis/LoopInfo.cpp
+++ b/llvm/lib/Analysis/LoopInfo.cpp
@@ -112,8 +112,10 @@ 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.
+  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}
+;.

[nikic]

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.

I don't think this is generally true, because the instruction might be guarded by control flow. Just to give the most obvious example, if you have:

loop {
  if (c) {
    c ? x : a;
  }
}

Then at the current position, the branch probability of the select is 1, while outside the loop it is unknown.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[boomanaiden154]

I don't think this is generally true, because the instruction might be guarded by control flow. Just to give the most obvious example, if you have:

Good point. I've updated the patch to account for that. Now we only propagate profile information if the instruction is in either the header or in the latch (and there is a single latch), effectively excluding any control flow in the loop body. When tested compiling LLVM, this gives works pretty well (~85% of the time), so I think it's worth the slight additional complexity.

[nikic]

This is right for the header, but I'm not sure I understand why it's right for a unique latch.

A variation on my previous example would be:

loop {
  if (c) {
    break;
  }
  c ? x : a;
}

Here the select is in a unique latch block, but the condition is correlated with an early exit.

[boomanaiden154]

In that case it would not be. I for some reason never thought of the case with an early exit and was thinking a unique latch would be the only way out of the loop. That is obviously wrong in hindsight. Patch updated.