[LoopFlatten] Update MemorySSA state

[SjoerdMeijer]

Summary

I would like to move LoopFlatten from LoopPass Manager LPM2 to LPM1 (in D116612), but that is a LPM that is using MemorySSA and so LoopFlatten needs to preserve MemorySSA and this adds that. More specifically, LoopFlatten restructures the CFG and with this change the MSSA state is updated accordingly, where we also update the DomTree. LoopFlatten doesn't rewrite/optimise/delete load or store instructions, so I have not added any MSSA updates for that.

[SjoerdMeijer]

Ah, forgot to mention tests. Without the update MSSAU->removeEdge(), I was expecting AR.MSSA->verifyMemorySSA() to run in an assert (with the expensive checks enabled in my build), but somewhat surprisingly to me it didn't.
So not sure yet if we need tests for this, but am looking into this.

[fhahn]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

824

Could MSSAU be created unconditionally like with the legacy pass manager?

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

824

Thanks for taking a look.

I borrowed this boilerplate code from SimpleLoopUnswitch, but sure, I guess so and will add that.

[SjoerdMeijer]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

824

Hmm, I see this pattern used in most passes, like LoopRotation, LoopSimplifyCFG (SimpleLoopUnswitch).

For my understanding, what would be the benefit of doing this unconditionally?

[fhahn]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

824

slightly less & more straight-forward code?

[SjoerdMeijer]

slightly less & more straight-forward code?

Fair enough. :-) Was wondering if there was more to it.
Now getting MSSAU unconditionally.

[SjoerdMeijer]

I somehow missed the failing regression test.
Turns out I do need to check AR.MSSA for passing on the MSSA object (or a nullptr).

[aeubanks]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

824

with the legacy PM it looks like we're unconditionally requesting MSSA (constructing it if it doesn't yet exist) and then updating it, which is wasted work if MSSA wasn't already available. we only want to keep MSSA up to date if it's already cached

[SjoerdMeijer]

Ping.

[nikic]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

892

Shouldn't this be using getAnalysisIfAvailable? MSSA is not a required analysis here. Also, getAnalysisUsage() should be adjusted to mark MSSA as preserved.

[SjoerdMeijer]

Thanks for the review, sorry for the delay.
Addressed comments and have added that to the LPM.

[nikic]

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

892

This doesn't look right ... won't this crash if getAnalysisIfAvailable() returns nullptr?

[SjoerdMeijer]

Sorry, my bad, should be fixed now.

[asbirlea]

Changes LGTM.
Nit: clang-format.
A simpler way to check all verification is now to pass VerificationLevel::Full to verifyMemorySSA(); I'm mentioning this for your local testing, in case your build doesn't have EXPENSIVE_CHECKS on already, not to include in this patch.

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

824

Nit: I'd stick with the Optional<MemorySSAUpdater> MSSAU; to avoid the object allocation when MSSA is not present.

[SjoerdMeijer]

In D116660#3252608, @asbirlea wrote:

Changes LGTM.
Nit: clang-format.
A simpler way to check all verification is now to pass VerificationLevel::Full to verifyMemorySSA(); I'm mentioning this for your local testing, in case your build doesn't have EXPENSIVE_CHECKS on already, not to include in this patch.

I have applied clang-format, double checked locally VerificationLevel::Full (but was using an expensive checks build), and restored optionally allocating MSSAU.

Thanks for reviewing!

llvm/lib/Transforms/Scalar/LoopFlatten.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Scalar/LoopFlatten.h"
 
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"

 
 static cl::opt<bool>
     AssumeNoOverflow("loop-flatten-assume-no-overflow", cl::Hidden,
                      cl::init(false),
                      cl::desc("Assume that the product of the two iteration "
                               "trip counts will never overflow"));
 
 static cl::opt<bool>
