[GuardWidening] Preserve MemorySSA

[nikic]

Summary

As reported on https://bugs.llvm.org/show_bug.cgi?id=51020, the guard widening loop pass doesn't preserve MemorySSA, so it can no longer be scheduled in the same loop pass manager as LICM. However, the loop-schedule.ll test indicates that this is supposed to work.

Fix this by preserving MemorySSA if available, as this seems to be trivial in this case (we only need to drop the memory access for the guards).

[nikic]

Mark as preserved in NPM (it's not part of getLoopPassPreservedAnalyses).

[asbirlea]

llvm/lib/Transforms/Scalar/GuardWidening.cpp

772–773

Does it make sense to get the cached analysis in the Function pass if it exists and preserve it too in that case?

823–824

Why not also preserve in the Function pass if MemorySSA is available?

[reames]

LGTM for the loop passes. As with Alina, I'd encourage you to do the function passes too for completeness. Feel free to either do a separate change for that, or roll it into this commit.

[nikic]

Preserve MSSA in function pass as well.

[asbirlea]

Thanks, LG with one comment to address.

llvm/lib/Transforms/Scalar/GuardWidening.cpp

783–784

PA.preserve<MemorySSAAnalysis>();

llvm/lib/Transforms/Scalar/GuardWidening.cpp

 #include "llvm/Transforms/Scalar/GuardWidening.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/GuardUtils.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/MemorySSAUpdater.h"
 #include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
assert(GI->getIntrinsicID() == Intrinsic::experimental_guard &&
            "Bad guard intrinsic?");
     GI->setArgOperand(0, NewCond);
     return;
   }
   cast<BranchInst>(I)->setCondition(NewCond);
 }
 
 // Eliminates the guard instruction properly.
-static void eliminateGuard(Instruction *GuardInst) {
+static void eliminateGuard(Instruction *GuardInst, MemorySSAUpdater *MSSAU) {
   GuardInst->eraseFromParent();
+  if (MSSAU)
+    MSSAU->removeMemoryAccess(GuardInst);
   ++GuardsEliminated;
 }
 
 class GuardWideningImpl {
   DominatorTree &DT;
   PostDominatorTree *PDT;
   LoopInfo &LI;
+  MemorySSAUpdater *MSSAU;
 
   /// Together, these describe the region of interest.  This might be all of
   /// the blocks within a function, or only a given loop's blocks and preheader.
   DomTreeNode *Root;
   std::function<bool(BasicBlock*)> BlockFilter;
 
   /// The set of guards and conditional branches whose conditions have been
   /// widened into dominating guards.
void widenGuard(Instruction *ToWiden, Value *NewCondition,
     if (isGuardAsWidenableBranch(ToWiden)) {
       setWidenableBranchCond(cast<BranchInst>(ToWiden), Result);
       return;
     }
     setCondition(ToWiden, Result);
   }
 
 public:
-
   explicit GuardWideningImpl(DominatorTree &DT, PostDominatorTree *PDT,
-                             LoopInfo &LI, DomTreeNode *Root,
+                             LoopInfo &LI, MemorySSAUpdater *MSSAU,
+                             DomTreeNode *Root,
                              std::function<bool(BasicBlock*)> BlockFilter)
-    : DT(DT), PDT(PDT), LI(LI), Root(Root), BlockFilter(BlockFilter)
-        {}
+      : DT(DT), PDT(PDT), LI(LI), MSSAU(MSSAU), Root(Root),
+        BlockFilter(BlockFilter) {}
 
   /// The entry point for this pass.
   bool run();
 };
 }
 
 static bool isSupportedGuardInstruction(const Instruction *Insn) {
   if (isGuard(Insn))
for (auto *II : CurrentList)
       Changed |= eliminateInstrViaWidening(II, DFI, GuardsInBlock);
   }
 
   assert(EliminatedGuardsAndBranches.empty() || Changed);
   for (auto *I : EliminatedGuardsAndBranches)
     if (!WidenedGuards.count(I)) {
       assert(isa<ConstantInt>(getCondition(I)) && "Should be!");
       if (isSupportedGuardInstruction(I))
-        eliminateGuard(I);
+        eliminateGuard(I, MSSAU);
       else {
         assert(isa<BranchInst>(I) &&
                "Eliminated something other than guard or branch?");
         ++CondBranchEliminated;
       }
     }
 
   return Changed;
