From 6bfcac612fa6d2093fc4bcde8257700f2594c0f1 Mon Sep 17 00:00:00 2001
From: Max Kazantsev <mkazantsev@azul.com>
Date: Wed, 12 Oct 2022 13:17:28 +0700
Subject: [PATCH] [SimpleLoopUnswitch][NFC] Separate legality checks from cost
 computation

These are semantically two different stages, but were entwined in the
old implementation. Now cost computation does not do legality checks,
and they all are done beforehead.
---
 .../Transforms/Scalar/SimpleLoopUnswitch.cpp  | 55 ++++++++++---------
 1 file changed, 30 insertions(+), 25 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index fc3f2769d01de..fa570addeb628 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -2827,8 +2827,23 @@ struct NonTrivialUnswitchCandidate {
 };
 } // end anonymous namespace.
 
-static Optional<NonTrivialUnswitchCandidate>
-findBestNonTrivialUnswitchCandidate(
+static bool isSafeToClone(const Loop &L) {
+  if (!L.isSafeToClone())
+    return false;
+  for (auto *BB : L.blocks())
+    for (auto &I : *BB) {
+      if (I.getType()->isTokenTy() && I.isUsedOutsideOfBlock(BB))
+        return false;
+      if (auto *CB = dyn_cast<CallBase>(&I)) {
+        assert(!CB->cannotDuplicate() && "Checked by L.isSafeToClone().");
+        if (CB->isConvergent())
+          return false;
+      }
+    }
+  return true;
+}
+
+static NonTrivialUnswitchCandidate findBestNonTrivialUnswitchCandidate(
     ArrayRef<std::pair<Instruction *, TinyPtrVector<Value *> > >
         UnswitchCandidates, const Loop &L, const DominatorTree &DT,
     const LoopInfo &LI, AssumptionCache &AC, const TargetTransformInfo &TTI,
@@ -2857,13 +2872,6 @@ findBestNonTrivialUnswitchCandidate(
     for (auto &I : *BB) {
       if (EphValues.count(&I))
         continue;
-
-      if (I.getType()->isTokenTy() && I.isUsedOutsideOfBlock(BB))
-        return None;
-      if (auto *CB = dyn_cast<CallBase>(&I))
-        if (CB->isConvergent() || CB->cannotDuplicate())
-          return None;
-
       Cost += TTI.getInstructionCost(&I, CostKind);
     }
     assert(Cost >= 0 && "Must not have negative costs!");
@@ -3030,31 +3038,28 @@ static bool unswitchBestCondition(
       dbgs() << "Considering " << UnswitchCandidates.size()
              << " non-trivial loop invariant conditions for unswitching.\n");
 
-  Optional<NonTrivialUnswitchCandidate> Best =
-      findBestNonTrivialUnswitchCandidate(UnswitchCandidates, L, DT, LI, AC,
-                                          TTI, PartialIVInfo);
-  if (!Best)
-    return false;
+  NonTrivialUnswitchCandidate Best = findBestNonTrivialUnswitchCandidate(
+      UnswitchCandidates, L, DT, LI, AC, TTI, PartialIVInfo);
 
-  assert(Best->TI && "Failed to find loop unswitch candidate");
+  assert(Best.TI && "Failed to find loop unswitch candidate");
 
-  if (Best->Cost >= UnswitchThreshold) {
-    LLVM_DEBUG(dbgs() << "Cannot unswitch, lowest cost found: " << Best->Cost
+  if (Best.Cost >= UnswitchThreshold) {
+    LLVM_DEBUG(dbgs() << "Cannot unswitch, lowest cost found: " << Best.Cost
                       << "\n");
     return false;
   }
 
-  if (Best->TI != PartialIVCondBranch)
+  if (Best.TI != PartialIVCondBranch)
     PartialIVInfo.InstToDuplicate.clear();
 
   // If the best candidate is a guard, turn it into a branch.
-  if (isGuard(Best->TI))
-    Best->TI = turnGuardIntoBranch(cast<IntrinsicInst>(Best->TI), L, ExitBlocks,
-                                   DT, LI, MSSAU);
+  if (isGuard(Best.TI))
+    Best.TI = turnGuardIntoBranch(cast<IntrinsicInst>(Best.TI), L, ExitBlocks,
+                                  DT, LI, MSSAU);
 
-  LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = " << Best->Cost
-                    << ") terminator: " << *Best->TI << "\n");
-  unswitchNontrivialInvariants(L, *Best->TI, Best->Invariants, ExitBlocks,
+  LLVM_DEBUG(dbgs() << "  Unswitching non-trivial (cost = " << Best.Cost
+                    << ") terminator: " << *Best.TI << "\n");
+  unswitchNontrivialInvariants(L, *Best.TI, Best.Invariants, ExitBlocks,
                                PartialIVInfo, DT, LI, AC, UnswitchCB, SE, MSSAU,
                                DestroyLoopCB);
   return true;
@@ -3133,7 +3138,7 @@ unswitchLoop(Loop &L, DominatorTree &DT, LoopInfo &LI, AssumptionCache &AC,
   }
 
   // Skip non-trivial unswitching for loops that cannot be cloned.
-  if (!L.isSafeToClone())
+  if (!isSafeToClone(L))
     return false;
 
   // For non-trivial unswitching, because it often creates new loops, we rely on