[LoopIdiom] let the pass deal with runtime memset size

[eopXD]

Summary

The current LIR does not deal with runtime-determined memset-size. This patch
utilizes SCEV and check if the PointerStrideSCEV and the MemsetSizeSCEV are equal.
Before comparison the pass would try to fold the expression that is already
protected by the loop guard.

Testcase file memset-runtime.ll, memset-runtime-debug.ll added.

This patch deals with proper loop-idiom. Proceeding patch wants to deal with SCEV-s
that are inequal after folding with the loop guards.

[eopXD]

Fix clang-format warnings.

[eopXD]

Fix more clang-format warning.

[lebedev.ri]

Thanks.
I would also like to ask to split this up further.
In particular, SCEVExprConverter looks scary.
Can we handle PositiveStrideSCEV != MemsetSizeSCEV case later?

[eopXD]

Address comments.

Abort optimization if PositiveStrideSCEV != MemsetSizeSCEV.
Added TODO and will deal with it later.

[eopXD]

In D107353#2922787, @lebedev.ri wrote:

In particular, SCEVExprConverter looks scary.

I agree that adding a bug chunk is quite scary.
Its purpose is to reduce SCEV-s for comparison (convert sext to zext or fold smax constants)
I think most expression cases in the converter is not touched (as LIR is dealing with loops here).

However I currently can't think of a more simpler structure though, do you have any suggestions?
(Or is it acceptable since its a local structure in LIR?)

[lebedev.ri]

It might be good do change the function signatures (to take the ScalarEvolution argument, to take a SCEV instead of unsigned) in a preparatory NFC patch.
Some more catch-all inline comments..

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

220

Can this be changed in a preparatory NFC patch?

785–787

Can this be changed in a preparatory NFC patch?

917–918

Can you make this change globally trouough the file in a preparatory NFC patch?

(Though, somehow it feels conceptually wrong to pass it into processLoopStridedStore().)

949

Negative test missing

953

Negative test missing

954

Negative test missing

961

Test missing

968

Test missing

993

Can this be changed in a preparatory NFC patch?

[eopXD]

NFC patch for changing StoreSize from unsigned to SCEV: D104595

[eopXD]

Rebase to main HEAD.

[eopXD]

Landed D104595, rebase.

[eopXD]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

917–918

NFC patch: D107570.

[eopXD]

Landed D107570, rebase.

[eopXD]

Rebase.

[eopXD]

Address comments.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

949

isLegalStore already blocks non-zero address spaces. So this guard will never trigger and should be deleted.

953

Added memset-runtime-negative-cases.ll - @MemsetSize_LoopVariant.

954

isLegalStore already blocks non-affine SCEV for the pointer in the memset instruction. So this guard will never trigger and should be deleted.

961

Added memset-runtime.ll - @For_NegativeStride

968

Added memset-runtime-negative-cases.ll - @MemsetSize_Stride_Mismatch.

[eopXD]

isLegalStore should be kept the same.

[eopXD]

processLoopMemset does not goes through isLegalStore.
Checks and test cases added.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

949

Added memset-runtime-debug.ll - @NonZeroAddressSpace.

954

Added memset-runtime-debug.ll - @NonAffinePointer.

[Whitney]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

920–973

initialize the variable?

924

Add a period at the end of a sentence.

943

Add a period at the end of a sentence.

950

// If the original StrideSCEV and MemsetSizeSCEV does not match, the pass
// will fold expressions that is covered by the loop guard at loop entry.
// The pass will compare again after the folding and proceed if equal.

That should be added in the future patch that implement that.

llvm/test/Transforms/LoopIdiom/memset-runtime-debug.ll

1 ↗

