[NFCI] [LoopIdiom] Let processLoopStridedStore take StoreSize as SCEV instead of unsigned

[eopXD]

Summary

Letting it take SCEV allows further modification on the function to optimize
if the StoreSize / Stride is runtime determined.

This is a preceeding of D107353.
The big picture is to let LoopIdiom deal with runtime-determined sizes.

[Whitney]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

952

Do you think is beneficial to make this a forwarder in the meantime? like

const SCEV *StoreSizeSCEV = SE->getSCEV(StoreSize);
mayLoopAccessLocation(Ptr, Access, L, BECount, StoreSizeSCEV, AA, IgnoredStores);

982

If you plan to fully replace the original version, can you copy all the comments as well.

1068

Initialize to nullptr.

[eopXD]

Address comments.

[eopXD]

Closing this revision because D107353 is addressing these changes and additional ones that indicates its usage - Let the LoopIdiomRecognizePass deal with runtime-determined memset size.

[eopXD]

Reopen this patch since a NFC can be done prior to the main change (D107353) on LoopIdiom.

[eopXD]

Rebase patch to latest main HEAD.

[eopXD]

Fix clang-format warnings.

[eopXD]

Fix more clang-format warnings.

[Whitney]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

916

Is it better to have line 916-918 after 921-923?

In what case, !MemsetSizeSCEV? Do we early exit more often than before?

982

comments start with uppercase.

1142

Is this really NFC?

[eopXD]

Address comments.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

916

The variable is not used other than forwarding the store size of the memset as SCEV into processLoopStridedStore. So I deleted this variable. Also the if-statement will never be triggered and is misleading.

1142

Right here the code does the exact same thing as above except no conversion on StoreSize is needed. (since its now a SCEV).

Created another getNumBytes function that takes StoreSize as SCEV, which better shows that this is a NFC.

[Whitney]

Please make sure clang-format is good for the new lines you changed.
Otherwise, LGTM.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

941

probably need to fix clang-format here.

[lebedev.ri]

This does seem like a NFC refactor.
Thanks.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

984

I think this should just be inlined into users.

1001–1002

1012

Same, i think this should just be inlined into users.

[eopXD]

Address comments.

[eopXD]

Delete redundent brackets added.

[eopXD]

@lebedev.ri
If no further modification is required, please also accept this patch so that I can land it and move on to D107353.

[lebedev.ri]

This looks fine i guess.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

private:
   template <typename MemInst>
   bool processLoopMemIntrinsic(
       BasicBlock *BB,
       bool (LoopIdiomRecognize::*Processor)(MemInst *, const SCEV *),
       const SCEV *BECount);
   bool processLoopMemCpy(MemCpyInst *MCI, const SCEV *BECount);
   bool processLoopMemSet(MemSetInst *MSI, const SCEV *BECount);
 
