[LoopLoadElim] Form simplify form by demand

[mkazantsev]

Summary

The LoopLoadElim pass effectively didn't work in the new PM,
because unlike the old pass manager, it never bothered to simplify loops
before it starts working. As result, at some point it tried to get a loop
preheader, and it might not exist, which lead to crash (fixed as D77775).
Old pass manager, on other hand, was requesting simplification of all
loops before the run, potentially doing redundant work.

This patch unifies operation in both pass managers. Now neither of them
requests simplification of all loops before the run. Instead, it simplifies
loop by itself when it thinks it's necessary.

It also enables all tests run in both pipelines and makes sure they
produce equivalent results.

[fedor.sergeev]

Looks good to me, though I would like somebody else to check if running loopSimplify instead of simply bailing out is appropriate here.

[mkazantsev]

I'm not sure either. I can make 1st patch that just fixes the crash with bailout and then rebase this one on top of it.

[mkazantsev]

I made conservative fix D77775. Let's merge it and then consider this one as improvement.

[mkazantsev]

Rebased.

[mkazantsev]

Ping?

llvm/lib/Transforms/Scalar/LoopLoadElimination.cpp

 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/LazyBlockFrequencyInfo.h"
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/MemorySSA.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpander.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Value.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Utils.h"
+#include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopVersioning.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
 #include <algorithm>
 #include <cassert>
 #include <forward_list>
 #include <set>
 #include <tuple>
 #include <utility>

 }
 
 namespace {
 
 /// The per-loop class that does most of the work.
 class LoadEliminationForLoop {
 public:
   LoadEliminationForLoop(Loop *L, LoopInfo *LI, const LoopAccessInfo &LAI,
-                         DominatorTree *DT, BlockFrequencyInfo *BFI,
-                         ProfileSummaryInfo* PSI)
-      : L(L), LI(LI), LAI(LAI), DT(DT), BFI(BFI), PSI(PSI), PSE(LAI.getPSE()) {}
+                         DominatorTree *DT, AssumptionCache *AC,
+                         MemorySSA *MSSA, BlockFrequencyInfo *BFI,
+                         ProfileSummaryInfo *PSI)
+      : L(L), LI(LI), LAI(LAI), DT(DT), AC(AC), BFI(BFI), PSI(PSI),
+        PSE(LAI.getPSE()) {
+    if (MSSA)
+      MSSAU = std::make_unique<MemorySSAUpdater>(MSSA);
+  }
 
   /// Look through the loop-carried and loop-independent dependences in
   /// this loop and find store->load dependences.
   ///
   /// Note that no candidate is returned if LAA has failed to analyze the loop
   /// (e.g. if it's not bottom-tested, contains volatile memops, etc.)
   std::forward_list<StoreToLoadForwardingCandidate>
   findStoreToLoadDependences(const LoopAccessInfo &LAI) {
bool processLoop() {
     }
 
     if (LAI.getPSE().getUnionPredicate().getComplexity() >
         LoadElimSCEVCheckThreshold) {
       LLVM_DEBUG(dbgs() << "Too many SCEV run-time checks needed.\n");
       return false;
     }
 
-    if (!L->isLoopSimplifyForm()) {
-      LLVM_DEBUG(dbgs() << "Loop is not is loop-simplify form");
-      return false;
-    }
-
     if (!Checks.empty() || !LAI.getPSE().getUnionPredicate().isAlwaysTrue()) {
       if (LAI.hasConvergentOp()) {
         LLVM_DEBUG(dbgs() << "Versioning is needed but not allowed with "
                              "convergent calls\n");
         return false;
       }
 
       auto *HeaderBB = L->getHeader();
       auto *F = HeaderBB->getParent();
       bool OptForSize = F->hasOptSize() ||
                         llvm::shouldOptimizeForSize(HeaderBB, PSI, BFI,
                                                     PGSOQueryType::IRPass);
       if (OptForSize) {
         LLVM_DEBUG(
             dbgs() << "Versioning is needed but not allowed when optimizing "
                       "for size.\n");
         return false;
       }
 
       // Point of no-return, start the transformation.  First, version the loop
       // if necessary.
+      // TODO: Is simplification really necessary here?
+      if (!L->isLoopSimplifyForm())
+        simplifyLoop(L, DT, LI, PSE.getSE(), AC, MSSAU.get(),
+                     /*PreserveLCSSA*/ false);
 
       LoopVersioning LV(LAI, L, LI, DT, PSE.getSE(), false);
       LV.setAliasChecks(std::move(Checks));
       LV.setSCEVChecks(LAI.getPSE().getUnionPredicate());
       LV.versionLoop();
     }
 
+    // Make sure the loop is in simplify form. We'll need preheader further.
+    if (!L->isLoopSimplifyForm())
+      simplifyLoop(L, DT, LI, PSE.getSE(), AC, MSSAU.get(),
+                   /*PreserveLCSSA*/ false);
+
     // Next, propagate the value stored by the store to the users of the load.
     // Also for the first iteration, generate the initial value of the load.
     SCEVExpander SEE(*PSE.getSE(), L->getHeader()->getModule()->getDataLayout(),
                      "storeforward");
     for (const auto &Cand : Candidates)
       propagateStoredValueToLoadUsers(Cand, SEE);
     NumLoopLoadEliminted += NumForwarding;
 
     return true;
   }
 
 private:
   Loop *L;
 
   /// Maps the load/store instructions to their index according to
   /// program order.
   DenseMap<Instruction *, unsigned> InstOrder;
 
   // Analyses used.
   LoopInfo *LI;
   const LoopAccessInfo &LAI;
   DominatorTree *DT;
+  AssumptionCache *AC;
   BlockFrequencyInfo *BFI;
   ProfileSummaryInfo *PSI;
   PredicatedScalarEvolution PSE;
+  std::unique_ptr<MemorySSAUpdater> MSSAU;
 };
 
 } // end anonymous namespace
 
 static bool
 eliminateLoadsAcrossLoops(Function &F, LoopInfo &LI, DominatorTree &DT,
-                          BlockFrequencyInfo *BFI, ProfileSummaryInfo *PSI,
+                          ScalarEvolution &SE, AssumptionCache *AC,
+                          MemorySSA *MSSA, BlockFrequencyInfo *BFI,
+                          ProfileSummaryInfo *PSI,
                           function_ref<const LoopAccessInfo &(Loop &)> GetLAI) {
   // Build up a worklist of inner-loops to transform to avoid iterator
   // invalidation.
   // FIXME: This logic comes from other passes that actually change the loop
   // nest structure. It isn't clear this is necessary (or useful) for a pass
   // which merely optimizes the use of loads in a loop.
   SmallVector<Loop *, 8> Worklist;
 
   for (Loop *TopLevelLoop : LI)
     for (Loop *L : depth_first(TopLevelLoop))
       // We only handle inner-most loops.
       if (L->empty())
         Worklist.push_back(L);
 
   // Now walk the identified inner loops.
   bool Changed = false;
   for (Loop *L : Worklist) {
     // The actual work is performed by LoadEliminationForLoop.
-    LoadEliminationForLoop LEL(L, &LI, GetLAI(*L), &DT, BFI, PSI);
+    LoadEliminationForLoop LEL(L, &LI, GetLAI(*L), &DT, AC, MSSA, BFI, PSI);
     Changed |= LEL.processLoop();
   }
   return Changed;
 }
 
 namespace {
 
 /// The pass.  Most of the work is delegated to the per-loop
bool runOnFunction(Function &F) override {
 
     auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     auto &LAA = getAnalysis<LoopAccessLegacyAnalysis>();
     auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     auto *PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
     auto *BFI = (PSI && PSI->hasProfileSummary()) ?
                 &getAnalysis<LazyBlockFrequencyInfoPass>().getBFI() :
                 nullptr;
+    auto &SE = getAnalysis<ScalarEvolutionWrapperPass>().getSE();
 
     // Process each loop nest in the function.
-    return eliminateLoadsAcrossLoops(
-        F, LI, DT, BFI, PSI,
-        [&LAA](Loop &L) -> const LoopAccessInfo & { return LAA.getInfo(&L); });
+    return eliminateLoadsAcrossLoops(F, LI, DT, SE, nullptr, nullptr, BFI, PSI,
+                                     [&LAA](Loop & L)->const LoopAccessInfo &
+    { return LAA.getInfo(&L); });
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequiredID(LoopSimplifyID);
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addPreserved<LoopInfoWrapperPass>();
     AU.addRequired<LoopAccessLegacyAnalysis>();
     AU.addRequired<ScalarEvolutionWrapperPass>();

 
 static const char LLE_name[] = "Loop Load Elimination";
 
 INITIALIZE_PASS_BEGIN(LoopLoadElimination, LLE_OPTION, LLE_name, false, false)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LoopAccessLegacyAnalysis)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
-INITIALIZE_PASS_DEPENDENCY(LoopSimplify)
 INITIALIZE_PASS_DEPENDENCY(ProfileSummaryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LazyBlockFrequencyInfoPass)
 INITIALIZE_PASS_END(LoopLoadElimination, LLE_OPTION, LLE_name, false, false)
 
 FunctionPass *llvm::createLoopLoadEliminationPass() {
   return new LoopLoadElimination();
 }
 
PreservedAnalyses LoopLoadEliminationPass::run(Function &F,
   auto *PSI = MAM.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
   auto *BFI = (PSI && PSI->hasProfileSummary()) ?
       &AM.getResult<BlockFrequencyAnalysis>(F) : nullptr;
   MemorySSA *MSSA = EnableMSSALoopDependency
                         ? &AM.getResult<MemorySSAAnalysis>(F).getMSSA()
                         : nullptr;
 
   auto &LAM = AM.getResult<LoopAnalysisManagerFunctionProxy>(F).getManager();
-  bool Changed = eliminateLoadsAcrossLoops(
-      F, LI, DT, BFI, PSI, [&](Loop &L) -> const LoopAccessInfo & {
+  bool Changed =
+      eliminateLoadsAcrossLoops(F, LI, DT, SE, &AC, MSSA, BFI, PSI,
+                                [&](Loop & L)->const LoopAccessInfo & {
         LoopStandardAnalysisResults AR = {AA, AC, DT, LI, SE, TLI, TTI, MSSA};
         return LAM.getResult<LoopAccessAnalysis>(L, AR);
       });
 
   if (!Changed)
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   return PA;
 }

llvm/test/Transforms/LoopLoadElim/cond-load.ll

 ; RUN: opt -S -loop-load-elim < %s | FileCheck %s
+; RUN: opt -S -passes=loop-load-elim < %s | FileCheck %s
 
 ; We can't hoist conditional loads to the preheader for the initial value.
 ; E.g. in the loop below we'd access array[-1] if we did:
 ;
 ;   for(int i = 0 ; i < n ; i++ )
 ;     array[i] = ( i > 0 ? array[i - 1] : 0 ) + 4;
 
 target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"

llvm/test/Transforms/LoopLoadElim/def-store-before-load.ll

 ; RUN: opt -loop-load-elim -S < %s | FileCheck %s
+; RUN: opt -passes=loop-load-elim -S < %s | FileCheck %s
 
 ;  No loop-carried forwarding: The intervening store to A[i] kills the stored
 ;  value from the previous iteration.
 ;
 ;   for (unsigned i = 0; i < 100; i++) {
 ;     A[i] = 1;
 ;     A[i+1] = A[i] + B[i];
 ;   }

llvm/test/Transforms/LoopLoadElim/loop-simplify-dep.ll

 ; RUN: opt -loop-load-elim -S < %s | FileCheck %s
+; RUN: opt -passes=loop-load-elim -S < %s | FileCheck %s
 
 ; Make sure we create a preheader if we dont' have one.
 
 target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 
 define void @f(i32* noalias nocapture %A, i32* noalias nocapture readonly %B, i64 %N, i1 %C) {
 entry:
   br i1 %C, label %for.body, label %for.end

llvm/test/Transforms/LoopLoadElim/multiple-stores-same-block.ll

 ; RUN: opt -basicaa -loop-load-elim -S < %s | FileCheck %s
-
+; RUN: opt -aa-pipeline=basic-aa -passes='require<aa>,loop-load-elim' -S < %s | FileCheck %s
 ; In this case the later store forward to the load:
 ;
 ;   for (unsigned i = 0; i < 100; i++) {
 ;     B[i] = A[i] + 1;
 ;     A[i+1] = C[i] + 2;
 ;     A[i+1] = D[i] + 3;
 ;   }
 

llvm/test/Transforms/LoopLoadElim/new-pm-preheader-problem.ll

+; RUN: opt -loop-load-elim -S < %s | FileCheck %s
 ; RUN: opt -passes=loop-load-elim -S < %s | FileCheck %s
 ; Make sure it doesn't crash in new pass manager due to missing preheader.
 target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 
 define void @test(i32* noalias nocapture %A, i32* noalias nocapture readonly %B, i64 %N, i1 %C) {
 entry:
   br i1 %C, label %for.body, label %for.end
 
 ; CHECK: test
+; CHECK: for.body.preheader:
+; CHECK-NEXT: %load_initial = load i32, i32* %A
+; CHECK-NEXT: br label %for.body
+
+; CHECK: for.body:
 for.body:
+; CHECK-NEXT: %store_forwarded = phi i32 [ %load_initial, %for.body.preheader ], [ %add, %for.body ]
   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
   %arrayidx = getelementptr inbounds i32, i32* %A, i64 %indvars.iv
   %load = load i32, i32* %arrayidx, align 4
   %arrayidx2 = getelementptr inbounds i32, i32* %B, i64 %indvars.iv
   %load_1 = load i32, i32* %arrayidx2, align 4
+; CHECK: %add = add i32 %load_1, %store_forwarded
   %add = add i32 %load_1, %load
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %arrayidx_next = getelementptr inbounds i32, i32* %A, i64 %indvars.iv.next
   store i32 %add, i32* %arrayidx_next, align 4
   %exitcond = icmp eq i64 %indvars.iv.next, %N
   br i1 %exitcond, label %for.end, label %for.body
 
 for.end:
   ret void
 }

llvm/test/Transforms/LoopLoadElim/non-consecutive.ll

 ; RUN: opt -loop-load-elim -S < %s | FileCheck %s
+; RUN: opt -passes=loop-load-elim -S < %s | FileCheck %s
 
 ; The accesses to A are independent here but LAA reports it as a loop-carried
 ; forward dependence.  Check that we don't perform st->ld forwarding between
 ; them.
 ;
 ;   for (unsigned i = 0; i < 100; i++) {
 ;     A[i][1] = B[i] + 2;
 ;     C[i] = A[i][0] * 2;

llvm/test/Transforms/LoopLoadElim/opt-size.ll

 ; RUN: opt -basicaa -loop-load-elim -S < %s | FileCheck %s
 ; RUN: opt -basicaa -loop-load-elim -pgso -S < %s | FileCheck %s -check-prefix=PGSO
 ; RUN: opt -basicaa -loop-load-elim -pgso=false -S < %s | FileCheck %s -check-prefix=NPGSO
+; RUN: opt -aa-pipeline=basic-aa -passes='require<aa>,loop-load-elim' -S < %s | FileCheck %s
+; TODO: Enable tests with PGSO in new pipeline.
 
 ; When optimizing for size don't eliminate in this loop because the loop would
 ; have to be versioned first because A and C may alias.
 ;
 ;   for (unsigned i = 0; i < 100; i++) {
 ;     A[i+1] = B[i] + 2;
 ;     C[i] = A[i] * 2;
 ;   }

llvm/test/Transforms/LoopLoadElim/type-mismatch.ll

 ; RUN: opt -loop-load-elim -S < %s | FileCheck %s
+; RUN: opt -passes=loop-load-elim -S < %s | FileCheck %s
 
 ; Don't crash if the store and the load use different types.
 ;
 ;   for (unsigned i = 0; i < 100; i++) {
 ;     A[i+1] = B[i] + 2;
 ;     C[i] = ((float*)A)[i] * 2;
 ;   }
 

llvm/test/Transforms/LoopLoadElim/unknown-dep.ll

 ; RUN: opt -basicaa -loop-load-elim -S < %s | FileCheck %s
-
+; RUN: opt -aa-pipeline=basic-aa -passes='require<aa>,loop-load-elim' -S < %s | FileCheck %s
 target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
 
 ; Give up in the presence of unknown deps. Here, the different strides result
 ; in unknown dependence:
 ;
 ;   for (unsigned i = 0; i < 100; i++) {
 ;     A[i+1] = B[i] + 2;
 ;     A[2*i] = C[i] + 2;