LICM: Hoist LOAD without STORE

[djtodoro]

Summary

When doing load/store promotion within LICM, if we cannot prove that it is safe to sink the store we won't hoist the load, even though we can prove the load could be dereferenced and moved outside the loop. This patch implements the load promotion by moving it in the loop preheader by inserting proper PHI in the loop. The store is kept as is in the loop. By doing this, we avoid doing the load from a memory location in each iteration.

Please consider this small example:

loop {
  var = *ptr;
  if (var) break;
  *ptr= var + 1;
}

After this patch, it will be:

var0 = *ptr;
loop {
  var1 = phi (var0, var2);
  if (var1) break;
  var2 = var1 + 1;
  *ptr = var2;
}

This addresses some problems from [0].

[0] https://bugs.llvm.org/show_bug.cgi?id=51193

[lebedev.ri]

This seems to be in-line with our recent discussions about whether on not LICM is a general canonicalization pass,
or should it have target-specific configurations. @reames thoughts? (to be noted, this more applies to D113290, if it was done without a target hook it would not be a question)

[efriedma]

llvm/lib/Transforms/Scalar/LICM.cpp

2012

Maybe move this explanation to the definition of the variable "ShouldTryLoadHoistingOnly", since that's where it's relevant.

2313

Not sure I understand what you're doing with SavedInsts here.

If I'm following correctly, SavedInsts is empty the first time this is called from LoopInvariantCodeMotion::runOnLoop. On subsequent iterations, it contains the stores that correspond to loads that were optimized.

If you're trying to prevent an infinite loop, this isn't the way to do it; promoteLoopAccessesToScalars should just return false if there aren't any loads to promote. If you're just trying to save a little work in subsequent iterations, please leave it out for now; if there's actually some significant compile-time gain here, we can look into it as a followup.

llvm/test/Transforms/LICM/AArch64/hoist-load-without-store.ll

20

I'd like to see check lines for the PHI nodes.

Do we want any other tests here? Doesn't feel right to have just one test. I guess the cases where this transforms fails are generally covered by other tests, though, so not sure what we'd need to cover.

[lebedev.ri]

llvm/lib/Transforms/Scalar/LICM.cpp

120–125

This is really ambiguous to me. Is this talking about
*only* about hoisting loads, and preventing hoisting stores,
or hoisting loads even if we can't hoist the store?

[reames]

This needs much clearer motivation - specifically a clear explanation of what you're trying to achieve and why. Please consider that motivation blocking for further review.

[lebedev.ri]

In D113289#3135476, @reames wrote:

This needs much clearer motivation - specifically a clear explanation of what you're trying to achieve and why. Please consider that motivation blocking for further review.

FWIW, i think the general idea behind this is *really* important, "partial" mem2reg is still a win ; D113520 is somewhat related.

[djtodoro]

Thanks a lot for your comments!

In D113289#3135476, @reames wrote:

This needs much clearer motivation - specifically a clear explanation of what you're trying to achieve and why. Please consider that motivation blocking for further review.

I've updated the summary as a first step.

llvm/lib/Transforms/Scalar/LICM.cpp

120–125

Oh, I agree...

2012

Makes sense.

2313

Hmm, actually it doesn't improve anything, we should simply return false if there is no load to promote.

llvm/test/Transforms/LICM/AArch64/hoist-load-without-store.ll

20

The more tests the better. I'll try to make some other tests as well. As a starting point, I'll update the CHECKS for this one.

[djtodoro]

• Addressing comments
• Improving && reducing the test

[efriedma]

The new summary describing the optimization is a bit confusing to try to parse. I'd suggest splitting into three paragraphs:

1. What LLVM currently does on the testcase.
2. How the current patch changes the code generation, and why it's an improvement.
3. Potential future directions (fallow-store-data-races etc.)

[reames]

Ok, let me describe back to you what you're trying to do. Give a case like:

loop {
  a = *g;
  if (a) break;
  *g = a + 1
}

Where g is some loop invariant address, you want to produce a form which looks like this:

a0 = *g
loop {
  a1 = phi (a0, a2)
  if (a1) break;
  a2 = a1 + 1
  *g = a2;
}

