[MemorySSA] Support lazy use optimization

[nikic]

Summary

This changes MemorySSA to be constructed in unoptimized form. MemorySSA::ensureOptimized() can be called to optimize all uses (once). This should be done by passes where having optimized uses is beneficial, either because we're going to query all uses anyway, or because we're doing def-use walks.

This should help reduce the compile-time impact of MemorySSA for some use cases (the reason why I started looking into this is D117926), which can avoid optimizing all uses upfront, and instead only optimize those that are actually queried.

Actually, we have an existing use-case for this, which is EarlyCSE. Disabling eager use optimization there gives a significant compile-time improvement: http://llvm-compile-time-tracker.com/compare.php?from=59191057243e34d85b644716ef2811bfea8efd1e&to=1dd3499045716e71bad61ccd70700e734a74d350&stat=instructions This is because EarlyCSE will generally only query clobbers for a subset of all uses.

To be more conservative I haven't included the EarlyCSE change here, but I can add it.

[asbirlea]

Thank you for sending this out!

Broadly the changes look good. The patch as is can also give some compile-time improvements for cases where MSSA is built and preserved in loop pipelines and need only be optimized before LICM, not before the first loop pass in the pipeline. This also means results may change based on the delayed optimization.
This should hold true for both Legacy and NPM, but I see in practice there's a difference only for LPM (http://llvm-compile-time-tracker.com/compare.php?from=067c035012fc061ad6378458774ac2df117283c6&to=59191057243e34d85b644716ef2811bfea8efd1e&stat=instructions)

Let me test this out as is first. Then separately with the EarlyCSE ensureOptimized() removed to check if the gains in compile-time are not coupled with run time regressions.

[nikic]

In D121381#3373615, @asbirlea wrote:

Thank you for sending this out!

Broadly the changes look good. The patch as is can also give some compile-time improvements for cases where MSSA is built and preserved in loop pipelines and need only be optimized before LICM, not before the first loop pass in the pipeline. This also means results may change based on the delayed optimization.
This should hold true for both Legacy and NPM, but I see in practice there's a difference only for LPM (http://llvm-compile-time-tracker.com/compare.php?from=067c035012fc061ad6378458774ac2df117283c6&to=59191057243e34d85b644716ef2811bfea8efd1e&stat=instructions)

I think this is because with the LegacyPM, loop analyses (including MSSA) will get scheduled unconditionally, even if there are no loops, so we sometimes compute MSSA unnecessarily. With the NewPM, analyses are only computed if there are loops, so we generally don't perform unnecessary MSSA constructions there.

Let me test this out as is first. Then separately with the EarlyCSE ensureOptimized() removed to check if the gains in compile-time are not coupled with run time regressions.

I should mention here that this does impact EarlyCSE results due to this fallback: https://github.com/llvm/llvm-project/blob/54d7fde46e8a0e425245e18732c2a78e64fa7b35/llvm/lib/Transforms/Scalar/EarlyCSE.cpp#L1030-L1034 Without eagerly optimizing uses, the getDefiningAccess() fallback will be unoptimized for uses. But I think we can compensate this by increasing the EarlyCSEMssaOptCap if necessary.

[reames]

Conceptually, I think this makes a lot of sense.

Biggest thing I see in this patch is a strong need to update comments and APIs. This introduces the possibility of clients seeing unoptimized uses which is new, and not currently reflected in the API documentation. Please address before landing.

llvm/include/llvm/Analysis/MemorySSA.h

805

Strong comment: ensureOptimizedLoads.

This does *NOT* optimize stores, and that's important to distinguish in comments and naming.

[asbirlea]

LGTM for this patch on the impact side, testing looks good.

+1 on improving documentation for future uses, to include that by default uses are no longer optimized. Also rename the new method to ensureOptimizedUses() per @reames's suggestion.
(docs can be landed separately as long as they come shortly after)

On the impact for the EarlyCSE follow-up I'm seeing a couple of small regressions in fdo configuration that I would not block on. I'd go with landing the EarlyCSE change separately from this and evaluate if increasing the cap makes sense or not based on if other folks notice any performance impact from the change.

