[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

348

Can this be changed in a preparatory NFC patch?

913–915

Can this be changed in a preparatory NFC patch?

1045–1046

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().)

1077

Negative test missing

1081

Negative test missing

1082

Negative test missing

1089

Test missing

1096

Test missing

1132

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

1045–1046

NFC patch: D107570.

[eopXD]

Landed D107570, rebase.

[eopXD]

Rebase.

[eopXD]

Address comments.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1077

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

1081

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

1082

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

1089

Added memset-runtime.ll - @For_NegativeStride

1096

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

1077

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

1082

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

[Whitney]

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

1048–1112

initialize the variable?

1052

Add a period at the end of a sentence.

1071

Add a period at the end of a sentence.

1078

// 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

1073

isn't -> aren't

1074

handle -> handles

1077

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

1023

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

1088

"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

1053–1057

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

1256–1263

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

1053

Comment needs updating?

[eopXD]

Rebase to main.

[eopXD]

Address comments.

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

 static cl::opt<bool> UseLIRCodeSizeHeurs(
     "use-lir-code-size-heurs",
     cl::desc("Use loop idiom recognition code size heuristics when compiling"
              "with -Os/-Oz"),
     cl::init(true), cl::Hidden);
 
 namespace {
 
+/// A helper class to do the following SCEV expression conversions.
+///   1) "sext %val" to "zext %val"
+///   2) "SOME_CONSTANT_VALUE smax %val" to "%val"
+/// The converter would use ScalarEvolution::isLoopEntryGuardedByCond to see if
+/// the expression can be folded.
+class SCEVExprConverter {
+public:
+  Loop *CurLoop;
+  ScalarEvolution *SE;
+
+  SCEVExprConverter(ScalarEvolution *SE) : SE(SE) {
+    assert(SE != nullptr && "expect SE provided");
+  }
+
+  const SCEV *convertSCEV(const SCEV *Expr);
+};
+
+/// Tries to fold the SCEV with regard to loop guards of CurLoop
+const SCEV *SCEVExprConverter::convertSCEV(const SCEV *Expr) {
+  switch (Expr->getSCEVType()) {
+  case scConstant:
+  case scUnknown:
+  case scCouldNotCompute:
+    return Expr;
+  case scTruncate: {
+    const SCEVTruncateExpr *Trunc = cast<SCEVTruncateExpr>(Expr);
+    Type *Ty = Trunc->getType();
+    const SCEV *NewTrunc = convertSCEV(Trunc->getOperand());
+    return SE->getTruncateExpr(NewTrunc, Ty);
+  }
+  case scZeroExtend: {
+    const SCEVZeroExtendExpr *Zext = cast<SCEVZeroExtendExpr>(Expr);
+    Type *Ty = Zext->getType();
+    const SCEV *NewZext = convertSCEV(Zext->getOperand());
+    return SE->getZeroExtendExpr(NewZext, Ty);
+  }
+  case scSignExtend: {
+    // Return original SCEV if expression is not guarded by loop with Zero
+    // Otherwise fold the expression
+    const SCEVSignExtendExpr *Sext = cast<SCEVSignExtendExpr>(Expr);
+    if (SE->isLoopEntryGuardedByCond(CurLoop, ICmpInst::ICMP_SGE, Sext, SE->getZero(Sext->getType())) == false)
+      return Sext;
+    const SCEV *NewZext = convertSCEV(Sext->getOperand());
+    return SE->getZeroExtendExpr(NewZext, Sext->getType());
+  }
+  case scAddExpr: {
+    const SCEVAddExpr *Add = cast<SCEVAddExpr>(Expr);
+    const SCEV *NewAdd = convertSCEV(Add->getOperand(0));
+    for (int I = 1, E = Add->getNumOperands(); I != E; ++I) {
+      NewAdd = SE->getAddExpr(NewAdd, convertSCEV(Add->getOperand(I)));
+    }
+    return NewAdd;
+  }
+  case scMulExpr: {
+    const SCEVMulExpr *Mul = cast<SCEVMulExpr>(Expr);
+    const SCEV *NewMul = convertSCEV(Mul->getOperand(0));
+    for (int I = 1, E = Mul->getNumOperands(); I != E; ++I) {
+      NewMul = SE->getMulExpr(NewMul, convertSCEV(Mul->getOperand(I)));
+    }
+    return NewMul;
+  }
+  case scUDivExpr: {
+    const SCEVUDivExpr *UDiv = cast<SCEVUDivExpr>(Expr);
+    const SCEV *NewLHS = convertSCEV(UDiv->getLHS());
+    const SCEV *NewRHS = convertSCEV(UDiv->getRHS());
+    return SE->getUDivExpr(NewLHS, NewRHS);
+  }
+  case scAddRecExpr:
+    assert(false && "Do not expect AddRec here!");
+  case scUMaxExpr: {
+    const SCEVUMaxExpr *UMax = cast<SCEVUMaxExpr>(Expr);
+    const SCEV *NewUMax = convertSCEV(UMax->getOperand(0));
+    for (int I = 1, E = UMax->getNumOperands(); I != E; ++I) {
+      NewUMax = SE->getUMaxExpr(NewUMax, convertSCEV(UMax->getOperand(I)));
+    }
+    return NewUMax;
+  }
+  case scSMaxExpr: {
+    // Return original SCEV if expression is not guarded by loop with Zero
+    // Otherwise fold the expression
+    const SCEVSMaxExpr *SMax = cast<SCEVSMaxExpr>(Expr);
+    const int NumOfOps = SMax->getNumOperands();
+    bool Fold = false;
+    // If an operand is constant, it will be the first operand.
+    const SCEV *SMaxOp0 = SMax->getOperand(0);
+    const SCEVConstant *Cst = dyn_cast<SCEVConstant>(SMaxOp0);
+
+    if (Cst) {
+      // check if the operand is guarded to the constant
+      // if not, return orignal expression
+      Fold = true;
+      for (int I = 1, E = NumOfOps; I != E; ++I) {
+        const SCEV *Ev = SMax->getOperand(I);
+        if (SE->isLoopEntryGuardedByCond(CurLoop, ICmpInst::ICMP_SGE, Ev, Cst) == false)
+          return SMax;
+      }
+    }
+
+    const int StartIdx = Fold ? 1 : 0;
+    const SCEV *NewSMax = convertSCEV(SMax->getOperand(StartIdx));
+    for (int I = StartIdx + 1, E = NumOfOps; I != E; ++I) {
+      NewSMax = SE->getSMaxExpr(NewSMax, convertSCEV(SMax->getOperand(I)));
+    }
+    return NewSMax;
+  }
+  case scUMinExpr: {
+    const SCEVUMinExpr *UMin = cast<SCEVUMinExpr>(Expr);
+    const SCEV *NewUMin = convertSCEV(UMin->getOperand(0));
+    for (int I = 1, E = UMin->getNumOperands(); I != E; ++I) {
+      NewUMin = SE->getUMinExpr(NewUMin, convertSCEV(UMin->getOperand(I)));
+    }
+    return NewUMin;
+  }
+  case scSMinExpr: {
+    const SCEVSMinExpr *SMin = cast<SCEVSMinExpr>(Expr);
+    const SCEV *NewSMin = convertSCEV(SMin->getOperand(0));
+    for (int I = 1, E = SMin->getNumOperands(); I != E; ++I) {
+      NewSMin = SE->getSMinExpr(NewSMin, convertSCEV(SMin->getOperand(I)));
+    }
+    return NewSMin;
+  }
+  default:
+    llvm_unreachable("Unexpected SCEV expression!");
+  }
+}
+
 class LoopIdiomRecognize {
   Loop *CurLoop = nullptr;
   AliasAnalysis *AA;
   DominatorTree *DT;
   LoopInfo *LI;
   ScalarEvolution *SE;
   TargetLibraryInfo *TLI;
   const TargetTransformInfo *TTI;
   const DataLayout *DL;
   OptimizationRemarkEmitter &ORE;
   bool ApplyCodeSizeHeuristics;
   std::unique_ptr<MemorySSAUpdater> MSSAU;
+  SCEVExprConverter Converter;
 
 public:
   explicit LoopIdiomRecognize(AliasAnalysis *AA, DominatorTree *DT,
                               LoopInfo *LI, ScalarEvolution *SE,
                               TargetLibraryInfo *TLI,
                               const TargetTransformInfo *TTI, MemorySSA *MSSA,
                               const DataLayout *DL,
                               OptimizationRemarkEmitter &ORE)
-      : AA(AA), DT(DT), LI(LI), SE(SE), TLI(TLI), TTI(TTI), DL(DL), ORE(ORE) {
+      : AA(AA), DT(DT), LI(LI), SE(SE), TLI(TLI), TTI(TTI), DL(DL), ORE(ORE),
+        Converter(SE) {
     if (MSSA)
       MSSAU = std::make_unique<MemorySSAUpdater>(MSSA);
   }
 
   bool runOnLoop(Loop *L);
 
 private:
   using StoreList = SmallVector<StoreInst *, 8>;
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,
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;
 
-    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;
     }
   }
 
   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) {
+      Converter.CurLoop = CurLoop;
+      const SCEV *FoldedPositiveStride =
+          Converter.convertSCEV(PositiveStrideSCEV);
+      const SCEV *FoldedMemsetSize = Converter.convertSCEV(MemsetSizeSCEV);
+      LLVM_DEBUG(
+          dbgs() << "  Try to fold SCEV expression covered by loop guard\n"
+                 << "    FoldedMemsetSCEV: " << *FoldedMemsetSize << "\n"
+                 << "    FoldedPositiveStrideSCEV: "
+                 << *FoldedPositiveStride << "\n");
+
+      if (FoldedPositiveStride != FoldedMemsetSize) {
+        LLVM_DEBUG(dbgs() << "  folded SCEV unmatch, 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, (unsigned)SizeInBytes, MaybeAlign(MSI->getDestAlignment()),
-      SplatValue, MSI, MSIs, Ev, BECount, NegStride, /*IsLoopMemset=*/true);
+      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, 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);
 }
 