(On Diff #364966)

is that true? Or you manually generated this test case?
If not, please remove this line.

47 ↗

(On Diff #364966)

That's not ar[i]

llvm/test/Transforms/LoopIdiom/memset-runtime.ll

3

Why do we need assert for this test case?

[bmahjour]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

945

isn't -> aren't

946

handle -> handles

949

Is this part of the comment talking about what you plan to do in the future? If so please either move it to the TODO section on line 977, or put a TODO marker here.

llvm/test/Transforms/LoopIdiom/memset-runtime-debug.ll

43 ↗

(On Diff #364966)

int *ar is actually the first parameter in the IR below.

114 ↗

(On Diff #364966)

Would be useful to have the psudo-C comment like other tests.

llvm/test/Transforms/LoopIdiom/memset-runtime.ll

3

This one doesn't require asserts.

52

ar is the first parameter in the IR below

60

could you rename %1 and %2 to %n and %m and update the CHECKS?

[eopXD]

Address comments.

[Whitney]

LGTM

[qianzhen]

Added some minor NFC comments.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

895

"with a constant size" can be removed, since the new code is trying to handle a non-constant size.

960

"MemsizeSizeSCEV" -> "MemsetSizeSCEV"

[eopXD]

Address comments.

[eopXD]

Thank you @lebedev.ri for guiding my huge code chunk in D104636 into this patch that is more concise.
Also thank you @Whitney, @bmahjour, @qianzhen for reviewing my code.
I am now going to commit it.

Next I would work on the case when MemsetSizeSCEV != PointerStrideSCEV. In this case, the SCEV-s can be folded given the loop entry is already protected by loop guard.

[lebedev.ri]

Thank you, i think this is about right for this chunk.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

925–929

Since non-constant case has to deal with this problem also, i suspect this should go away

llvm/test/Transforms/LoopIdiom/memset-runtime-debug.ll

3 ↗

(On Diff #366304)

It's not great that we don't check what the IR is.
At the very least i think you want FileCheck --implicit-check-not=memset %s

[eopXD]

Address comments.

Put debug output checks at very front, then IRs are checked.

[lebedev.ri]

Actually, we've talked about not needing the overflow check, but IIRC just saturating the number of bytes.
Can you point me at the code that does that?

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1069–1095

This doesn't look right.

1. There absolutely should not be two almost identical functions
2. Shouldn't this be saturating unsigned multiplication?

[eopXD]

In D107353#2944129, @lebedev.ri wrote:

Actually, we've talked about not needing the overflow check, but IIRC just saturating the number of bytes.

May I ask what does IIRC mean?

[lebedev.ri]

In D107353#2944137, @eopXD wrote:

In D107353#2944129, @lebedev.ri wrote:

Actually, we've talked about not needing the overflow check, but IIRC just saturating the number of bytes.

May I ask what does IIRC mean?

If I Remember Correctly

[eopXD]

In D107353#2944129, @lebedev.ri wrote:

Actually, we've talked about not needing the overflow check, but IIRC just saturating the number of bytes.
Can you point me at the code that does that?

No code for saturation here.
I thought restricting to address space zero prevents the need of saturation.

[eopXD]

Address comment.
There is only one getNumBytes now.

[lebedev.ri]

In D107353#2944144, @eopXD wrote:

In D107353#2944129, @lebedev.ri wrote:

Actually, we've talked about not needing the overflow check, but IIRC just saturating the number of bytes.
Can you point me at the code that does that?

No code for saturation here.
I thought restricting to address space zero prevents the need of saturation.

@efriedma does that sound right?

[efriedma]

In D107353#2944152, @lebedev.ri wrote:

In D107353#2944144, @eopXD wrote:

No code for saturation here.
I thought restricting to address space zero prevents the need of saturation.

@efriedma does that sound right?

I think I mentioned this on the other bug... but sounds right to me; overflow implies undefined behavior. If the amount does overflow, that means the memset is writing at least once to every byte in the address space. Among other things, that implies a store to null.

[lebedev.ri]

In D107353#2944315, @efriedma wrote:

In D107353#2944152, @lebedev.ri wrote:

In D107353#2944144, @eopXD wrote:

No code for saturation here.
I thought restricting to address space zero prevents the need of saturation.

@efriedma does that sound right?

I think I mentioned this on the other bug... but sounds right to me; overflow implies undefined behavior. If the amount does overflow, that means the memset is writing at least once to every byte in the address space. Among other things, that implies a store to null.

Well, alrighy then.
@efriedma feel like stamping this too then?

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

925

Comment needs updating?

[eopXD]

Rebase to main.

[eopXD]

Address comments.

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, 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,
LoopIdiomRecognize::isLegalStore(StoreInst *SI) {
   // 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 *StoreEv =
       dyn_cast<SCEVAddRecExpr>(SE->getSCEV(StorePtr));
   if (!StoreEv || StoreEv->getLoop() != CurLoop || !StoreEv->isAffine())
     return LegalStoreKind::None;
 
-  // Check to see if we have a constant stride.
-  if (!isa<SCEVConstant>(StoreEv->getOperand(1)))
-    return LegalStoreKind::None;
-
   // See if the store can be turned into a memset.
 
   // If the stored value is a byte-wise value (like i32 -1), then it may be
   // turned into a memset of i8 -1, assuming that all the consecutive bytes
   // are stored.  A store of i32 0x01020304 can never be turned into a memset,
   // but it can be turned into memset_pattern if the target supports it.
   Value *SplatValue = isBytewiseValue(StoredVal, *DL);
 
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;
 
     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;
     }
   }
 
   return Changed;
bool LoopIdiomRecognize::processLoopMemCpy(MemCpyInst *MCI,
   return processLoopStoreOfLoopLoad(Dest, Source, (unsigned)SizeInBytes,
                                     MCI->getDestAlign(), MCI->getSourceAlign(),
                                     MCI, MCI, StoreEv, LoadEv, BECount);
 }
 
 /// processLoopMemSet - See if this memset can be promoted to a large memset.
 bool LoopIdiomRecognize::processLoopMemSet(MemSetInst *MSI,
                                            const SCEV *BECount) {
   // We can only handle non-volatile memsets with a constant size.
-  if (MSI->isVolatile() || !isa<ConstantInt>(MSI->getLength()))
+  if (MSI->isVolatile())
     return false;
 
   // If we're not allowed to hack on memset, we fail.
   if (!HasMemset || DisableLIRP::Memset)
     return false;
 
   Value *Pointer = MSI->getDest();
 
   // 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;
 
+  const SCEV *StrideSCEV = Ev->getOperand(1);
+  const SCEV *MemsetSizeSCEV = SE->getSCEV(MSI->getLength());
+  if (!StrideSCEV || !MemsetSizeSCEV)
+    return false;
+
+  bool NegStride;
+  const bool IsConstantSize = isa<ConstantInt>(MSI->getLength());
+
+  if (IsConstantSize) {
+    // Memset size is constant
     // Reject memsets that are so large that they overflow an unsigned.
+    LLVM_DEBUG(dbgs() << "  memset size is constant\n");
     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.
+    // 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;
 
+    NegStride = SizeInBytes == -Stride;
+  } else {
+    // Memset size is non-constant
+    // Check if the pointer stride matches the memset size.
+    // Also, the pass only handle memset length and stride that are invariant
+    // for the top level loop. To be conservative, the pass would not promote
+    // pointers that isn't in address space zero.
+    // If the original StrideSCEV and MemsetSizeSCEV does not match, the pass
+    // will fold expressions that is covered by the loop guard at loop entry.
+    // The pass will compare again after the folding and proceed if equal.
+    LLVM_DEBUG(dbgs() << "  memset size is non-constant\n");
+    if (Pointer->getType()->getPointerAddressSpace() != 0) {
+      LLVM_DEBUG(dbgs() << "  pointer is not in address space zero\n");
+      return false;
+    }
+    if (!SE->isLoopInvariant(MemsetSizeSCEV, CurLoop) ||
+        !SE->isLoopInvariant(StrideSCEV, CurLoop)) {
+      LLVM_DEBUG(dbgs() << "  memset size or stride is not a loop-invariant, "
+                        << "abort\n");
+      return false;
+    }
+
+    // compare positive direction strideSCEV with MemsizeSizeSCEV
+    NegStride = StrideSCEV->isNonConstantNegative();
+    const SCEV *PositiveStrideSCEV =
+        NegStride ? SE->getNegativeSCEV(StrideSCEV) : StrideSCEV;
+    LLVM_DEBUG(dbgs() << "  MemsetSizeSCEV: " << *MemsetSizeSCEV << "\n"
+                      << "  PositiveStrideSCEV: " << *PositiveStrideSCEV
+                      << "\n");
+
+    if (PositiveStrideSCEV != MemsetSizeSCEV) {
+      // TODO: folding can be done to the SCEVs
+      LLVM_DEBUG(dbgs() << "  SCEV don't match, abort\n");
+      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, SE->getSCEV(MSI->getLength()),
-                                 MaybeAlign(MSI->getDestAlignment()),
-                                 SplatValue, MSI, MSIs, Ev, BECount, NegStride,
-                                 /*IsLoopMemset=*/true);
+  return processLoopStridedStore(
+      Pointer, MemsetSizeSCEV, 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, 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
static const SCEV *getTripCount(const SCEV *BECount, Type *IntPtr,
   } else {
     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) {
     return SE->getMulExpr(TripCountSCEV, SE->getConstant(IntPtr, StoreSize),
                           SCEV::FlagNUW);
   }
   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, const SCEV *StoreSizeSCEV, MaybeAlign StoreAlignment,
     Value *StoredVal, Instruction *TheStore,
     SmallPtrSetImpl<Instruction *> &Stores, const SCEVAddRecExpr *Ev,
     const SCEV *BECount, bool NegStride, bool IsLoopMemset) {

llvm/test/Transforms/LoopIdiom/memset-runtime.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes="function(loop(loop-idiom,loop-deletion),simplifycfg)" < %s -S | FileCheck %s
+; The C code to generate this testcase:
+; void test(int n, int m, int *ar)
+; {
+;   long i;
+;   for (i=0; i<n; ++i) {
+;     int *arr = ar + i * m; // ar[i];
+;     memset(arr, 0, m * sizeof(int));
+;   }
+; }
+; The optimized IR should be similar to the following:
+; void test(int n, int m, int *ar)
+; {
+;   memset(ar, 0, m * n * sizeof(int));
+; }
+define void @test_simple_for(i32* nocapture %ar, i64 %n, i64 %m) {
+; CHECK-LABEL: @test_simple_for(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AR1:%.*]] = bitcast i32* [[AR:%.*]] to i8*
+; CHECK-NEXT:    [[TMP0:%.*]] = shl nuw i64 [[M:%.*]], 2
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i64 [[M]], [[N:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = shl i64 [[TMP1]], 2
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 4 [[AR1]], i8 0, i64 [[TMP2]], i1 false)
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = shl nuw i64 %m, 2
+  br label %for.cond1.preheader
+
+for.cond1.preheader:                              ; preds = %for.inc4, %entry
+  %i.017 = phi i64 [ 0, %entry ], [ %inc5, %for.inc4 ]
+  %1 = mul i64 %m, %i.017
+  %scevgep = getelementptr i32, i32* %ar, i64 %1
+  %scevgep1 = bitcast i32* %scevgep to i8*
+  %mul = mul nsw i64 %i.017, %m
+  call void @llvm.memset.p0i8.i64(i8* align 4 %scevgep1, i8 0, i64 %0, i1 false)
+  br label %for.inc4
+
+for.inc4:                                         ; preds = %for.cond1.preheader
+  %inc5 = add nuw nsw i64 %i.017, 1
+  %exitcond18.not = icmp eq i64 %inc5, %n
+  br i1 %exitcond18.not, label %for.end6, label %for.cond1.preheader
+
+for.end6:                                         ; preds = %for.inc4
+  ret void
+}
+
+; Function Attrs: argmemonly nofree nounwind willreturn writeonly
+declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i1 immarg)