[asbirlea]

Sent out https://reviews.llvm.org/D121740.

llvm/include/llvm/Analysis/MemorySSA.h

804

Please expand the doc for this API here.

[nikic]

Rename to ensureOptimizedUses(), some comment updates.

[fhahn]

Thanks, this looks like a nice improvement in combination with D121740!

[asbirlea]

LGTM with a couple nits.

llvm/include/llvm/Analysis/MemorySSA.h

350

This is not necessarily true due to the cap set by MaxCheckLimit (memssa-check-limit).
Also not true if a MemoryDef was added that invalidated optimizations (OptimizedID will no longer match).

806

Add: By default, during MemorySSA build, uses are *not* optimized

[reames]

I noticed by accident that MemCpyOpt does not call ensureOptimizedUses, and thus changes behavior after this change. Not sure if that's good or bad, but it doesn't seem to have been explicitly noted in the review. For call slot optimization, we do perform a lot of clobber walks from loads.

(Separately, we could recognize the code to perform far fewer clobber walks. - e.g. there's no point in searching for the clobber for any non-alloca location.)

[nikic]

In D121381#3410466, @reames wrote:

I noticed by accident that MemCpyOpt does not call ensureOptimizedUses, and thus changes behavior after this change. Not sure if that's good or bad, but it doesn't seem to have been explicitly noted in the review. For call slot optimization, we do perform a lot of clobber walks from loads.

I'm surprised that this would cause a change in behavior. MemCpyOpt always does clobber walks (without any limit that falls back to getDefiningAccess), so it really shouldn't care whether uses are optimized or not. Possibly our optimization cutoffs work slightly different between the eager use optimization and the clobber walk, and that can cause differences?

[reames]

In D121381#3410717, @nikic wrote:

In D121381#3410466, @reames wrote:

I noticed by accident that MemCpyOpt does not call ensureOptimizedUses, and thus changes behavior after this change. Not sure if that's good or bad, but it doesn't seem to have been explicitly noted in the review. For call slot optimization, we do perform a lot of clobber walks from loads.

I'm surprised that this would cause a change in behavior. MemCpyOpt always does clobber walks (without any limit that falls back to getDefiningAccess), so it really shouldn't care whether uses are optimized or not. Possibly our optimization cutoffs work slightly different between the eager use optimization and the clobber walk, and that can cause differences?

Er, sorry, was apparently sloppy in my wording.

I did *not* see a functional difference. I saw a compile time difference. On at least one example, not having the eager use optimization caused a significant shift in where we spend time. I'm not even saying it slowed down (the results were too noisy to tell), just that it was different and that doesn't seem obviously expected from the review.

llvm/include/llvm/Analysis/MemorySSA.h

public:
 
   void print(raw_ostream &OS) const;
 
   void setOptimized(MemoryAccess *DMA) {
     OptimizedID = DMA->getID();
     setOperand(0, DMA);
   }
 
-  /// The defining access of a MemoryUses are always optimized if queried from
-  /// outside MSSA construction itself.  This result is only useful inside
-  /// the MSSA implementation.
+  /// Whether the MemoryUse is optimized. If ensureOptimizedUses() was called,
+  /// uses will usually be optimized, but this is not guaranteed (e.g. due to
+  /// invalidation and optimization limits.)
   bool isOptimized() const {
     return getDefiningAccess() && OptimizedID == getDefiningAccess()->getID();
   }
 
   MemoryAccess *getOptimized() const {
     return getDefiningAccess();
   }
 
public:
   /// Verify that MemorySSA is self consistent (IE definitions dominate
   /// all uses, uses appear in the right places).  This is used by unit tests.
   void verifyMemorySSA(VerificationLevel = VerificationLevel::Fast) const;
 
   /// Used in various insertion functions to specify whether we are talking
   /// about the beginning or end of a block.
   enum InsertionPlace { Beginning, End, BeforeTerminator };
 
+  /// By default, uses are *not* optimized during MemorySSA construction.
+  /// Calling this method will attempt to optimize all MemoryUses, if this has
+  /// not happened yet for this MemorySSA instance. This should be done if you
+  /// plan to query the clobbering access for most uses, or if you walk the
+  /// def-use chain of uses.
+  void ensureOptimizedUses();
+
 protected:
   // Used by Memory SSA dumpers and wrapper pass
   friend class MemorySSAPrinterLegacyPass;
   friend class MemorySSAUpdater;
 
   void verifyOrderingDominationAndDefUses(
       Function &F, VerificationLevel = VerificationLevel::Fast) const;
   void verifyDominationNumbers(const Function &F) const;
private:
   mutable SmallPtrSet<const BasicBlock *, 16> BlockNumberingValid;
   mutable DenseMap<const MemoryAccess *, unsigned long> BlockNumbering;
 
   // Memory SSA building info
   std::unique_ptr<ClobberWalkerBase<AliasAnalysis>> WalkerBase;
   std::unique_ptr<CachingWalker<AliasAnalysis>> Walker;
   std::unique_ptr<SkipSelfWalker<AliasAnalysis>> SkipWalker;
   unsigned NextID = 0;
+  bool IsOptimized = false;
 };
 
 /// Enables verification of MemorySSA.
 ///
 /// The checks which this flag enables is exensive and disabled by default
 /// unless `EXPENSIVE_CHECKS` is defined.  The flag `-verify-memoryssa` can be
 /// used to selectively enable the verification without re-compilation.
 extern bool VerifyMemorySSA;

