Add FuncPass Implementation for non-trivial LoopUnswitching

[drcut]

Summary

Implement non-trivial LoopUnswitching with FunctionPass. Since non-trivial LoopUnswitching may rely on Function Analysis, it's better be implemented based on FuncPass.

[sameerds]

It seems the loop pass in the new PM already has an option to perform non-trivial unswitching, the purpose of this pass is not clear. Doesn't the core function unswitchLoop() use the same analyses in both cases? Does the function pass use an analysis that is not in the LoopStandardAnalysisResults?

llvm/lib/Passes/PassRegistry.def

362

Is there a better way to name this option? The function pass feels like an implementation detail. Is there something else that conveys the significance?

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

3219–3226

The body of this pass closely resembles the loop pass, but the state is too mixed up to refactor. But at least this part with the string literals deserves to be commoned out.

[aeubanks]

this was done for https://discourse.llvm.org/t/rfc-re-implement-non-trivial-simpleloopunswitch-with-functionpass/62937, which we're no longer pursuing
we should probably abandon this patch

[drcut]

This patch implements SimpleLoopUnswitch using FuncPass, which can be used to replace the LoopPass implementation currently used in the optimization pipeline.
This patch has been verified and can be re-opened once we find using FuncPass for SLU is necessary/beneficial.

llvm/include/llvm/Transforms/Scalar/SimpleLoopUnswitch.h

public:
 
   PreservedAnalyses run(Loop &L, LoopAnalysisManager &AM,
                         LoopStandardAnalysisResults &AR, LPMUpdater &U);
 
   void printPipeline(raw_ostream &OS,
                      function_ref<StringRef(StringRef)> MapClassName2PassName);
 };
 
+// (TODO): use this FunctionPass for all non-trivial SLU
+class FuncSimpleLoopUnswitchPass
+    : public PassInfoMixin<FuncSimpleLoopUnswitchPass> {
+public:
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+};
+
 /// Create the legacy pass object for the simple loop unswitcher.
 ///
 /// See the documentaion for `SimpleLoopUnswitchPass` for details.
 Pass *createSimpleLoopUnswitchLegacyPass(bool NonTrivial = false);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_SIMPLELOOPUNSWITCH_H

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("print-memderefs", MemDerefPrinterPass(dbgs()))
 FUNCTION_PASS("reassociate", ReassociatePass())
 FUNCTION_PASS("redundant-dbg-inst-elim", RedundantDbgInstEliminationPass())
 FUNCTION_PASS("reg2mem", RegToMemPass())
 FUNCTION_PASS("scalarize-masked-mem-intrin", ScalarizeMaskedMemIntrinPass())
 FUNCTION_PASS("scalarizer", ScalarizerPass())
 FUNCTION_PASS("separate-const-offset-from-gep", SeparateConstOffsetFromGEPPass())
 FUNCTION_PASS("sccp", SCCPPass())
+FUNCTION_PASS("simple-loop-unswitch-func", FuncSimpleLoopUnswitchPass())
 FUNCTION_PASS("sink", SinkingPass())
 FUNCTION_PASS("slp-vectorizer", SLPVectorizerPass())
 FUNCTION_PASS("slsr", StraightLineStrengthReducePass())
 FUNCTION_PASS("speculative-execution", SpeculativeExecutionPass())
 FUNCTION_PASS("sroa", SROAPass())
 FUNCTION_PASS("strip-gc-relocates", StripGCRelocates())
 FUNCTION_PASS("structurizecfg", StructurizeCFGPass())
 FUNCTION_PASS("tailcallelim", TailCallElimPass())

llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp

PreservedAnalyses SimpleLoopUnswitchPass::run(Loop &L, LoopAnalysisManager &AM,
   assert(AR.DT.verify(DominatorTree::VerificationLevel::Fast));
 
   auto PA = getLoopPassPreservedAnalyses();
   if (AR.MSSA)
     PA.preserve<MemorySSAAnalysis>();
   return PA;
 }
 