-  bool processLoopStridedStore(Value *DestPtr, unsigned StoreSize,
+  bool processLoopStridedStore(Value *DestPtr, const SCEV *StoreSizeSCEV,
                                MaybeAlign StoreAlignment, Value *StoredVal,
                                Instruction *TheStore,
                                SmallPtrSetImpl<Instruction *> &Stores,
                                const SCEVAddRecExpr *Ev, const SCEV *BECount,
                                bool NegStride, bool IsLoopMemset = false);
   bool processLoopStoreOfLoopLoad(StoreInst *SI, const SCEV *BECount);
   bool processLoopStoreOfLoopLoad(Value *DestPtr, Value *SourcePtr,
                                   unsigned StoreSize, MaybeAlign StoreAlign,
for (SetVector<StoreInst *>::iterator it = Heads.begin(), e = Heads.end();
 
     // Check to see if the stride matches the size of the stores.  If so, then
     // we know that every byte is touched in the loop.
     if (StoreSize != Stride && StoreSize != -Stride)
       continue;
 
     bool NegStride = StoreSize == -Stride;
 
-    if (processLoopStridedStore(StorePtr, StoreSize,
+    const SCEV *StoreSizeSCEV = SE->getConstant(BECount->getType(), StoreSize);
+    if (processLoopStridedStore(StorePtr, StoreSizeSCEV,
                                 MaybeAlign(HeadStore->getAlignment()),
                                 StoredVal, HeadStore, AdjacentStores, StoreEv,
                                 BECount, NegStride)) {
       TransformedStores.insert(AdjacentStores.begin(), AdjacentStores.end());
       Changed = true;
     }
   }
 
bool LoopIdiomRecognize::processLoopMemSet(MemSetInst *MSI,
 
   // See if the pointer expression is an AddRec like {base,+,1} on the current
   // loop, which indicates a strided store.  If we have something else, it's a
   // random store we can't handle.
   const SCEVAddRecExpr *Ev = dyn_cast<SCEVAddRecExpr>(SE->getSCEV(Pointer));
   if (!Ev || Ev->getLoop() != CurLoop || !Ev->isAffine())
     return false;
 
   // Reject memsets that are so large that they overflow an unsigned.
   uint64_t SizeInBytes = cast<ConstantInt>(MSI->getLength())->getZExtValue();
   if ((SizeInBytes >> 32) != 0)
     return false;
 
   // Check to see if the stride matches the size of the memset.  If so, then we
   // know that every byte is touched in the loop.
   const SCEVConstant *ConstStride = dyn_cast<SCEVConstant>(Ev->getOperand(1));
   if (!ConstStride)
     return false;
 
   APInt Stride = ConstStride->getAPInt();
   if (SizeInBytes != Stride && SizeInBytes != -Stride)
     return false;
 
   // Verify that the memset value is loop invariant.  If not, we can't promote
   // the memset.
   Value *SplatValue = MSI->getValue();
   if (!SplatValue || !CurLoop->isLoopInvariant(SplatValue))
     return false;
 
   SmallPtrSet<Instruction *, 1> MSIs;
   MSIs.insert(MSI);
   bool NegStride = SizeInBytes == -Stride;
   return processLoopStridedStore(
-      Pointer, (unsigned)SizeInBytes, MaybeAlign(MSI->getDestAlignment()),
-      SplatValue, MSI, MSIs, Ev, BECount, NegStride, /*IsLoopMemset=*/true);
+      Pointer, SE->getSCEV(MSI->getLength()), MaybeAlign(MSI->getDestAlignment()), SplatValue,
+      MSI, MSIs, Ev, BECount, NegStride, /*IsLoopMemset=*/true);
 }
 
 /// mayLoopAccessLocation - Return true if the specified loop might access the
 /// specified pointer location, which is a loop-strided access.  The 'Access'
 /// argument specifies what the verboten forms of access are (read or write).
 static bool
 mayLoopAccessLocation(Value *Ptr, ModRefInfo Access, Loop *L,
-                      const SCEV *BECount, unsigned StoreSize,
+                      const SCEV *BECount, const SCEV *StoreSizeSCEV,
                       AliasAnalysis &AA,
                       SmallPtrSetImpl<Instruction *> &IgnoredStores) {
   // Get the location that may be stored across the loop.  Since the access is
   // strided positively through memory, we say that the modified location starts
   // at the pointer and has infinite size.
   LocationSize AccessSize = LocationSize::afterPointer();
 
   // If the loop iterates a fixed number of times, we can refine the access size
   // to be exactly the size of the memset, which is (BECount+1)*StoreSize
-  if (const SCEVConstant *BECst = dyn_cast<SCEVConstant>(BECount))
+  const SCEVConstant *BECst = dyn_cast<SCEVConstant>(BECount);
+  const SCEVConstant *ConstSize = dyn_cast<SCEVConstant>(StoreSizeSCEV);
+  if (BECst && ConstSize)
     AccessSize = LocationSize::precise((BECst->getValue()->getZExtValue() + 1) *
-                                       StoreSize);
+                                       ConstSize->getValue()->getZExtValue());
 
   // TODO: For this to be really effective, we have to dive into the pointer
   // operand in the store.  Store to &A[i] of 100 will always return may alias
   // with store of &A[100], we need to StoreLoc to be "A" with size of 100,
   // which will then no-alias a store to &A[100].
   MemoryLocation StoreLoc(Ptr, AccessSize);
 
   for (Loop::block_iterator BI = L->block_begin(), E = L->block_end(); BI != E;
        ++BI)
     for (Instruction &I : **BI)
       if (IgnoredStores.count(&I) == 0 &&
           isModOrRefSet(
               intersectModRef(AA.getModRefInfo(&I, StoreLoc), Access)))
         return true;
-
   return false;
 }
 
+// Forwards StoreSize as SCEV, aim to replace usage of this prototype to
+// let LIR deal with runtime-determined store size.
+static bool mayLoopAccessLocation(Value *Ptr, ModRefInfo Access, Loop *L,
+                                  const SCEV *BECount, unsigned StoreSize,
+                                  AliasAnalysis &AA,
+                                  SmallPtrSetImpl<Instruction *> &IgnoredStores,
+                                  ScalarEvolution *SE) {
+  const SCEV *StoreSizeSCEV = SE->getConstant(BECount->getType(), StoreSize);
+  return mayLoopAccessLocation(Ptr, Access, L, BECount, StoreSizeSCEV, AA,
+                               IgnoredStores);
+}
+
 // If we have a negative stride, Start refers to the end of the memory location
 // we're trying to memset.  Therefore, we need to recompute the base pointer,
 // which is just Start - BECount*Size.
 static const SCEV *getStartForNegStride(const SCEV *Start, const SCEV *BECount,
-                                        Type *IntPtr, unsigned StoreSize,
+                                        Type *IntPtr, const SCEV *StoreSizeSCEV,
                                         ScalarEvolution *SE) {
   const SCEV *Index = SE->getTruncateOrZeroExtend(BECount, IntPtr);
-  if (StoreSize != 1)
-    Index = SE->getMulExpr(Index, SE->getConstant(IntPtr, StoreSize),
+  const SCEVConstant *ConstSize = dyn_cast<SCEVConstant>(StoreSizeSCEV);
+  if (!ConstSize || (ConstSize->getAPInt() != 1)) {
+    // index = back edge count * store size
+    Index = SE->getMulExpr(Index,
+                           SE->getTruncateOrZeroExtend(StoreSizeSCEV, IntPtr),
                            SCEV::FlagNUW);
+  }
+  // base pointer = start - index * store size
   return SE->getMinusSCEV(Start, Index);
 }
 
-/// Compute the number of bytes as a SCEV from the backedge taken count.
-///
-/// This also maps the SCEV into the provided type and tries to handle the
-/// computation in a way that will fold cleanly.
-static const SCEV *getNumBytes(const SCEV *BECount, Type *IntPtr,
-                               unsigned StoreSize, Loop *CurLoop,
-                               const DataLayout *DL, ScalarEvolution *SE) {
-  const SCEV *NumBytesS;
-  // The # stored bytes is (BECount+1)*Size.  Expand the trip count out to
+// Forwards StoreSize as SCEV, aim to replace usage of this prototype to
+// let LIR deal with runtime-determined store size.
+static const SCEV *getStartForNegStride(const SCEV *Start, const SCEV *BECount,
+                                        Type *IntPtr, unsigned StoreSize,
+                                        ScalarEvolution *SE) {
+  const SCEV *StoreSizeSCEV = SE->getConstant(BECount->getType(), StoreSize);
+  return getStartForNegStride(Start, BECount, IntPtr, StoreSizeSCEV, SE);
+}
+
+/// Compute trip count from the backedge taken count.
+static const SCEV *getTripCount(const SCEV *BECount, Type *IntPtr,
+                                Loop *CurLoop, const DataLayout *DL,
+                                ScalarEvolution *SE) {
+  const SCEV *TripCountS = nullptr;
+  // The # stored bytes is (BECount+1).  Expand the trip count out to
   // pointer size if it isn't already.
   //
   // If we're going to need to zero extend the BE count, check if we can add
   // one to it prior to zero extending without overflow. Provided this is safe,
   // it allows better simplification of the +1.
-  if (DL->getTypeSizeInBits(BECount->getType()).getFixedSize() <
-          DL->getTypeSizeInBits(IntPtr).getFixedSize() &&
+  if (DL->getTypeSizeInBits(BECount->getType()) <
+          DL->getTypeSizeInBits(IntPtr) &&
       SE->isLoopEntryGuardedByCond(
           CurLoop, ICmpInst::ICMP_NE, BECount,
           SE->getNegativeSCEV(SE->getOne(BECount->getType())))) {
-    NumBytesS = SE->getZeroExtendExpr(
+    TripCountS = SE->getZeroExtendExpr(
         SE->getAddExpr(BECount, SE->getOne(BECount->getType()), SCEV::FlagNUW),
         IntPtr);
   } else {
-    NumBytesS = SE->getAddExpr(SE->getTruncateOrZeroExtend(BECount, IntPtr),
+    TripCountS = SE->getAddExpr(SE->getTruncateOrZeroExtend(BECount, IntPtr),
                                 SE->getOne(IntPtr), SCEV::FlagNUW);
   }
 
+  return TripCountS;
+}
+
+/// Compute the number of bytes as a SCEV from the backedge taken count.
+///
+/// This also maps the SCEV into the provided type and tries to handle the
+/// computation in a way that will fold cleanly.
+static const SCEV *getNumBytes(const SCEV *BECount, Type *IntPtr,
+                               unsigned StoreSize, Loop *CurLoop,
+                               const DataLayout *DL, ScalarEvolution *SE) {
+  const SCEV *TripCountSCEV = getTripCount(BECount, IntPtr, CurLoop, DL, SE);
+
   // And scale it based on the store size.
   if (StoreSize != 1) {
-    NumBytesS = SE->getMulExpr(NumBytesS, SE->getConstant(IntPtr, StoreSize),
+    return SE->getMulExpr(TripCountSCEV, SE->getConstant(IntPtr, StoreSize),
                                SCEV::FlagNUW);
   }
-  return NumBytesS;
+  return TripCountSCEV;
+}
+
+/// getNumBytes that takes StoreSize as a SCEV
+static const SCEV *getNumBytes(const SCEV *BECount, Type *IntPtr,
+                               const SCEV *StoreSizeSCEV, Loop *CurLoop,
+                               const DataLayout *DL, ScalarEvolution *SE) {
+  const SCEV *TripCountSCEV = getTripCount(BECount, IntPtr, CurLoop, DL, SE);
+
+  return SE->getMulExpr(
+      TripCountSCEV, SE->getTruncateOrZeroExtend(StoreSizeSCEV, IntPtr),
+      SCEV::FlagNUW);
 }
 
+
 /// processLoopStridedStore - We see a strided store of some value.  If we can
 /// transform this into a memset or memset_pattern in the loop preheader, do so.
 bool LoopIdiomRecognize::processLoopStridedStore(
-    Value *DestPtr, unsigned StoreSize, MaybeAlign StoreAlignment,
+    Value *DestPtr, const SCEV *StoreSizeSCEV, MaybeAlign StoreAlignment,
     Value *StoredVal, Instruction *TheStore,
     SmallPtrSetImpl<Instruction *> &Stores, const SCEVAddRecExpr *Ev,
     const SCEV *BECount, bool NegStride, bool IsLoopMemset) {
   Value *SplatValue = isBytewiseValue(StoredVal, *DL);
   Constant *PatternValue = nullptr;
 
   if (!SplatValue)
     PatternValue = getMemSetPatternValue(StoredVal, DL);
bool LoopIdiomRecognize::processLoopStridedStore(
 
   Type *DestInt8PtrTy = Builder.getInt8PtrTy(DestAS);
   Type *IntIdxTy = DL->getIndexType(DestPtr->getType());
 
   bool Changed = false;
   const SCEV *Start = Ev->getStart();
   // Handle negative strided loops.
   if (NegStride)
-    Start = getStartForNegStride(Start, BECount, IntIdxTy, StoreSize, SE);
+    Start = getStartForNegStride(Start, BECount, IntIdxTy, StoreSizeSCEV, SE);
 
   // TODO: ideally we should still be able to generate memset if SCEV expander
   // is taught to generate the dependencies at the latest point.
   if (!isSafeToExpand(Start, *SE))
     return Changed;
 
   // Okay, we have a strided store "p[i]" of a splattable value.  We can turn
   // this into a memset in the loop preheader now if we want.  However, this
   // would be unsafe to do if there is anything else in the loop that may read
   // or write to the aliased location.  Check for any overlap by generating the
   // base pointer and checking the region.
   Value *BasePtr =
       Expander.expandCodeFor(Start, DestInt8PtrTy, Preheader->getTerminator());
 
   // From here on out, conservatively report to the pass manager that we've
   // changed the IR, even if we later clean up these added instructions. There
   // may be structural differences e.g. in the order of use lists not accounted
   // for in just a textual dump of the IR. This is written as a variable, even
   // though statically all the places this dominates could be replaced with
   // 'true', with the hope that anyone trying to be clever / "more precise" with
   // the return value will read this comment, and leave them alone.
   Changed = true;
 
   if (mayLoopAccessLocation(BasePtr, ModRefInfo::ModRef, CurLoop, BECount,
-                            StoreSize, *AA, Stores))
+                            StoreSizeSCEV, *AA, Stores))
     return Changed;
 
   if (avoidLIRForMultiBlockLoop(/*IsMemset=*/true, IsLoopMemset))
     return Changed;
 
   // Okay, everything looks good, insert the memset.
 
-  const SCEV *NumBytesS =
-      getNumBytes(BECount, IntIdxTy, StoreSize, CurLoop, DL, SE);
+  const SCEV *NumBytesS = getNumBytes(BECount, IntIdxTy,StoreSizeSCEV, CurLoop, DL, SE);
 
   // TODO: ideally we should still be able to generate memset if SCEV expander
   // is taught to generate the dependencies at the latest point.
   if (!isSafeToExpand(NumBytesS, *SE))
     return Changed;
 
   Value *NumBytes =
       Expander.expandCodeFor(NumBytesS, IntIdxTy, Preheader->getTerminator());
bool LoopIdiomRecognize::processLoopStoreOfLoopLoad(
   SmallPtrSet<Instruction *, 2> Stores;
   Stores.insert(TheStore);
 
   bool IsMemCpy = isa<MemCpyInst>(TheStore);
   const StringRef InstRemark = IsMemCpy ? "memcpy" : "load and store";
 
   bool UseMemMove =
       mayLoopAccessLocation(StoreBasePtr, ModRefInfo::ModRef, CurLoop, BECount,
-                            StoreSize, *AA, Stores);
+                            StoreSize, *AA, Stores, SE);
   if (UseMemMove) {
     Stores.insert(TheLoad);
     if (mayLoopAccessLocation(StoreBasePtr, ModRefInfo::ModRef, CurLoop,
-                              BECount, StoreSize, *AA, Stores)) {
+                              BECount, StoreSize, *AA, Stores, SE)) {
       ORE.emit([&]() {
         return OptimizationRemarkMissed(DEBUG_TYPE, "LoopMayAccessStore",
                                         TheStore)
                << ore::NV("Inst", InstRemark) << " in "
                << ore::NV("Function", TheStore->getFunction())
                << " function will not be hoisted: "
                << ore::NV("Reason", "The loop may access store location");
       });
bool LoopIdiomRecognize::processLoopStoreOfLoopLoad(
   Value *LoadBasePtr = Expander.expandCodeFor(
       LdStart, Builder.getInt8PtrTy(LdAS), Preheader->getTerminator());
 
   // If the store is a memcpy instruction, we must check if it will write to
   // the load memory locations. So remove it from the ignored stores.
   if (IsMemCpy)
     Stores.erase(TheStore);
   if (mayLoopAccessLocation(LoadBasePtr, ModRefInfo::Mod, CurLoop, BECount,
-                            StoreSize, *AA, Stores)) {
+                            StoreSize, *AA, Stores, SE)) {
     ORE.emit([&]() {
       return OptimizationRemarkMissed(DEBUG_TYPE, "LoopMayAccessLoad", TheLoad)
              << ore::NV("Inst", InstRemark) << " in "
              << ore::NV("Function", TheStore->getFunction())
              << " function will not be hoisted: "
              << ore::NV("Reason", "The loop may access load location");
     });
     return Changed;