-    WidenIV("loop-flatten-widen-iv", cl::Hidden,
-            cl::init(true),
+    WidenIV("loop-flatten-widen-iv", cl::Hidden, cl::init(true),
             cl::desc("Widen the loop induction variables, if possible, so "
                      "overflow checks won't reject flattening"));
 
 struct FlattenInfo {
   Loop *OuterLoop = nullptr;
   Loop *InnerLoop = nullptr;
   // These PHINodes correspond to loop induction variables, which are expected
   // to start at zero and increment by one on each loop.
struct FlattenInfo {
   bool Widened = false;
 
   // Holds the old/narrow induction phis, i.e. the Phis before IV widening has
   // been applied. This bookkeeping is used so we can skip some checks on these
   // phi nodes.
   PHINode *NarrowInnerInductionPHI = nullptr;
   PHINode *NarrowOuterInductionPHI = nullptr;
 
-  FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL) {};
+  FlattenInfo(Loop *OL, Loop *IL) : OuterLoop(OL), InnerLoop(IL){};
 
   bool isNarrowInductionPhi(PHINode *Phi) {
     // This can't be the narrow phi if we haven't widened the IV first.
     if (!Widened)
       return false;
     return NarrowInnerInductionPHI == Phi || NarrowOuterInductionPHI == Phi;
   }
 };
if (isa<SCEVCouldNotCompute>(BackedgeTakenCount)) {
     LLVM_DEBUG(dbgs() << "Backedge-taken count is not predictable\n");
     return false;
   }
   // The use of the Extend=false flag on getTripCountFromExitCount was added
   // during a refactoring to preserve existing behavior.  However, there's
   // nothing obvious in the surrounding code when handles the overflow case.
   // FIXME: audit code to establish whether there's a latent bug here.
   const SCEV *SCEVTripCount =
       SE->getTripCountFromExitCount(BackedgeTakenCount, false);
   const SCEV *SCEVRHS = SE->getSCEV(RHS);
   if (SCEVRHS == SCEVTripCount)
     return setLoopComponents(RHS, TripCount, Increment, IterationInstructions);
   ConstantInt *ConstantRHS = dyn_cast<ConstantInt>(RHS);
   if (ConstantRHS) {
     const SCEV *BackedgeTCExt = nullptr;
     if (IsWidened) {
       const SCEV *SCEVTripCountExt;
if (!checkIVUsers(FI))
     return false;
 
   LLVM_DEBUG(dbgs() << "CanFlattenLoopPair: OK\n");
   return true;
 }
 
 static bool DoFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                               ScalarEvolution *SE, AssumptionCache *AC,
-                              const TargetTransformInfo *TTI, LPMUpdater *U) {
+                              const TargetTransformInfo *TTI, LPMUpdater *U,
+                              MemorySSAUpdater *MSSAU) {
   Function *F = FI.OuterLoop->getHeader()->getParent();
   LLVM_DEBUG(dbgs() << "Checks all passed, doing the transformation\n");
   {
     using namespace ore;
     OptimizationRemark Remark(DEBUG_TYPE, "Flattened", FI.InnerLoop->getStartLoc(),
                               FI.InnerLoop->getHeader());
     OptimizationRemarkEmitter ORE(F);
     Remark << "Flattened into outer loop";
static bool DoFlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   // trip counts.
   cast<User>(FI.OuterBranch->getCondition())->setOperand(1, NewTripCount);
 
   // Replace the inner loop backedge with an unconditional branch to the exit.
   BasicBlock *InnerExitBlock = FI.InnerLoop->getExitBlock();
   BasicBlock *InnerExitingBlock = FI.InnerLoop->getExitingBlock();
   InnerExitingBlock->getTerminator()->eraseFromParent();
   BranchInst::Create(InnerExitBlock, InnerExitingBlock);
+
+  // Update the DomTree and MemorySSA.
   DT->deleteEdge(InnerExitingBlock, FI.InnerLoop->getHeader());
+  if (MSSAU)
+    MSSAU->removeEdge(InnerExitingBlock, FI.InnerLoop->getHeader());
 
   // Replace all uses of the polynomial calculated from the two induction
   // variables with the one new one.
   IRBuilder<> Builder(FI.OuterInductionPHI->getParent()->getTerminator());
   for (Value *V : FI.LinearIVUses) {
     Value *OuterValue = FI.OuterInductionPHI;
     if (FI.Widened)
       OuterValue = Builder.CreateTrunc(FI.OuterInductionPHI, V->getType(),
                                        "flatten.trunciv");
 
-    LLVM_DEBUG(dbgs() << "Replacing: "; V->dump();
-               dbgs() << "with:      "; OuterValue->dump());
+    LLVM_DEBUG(dbgs() << "Replacing: "; V->dump(); dbgs() << "with:      ";
+               OuterValue->dump());
     V->replaceAllUsesWith(OuterValue);
   }
 
   // Tell LoopInfo, SCEV and the pass manager that the inner loop has been
   // deleted, and any information that have about the outer loop invalidated.
   SE->forgetLoop(FI.OuterLoop);
   SE->forgetLoop(FI.InnerLoop);
   if (U)
static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   unsigned MaxLegalSize = DL.getLargestLegalIntTypeSizeInBits();
   auto *MaxLegalType = DL.getLargestLegalIntType(M->getContext());
 
   // If both induction types are less than the maximum legal integer width,
   // promote both to the widest type available so we know calculating
   // (OuterTripCount * InnerTripCount) as the new trip count is safe.
   if (InnerType != OuterType ||
       InnerType->getScalarSizeInBits() >= MaxLegalSize ||
-      MaxLegalType->getScalarSizeInBits() < InnerType->getScalarSizeInBits() * 2) {
+      MaxLegalType->getScalarSizeInBits() <
+          InnerType->getScalarSizeInBits() * 2) {
     LLVM_DEBUG(dbgs() << "Can't widen the IV\n");
     return false;
   }
 
   SCEVExpander Rewriter(*SE, DL, "loopflatten");
   SmallVector<WeakTrackingVH, 4> DeadInsts;
   unsigned ElimExt = 0;
   unsigned Widened = 0;
 
-  auto CreateWideIV = [&] (WideIVInfo WideIV, bool &Deleted) -> bool {
-    PHINode *WidePhi = createWideIV(WideIV, LI, SE, Rewriter, DT, DeadInsts,
-                                    ElimExt, Widened, true /* HasGuards */,
-                                    true /* UsePostIncrementRanges */);
+  auto CreateWideIV = [&](WideIVInfo WideIV, bool &Deleted) -> bool {
+    PHINode *WidePhi =
+        createWideIV(WideIV, LI, SE, Rewriter, DT, DeadInsts, ElimExt, Widened,
+                     true /* HasGuards */, true /* UsePostIncrementRanges */);
     if (!WidePhi)
       return false;
     LLVM_DEBUG(dbgs() << "Created wide phi: "; WidePhi->dump());
     LLVM_DEBUG(dbgs() << "Deleting old phi: "; WideIV.NarrowIV->dump());
     Deleted = RecursivelyDeleteDeadPHINode(WideIV.NarrowIV);
     return true;
   };
 
   bool Deleted;
-  if (!CreateWideIV({FI.InnerInductionPHI, MaxLegalType, false }, Deleted))
+  if (!CreateWideIV({FI.InnerInductionPHI, MaxLegalType, false}, Deleted))
     return false;
   // Add the narrow phi to list, so that it will be adjusted later when the
   // the transformation is performed.
   if (!Deleted)
     FI.InnerPHIsToTransform.insert(FI.InnerInductionPHI);
 
-  if (!CreateWideIV({FI.OuterInductionPHI, MaxLegalType, false }, Deleted))
+  if (!CreateWideIV({FI.OuterInductionPHI, MaxLegalType, false}, Deleted))
     return false;
 
   assert(Widened && "Widened IV expected");
   FI.Widened = true;
 
   // Save the old/narrow induction phis, which we need to ignore in CheckPHIs.
   FI.NarrowInnerInductionPHI = FI.InnerInductionPHI;
   FI.NarrowOuterInductionPHI = FI.OuterInductionPHI;
 
   // After widening, rediscover all the loop components.
   return CanFlattenLoopPair(FI, DT, LI, SE, AC, TTI);
 }
 
 static bool FlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
                             ScalarEvolution *SE, AssumptionCache *AC,
-                            const TargetTransformInfo *TTI, LPMUpdater *U) {
+                            const TargetTransformInfo *TTI, LPMUpdater *U,
+                            MemorySSAUpdater *MSSAU) {
   LLVM_DEBUG(
       dbgs() << "Loop flattening running on outer loop "
              << FI.OuterLoop->getHeader()->getName() << " and inner loop "
              << FI.InnerLoop->getHeader()->getName() << " in "
              << FI.OuterLoop->getHeader()->getParent()->getName() << "\n");
 
   if (!CanFlattenLoopPair(FI, DT, LI, SE, AC, TTI))
     return false;
static bool FlattenLoopPair(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   // 'RepeatedInstructionThreshold' is set to only 2, which can probably be
   // relaxed. Because this is making a code change (the IV widening, but not
   // the flattening), we return true here.
   if (FI.Widened && !CanFlatten)
     return true;
 
   // If we have widened and can perform the transformation, do that here.
   if (CanFlatten)
-    return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI, U);
+    return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI, U, MSSAU);
 
   // Otherwise, if we haven't widened the IV, check if the new iteration
   // variable might overflow. In this case, we need to version the loop, and
   // select the original version at runtime if the iteration space is too
   // large.
   // TODO: We currently don't version the loop.
   OverflowResult OR = checkOverflow(FI, DT, AC);
   if (OR == OverflowResult::AlwaysOverflowsHigh ||
       OR == OverflowResult::AlwaysOverflowsLow) {
     LLVM_DEBUG(dbgs() << "Multiply would always overflow, so not profitable\n");
     return false;
   } else if (OR == OverflowResult::MayOverflow) {
     LLVM_DEBUG(dbgs() << "Multiply might overflow, not flattening\n");
     return false;
   }
 
   LLVM_DEBUG(dbgs() << "Multiply cannot overflow, modifying loop in-place\n");
-  return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI, U);
+  return DoFlattenLoopPair(FI, DT, LI, SE, AC, TTI, U, MSSAU);
 }
 
 bool Flatten(LoopNest &LN, DominatorTree *DT, LoopInfo *LI, ScalarEvolution *SE,
-             AssumptionCache *AC, TargetTransformInfo *TTI, LPMUpdater *U) {
+             AssumptionCache *AC, TargetTransformInfo *TTI, LPMUpdater *U,
+             MemorySSAUpdater *MSSAU) {
   bool Changed = false;
   for (Loop *InnerLoop : LN.getLoops()) {
     auto *OuterLoop = InnerLoop->getParentLoop();
     if (!OuterLoop)
       continue;
     FlattenInfo FI(OuterLoop, InnerLoop);
-    Changed |= FlattenLoopPair(FI, DT, LI, SE, AC, TTI, U);
+    Changed |= FlattenLoopPair(FI, DT, LI, SE, AC, TTI, U, MSSAU);
   }
   return Changed;
 }
 
 PreservedAnalyses LoopFlattenPass::run(LoopNest &LN, LoopAnalysisManager &LAM,
                                        LoopStandardAnalysisResults &AR,
                                        LPMUpdater &U) {
 
   bool Changed = false;
 
+  Optional<MemorySSAUpdater> MSSAU;
+  if (AR.MSSA) {
+    MSSAU = MemorySSAUpdater(AR.MSSA);
+    if (VerifyMemorySSA)
+      AR.MSSA->verifyMemorySSA();
+  }
+
   // The loop flattening pass requires loops to be
   // in simplified form, and also needs LCSSA. Running
   // this pass will simplify all loops that contain inner loops,
   // regardless of whether anything ends up being flattened.
-  Changed |= Flatten(LN, &AR.DT, &AR.LI, &AR.SE, &AR.AC, &AR.TTI, &U);
+  Changed |= Flatten(LN, &AR.DT, &AR.LI, &AR.SE, &AR.AC, &AR.TTI, &U,
+                     MSSAU.hasValue() ? MSSAU.getPointer() : nullptr);
 
   if (!Changed)
     return PreservedAnalyses::all();
 
-  return getLoopPassPreservedAnalyses();
+  if (AR.MSSA && VerifyMemorySSA)
+    AR.MSSA->verifyMemorySSA();
+
+  auto PA = getLoopPassPreservedAnalyses();
+  if (AR.MSSA)
+    PA.preserve<MemorySSAAnalysis>();
+  return PA;
 }
 
 namespace {
 class LoopFlattenLegacyPass : public FunctionPass {
 public:
   static char ID; // Pass ID, replacement for typeid
   LoopFlattenLegacyPass() : FunctionPass(ID) {
     initializeLoopFlattenLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   // Possibly flatten loop L into its child.
   bool runOnFunction(Function &F) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     getLoopAnalysisUsage(AU);
     AU.addRequired<TargetTransformInfoWrapperPass>();
     AU.addPreserved<TargetTransformInfoWrapperPass>();
     AU.addRequired<AssumptionCacheTracker>();
     AU.addPreserved<AssumptionCacheTracker>();
+    AU.addPreserved<MemorySSAWrapperPass>();
   }
 };
 } // namespace
 
 char LoopFlattenLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(LoopFlattenLegacyPass, "loop-flatten", "Flattens loops",
                       false, false)
 INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_END(LoopFlattenLegacyPass, "loop-flatten", "Flattens loops",
                     false, false)
 