+PreservedAnalyses FuncSimpleLoopUnswitchPass::run(Function &F,
+                                                  FunctionAnalysisManager &AM) {
+  LoopInfo &LI = AM.getResult<LoopAnalysis>(F);
+  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
+  auto &SE = AM.getResult<ScalarEvolutionAnalysis>(F);
+  auto MSSA = &AM.getResult<MemorySSAAnalysis>(F).getMSSA();
+  Optional<MemorySSAUpdater> MSSAU;
+  if (MSSA) {
+    MSSAU = MemorySSAUpdater(MSSA);
+    if (VerifyMemorySSA)
+      MSSA->verifyMemorySSA();
+  }
+  bool Changed = false;
+  for (const auto &L : LI) {
+    Changed |= simplifyLoop(L, &DT, &LI, &SE, nullptr, MSSAU.getPointer(),
+                            /*PreserveLCSSA=*/false);
+    Changed |= formLCSSARecursively(*L, DT, &LI, &SE);
+  }
+
+  SmallPriorityWorklist<Loop *, 4> Worklist;
+  appendLoopsToWorklist(LI, Worklist);
+  while (!Worklist.empty()) {
+    Loop &L = *Worklist.pop_back_val();
+    LLVM_DEBUG(dbgs() << "Unswitching loop in " << F.getName() << ": " << L
+                      << "\n");
+
+    // Save the current loop name in a variable so that we can report it even
+    // after it has been deleted.
+    std::string LoopName = std::string(L.getName());
+
+    auto &LAM = AM.getResult<LoopAnalysisManagerFunctionProxy>(F).getManager();
+
+    auto UnswitchCB = [&L, &LAM, &LoopName, &Worklist](
+                          bool CurrentLoopValid, bool PartiallyInvariant,
+                          ArrayRef<Loop *> NewLoops) {
+      // If we did a non-trivial unswitch, we have added new (cloned) loops.
+      if (!NewLoops.empty())
+        appendLoopsToWorklist(NewLoops, Worklist);
+
+      // If the current loop remains valid, we should revisit it to catch any
+      // other unswitch opportunities. Otherwise, we need to mark it as deleted.
+      if (CurrentLoopValid) {
+        if (PartiallyInvariant) {
+          // Mark the new loop as partially unswitched, to avoid unswitching on
+          // the same condition again.
+          auto &Context = L.getHeader()->getContext();
+          MDNode *DisableUnswitchMD = MDNode::get(
+              Context,
+              MDString::get(Context, "llvm.loop.unswitch.partial.disable"));
+          MDNode *NewLoopID = makePostTransformationMetadata(
+              Context, L.getLoopID(), {"llvm.loop.unswitch.partial"},
+              {DisableUnswitchMD});
+          L.setLoopID(NewLoopID);
+        } else
+          Worklist.insert(&L);
+      } else
+        LAM.clear(L, LoopName);
+    };
+
+    auto DestroyLoopCB = [&LAM](Loop &L, StringRef Name) {
+      LAM.clear(L, Name);
+    };
+    Changed |= unswitchLoop(
+        L, AM.getResult<DominatorTreeAnalysis>(F),
+        AM.getResult<LoopAnalysis>(F), AM.getResult<AssumptionAnalysis>(F),
+        AM.getResult<AAManager>(F), AM.getResult<TargetIRAnalysis>(F),
+        /*Trivial=*/false, /*NonTrivial=*/true,
+        UnswitchCB, &AM.getResult<ScalarEvolutionAnalysis>(F),
+        MSSAU.hasValue() ? MSSAU.getPointer() : nullptr, DestroyLoopCB);
+  }
+  if (!Changed)
+    return PreservedAnalyses::all();
+
+  if (MSSA && VerifyMemorySSA)
+    MSSA->verifyMemorySSA();
+
+  // Historically this pass has had issues with the dominator tree so verify it
+  // in asserts builds.
+  assert(AM.getResult<DominatorTreeAnalysis>(F).verify(
+      DominatorTree::VerificationLevel::Fast));
+
+  auto PA = getLoopPassPreservedAnalyses();
+  if (MSSA)
+    PA.preserve<MemorySSAAnalysis>();
+  return PA;
+}
+
 void SimpleLoopUnswitchPass::printPipeline(
     raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
   static_cast<PassInfoMixin<SimpleLoopUnswitchPass> *>(this)->printPipeline(
       OS, MapClassName2PassName);
 
   OS << "<";
   OS << (NonTrivial ? "" : "no-") << "nontrivial;";
   OS << (Trivial ? "" : "no-") << "trivial";

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-cost.ll

 ; Specifically exercise the cost modeling for non-trivial loop unswitching.
 ;
 ; RUN: opt -passes='loop(simple-loop-unswitch<nontrivial>),verify<loops>' -unswitch-threshold=5 -S < %s | FileCheck %s
 ; RUN: opt -passes='loop-mssa(simple-loop-unswitch<nontrivial>),verify<loops>' -unswitch-threshold=5 -S < %s | FileCheck %s
 ; RUN: opt -simple-loop-unswitch -enable-nontrivial-unswitch -unswitch-threshold=5 -verify-memoryssa -S < %s | FileCheck %s
+; RUN: opt -passes='simple-loop-unswitch-func,verify<loops>' -unswitch-threshold=5 -S < %s | FileCheck %s
 
 declare void @a()
 declare void @b()
 declare void @x()
 
 ; First establish enough code size in the duplicated 'loop_begin' block to
 ; suppress unswitching.
 define void @test_no_unswitch(i1* %ptr, i1 %cond) {

llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch-freeze.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -freeze-loop-unswitch-cond -passes='loop(simple-loop-unswitch<nontrivial>),verify<loops>' -S < %s | FileCheck %s
 ; RUN: opt -freeze-loop-unswitch-cond -passes='loop-mssa(simple-loop-unswitch<nontrivial>),verify<loops>' -S < %s | FileCheck %s
 ; RUN: opt -freeze-loop-unswitch-cond -simple-loop-unswitch -enable-nontrivial-unswitch -verify-memoryssa -S < %s | FileCheck %s
+; RUN: opt -freeze-loop-unswitch-cond -passes='simple-loop-unswitch-func,verify<loops>' -S < %s | FileCheck %s
 
 declare i32 @a()
 declare i32 @b()
 declare i32 @c()
 declare i32 @d()
 
 declare void @sink1(i32)
 declare void @sink2(i32)