[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

887

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

887

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"
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(),
static bool CanWidenIV(FlattenInfo &FI, DominatorTree *DT, LoopInfo *LI,
   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;
 
+  MemorySSAUpdater MSSAU = MemorySSAUpdater(AR.MSSA);
+  if (AR.MSSA && 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,
+                     AR.MSSA ? &MSSAU : 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(); }
 
 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;
 }