Assuming I've correctly described that, I agree this is valuable. I'm honestly a bit shocked we didn't already handle it. However, this should absolutely not be target dependent. It should simply be always enabled.

llvm/lib/Transforms/Scalar/LICM.cpp

1872

Please name this CanInsertStoresInExitBlocks and invert the code appropriately.

2136

Please move this above the if block. Next to the SawAtomic/SawNotAtomic is a good place.

2231–2232

This block simplifies with the removal of the flag.

llvm/lib/Transforms/Utils/SSAUpdater.cpp

446

I see what you're doing here, but this is really ugly. I think we need to separate the set of uses being visited from the set of uses being *deleted*.

Please add a "shouldDelete(Inst)" callback on the LoadAndStorePromoter, default it to answering yes, then implement that in the sub-class. Still not great, but at least a bit better.

llvm/test/Transforms/LICM/AArch64/hoist-load-without-store.ll

32

Please use the simplest test possible. An IR form of my example from the review is probably good. You must strip all irrelevant IR details (e.g. metadata arguments, globals, etc..).

[djtodoro]

In D113289#3147618, @reames wrote:

Ok, let me describe back to you what you're trying to do. Give a case like:

loop {
  a = *g;
  if (a) break;
  *g = a + 1
}

Where g is some loop invariant address, you want to produce a form which looks like this:

a0 = *g
loop {
  a1 = phi (a0, a2)
  if (a1) break;
  a2 = a1 + 1
  *g = a2;
}

Yep, that's it!

Assuming I've correctly described that, I agree this is valuable. I'm honestly a bit shocked we didn't already handle it. However, this should absolutely not be target dependent. It should simply be always enabled.

Thanks! I totally agree, so I removed the flag and abandoned the second patch from the stack that was target-dependent.

[djtodoro]

• removed the option (so this is enabled always)
• reduced the test
• refactored the code a bit

[reames]

Heading in the right direction.

In addition to the inline comments, please update and simplify the submission comment. The comment should not explain the motivating workload (except maybe as a link to a bug), and should instead focus on explaining the effect of the transform. i.e. use a more minimal example, and explain *why* we're implementing the transform.

llvm/lib/Transforms/Scalar/LICM.cpp

1872

Comment is stale.

1963

The whole preserveduses mechanism seems like overkill. The preserveuses set should simply be every store in the use set when CanInstertStoreInExistingBlocks is false. If it's not, I'm missing something and we probably need some comments/tests.

2234

You've already proven dereferenceability (i.e. load legality) above. This check should be:
if (!SafeToInsertStore && !FoundLoadToPromote)

return false;

llvm/lib/Transforms/Utils/SSAUpdater.cpp

453

Drop isa<StoreInst> from this check. The shouldDelete interface should handle this.

[djtodoro]

In D113289#3152152, @reames wrote:

Heading in the right direction.

In addition to the inline comments, please update and simplify the submission comment. The comment should not explain the motivating workload (except maybe as a link to a bug), and should instead focus on explaining the effect of the transform. i.e. use a more minimal example, and explain *why* we're implementing the transform.

It makes sense to me.

llvm/lib/Transforms/Scalar/LICM.cpp

1963

Yes... I've had some tests failing due to some other reason, but this shouldn't have been included. Thanks.

[djtodoro]

• Simplify the code a bit
• Address comments

[reames]

LGTM w/required changes before submission.

llvm/include/llvm/Transforms/Utils/SSAUpdater.h

173

Make this comment generic.

e.g. Will return true if a sub-class wants to keep one of the loads or stores after SSA construction.

llvm/lib/Transforms/Scalar/LICM.cpp

478–479

Whitespace only change - please remove.

llvm/test/Transforms/LICM/scalar-promote.ll

1 ↗