+// 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);
+}
+
+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()) <
+          DL->getTypeSizeInBits(IntPtr) &&
+      SE->isLoopEntryGuardedByCond(
+          CurLoop, ICmpInst::ICMP_NE, BECount,
+          SE->getNegativeSCEV(SE->getOne(BECount->getType())))) {
+    TripCountS = SE->getZeroExtendExpr(
+        SE->getAddExpr(BECount, SE->getOne(BECount->getType()), SCEV::FlagNUW),
+        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 *NumBytesS;
if (DL->getTypeSizeInBits(BECount->getType()).getFixedSize() <
           SE->getNegativeSCEV(SE->getOne(BECount->getType())))) {
     NumBytesS = SE->getZeroExtendExpr(
         SE->getAddExpr(BECount, SE->getOne(BECount->getType()), SCEV::FlagNUW),
         IntPtr);
   } else {
     NumBytesS = SE->getAddExpr(SE->getTruncateOrZeroExtend(BECount, IntPtr),
                                SE->getOne(IntPtr), SCEV::FlagNUW);
   }
 
   // And scale it based on the store size.
   if (StoreSize != 1) {
     NumBytesS = SE->getMulExpr(NumBytesS, SE->getConstant(IntPtr, StoreSize),
                                SCEV::FlagNUW);
   }
   return NumBytesS;
 }
 
 /// 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 *TripCountS = getTripCount(BECount, IntIdxTy, CurLoop, DL, SE);
