[MemorySSA] Always perform MemoryUses liveOnEntry optimization on MSSA construction

[aeubanks]

Summary

Fixes invariant memory regressions in future DSE patches.

Also add a flag to print<memoryssa> to not ensure optimized uses to test this.

Noticeable compile time regression [1], but a future DSE change that depends on this more than makes up for it.

[1] https://llvm-compile-time-tracker.com/compare.php?from=9d5466849a770eeab222d5a5890376d3596e8ad6&to=95682dbe11d76a3342870437377216e96b167504&stat=instructions:u

[asbirlea]

llvm/lib/Analysis/MemorySSA.cpp

368

Can you check if the the callsite in optimizeUses can be removed with this change?
I think the one in getClobbering needs to stay.

llvm/lib/Passes/PassBuilder.cpp

1058

nit: Simpify the test option to "no-optimized-uses".

[aeubanks]

address comments
change MemorySSAPrinterPass option to not be negative

llvm/lib/Analysis/MemorySSA.cpp

368

removed, and yeah the one in getClobbering needs to stay when you invalidate an optimized access of a MemoryUse and then try and optimize it again

llvm/lib/Passes/PassBuilder.cpp

1058

that's can be misleading because if memoryssa was already optimized previously then it'll still be optimized

[nikic]

llvm/include/llvm/Analysis/MemorySSA.h

953

Mismatch in names

[aeubanks]

fix mismatched names

[asbirlea]

lgtm

llvm/include/llvm/Analysis/MemorySSA.h

public:
   };
 
   Result run(Function &F, FunctionAnalysisManager &AM);
 };
 
 /// Printer pass for \c MemorySSA.
 class MemorySSAPrinterPass : public PassInfoMixin<MemorySSAPrinterPass> {
   raw_ostream &OS;
+  bool EnsureOptimizedUses;
 
 public:
-  explicit MemorySSAPrinterPass(raw_ostream &OS) : OS(OS) {}
+  explicit MemorySSAPrinterPass(raw_ostream &OS, bool EnsureOptimizedUses)
+      : OS(OS), EnsureOptimizedUses(EnsureOptimizedUses) {}
 
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 /// Printer pass for \c MemorySSA via the walker.
 class MemorySSAWalkerPrinterPass
     : public PassInfoMixin<MemorySSAWalkerPrinterPass> {
   raw_ostream &OS;

llvm/lib/Analysis/MemorySSA.cpp

UpwardsMemoryQuery(const Instruction *Inst, const MemoryAccess *Access)
       : IsCall(isa<CallBase>(Inst)), Inst(Inst), OriginalAccess(Access) {
     if (!IsCall)
       StartingLoc = MemoryLocation::get(Inst);
   }
 };
 
 } // end anonymous namespace
 
-static bool isUseTriviallyOptimizableToLiveOnEntry(BatchAAResults &AA,
+template <typename AliasAnalysisType>
+static bool isUseTriviallyOptimizableToLiveOnEntry(AliasAnalysisType &AA,
                                                    const Instruction *I) {
   // If the memory can't be changed, then loads of the memory can't be
   // clobbered.
   if (auto *LI = dyn_cast<LoadInst>(I)) {
     return I->hasMetadata(LLVMContext::MD_invariant_load) ||
            !isModSet(AA.getModRefInfoMask(MemoryLocation::get(LI)));
   }
   return false;
if (!MU) {
       VersionStack.push_back(&MA);
       ++StackEpoch;
       continue;
     }
 
     if (MU->isOptimized())
       continue;
 
-    if (isUseTriviallyOptimizableToLiveOnEntry(*AA, MU->getMemoryInst())) {
-      MU->setDefiningAccess(MSSA->getLiveOnEntryDef(), true);
-      continue;
-    }
-
     MemoryLocOrCall UseMLOC(MU);
     auto &LocInfo = LocStackInfo[UseMLOC];
     // If the pop epoch changed, it means we've removed stuff from top of
     // stack due to changing blocks. We may have to reset the lower bound or
     // last kill info.
     if (LocInfo.PopEpoch != PopEpoch) {
       LocInfo.PopEpoch = PopEpoch;
       LocInfo.StackEpoch = StackEpoch;
#endif
   }
 
   // It's possible for an instruction to not modify memory at all. During
   // construction, we ignore them.
   if (!Def && !Use)
     return nullptr;
 
   MemoryUseOrDef *MUD;
-  if (Def)
+  if (Def) {
     MUD = new MemoryDef(I->getContext(), nullptr, I, I->getParent(), NextID++);
-  else
+  } else {
     MUD = new MemoryUse(I->getContext(), nullptr, I, I->getParent());
+    if (isUseTriviallyOptimizableToLiveOnEntry(*AAP, I)) {
+      MemoryAccess *LiveOnEntry = getLiveOnEntryDef();
+      MUD->setOptimized(LiveOnEntry);
+    }
+  }
   ValueToMemoryAccess[I] = MUD;
   return MUD;
 }
 
 /// Properly remove \p MA from all of MemorySSA's lookup tables.
 void MemorySSA::removeFromLookups(MemoryAccess *MA) {
   assert(MA->use_empty() &&
          "Trying to remove memory access that still has uses");
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();
+  if (EnsureOptimizedUses)
     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/Passes/PassBuilder.cpp

 }
 
 Expected<bool> parseSeparateConstOffsetFromGEPPassOptions(StringRef Params) {
   return parseSinglePassOption(Params, "lower-gep",
                                "SeparateConstOffsetFromGEP");
 }
 
 Expected<OptimizationLevel>
 parseFunctionSimplificationPipelineOptions(StringRef Params) {
   std::optional<OptimizationLevel> L = parseOptLevel(Params);
   if (!L || *L == OptimizationLevel::O0) {
     return make_error<StringError>(
         formatv("invalid function-simplification parameter '{0}' ", Params)
             .str(),
         inconvertibleErrorCode());
   };
   return *L;
 }
 
+Expected<bool> parseMemorySSAPrinterPassOptions(StringRef Params) {
+  return parseSinglePassOption(Params, "no-ensure-optimized-uses",
+                               "MemorySSAPrinterPass");
+}
+
 } // namespace
 
 /// Tests whether a pass name starts with a valid prefix for a default pipeline
 /// alias.
 static bool startsWithDefaultPipelineAliasPrefix(StringRef Name) {
   return Name.startswith("default") || Name.startswith("thinlto") ||
          Name.startswith("lto");
 }