(On Diff #389963)

Please autogen this in a separate commit, then rebase your change over it so that diff is visible in commit.

[djtodoro]

@reames Thanks a lot for the comments!

llvm/include/llvm/Transforms/Utils/SSAUpdater.h

173

Actually opposite, but it makes sense.

llvm/test/Transforms/LICM/scalar-promote.ll

1 ↗

(On Diff #389963)

Sure. @nikic already addressed that with eee035235eb.

[djtodoro]

• Update the llvm/test/Transforms/LICM/scalar-promote-opaque-ptrs.ll

@lebedev.ri is this ok to go as is?

[lebedev.ri]

llvm/test/Transforms/LICM/hoist-load-without-store.ll

1 ↗

(On Diff #391009)

I mean, can you just commit this test file first, with proper check lines,
so that they are only changed by the following patch?

[djtodoro]

llvm/test/Transforms/LICM/hoist-load-without-store.ll

1 ↗

(On Diff #391009)

Got it! Sure. I'll do that.

I thought you wanted to double check the new changes in the failing tests.

[nikic]

llvm/test/Transforms/LICM/scalar-promote.ll

312 ↗

(On Diff #391009)

Comment should probably be updated.

[djtodoro]

llvm/test/Transforms/LICM/scalar-promote.ll

312 ↗

(On Diff #391009)

I'll do that before commiting. Thanks!

[uabelho]

I wrote an issue about LICM behaving differently with/without debug information with this patch:
https://github.com/llvm/llvm-project/issues/55915

llvm/include/llvm/Transforms/Utils/LoopUtils.h

 /// LoopInfo, DominatorTree, Loop, AliasSet information for all instructions
 /// of the loop and loop safety information as arguments.
 /// Diagnostics is emitted via \p ORE. It returns changed status.
 bool promoteLoopAccessesToScalars(
     const SmallSetVector<Value *, 8> &, SmallVectorImpl<BasicBlock *> &,
     SmallVectorImpl<Instruction *> &, SmallVectorImpl<MemoryAccess *> &,
     PredIteratorCache &, LoopInfo *, DominatorTree *, const TargetLibraryInfo *,
     Loop *, MemorySSAUpdater *, ICFLoopSafetyInfo *,
-    OptimizationRemarkEmitter *);
+    OptimizationRemarkEmitter *, SmallPtrSetImpl<Instruction *> &);
 
 /// Does a BFS from a given node to all of its children inside a given loop.
 /// The returned vector of nodes includes the starting point.
 SmallVector<DomTreeNode *, 16> collectChildrenInLoop(DomTreeNode *N,
                                                      const Loop *CurLoop);
 
 /// Returns the instructions that use values defined in the loop.
 SmallVector<Instruction *, 8> findDefsUsedOutsideOfLoop(Loop *L);

llvm/include/llvm/Transforms/Utils/SSAUpdater.h

public:
   /// a load is RAUW'd another value.
   virtual void replaceLoadWithValue(LoadInst *LI, Value *V) const {}
 
   /// Called before each instruction is deleted.
   virtual void instructionDeleted(Instruction *I) const {}
 
   /// Called to update debug info associated with the instruction.
   virtual void updateDebugInfo(Instruction *I) const {}
+
+  /// Used to indicate if we want to keep stores by hoisting loads only.
+  virtual bool shouldPreserveStore() const { return false; }
+
+  /// Use preserved.
+  virtual void instructionPreserved(Instruction *I) const {}
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_SSAUPDATER_H

llvm/lib/Transforms/Scalar/LICM.cpp

static cl::opt<unsigned> HoistSinkColdnessThreshold(
     cl::desc("Relative coldness Threshold of hoisting/sinking destination "
              "block for LICM to be considered beneficial"));
 
 static cl::opt<uint32_t> MaxNumUsesTraversed(
     "licm-max-num-uses-traversed", cl::Hidden, cl::init(8),
     cl::desc("Max num uses visited for identifying load "
              "invariance in loop using invariant start (default = 8)"));
 
+// Hoist the load without hoisting the store. Some targets (such as aarch64)
+// will gain benefits out of this, since there will be better register
+// promotion.
+static cl::opt<bool> EnableLoadHoistOnly(
+    "licm-load-hoisting-only", cl::Hidden, cl::init(false),
+    cl::desc("Enable hoisting loads without hoisting the stores"));
+
 // Experimental option to allow imprecision in LICM in pathological cases, in
 // exchange for faster compile. This is to be removed if MemorySSA starts to
 // address the same issue. This flag applies only when LICM uses MemorySSA
 // instead on AliasSetTracker. LICM calls MemorySSAWalker's
 // getClobberingMemoryAccess, up to the value of the Cap, getting perfect
 // accuracy. Afterwards, LICM will call into MemorySSA's getDefiningAccess,
 // which may not be precise, since optimizeUses is capped. The result is
 // correct, but we may not get as "far up" as possible to get which access is

 
 static void moveInstructionBefore(Instruction &I, Instruction &Dest,
                                   ICFLoopSafetyInfo &SafetyInfo,
                                   MemorySSAUpdater *MSSAU, ScalarEvolution *SE);
 
 static void foreachMemoryAccess(MemorySSA *MSSA, Loop *L,
                                 function_ref<void(Instruction *)> Fn);
 static SmallVector<SmallSetVector<Value *, 8>, 0>
-collectPromotionCandidates(MemorySSA *MSSA, AliasAnalysis *AA, Loop *L);
+collectPromotionCandidates(MemorySSA *MSSA, AliasAnalysis *AA, Loop *L,
+                           SmallPtrSetImpl<Instruction *> &SavedInsts);
 
 namespace {
 struct LoopInvariantCodeMotion {
   bool runOnLoop(Loop *L, AAResults *AA, LoopInfo *LI, DominatorTree *DT,
                  BlockFrequencyInfo *BFI, TargetLibraryInfo *TLI,
                  TargetTransformInfo *TTI, ScalarEvolution *SE, MemorySSA *MSSA,
                  OptimizationRemarkEmitter *ORE, bool LoopNestMode = false);
 
bool LoopInvariantCodeMotion::runOnLoop(
 
   // Get the preheader block to move instructions into...
   BasicBlock *Preheader = L->getLoopPreheader();
 
   // Compute loop safety information.
   ICFLoopSafetyInfo SafetyInfo;
   SafetyInfo.computeLoopSafetyInfo(L);
 
+  llvm::SmallPtrSet<Instruction *, 32> SavedInstructions;
+
   // We want to visit all of the instructions in this loop... that are not parts
   // of our subloops (they have already had their invariants hoisted out of
   // their loop, into this loop, so there is no need to process the BODIES of
   // the subloops).
   //
   // Traverse the body of the loop in depth first order on the dominator tree so
   // that we are guaranteed to see definitions before we see uses.  This allows
   // us to sink instructions in one pass, without iteration.  After sinking
if (!HasCatchSwitch) {
       // Promoting one set of accesses may make the pointers for another set
       // loop invariant, so run this in a loop (with the MaybePromotable set
       // decreasing in size over time).
       bool Promoted = false;
       bool LocalPromoted;
       do {
         LocalPromoted = false;
         for (const SmallSetVector<Value *, 8> &PointerMustAliases :
-             collectPromotionCandidates(MSSA, AA, L)) {
+             collectPromotionCandidates(MSSA, AA, L, SavedInstructions)) {
           LocalPromoted |= promoteLoopAccessesToScalars(
-              PointerMustAliases, ExitBlocks, InsertPts, MSSAInsertPts, PIC,
-              LI, DT, TLI, L, &MSSAU, &SafetyInfo, ORE);
+              PointerMustAliases, ExitBlocks, InsertPts, MSSAInsertPts, PIC, LI,
+              DT, TLI, L, &MSSAU, &SafetyInfo, ORE, SavedInstructions);
         }
         Promoted |= LocalPromoted;
       } while (LocalPromoted);
 
       // Once we have promoted values across the loop body we have to
       // recursively reform LCSSA as any nested loop may now have values defined
       // within the loop used in the outer loop.
       // FIXME: This is really heavy handed. It would be a bit better to use an
bool LoopInvariantCodeMotion::runOnLoop(
   assert((L->isOutermost() || L->getParentLoop()->isLCSSAForm(*DT)) &&
          "Parent loop not left in LCSSA form after LICM!");
 
   if (VerifyMemorySSA)
     MSSA->verifyMemorySSA();
 
   if (Changed && SE)
     SE->forgetLoopDispositions(L);
+
   return Changed;
 }
 
 /// Walk the specified region of the CFG (defined by all blocks dominated by
 /// the specified block, and that are in the current loop) in reverse depth
 /// first order w.r.t the DominatorTree.  This allows us to visit uses before
 /// definitions, allowing us to sink a loop body in one pass without iteration.
 ///
class LoopPromoter : public LoadAndStorePromoter {
   PredIteratorCache &PredCache;
   MemorySSAUpdater *MSSAU;
   LoopInfo &LI;
   DebugLoc DL;
   int Alignment;
   bool UnorderedAtomic;
   AAMDNodes AATags;
   ICFLoopSafetyInfo &SafetyInfo;
+  bool HoistLoadOnly;
+  SmallPtrSetImpl<Instruction *> &SavedInsts;
 
   // We're about to add a use of V in a loop exit block.  Insert an LCSSA phi
   // (if legal) if doing so would add an out-of-loop use to an instruction
   // defined in-loop.
   Value *maybeInsertLCSSAPHI(Value *V, BasicBlock *BB) const {
     if (!LI.wouldBeOutOfLoopUseRequiringLCSSA(V, BB))
       return V;
 

 public:
   LoopPromoter(Value *SP, ArrayRef<const Instruction *> Insts, SSAUpdater &S,
                const SmallSetVector<Value *, 8> &PMA,
                SmallVectorImpl<BasicBlock *> &LEB,
                SmallVectorImpl<Instruction *> &LIP,
                SmallVectorImpl<MemoryAccess *> &MSSAIP, PredIteratorCache &PIC,
                MemorySSAUpdater *MSSAU, LoopInfo &li, DebugLoc dl,
                int alignment, bool UnorderedAtomic, const AAMDNodes &AATags,
-               ICFLoopSafetyInfo &SafetyInfo)
+               ICFLoopSafetyInfo &SafetyInfo, bool HoistLoadOnly,
+               SmallPtrSetImpl<Instruction *> &SavedInsts)
       : LoadAndStorePromoter(Insts, S), SomePtr(SP), PointerMustAliases(PMA),
         LoopExitBlocks(LEB), LoopInsertPts(LIP), MSSAInsertPts(MSSAIP),
         PredCache(PIC), MSSAU(MSSAU), LI(li), DL(std::move(dl)),
         Alignment(alignment), UnorderedAtomic(UnorderedAtomic), AATags(AATags),
-        SafetyInfo(SafetyInfo) {}
+        SafetyInfo(SafetyInfo), HoistLoadOnly(HoistLoadOnly),
+        SavedInsts(SavedInsts) {}
 
   bool isInstInList(Instruction *I,
                     const SmallVectorImpl<Instruction *> &) const override {
     Value *Ptr;
     if (LoadInst *LI = dyn_cast<LoadInst>(I))
       Ptr = LI->getOperand(0);
     else
       Ptr = cast<StoreInst>(I)->getPointerOperand();
     return PointerMustAliases.count(Ptr);
   }
 
-  void doExtraRewritesBeforeFinalDeletion() override {
+  void insertStoresInLoopExitBlocks() {
     // Insert stores after in the loop exit blocks.  Each exit block gets a
     // store of the live-out values that feed them.  Since we've already told
     // the SSA updater about the defs in the loop and the preheader
     // definition, it is all set and we can start using it.
     for (unsigned i = 0, e = LoopExitBlocks.size(); i != e; ++i) {
       BasicBlock *ExitBlock = LoopExitBlocks[i];
       Value *LiveInValue = SSA.GetValueInMiddleOfBlock(ExitBlock);
       LiveInValue = maybeInsertLCSSAPHI(LiveInValue, ExitBlock);
for (unsigned i = 0, e = LoopExitBlocks.size(); i != e; ++i) {
             MSSAU->createMemoryAccessAfter(NewSI, nullptr, MSSAInsertPoint);
       }
       MSSAInsertPts[i] = NewMemAcc;
       MSSAU->insertDef(cast<MemoryDef>(NewMemAcc), true);
       // FIXME: true for safety, false may still be correct.
     }
   }
 
+  void doExtraRewritesBeforeFinalDeletion() override {
+    if (!HoistLoadOnly)
+      insertStoresInLoopExitBlocks();
+  }
+
   void instructionDeleted(Instruction *I) const override {
     SafetyInfo.removeInstruction(I);
     MSSAU->removeMemoryAccess(I);
   }
+
+  void instructionPreserved(Instruction *I) const override {
+    SavedInsts.insert(I);
+  }
+
+  bool shouldPreserveStore() const override { return HoistLoadOnly; }
 };
 
 bool isNotCapturedBeforeOrInLoop(const Value *V, const Loop *L,
                                  DominatorTree *DT) {
   // We can perform the captured-before check against any instruction in the
   // loop header, as the loop header is reachable from any instruction inside
   // the loop.
   // TODO: ReturnCaptures=true shouldn't be necessary here.

 }
 
 } // namespace
 
 /// Try to promote memory values to scalars by sinking stores out of the
 /// loop and moving loads to before the loop.  We do this by looping over
 /// the stores in the loop, looking for stores to Must pointers which are
 /// loop invariant.
+/// There is a special case where we move the load before the loop by
+/// preserving the stores as is, since it may improve register promotion
+/// for some targets. This is controled by the EnableLoadHoistOnly flag.
 ///
 bool llvm::promoteLoopAccessesToScalars(
     const SmallSetVector<Value *, 8> &PointerMustAliases,
     SmallVectorImpl<BasicBlock *> &ExitBlocks,
     SmallVectorImpl<Instruction *> &InsertPts,
     SmallVectorImpl<MemoryAccess *> &MSSAInsertPts, PredIteratorCache &PIC,
     LoopInfo *LI, DominatorTree *DT, const TargetLibraryInfo *TLI,
     Loop *CurLoop, MemorySSAUpdater *MSSAU, ICFLoopSafetyInfo *SafetyInfo,
-    OptimizationRemarkEmitter *ORE) {
+    OptimizationRemarkEmitter *ORE,
+    SmallPtrSetImpl<Instruction *> &SavedInsts) {
   // Verify inputs.
   assert(LI != nullptr && DT != nullptr && CurLoop != nullptr &&
          SafetyInfo != nullptr &&
          "Unexpected Input to promoteLoopAccessesToScalars");
 
   Value *SomePtr = *PointerMustAliases.begin();
   BasicBlock *Preheader = CurLoop->getLoopPreheader();
 
for (User *U : ASIV->users()) {
         // alignment as well.
         if (!DereferenceableInPH || (InstAlignment > Alignment))
           if (isSafeToExecuteUnconditionally(*Load, DT, TLI, CurLoop,
                                              SafetyInfo, ORE,
                                              Preheader->getTerminator())) {
             DereferenceableInPH = true;
             Alignment = std::max(Alignment, InstAlignment);
           }
       } else if (const StoreInst *Store = dyn_cast<StoreInst>(UI)) {
         // Stores *of* the pointer are not interesting, only stores *to* the
         // pointer.
         if (UI->getOperand(1) != ASIV)
           continue;
         if (!Store->isUnordered())
           return false;
 
         SawUnorderedAtomic |= Store->isAtomic();
if (!SafeToInsertStore) {
     else {
       Value *Object = getUnderlyingObject(SomePtr);
       SafeToInsertStore =
           (isAllocLikeFn(Object, TLI) || isa<AllocaInst>(Object)) &&
           isNotCapturedBeforeOrInLoop(Object, CurLoop, DT);
     }
   }
 
+  bool ShouldTryLoadHoistingOnly = false;
   // If we've still failed to prove we can sink the store, give up.
-  if (!SafeToInsertStore)
+  if (!SafeToInsertStore) {
+    if (DereferenceableInPH && EnableLoadHoistOnly)
+      ShouldTryLoadHoistingOnly = true;
+    else
       return false;
+  }
 
   // Otherwise, this is safe to promote, lets do it!
   LLVM_DEBUG(dbgs() << "LICM: Promoting value stored to in loop: " << *SomePtr
                     << '\n');
   ORE->emit([&]() {
     return OptimizationRemark(DEBUG_TYPE, "PromoteLoopAccessesToScalar",
                               LoopUses[0])
            << "Moving accesses to memory location out of the loop";
   });
   ++NumPromoted;
 
   // Look at all the loop uses, and try to merge their locations.
   std::vector<const DILocation *> LoopUsesLocs;
   for (auto U : LoopUses)
     LoopUsesLocs.push_back(U->getDebugLoc().get());
   auto DL = DebugLoc(DILocation::getMergedLocations(LoopUsesLocs));
 
   // We use the SSAUpdater interface to insert phi nodes as required.
   SmallVector<PHINode *, 16> NewPHIs;
   SSAUpdater SSA(&NewPHIs);
   LoopPromoter Promoter(SomePtr, LoopUses, SSA, PointerMustAliases, ExitBlocks,
                         InsertPts, MSSAInsertPts, PIC, MSSAU, *LI, DL,
                         Alignment.value(), SawUnorderedAtomic, AATags,
-                        *SafetyInfo);
+                        *SafetyInfo, ShouldTryLoadHoistingOnly, SavedInsts);
 
   // Set up the preheader to have a definition of the value.  It is the live-out
   // value from the preheader that uses in the loop will use.
   LoadInst *PreheaderLoad = new LoadInst(
       SomePtr->getType()->getPointerElementType(), SomePtr,
       SomePtr->getName() + ".promoted", Preheader->getTerminator());
   if (SawUnorderedAtomic)
     PreheaderLoad->setOrdering(AtomicOrdering::Unordered);
static void foreachMemoryAccess(MemorySSA *MSSA, Loop *L,
   for (const BasicBlock *BB : L->blocks())
     if (const auto *Accesses = MSSA->getBlockAccesses(BB))
       for (const auto &Access : *Accesses)
         if (const auto *MUD = dyn_cast<MemoryUseOrDef>(&Access))
           Fn(MUD->getMemoryInst());
 }
 
 static SmallVector<SmallSetVector<Value *, 8>, 0>
-collectPromotionCandidates(MemorySSA *MSSA, AliasAnalysis *AA, Loop *L) {
+collectPromotionCandidates(MemorySSA *MSSA, AliasAnalysis *AA, Loop *L,
+                           SmallPtrSetImpl<Instruction *> &SavedInsts) {
   AliasSetTracker AST(*AA);
 
-  auto IsPotentiallyPromotable = [L](const Instruction *I) {
-    if (const auto *SI = dyn_cast<StoreInst>(I))
+  auto IsPotentiallyPromotable = [L, &SavedInsts](const Instruction *I) {
+    if (const auto *SI = dyn_cast<StoreInst>(I)) {
+      if (SavedInsts.count(I))
+        return false;
       return L->isLoopInvariant(SI->getPointerOperand());
+    }
     if (const auto *LI = dyn_cast<LoadInst>(I))
       return L->isLoopInvariant(LI->getPointerOperand());
     return false;
   };
 
   // Populate AST with potentially promotable accesses and remove them from
   // MaybePromotable, so they will not be checked again on the next iteration.
   SmallPtrSet<Value *, 16> AttemptingPromotion;

llvm/lib/Transforms/Utils/SSAUpdater.cpp

for (LoadInst *ALoad : LiveInLoads) {
     if (NewVal == ALoad) NewVal = UndefValue::get(NewVal->getType());
     ALoad->replaceAllUsesWith(NewVal);
     ReplacedLoads[ALoad] = NewVal;
   }
 
   // Allow the client to do stuff before we start nuking things.
   doExtraRewritesBeforeFinalDeletion();
 
+  // Hoist the loads only if we didn't manage to sink the stores safely.
+  bool PreserveTheStores = shouldPreserveStore();
+
   // Now that everything is rewritten, delete the old instructions from the
   // function.  They should all be dead now.
   for (Instruction *User : Insts) {
+    // Keep the store since the load is moved above only.
+    if (PreserveTheStores && isa<StoreInst>(User)) {
+      // We pretend as if this was deleted, since we don't want
+      // to track this mem access here any more.
+      instructionPreserved(User);
+      continue;
+    }
+
     // If this is a load that still has uses, then the load must have been added
     // as a live value in the SSAUpdate data structure for a block (e.g. because
     // the loaded value was stored later).  In this case, we need to recursively
     // propagate the updates until we get to the real value.
     if (!User->use_empty()) {
       Value *NewVal = ReplacedLoads[User];
       assert(NewVal && "not a replaced load?");
 

llvm/test/Transforms/LICM/AArch64/hoist-load-without-store.ll

+;; C reproducer:
+;; int u, v;
+;;
+;; void f(int a[restrict], int b[restrict], int n) {
+;;     for (int i = 0; i < n; ++i) {
+;;         if (a[i]) {
+;;             ++u;
+;;             break;
+;;         }
+;;         ++u;
+;;         if (b[i])
+;;             ++v;
+;;     }
+;; }
+
+; RUN: opt -licm -licm-load-hoisting-only -mtriple aarch64-linux-gnu -S < %s | FileCheck %s
+
+; CHECK-LABEL: for.body.preheader:
+; CHECK: %u.promoted = load i32, i32* @u
+; CHECK: %v.promoted = load i32, i32* @v
+
+; ModuleID = './test.ll'
+source_filename = "test.c"
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux"
+
+@u = dso_local local_unnamed_addr global i32 0, align 4
+@v = dso_local local_unnamed_addr global i32 0, align 4
+
+; Function Attrs: nofree norecurse nosync nounwind uwtable
+define dso_local void @f(i32* noalias nocapture readonly %a, i32* noalias nocapture readonly %b, i32 %n) local_unnamed_addr #0 {
+entry:
+  %cmp5 = icmp sgt i32 %n, 0
+  br i1 %cmp5, label %for.body.preheader, label %cleanup
+
+for.body.preheader:                               ; preds = %entry
+  %wide.trip.count = zext i32 %n to i64
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.inc
+  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.inc ]
+  %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4, !tbaa !8
+  %tobool.not = icmp eq i32 %0, 0
+  %1 = load i32, i32* @u, align 4, !tbaa !8
+  %inc1 = add nsw i32 %1, 1
+  store i32 %inc1, i32* @u, align 4, !tbaa !8
+  br i1 %tobool.not, label %if.end, label %cleanup.loopexit
+
+if.end:                                           ; preds = %for.body
+  %arrayidx3 = getelementptr inbounds i32, i32* %b, i64 %indvars.iv
+  %2 = load i32, i32* %arrayidx3, align 4, !tbaa !8
+  %tobool4.not = icmp eq i32 %2, 0
+  br i1 %tobool4.not, label %for.inc, label %if.then5
+
+if.then5:                                         ; preds = %if.end
+  %3 = load i32, i32* @v, align 4, !tbaa !8
+  %inc6 = add nsw i32 %3, 1
+  store i32 %inc6, i32* @v, align 4, !tbaa !8
+  br label %for.inc
+
+for.inc:                                          ; preds = %if.end, %if.then5
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %cleanup.loopexit, label %for.body, !llvm.loop !12
+
+cleanup.loopexit:                                 ; preds = %for.body, %for.inc
+  br label %cleanup
+
+cleanup:                                          ; preds = %cleanup.loopexit, %entry
+  ret void
+}
+
+attributes #0 = { nofree norecurse nosync nounwind uwtable "frame-pointer"="non-leaf" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="generic" "target-features"="+neon" }
+
+!llvm.module.flags = !{!0, !1, !2, !3, !4, !5, !6}
+!llvm.ident = !{!7}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 1, !"branch-target-enforcement", i32 0}
+!2 = !{i32 1, !"sign-return-address", i32 0}
+!3 = !{i32 1, !"sign-return-address-all", i32 0}
+!4 = !{i32 1, !"sign-return-address-with-bkey", i32 0}
+!5 = !{i32 7, !"uwtable", i32 1}
+!6 = !{i32 7, !"frame-pointer", i32 1}
+!7 = !{!"clang version 14.0.0"}
+!8 = !{!9, !9, i64 0}
+!9 = !{!"int", !10, i64 0}
+!10 = !{!"omnipotent char", !11, i64 0}
+!11 = !{!"Simple C/C++ TBAA"}
+!12 = distinct !{!12, !13}
+!13 = !{!"llvm.loop.mustprogress"}
 No newline at end of file