+  const SCEVConstant *ConstSize = dyn_cast<SCEVConstant>(StoreSizeSCEV);
+  const SCEV *NumBytesS;
+  if (ConstSize && ConstSize->getAPInt() == 1) {
+    NumBytesS = TripCountS;
+    LLVM_DEBUG(dbgs() << "StoreSize = 1, NumbytesS: " << *NumBytesS << "\n");
+  } else {
+    NumBytesS = SE->getMulExpr(
+        TripCountS, SE->getTruncateOrZeroExtend(StoreSizeSCEV, IntIdxTy),
+        SCEV::FlagNUW);
+    LLVM_DEBUG(dbgs() << "  Calculate NumBytesS = TripCountS * StoreSizeSCEV\n"
+                      << "    TripCountS: " << *TripCountS << "\n"
+                      << "    StoreSizeSCEV: " << *StoreSizeSCEV << "\n"
+                      << "    NumBytesS: " << *NumBytesS << "\n");
+  }
 
   // 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;

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
+}
+; The C code to generate this testcase:
+; void test(int n, int m, int o, int *ar)
+; {
+;   for (int i=0; i<n; ++i) {
+;     for (int j=0; j<m; ++j) {
+;       int *arr = ar + i * m * o + j * o;
+;       memset(arr, 0, o * sizeof(int));
+;     }
+;   }
+; }
+; The optimized IR should be similar to the following:
+; void test(int n, int m, int o, int *ar)
+; {
+;   memset(ar, 0, m * n * o * sizeof(int));
+; }
+define void @test_nested_for(i32 %n, i32 %m, i32 %o, i32* %ar) {
+; CHECK-LABEL: @test_nested_for(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AR2:%.*]] = bitcast i32* [[AR:%.*]] to i8*
+; CHECK-NEXT:    [[CMP3:%.*]] = icmp slt i32 0, [[N:%.*]]
+; CHECK-NEXT:    br i1 [[CMP3]], label [[FOR_BODY_LR_PH:%.*]], label [[FOR_END11:%.*]]
+; CHECK:       for.body.lr.ph:
+; CHECK-NEXT:    [[CMP21:%.*]] = icmp slt i32 0, [[M:%.*]]
+; CHECK-NEXT:    [[CONV:%.*]] = sext i32 [[O:%.*]] to i64
+; CHECK-NEXT:    [[MUL8:%.*]] = mul i64 [[CONV]], 4
+; CHECK-NEXT:    br i1 [[CMP21]], label [[FOR_BODY_LR_PH_SPLIT_US:%.*]], label [[FOR_END11]]
+; CHECK:       for.body.lr.ph.split.us:
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[O]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[M]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = sext i32 [[O]] to i64
+; CHECK-NEXT:    [[WIDE_TRIP_COUNT10:%.*]] = zext i32 [[N]] to i64
+; CHECK-NEXT:    [[TMP3:%.*]] = mul i64 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[M]] to i64
+; CHECK-NEXT:    [[TMP5:%.*]] = mul i64 [[TMP0]], [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = shl i64 [[TMP5]], 2
+; CHECK-NEXT:    [[TMP7:%.*]] = mul i64 [[TMP5]], [[WIDE_TRIP_COUNT10]]
+; CHECK-NEXT:    [[TMP8:%.*]] = shl i64 [[TMP7]], 2
+; CHECK-NEXT:    call void @llvm.memset.p0i8.i64(i8* align 4 [[AR2]], i8 0, i64 [[TMP8]], i1 false)
+; CHECK-NEXT:    br label [[FOR_END11]]
+; CHECK:       for.end11:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp3 = icmp slt i32 0, %n
+  br i1 %cmp3, label %for.body.lr.ph, label %for.end11
+
+for.body.lr.ph:                                   ; preds = %entry
+  %cmp21 = icmp slt i32 0, %m
+  %conv = sext i32 %o to i64
+  %mul8 = mul i64 %conv, 4
+  br i1 %cmp21, label %for.body.lr.ph.split.us, label %for.body.lr.ph.split
+
+for.body.lr.ph.split.us:                          ; preds = %for.body.lr.ph
+  %0 = sext i32 %o to i64
+  %1 = sext i32 %m to i64
+  %2 = sext i32 %o to i64
+  %wide.trip.count10 = zext i32 %n to i64
+  br label %for.body.us
+
+for.body.us:                                      ; preds = %for.inc9.us, %for.body.lr.ph.split.us
+  %indvars.iv6 = phi i64 [ %indvars.iv.next7, %for.inc9.us ], [ 0, %for.body.lr.ph.split.us ]
+  br label %for.body3.lr.ph.us
+
+for.end.us:                                       ; preds = %for.cond1.for.end_crit_edge.us
+  br label %for.inc9.us
+
+for.inc9.us:                                      ; preds = %for.end.us
+  %indvars.iv.next7 = add nuw nsw i64 %indvars.iv6, 1
+  %exitcond11 = icmp ne i64 %indvars.iv.next7, %wide.trip.count10
+  br i1 %exitcond11, label %for.body.us, label %for.cond.for.end11_crit_edge.split.us
+
+for.body3.us:                                     ; preds = %for.body3.lr.ph.us, %for.inc.us
+  %indvars.iv = phi i64 [ 0, %for.body3.lr.ph.us ], [ %indvars.iv.next, %for.inc.us ]
+  %3 = mul nsw i64 %indvars.iv, %0
+  %add.ptr7.us = getelementptr inbounds i32, i32* %add.ptr.us, i64 %3
+  %4 = bitcast i32* %add.ptr7.us to i8*
+  call void @llvm.memset.p0i8.i64(i8* align 4 %4, i8 0, i64 %mul8, i1 false)
+  br label %for.inc.us
+
+for.inc.us:                                       ; preds = %for.body3.us
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp ne i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond, label %for.body3.us, label %for.cond1.for.end_crit_edge.us
+
+for.body3.lr.ph.us:                               ; preds = %for.body.us
+  %5 = mul nsw i64 %indvars.iv6, %1
+  %6 = mul nsw i64 %5, %2
+  %add.ptr.us = getelementptr inbounds i32, i32* %ar, i64 %6
+  %wide.trip.count = zext i32 %m to i64
+  br label %for.body3.us
+
+for.cond1.for.end_crit_edge.us:                   ; preds = %for.inc.us
+  br label %for.end.us
+
+for.cond.for.end11_crit_edge.split.us:            ; preds = %for.inc9.us
+  br label %for.cond.for.end11_crit_edge
+
+for.body.lr.ph.split:                             ; preds = %for.body.lr.ph
+  br label %for.cond.for.end11_crit_edge.split
+
+for.cond.for.end11_crit_edge.split:               ; preds = %for.body.lr.ph.split
+  br label %for.cond.for.end11_crit_edge
+
+for.cond.for.end11_crit_edge:                     ; preds = %for.cond.for.end11_crit_edge.split.us, %for.cond.for.end11_crit_edge.split
+  br label %for.end11
+
+for.end11:                                        ; preds = %for.cond.for.end11_crit_edge, %entry
+  ret void
+}
+
+; Function Attrs: argmemonly nofree nounwind willreturn writeonly
+declare void @llvm.memset.p0i8.i64(i8* nocapture writeonly, i8, i64, i1 immarg)