llvm/lib/Passes/PassRegistry.def

 FUNCTION_PASS("print<delinearization>", DelinearizationPrinterPass(dbgs()))
 FUNCTION_PASS("print<demanded-bits>", DemandedBitsPrinterPass(dbgs()))
 FUNCTION_PASS("print<domfrontier>", DominanceFrontierPrinterPass(dbgs()))
 FUNCTION_PASS("print<func-properties>", FunctionPropertiesPrinterPass(dbgs()))
 FUNCTION_PASS("print<inline-cost>", InlineCostAnnotationPrinterPass(dbgs()))
 FUNCTION_PASS("print<inliner-size-estimator>",
   InlineSizeEstimatorAnalysisPrinterPass(dbgs()))
 FUNCTION_PASS("print<loops>", LoopPrinterPass(dbgs()))
-FUNCTION_PASS("print<memoryssa>", MemorySSAPrinterPass(dbgs()))
 FUNCTION_PASS("print<memoryssa-walker>", MemorySSAWalkerPrinterPass(dbgs()))
 FUNCTION_PASS("print<phi-values>", PhiValuesPrinterPass(dbgs()))
 FUNCTION_PASS("print<regions>", RegionInfoPrinterPass(dbgs()))
 FUNCTION_PASS("print<scalar-evolution>", ScalarEvolutionPrinterPass(dbgs()))
 FUNCTION_PASS("print<stack-safety-local>", StackSafetyPrinterPass(dbgs()))
 FUNCTION_PASS("print<access-info>", LoopAccessInfoPrinterPass(dbgs()))
 // TODO: rename to print<foo> after NPM switch
 FUNCTION_PASS("print-alias-sets", AliasSetsPrinterPass(dbgs()))
FUNCTION_PASS_WITH_PARAMS("separate-const-offset-from-gep",
                           "lower-gep")
 FUNCTION_PASS_WITH_PARAMS("function-simplification",
                           "",
                            [this](OptimizationLevel OL) {
                              return buildFunctionSimplificationPipeline(OL, ThinOrFullLTOPhase::None);
                            },
                           parseFunctionSimplificationPipelineOptions,
                           "O1;O2;O3;Os;Oz")
+FUNCTION_PASS_WITH_PARAMS("print<memoryssa>",
+                          "MemorySSAPrinterPass",
+                           [](bool NoEnsureOptimizedUses) {
+                             return MemorySSAPrinterPass(dbgs(), !NoEnsureOptimizedUses);
+                           },
+                          parseMemorySSAPrinterPassOptions,
+                          "no-ensure-optimized-uses")
 #undef FUNCTION_PASS_WITH_PARAMS
 
 #ifndef LOOPNEST_PASS
 #define LOOPNEST_PASS(NAME, CREATE_PASS)
 #endif
 LOOPNEST_PASS("loop-flatten", LoopFlattenPass())
 LOOPNEST_PASS("loop-interchange", LoopInterchangePass())
 LOOPNEST_PASS("loop-unroll-and-jam", LoopUnrollAndJamPass())

llvm/test/Analysis/MemorySSA/tbaa.ll

+; RUN: opt -aa-pipeline=basic-aa,tbaa -passes='print<memoryssa><no-ensure-optimized-uses>' -disable-output < %s 2>&1 | FileCheck %s
+
+define i8 @test1_yes(ptr %a, ptr %b) {
+; CHECK: 1 = MemoryDef(liveOnEntry)
+  store i8 0, ptr %a, align 1
+; CHECK: MemoryUse(liveOnEntry)
+  %y = load i8, ptr %b, align 1, !tbaa !0
+  ret i8 %y
+}
+
+!0 = !{!1, !1, i64 0, i1 true}
+!1 = !{!"qux", !2}
+!2 = !{}