StringRef GuardWideningImpl::scoreTypeToString(WideningScore WS) {
   llvm_unreachable("Fully covered switch above!");
 }
 #endif
 
 PreservedAnalyses GuardWideningPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
   auto &LI = AM.getResult<LoopAnalysis>(F);
   auto &PDT = AM.getResult<PostDominatorTreeAnalysis>(F);
-  if (!GuardWideningImpl(DT, &PDT, LI, DT.getRootNode(),
-                         [](BasicBlock*) { return true; } ).run())
+  auto *MSSAA = AM.getCachedResult<MemorySSAAnalysis>(F);
+  std::unique_ptr<MemorySSAUpdater> MSSAU;
+  if (MSSAA)
+    MSSAU = std::make_unique<MemorySSAUpdater>(&MSSAA->getMSSA());
+  if (!GuardWideningImpl(DT, &PDT, LI, MSSAU ? MSSAU.get() : nullptr,
+                         DT.getRootNode(), [](BasicBlock *) { return true; })
+           .run())
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;
   PA.preserveSet<CFGAnalyses>();
+  PA.preserve<MemorySSAAnalysis>();
   return PA;
 }
 
 PreservedAnalyses GuardWideningPass::run(Loop &L, LoopAnalysisManager &AM,
                                          LoopStandardAnalysisResults &AR,
                                          LPMUpdater &U) {
   BasicBlock *RootBB = L.getLoopPredecessor();
   if (!RootBB)
     RootBB = L.getHeader();
   auto BlockFilter = [&](BasicBlock *BB) {
     return BB == RootBB || L.contains(BB);
   };
-  if (!GuardWideningImpl(AR.DT, nullptr, AR.LI, AR.DT.getNode(RootBB),
-                         BlockFilter).run())
+  std::unique_ptr<MemorySSAUpdater> MSSAU;
+  if (AR.MSSA)
+    MSSAU = std::make_unique<MemorySSAUpdater>(AR.MSSA);
+  if (!GuardWideningImpl(AR.DT, nullptr, AR.LI, MSSAU ? MSSAU.get() : nullptr,
+                         AR.DT.getNode(RootBB), BlockFilter).run())
     return PreservedAnalyses::all();
 
-  return getLoopPassPreservedAnalyses();
+  auto PA = getLoopPassPreservedAnalyses();
+  if (AR.MSSA)
+    PA.preserve<MemorySSAAnalysis>();
+  return PA;
 }
 
 namespace {
 struct GuardWideningLegacyPass : public FunctionPass {
   static char ID;
 
   GuardWideningLegacyPass() : FunctionPass(ID) {
     initializeGuardWideningLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override {
     if (skipFunction(F))
       return false;
     auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     auto &PDT = getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
-    return GuardWideningImpl(DT, &PDT, LI, DT.getRootNode(),
-                         [](BasicBlock*) { return true; } ).run();
+    auto *MSSAWP = getAnalysisIfAvailable<MemorySSAWrapperPass>();
+    std::unique_ptr<MemorySSAUpdater> MSSAU;
+    if (MSSAWP)
+      MSSAU = std::make_unique<MemorySSAUpdater>(&MSSAWP->getMSSA());
+    return GuardWideningImpl(DT, &PDT, LI, MSSAU ? MSSAU.get() : nullptr,
+                             DT.getRootNode(),
+                             [](BasicBlock *) { return true; })
+        .run();
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     AU.addRequired<DominatorTreeWrapperPass>();
     AU.addRequired<PostDominatorTreeWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
+    AU.addPreserved<MemorySSAWrapperPass>();
   }
 };
 
 /// Same as above, but restricted to a single loop at a time.  Can be
 /// scheduled with other loop passes w/o breaking out of LPM
 struct LoopGuardWideningLegacyPass : public LoopPass {
   static char ID;
 
   LoopGuardWideningLegacyPass() : LoopPass(ID) {
     initializeLoopGuardWideningLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnLoop(Loop *L, LPPassManager &LPM) override {
     if (skipLoop(L))
       return false;
     auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
     auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     auto *PDTWP = getAnalysisIfAvailable<PostDominatorTreeWrapperPass>();
     auto *PDT = PDTWP ? &PDTWP->getPostDomTree() : nullptr;
+    auto *MSSAWP = getAnalysisIfAvailable<MemorySSAWrapperPass>();
+    std::unique_ptr<MemorySSAUpdater> MSSAU;
+    if (MSSAWP)
+      MSSAU = std::make_unique<MemorySSAUpdater>(&MSSAWP->getMSSA());
+
     BasicBlock *RootBB = L->getLoopPredecessor();
     if (!RootBB)
       RootBB = L->getHeader();
     auto BlockFilter = [&](BasicBlock *BB) {
       return BB == RootBB || L->contains(BB);
     };
-    return GuardWideningImpl(DT, PDT, LI,
+    return GuardWideningImpl(DT, PDT, LI, MSSAU ? MSSAU.get() : nullptr,
                              DT.getNode(RootBB), BlockFilter).run();
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     getLoopAnalysisUsage(AU);
     AU.addPreserved<PostDominatorTreeWrapperPass>();
+    AU.addPreserved<MemorySSAWrapperPass>();
   }
 };
 }
 
 char GuardWideningLegacyPass::ID = 0;
 char LoopGuardWideningLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(GuardWideningLegacyPass, "guard-widening", "Widen guards",

llvm/test/Transforms/GuardWidening/basic-loop.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -S -loop-guard-widening -enable-new-pm=0 < %s | FileCheck %s
-; RUN: opt -S -passes="loop(guard-widening)" < %s | FileCheck %s
+; RUN: opt -S -loop-guard-widening -verify-memoryssa -enable-new-pm=0 < %s | FileCheck %s
+; RUN: opt -S -passes="loop-mssa(guard-widening)" -verify-memoryssa < %s | FileCheck %s
 
 declare void @llvm.experimental.guard(i1,...)
 
 @G = external global i32
 
 ; Show that we can widen into early checks within a loop, and in the process
 ; expose optimization oppurtunities.
 define void @widen_within_loop(i1 %cond_0, i1 %cond_1, i1 %cond_2) {

llvm/test/Transforms/GuardWidening/loop-schedule.ll

-; RUN: opt -S -licm -loop-guard-widening -licm -debug-pass=Structure -enable-new-pm=0 < %s 2>&1 | FileCheck %s --check-prefixes=LPM,CHECK
-; RUN: opt -S -passes='licm,guard-widening,licm' -debug-pass-manager < %s 2>&1 | FileCheck %s --check-prefixes=NPM,CHECK
+; RUN: opt -S -licm -loop-guard-widening -licm -verify-memoryssa -debug-pass=Structure -enable-new-pm=0 < %s 2>&1 | FileCheck %s --check-prefixes=LPM,CHECK
+; RUN: opt -S -passes='licm,guard-widening,licm' -verify-memoryssa -debug-pass-manager < %s 2>&1 | FileCheck %s --check-prefixes=NPM,CHECK
 
 ; Main point of this test is to check the scheduling -- there should be
 ; no analysis passes needed between LICM and LoopGuardWidening
 
 ; LPM: Loop Pass Manager
 ; LPM:   Loop Invariant Code Motion
-; LPM:   Widen guards (within a single loop, as a loop pass)
-; LPM:   Loop Invariant Code Motion
+; LPM-NEXT:   Widen guards (within a single loop, as a loop pass)
+; LPM-NEXT:   Loop Invariant Code Motion
 
 ; NPM: LICMPass
 ; NPM-NEXT: GuardWideningPass
 ; NPM-NEXT: LICMPass
 
 declare void @llvm.experimental.guard(i1,...)
 
 define void @iter(i32 %a, i32 %b, i1* %c_p) {