From d1b19da854fde9077a53400fabff486dc6dba022 Mon Sep 17 00:00:00 2001
From: Alina Sbirlea <asbirlea@google.com>
Date: Mon, 24 Oct 2022 15:23:16 -0700
Subject: [PATCH] [LoopPeeling] Add flag to disable support for peeling loops
 with non-latch exits

Add a flag to allow disabling the changes in
https://reviews.llvm.org/D134803.

Differential Revision: https://reviews.llvm.org/D136643
---
 llvm/lib/Transforms/Utils/LoopPeel.cpp             | 22 +++++++++++++++++++++-
 .../Transforms/LoopUnroll/peel-branch-weights.ll   | 18 +++++++++++++++++-
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/LoopPeel.cpp b/llvm/lib/Transforms/Utils/LoopPeel.cpp
index d4f14df..b7e3eb2 100644
--- a/llvm/lib/Transforms/Utils/LoopPeel.cpp
+++ b/llvm/lib/Transforms/Utils/LoopPeel.cpp
@@ -72,12 +72,32 @@ static cl::opt<unsigned> UnrollForcePeelCount(
     "unroll-force-peel-count", cl::init(0), cl::Hidden,
     cl::desc("Force a peel count regardless of profiling information."));
 
+static cl::opt<bool> DisableAdvancedPeeling(
+    "disable-advanced-peeling", cl::init(false), cl::Hidden,
+    cl::desc(
+        "Disable advance peeling. Issues for convergent targets (D134803)."));
+
 static const char *PeeledCountMetaData = "llvm.loop.peeled.count";
 
 // Check whether we are capable of peeling this loop.
 bool llvm::canPeel(Loop *L) {
   // Make sure the loop is in simplified form
-  return L->isLoopSimplifyForm();
+  if (!L->isLoopSimplifyForm())
+    return false;
+  if (!DisableAdvancedPeeling)
+    return true;
+
+  SmallVector<BasicBlock *, 4> Exits;
+  L->getUniqueNonLatchExitBlocks(Exits);
+  // The latch must either be the only exiting block or all non-latch exit
+  // blocks have either a deopt or unreachable terminator or compose a chain of
+  // blocks where the last one is either deopt or unreachable terminated. Both
+  // deopt and unreachable terminators are a strong indication they are not
+  // taken. Note that this is a profitability check, not a legality check. Also
+  // note that LoopPeeling currently can only update the branch weights of latch
+  // blocks and branch weights to blocks with deopt or unreachable do not need
+  // updating.
+  return llvm::all_of(Exits, IsBlockFollowedByDeoptOrUnreachable);
 }
 
 // This function calculates the number of iterations after which the given Phi
diff --git a/llvm/test/Transforms/LoopUnroll/peel-branch-weights.ll b/llvm/test/Transforms/LoopUnroll/peel-branch-weights.ll
index ddbb7dd..fade7a0 100644
--- a/llvm/test/Transforms/LoopUnroll/peel-branch-weights.ll
+++ b/llvm/test/Transforms/LoopUnroll/peel-branch-weights.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
 ; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=2 2>&1 | FileCheck %s
+; RUN: opt < %s -S -loop-unroll -unroll-force-peel-count=2 -disable-advanced-peeling 2>&1 | FileCheck %s --check-prefix=DISABLEADV
 
 declare i32 @get.x()
 
@@ -50,7 +51,22 @@ define void @test() {
 ; CHECK-NEXT:    br label [[LOOP_EXIT]]
 ; CHECK:       loop.exit:
 ; CHECK-NEXT:    ret void
-;
+
+; DISABLEADV-LABEL: @test()
+; DISABLEADV-NEXT: entry:
+; DISABLEADV-NEXT:  br label %loop
+; DISABLEADV: loop
+; DISABLEADV-NEXT:  %x = call i32 @get.x()
+; DISABLEADV-NEXT:  switch i32 %x, label %loop.latch [
+; DISABLEADV-NEXT:    i32 0, label %loop.latch
+; DISABLEADV-NEXT:    i32 1, label %loop.exit
+; DISABLEADV-NEXT:    i32 2, label %loop.exit
+; DISABLEADV-NEXT:  ], !prof !0
+; DISABLEADV: loop.latch:
+; DISABLEADV-NEXT:  br label %loop
+; DISABLEADV: loop.exit:
+; DISABLEADV-NEXT:  ret void
+
 entry:
   br label %loop
 
-- 
2.7.4