llvm/lib/Analysis/MemorySSA.cpp

void MemorySSA::OptimizeUses::optimizeUsesInBlock(
   for (MemoryAccess &MA : *Accesses) {
     auto *MU = dyn_cast<MemoryUse>(&MA);
     if (!MU) {
       VersionStack.push_back(&MA);
       ++StackEpoch;
       continue;
     }
 
+    if (MU->isOptimized())
+      continue;
+
     if (isUseTriviallyOptimizableToLiveOnEntry(*AA, MU->getMemoryInst())) {
       MU->setDefiningAccess(MSSA->getLiveOnEntryDef(), true, None);
       continue;
     }
 
     MemoryLocOrCall UseMLOC(MU);
     auto &LocInfo = LocStackInfo[UseMLOC];
     // If the pop epoch changed, it means we've removed stuff from top of
void MemorySSA::buildMemorySSA(BatchAAResults &BAA) {
   }
   placePHINodes(DefiningBlocks);
 
   // Now do regular SSA renaming on the MemoryDef/MemoryUse. Visited will get
   // filled in with all blocks.
   SmallPtrSet<BasicBlock *, 16> Visited;
   renamePass(DT->getRootNode(), LiveOnEntryDef.get(), Visited);
 
-  ClobberWalkerBase<BatchAAResults> WalkerBase(this, &BAA, DT);
-  CachingWalker<BatchAAResults> WalkerLocal(this, &WalkerBase);
-  OptimizeUses(this, &WalkerLocal, &BAA, DT).optimizeUses();
-
   // Mark the uses in unreachable blocks as live on entry, so that they go
   // somewhere.
   for (auto &BB : F)
     if (!Visited.count(&BB))
       markUnreachableAsLiveOnEntry(&BB);
 }
 
 MemorySSAWalker *MemorySSA::getWalker() { return getWalkerImpl(); }
if (UseBB != Dominator->getBlock())
       return DT->dominates(Dominator->getBlock(), UseBB);
     // If the UseBB and the DefBB are the same, compare locally.
     return locallyDominates(Dominator, cast<MemoryAccess>(Dominatee));
   }
   // If it's not a PHI node use, the normal dominates can already handle it.
   return dominates(Dominator, cast<MemoryAccess>(Dominatee.getUser()));
 }
 