-FunctionPass *llvm::createLoopFlattenPass() { return new LoopFlattenLegacyPass(); }
+FunctionPass *llvm::createLoopFlattenPass() {
+  return new LoopFlattenLegacyPass();
+}
 
 bool LoopFlattenLegacyPass::runOnFunction(Function &F) {
   ScalarEvolution *SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
   LoopInfo *LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>();
   DominatorTree *DT = DTWP ? &DTWP->getDomTree() : nullptr;
   auto &TTIP = getAnalysis<TargetTransformInfoWrapperPass>();
   auto *TTI = &TTIP.getTTI(F);
   auto *AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
+  auto *MSSA = getAnalysisIfAvailable<MemorySSAWrapperPass>();
+
+  Optional<MemorySSAUpdater> MSSAU;
+  if (MSSA)
+    MSSAU = MemorySSAUpdater(&MSSA->getMSSA());
+
   bool Changed = false;
   for (Loop *L : *LI) {
     auto LN = LoopNest::getLoopNest(*L, *SE);
-    Changed |= Flatten(*LN, DT, LI, SE, AC, TTI, nullptr);
+    Changed |= Flatten(*LN, DT, LI, SE, AC, TTI, nullptr,
+                       MSSAU.hasValue() ? MSSAU.getPointer() : nullptr);
   }
   return Changed;
 }