+void MemorySSA::ensureOptimizedUses() {
+  if (IsOptimized)
+    return;
+
+  BatchAAResults BatchAA(*AA);
+  ClobberWalkerBase<BatchAAResults> WalkerBase(this, &BatchAA, DT);
+  CachingWalker<BatchAAResults> WalkerLocal(this, &WalkerBase);
+  OptimizeUses(this, &WalkerLocal, &BatchAA, DT).optimizeUses();
+  IsOptimized = true;
+}
+
 void MemoryAccess::print(raw_ostream &OS) const {
   switch (getValueID()) {
   case MemoryPhiVal: return static_cast<const MemoryPhi *>(this)->print(OS);
   case MemoryDefVal: return static_cast<const MemoryDef *>(this)->print(OS);
   case MemoryUseVal: return static_cast<const MemoryUse *>(this)->print(OS);
   }
   llvm_unreachable("invalid value id");
 }
return getNodeLabel(Node, CFGInfo).find(';') != std::string::npos
                : "";
   }
 };
 
 } // namespace llvm
 
 bool MemorySSAPrinterLegacyPass::runOnFunction(Function &F) {
   auto &MSSA = getAnalysis<MemorySSAWrapperPass>().getMSSA();
+  MSSA.ensureOptimizedUses();
   if (DotCFGMSSA != "") {
     DOTFuncMSSAInfo CFGInfo(F, MSSA);
     WriteGraph(&CFGInfo, "", false, "MSSA", DotCFGMSSA);
   } else
     MSSA.print(dbgs());
 
   if (VerifyMemorySSA)
     MSSA.verifyMemorySSA();
bool MemorySSAAnalysis::Result::invalidate(
   return !(PAC.preserved() || PAC.preservedSet<AllAnalysesOn<Function>>()) ||
          Inv.invalidate<AAManager>(F, PA) ||
          Inv.invalidate<DominatorTreeAnalysis>(F, PA);
 }
 
 PreservedAnalyses MemorySSAPrinterPass::run(Function &F,
                                             FunctionAnalysisManager &AM) {
   auto &MSSA = AM.getResult<MemorySSAAnalysis>(F).getMSSA();
+  MSSA.ensureOptimizedUses();
   if (DotCFGMSSA != "") {
     DOTFuncMSSAInfo CFGInfo(F, MSSA);
     WriteGraph(&CFGInfo, "", false, "MSSA", DotCFGMSSA);
   } else {
     OS << "MemorySSA for function: " << F.getName() << "\n";
     MSSA.print(OS);
   }
 

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

 };
 
 static bool eliminateDeadStores(Function &F, AliasAnalysis &AA, MemorySSA &MSSA,
                                 DominatorTree &DT, PostDominatorTree &PDT,
                                 const TargetLibraryInfo &TLI,
                                 const LoopInfo &LI) {
   bool MadeChange = false;
 
+  MSSA.ensureOptimizedUses();
   DSEState State(F, AA, MSSA, DT, PDT, TLI, LI);
   // For each store:
   for (unsigned I = 0; I < State.MemDefs.size(); I++) {
     MemoryDef *KillingDef = State.MemDefs[I];
     if (State.SkipStores.count(KillingDef))
       continue;
     Instruction *KillingI = KillingDef->getMemoryInst();
 

llvm/lib/Transforms/Scalar/EarlyCSE.cpp

public:
   /// This is the current generation of the memory value.
   unsigned CurrentGeneration = 0;
 
   /// Set up the EarlyCSE runner for a particular function.
   EarlyCSE(const DataLayout &DL, const TargetLibraryInfo &TLI,
            const TargetTransformInfo &TTI, DominatorTree &DT,
            AssumptionCache &AC, MemorySSA *MSSA)
       : TLI(TLI), TTI(TTI), DT(DT), AC(AC), SQ(DL, &TLI, &DT, &AC), MSSA(MSSA),
-        MSSAUpdater(std::make_unique<MemorySSAUpdater>(MSSA)) {}
+        MSSAUpdater(std::make_unique<MemorySSAUpdater>(MSSA)) {
+    if (MSSA)
+      MSSA->ensureOptimizedUses();
+  }
 
   bool run();
 
 private:
   unsigned ClobberCounter = 0;
   // Almost a POD, but needs to call the constructors for the scoped hash
   // tables so that a new scope gets pushed on. These are RAII so that the
   // scope gets popped when the NodeScope is destroyed.

llvm/lib/Transforms/Scalar/GVNHoist.cpp

 // This pass hoists common computations across branches sharing common
 // dominator. The primary goal is to reduce the code size, and in some
 // cases reduce critical path (by exposing more ILP).
 class GVNHoist {
 public:
   GVNHoist(DominatorTree *DT, PostDominatorTree *PDT, AliasAnalysis *AA,
            MemoryDependenceResults *MD, MemorySSA *MSSA)
       : DT(DT), PDT(PDT), AA(AA), MD(MD), MSSA(MSSA),
-        MSSAUpdater(std::make_unique<MemorySSAUpdater>(MSSA)) {}
+        MSSAUpdater(std::make_unique<MemorySSAUpdater>(MSSA)) {
+    MSSA->ensureOptimizedUses();
+  }
 
   bool run(Function &F);
 
   // Copied from NewGVN.cpp
   // This function provides global ranking of operations so that we can place
   // them in a canonical order.  Note that rank alone is not necessarily enough
   // for a complete ordering, as constants all have the same rank.  However,
   // generally, we will simplify an operation with all constants so that it

llvm/lib/Transforms/Scalar/LICM.cpp

 bool LoopInvariantCodeMotion::runOnLoop(
     Loop *L, AAResults *AA, LoopInfo *LI, DominatorTree *DT,
     BlockFrequencyInfo *BFI, TargetLibraryInfo *TLI, TargetTransformInfo *TTI,
     ScalarEvolution *SE, MemorySSA *MSSA, OptimizationRemarkEmitter *ORE,
     bool LoopNestMode) {
   bool Changed = false;
 
   assert(L->isLCSSAForm(*DT) && "Loop is not in LCSSA form.");
+  MSSA->ensureOptimizedUses();
 
   // If this loop has metadata indicating that LICM is not to be performed then
   // just exit.
   if (hasDisableLICMTransformsHint(L)) {
     return false;
   }
 
   // Don't sink stores from loops with coroutine suspend instructions.

llvm/unittests/Analysis/MemorySSATest.cpp

TEST_F(MemorySSATest, TestLoadMustAlias) {
   B.CreateStore(ConstantInt::get(Int8, 1), AllocaA);
   // Check load from store/def
   LoadInst *LA3 = B.CreateLoad(Int8, AllocaA, "");
   // Check load alias cached for second load
   LoadInst *LA4 = B.CreateLoad(Int8, AllocaA, "");
 
   setupAnalyses();
   MemorySSA &MSSA = *Analyses->MSSA;
+  MSSA.ensureOptimizedUses();
 
   unsigned I = 0;
   for (LoadInst *V : {LA1, LA2}) {
     MemoryUse *MemUse = dyn_cast_or_null<MemoryUse>(MSSA.getMemoryAccess(V));
     EXPECT_EQ(MemUse->getOptimizedAccessType(), None)
         << "Load " << I << " doesn't have the correct alias information";
     // EXPECT_EQ expands such that if we increment I above, it won't get
     // incremented except when we try to print the error message.
TEST_F(MemorySSATest, TestLoadMayAlias) {
   LoadInst *LB1 = B.CreateLoad(Int8, PointerB, "");
   B.CreateStore(ConstantInt::get(Int8, 0), PointerA);
   LoadInst *LA2 = B.CreateLoad(Int8, PointerA, "");
   B.CreateStore(ConstantInt::get(Int8, 0), PointerB);
   LoadInst *LB2 = B.CreateLoad(Int8, PointerB, "");
 
   setupAnalyses();
   MemorySSA &MSSA = *Analyses->MSSA;
+  MSSA.ensureOptimizedUses();
 
   unsigned I = 0;
   for (LoadInst *V : {LA1, LB1}) {
     MemoryUse *MemUse = dyn_cast_or_null<MemoryUse>(MSSA.getMemoryAccess(V));
     EXPECT_EQ(MemUse->getOptimizedAccessType().getValue(),
               AliasResult::MayAlias)
         << "Load " << I << " doesn't have the correct alias information";
     // EXPECT_EQ expands such that if we increment I above, it won't get