[DA] runtime predicates for delinearization bounds checks

llvm/include/llvm/Analysis/Delinearization.h

@@ -26,6 +26,7 @@ class GetElementPtrInst;
 class Instruction;
 class ScalarEvolution;
 class SCEV;
+class SCEVPredicate;
 
 /// Compute the array dimensions Sizes from the set of Terms extracted from
 /// the memory access function of this SCEVAddRecExpr (second step of
@@ -144,11 +145,13 @@ bool delinearizeFixedSizeArray(ScalarEvolution &SE, const SCEV *Expr,
 /// Check that each subscript in \p Subscripts is within the corresponding size
 /// in \p Sizes. For the outermost dimension, the subscript being negative is
 /// allowed. If \p Ptr is not nullptr, it may be used to get information from
-/// the IR pointer value, which may help in the validation.
-bool validateDelinearizationResult(ScalarEvolution &SE,
-                                   ArrayRef<const SCEV *> Sizes,
-                                   ArrayRef<const SCEV *> Subscripts,
-                                   const Value *Ptr = nullptr);
+/// the IR pointer value, which may help in the validation. If \p Assume is not
+/// nullptr and a compile-time check fails, runtime predicates are added to
+/// \p Assume instead of returning false.
+bool validateDelinearizationResult(
+    ScalarEvolution &SE, ArrayRef<const SCEV *> Sizes,
+    ArrayRef<const SCEV *> Subscripts, const Value *Ptr = nullptr,
+    SmallVectorImpl<const SCEVPredicate *> *Assume = nullptr);
 
 /// Gathers the individual index expressions from a GEP instruction.
 ///

llvm/include/llvm/Analysis/DependenceAnalysis.h

@@ -754,15 +754,17 @@ class DependenceInfo {
   /// Given a linear access function, tries to recover subscripts
   /// for each dimension of the array element access.
   bool tryDelinearize(Instruction *Src, Instruction *Dst,
-                      SmallVectorImpl<Subscript> &Pair);
+                      SmallVectorImpl<Subscript> &Pair,
+                      SmallVectorImpl<const SCEVPredicate *> &Assume);
 
   /// Tries to delinearize \p Src and \p Dst access functions for a fixed size
   /// multi-dimensional array. Calls delinearizeFixedSizeArray() to delinearize
   /// \p Src and \p Dst separately,
   bool tryDelinearizeFixedSize(Instruction *Src, Instruction *Dst,
                                const SCEV *SrcAccessFn, const SCEV *DstAccessFn,
                                SmallVectorImpl<const SCEV *> &SrcSubscripts,
-                               SmallVectorImpl<const SCEV *> &DstSubscripts);
+                               SmallVectorImpl<const SCEV *> &DstSubscripts,
+                               SmallVectorImpl<const SCEVPredicate *> &Assume);
 
   /// Tries to delinearize access function for a multi-dimensional array with
   /// symbolic runtime sizes.
@@ -771,7 +773,8 @@ class DependenceInfo {
   tryDelinearizeParametricSize(Instruction *Src, Instruction *Dst,
                                const SCEV *SrcAccessFn, const SCEV *DstAccessFn,
                                SmallVectorImpl<const SCEV *> &SrcSubscripts,
-                               SmallVectorImpl<const SCEV *> &DstSubscripts);
+                               SmallVectorImpl<const SCEV *> &DstSubscripts,
+                               SmallVectorImpl<const SCEVPredicate *> &Assume);
 
   /// checkSubscript - Helper function for checkSrcSubscript and
   /// checkDstSubscript to avoid duplicate code

llvm/include/llvm/Analysis/ScalarEvolution.h

@@ -544,6 +544,14 @@ class ScalarEvolution {
                                 const SCEV *LHS, const SCEV *RHS,
                                 const Instruction *CtxI = nullptr);
 
+  /// Get the predicate that, if true at runtime, proves that the binary
+  /// operation \p BinOp between \p LHS and \p RHS does not have
+  /// signed/unsigned overflow (depending on \p Signed). Returns the
+  /// predicate, or nullptr if no-overflow is already provable at compile time.
+  LLVM_ABI const SCEVPredicate *
+  getNoOverflowPredicate(Instruction::BinaryOps BinOp, bool Signed,
+                         const SCEV *LHS, const SCEV *RHS);
+
   /// Parse NSW/NUW flags from add/sub/mul IR binary operation \p Op into
   /// SCEV no-wrap flags, and deduce flag[s] that aren't known yet.
   /// Does not mutate the original instruction. Returns std::nullopt if it could

llvm/lib/Analysis/Delinearization.cpp

@@ -753,37 +753,104 @@ static bool isKnownLessThan(ScalarEvolution *SE, const SCEV *S,
   return SE->isKnownNegative(LimitedBound);
 }
 
-bool llvm::validateDelinearizationResult(ScalarEvolution &SE,
-                                         ArrayRef<const SCEV *> Sizes,
-                                         ArrayRef<const SCEV *> Subscripts,
-                                         const Value *Ptr) {
+bool llvm::validateDelinearizationResult(
+    ScalarEvolution &SE, ArrayRef<const SCEV *> Sizes,
+    ArrayRef<const SCEV *> Subscripts, const Value *Ptr,
+    SmallVectorImpl<const SCEVPredicate *> *Assume) {
   // Sizes and Subscripts are as follows:
-  //
   //   Sizes:      [UNK][S_2]...[S_n]
   //   Subscripts: [I_1][I_2]...[I_n]
   //
   // where the size of the outermost dimension is unknown (UNK).
 
-  auto AddOverflow = [&](const SCEV *A, const SCEV *B) -> const SCEV * {
-    if (!SE.willNotOverflow(Instruction::Add, /*IsSigned=*/true, A, B))
-      return nullptr;
+  // Helper to add a predicate only if it's not already in the list.
+  // SCEV interns predicates, so pointer comparison is sufficient.
+  auto AddPredicate = [&](const SCEVPredicate *Pred) {
+    if (!Pred || !Assume)
+      return;
+    if (!llvm::is_contained(*Assume, Pred))
+      Assume->push_back(Pred);
+  };
+
+  // Unify types of two SCEVs to the wider type.
+  auto UnifyTypes =
+      [&](const SCEV *&A,
+          const SCEV *&B) -> std::pair<const SCEV *, const SCEV *> {
+    Type *WiderType = SE.getWiderType(A->getType(), B->getType());
+    return {SE.getNoopOrSignExtend(A, WiderType),
+            SE.getNoopOrSignExtend(B, WiderType)};
+  };
+
+  // Check if the result of A + B (signed) does not overflow. If it can be
+  // proven at compile-time, return the result. If it might overflow and Assume
+  // is provided, add a runtime equality predicate and return the result.
+  // Otherwise return nullptr.
+  auto AddNoOverflow = [&](const SCEV *A, const SCEV *B) -> const SCEV * {
+    std::tie(A, B) = UnifyTypes(A, B);
+    if (const auto *Pred = SE.getNoOverflowPredicate(Instruction::Add,
+                                                     /*Signed=*/true, A, B)) {
+      if (!Assume)
+        return nullptr;
+      AddPredicate(Pred);
+    }
     return SE.getAddExpr(A, B);
   };
 
-  auto MulOverflow = [&](const SCEV *A, const SCEV *B) -> const SCEV * {
-    if (!SE.willNotOverflow(Instruction::Mul, /*IsSigned=*/true, A, B))
-      return nullptr;
+  // Check if the result of A * B (signed) does not overflow. If it can be
+  // proven at compile-time, return the result. If it might overflow and Assume
+  // is provided, add a runtime equality predicate and return the result.
+  // Otherwise return nullptr.
+  auto MulNoOverflow = [&](const SCEV *A, const SCEV *B) -> const SCEV * {
+    std::tie(A, B) = UnifyTypes(A, B);
+    if (const auto *Pred = SE.getNoOverflowPredicate(Instruction::Mul,
+                                                     /*Signed=*/true, A, B)) {
+      if (!Assume)
+        return nullptr;
+      AddPredicate(Pred);
+    }
     return SE.getMulExpr(A, B);
   };
 
+  // Check if the result of A - B (signed) does not overflow. If it can be
+  // proven at compile-time or if Assume is provided (adding a runtime
+  // predicate), return true. Otherwise return false.
+  auto SubNoOverflow = [&](const SCEV *A, const SCEV *B) -> bool {
+    std::tie(A, B) = UnifyTypes(A, B);
+    if (const auto *Pred = SE.getNoOverflowPredicate(Instruction::Sub,
+                                                     /*Signed=*/true, A, B)) {
+      if (!Assume)
+        return false;
+      AddPredicate(Pred);
+    }
+    return true;
+  };
+
   // Range check: 0 <= I_k < S_k for k = 2..n.
   for (size_t I = 1; I < Sizes.size(); ++I) {
     const SCEV *Size = Sizes[I - 1];
     const SCEV *Subscript = Subscripts[I];
-    if (!isKnownNonNegative(&SE, Subscript, Ptr))
-      return false;
-    if (!isKnownLessThan(&SE, Subscript, Size))
-      return false;
+
+    // Check Subscript >= 0.
+    if (!isKnownNonNegative(&SE, Subscript, Ptr)) {
+      if (!Assume)
+        return false;
+      const SCEVPredicate *Pred = SE.getComparePredicate(
+          ICmpInst::ICMP_SGE, Subscript, SE.getZero(Subscript->getType()));
+      AddPredicate(Pred);
+    }
+
+    // Check Subscript < Size.
+    if (!isKnownLessThan(&SE, Subscript, Size)) {
+      if (!Assume)
+        return false;
+      // Need to unify types before creating the predicate.
+      Type *WiderType = SE.getWiderType(Subscript->getType(), Size->getType());
+      const SCEV *SubscriptExt = SE.getNoopOrSignExtend(Subscript, WiderType);
+      const SCEV *SizeExt = SE.getNoopOrSignExtend(Size, WiderType);
+      const SCEVPredicate *Pred =
+          SE.getComparePredicate(ICmpInst::ICMP_SLT, SubscriptExt, SizeExt);
+      AddPredicate(Pred);
+    }
   }
 
   // The offset computation is as follows:
@@ -811,21 +878,20 @@ bool llvm::validateDelinearizationResult(ScalarEvolution &SE,
   // NOTE: I_1 can be negative, so Min is not just 0.
   const SCEV *Prod = SE.getOne(Sizes[0]->getType());
   for (const SCEV *Size : Sizes) {
-    Prod = MulOverflow(Prod, Size);
+    Prod = MulNoOverflow(Prod, Size);
     if (!Prod)
       return false;
   }
-  const SCEV *Min = MulOverflow(Prod, Subscripts[0]);
+  const SCEV *Min = MulNoOverflow(Prod, Subscripts[0]);
   if (!Min)
     return false;
 
   // We have already checked that Min and Prod don't overflow, so it's enough
   // to check whether Min + Prod - 1 doesn't overflow.
-  const SCEV *MaxPlusOne = AddOverflow(Min, Prod);
+  const SCEV *MaxPlusOne = AddNoOverflow(Min, Prod);
   if (!MaxPlusOne)
     return false;
-  if (!SE.willNotOverflow(Instruction::Sub, /*IsSigned=*/true, MaxPlusOne,
-                          SE.getOne(MaxPlusOne->getType())))
+  if (!SubNoOverflow(MaxPlusOne, SE.getOne(MaxPlusOne->getType())))
     return false;
 
   return true;

llvm/lib/Analysis/DependenceAnalysis.cpp

@@ -3176,8 +3176,9 @@ const SCEV *DependenceInfo::getUpperBound(BoundInfo *Bound) const {
 /// source and destination array references are recurrences on a nested loop,
 /// this function flattens the nested recurrences into separate recurrences
 /// for each loop level.
-bool DependenceInfo::tryDelinearize(Instruction *Src, Instruction *Dst,
-                                    SmallVectorImpl<Subscript> &Pair) {
+bool DependenceInfo::tryDelinearize(
+    Instruction *Src, Instruction *Dst, SmallVectorImpl<Subscript> &Pair,
+    SmallVectorImpl<const SCEVPredicate *> &Assume) {
   assert(isLoadOrStore(Src) && "instruction is not load or store");
   assert(isLoadOrStore(Dst) && "instruction is not load or store");
   Value *SrcPtr = getLoadStorePointerOperand(Src);
@@ -3197,9 +3198,9 @@ bool DependenceInfo::tryDelinearize(Instruction *Src, Instruction *Dst,
   SmallVector<const SCEV *, 4> SrcSubscripts, DstSubscripts;
 
   if (!tryDelinearizeFixedSize(Src, Dst, SrcAccessFn, DstAccessFn,
-                               SrcSubscripts, DstSubscripts) &&
+                               SrcSubscripts, DstSubscripts, Assume) &&
       !tryDelinearizeParametricSize(Src, Dst, SrcAccessFn, DstAccessFn,
-                                    SrcSubscripts, DstSubscripts))
+                                    SrcSubscripts, DstSubscripts, Assume))
     return false;
 
   assert(isLoopInvariant(SrcBase, SrcLoop) &&
@@ -3245,7 +3246,8 @@ bool DependenceInfo::tryDelinearize(Instruction *Src, Instruction *Dst,
 bool DependenceInfo::tryDelinearizeFixedSize(
     Instruction *Src, Instruction *Dst, const SCEV *SrcAccessFn,
     const SCEV *DstAccessFn, SmallVectorImpl<const SCEV *> &SrcSubscripts,
-    SmallVectorImpl<const SCEV *> &DstSubscripts) {
+    SmallVectorImpl<const SCEV *> &DstSubscripts,
+    SmallVectorImpl<const SCEVPredicate *> &Assume) {
   LLVM_DEBUG({
     const SCEVUnknown *SrcBase =
         dyn_cast<SCEVUnknown>(SE->getPointerBase(SrcAccessFn));
@@ -3285,10 +3287,12 @@ bool DependenceInfo::tryDelinearizeFixedSize(
   // dimensions. For example some C language usage/interpretation make it
   // impossible to verify this at compile-time. As such we can only delinearize
   // iff the subscripts are positive and are less than the range of the
-  // dimension.
+  // dimension. If compile-time checks fail, add runtime predicates.
   if (!DisableDelinearizationChecks) {
-    if (!validateDelinearizationResult(*SE, SrcSizes, SrcSubscripts, SrcPtr) ||
-        !validateDelinearizationResult(*SE, DstSizes, DstSubscripts, DstPtr)) {
+    if (!validateDelinearizationResult(*SE, SrcSizes, SrcSubscripts, SrcPtr,
+                                       &Assume) ||
+        !validateDelinearizationResult(*SE, DstSizes, DstSubscripts, DstPtr,
+                                       &Assume)) {
       SrcSubscripts.clear();
       DstSubscripts.clear();
       return false;
@@ -3305,7 +3309,8 @@ bool DependenceInfo::tryDelinearizeFixedSize(
 bool DependenceInfo::tryDelinearizeParametricSize(
     Instruction *Src, Instruction *Dst, const SCEV *SrcAccessFn,
     const SCEV *DstAccessFn, SmallVectorImpl<const SCEV *> &SrcSubscripts,
-    SmallVectorImpl<const SCEV *> &DstSubscripts) {
+    SmallVectorImpl<const SCEV *> &DstSubscripts,
+    SmallVectorImpl<const SCEVPredicate *> &Assume) {
 
   Value *SrcPtr = getLoadStorePointerOperand(Src);
   Value *DstPtr = getLoadStorePointerOperand(Dst);
@@ -3346,15 +3351,13 @@ bool DependenceInfo::tryDelinearizeParametricSize(
       SrcSubscripts.size() != DstSubscripts.size())
     return false;
 
-  // Statically check that the array bounds are in-range. The first subscript we
-  // don't have a size for and it cannot overflow into another subscript, so is
-  // always safe. The others need to be 0 <= subscript[i] < bound, for both src
-  // and dst.
-  // FIXME: It may be better to record these sizes and add them as constraints
-  // to the dependency checks.
+  // Check that the array bounds are in-range. If compile-time checks fail,
+  // add runtime predicates.
   if (!DisableDelinearizationChecks)
-    if (!validateDelinearizationResult(*SE, Sizes, SrcSubscripts, SrcPtr) ||
-        !validateDelinearizationResult(*SE, Sizes, DstSubscripts, DstPtr))
+    if (!validateDelinearizationResult(*SE, Sizes, SrcSubscripts, SrcPtr,
+                                       &Assume) ||
+        !validateDelinearizationResult(*SE, Sizes, DstSubscripts, DstPtr,
+                                       &Assume))
       return false;
 
   return true;
@@ -3507,7 +3510,7 @@ DependenceInfo::depends(Instruction *Src, Instruction *Dst,
                                           SCEVUnionPredicate(Assume, *SE));
 
   if (Delinearize) {
-    if (tryDelinearize(Src, Dst, Pair)) {
+    if (tryDelinearize(Src, Dst, Pair, Assume)) {
       LLVM_DEBUG(dbgs() << "    delinearized\n");
       Pairs = Pair.size();
     }

llvm/lib/Analysis/ScalarEvolution.cpp

@@ -2376,6 +2376,49 @@ bool ScalarEvolution::willNotOverflow(Instruction::BinaryOps BinOp, bool Signed,
   }
 }
 
+const SCEVPredicate *
+ScalarEvolution::getNoOverflowPredicate(Instruction::BinaryOps BinOp,
+                                        bool Signed, const SCEV *LHS,
+                                        const SCEV *RHS) {
+  // First check if no-overflow can be proven at compile time.
+  if (willNotOverflow(BinOp, Signed, LHS, RHS))
+    return nullptr;
+
+  const SCEV *(ScalarEvolution::*Operation)(const SCEV *, const SCEV *,
+                                            SCEV::NoWrapFlags, unsigned);
+  switch (BinOp) {
+  default:
+    llvm_unreachable("Unsupported binary op");
+  case Instruction::Add:
+    Operation = &ScalarEvolution::getAddExpr;
+    break;
+  case Instruction::Sub:
+    Operation = &ScalarEvolution::getMinusSCEV;
+    break;
+  case Instruction::Mul:
+    Operation = &ScalarEvolution::getMulExpr;
+    break;
+  }
+
+  const SCEV *(ScalarEvolution::*Extension)(const SCEV *, Type *, unsigned) =
+      Signed ? &ScalarEvolution::getSignExtendExpr
+             : &ScalarEvolution::getZeroExtendExpr;
+
+  // Build predicate: ext(LHS op RHS) == ext(LHS) op ext(RHS)
+  auto *NarrowTy = cast<IntegerType>(LHS->getType());
+  auto *WideTy =
+      IntegerType::get(NarrowTy->getContext(), NarrowTy->getBitWidth() * 2);
+
+  const SCEV *A = (this->*Extension)(
+      (this->*Operation)(LHS, RHS, SCEV::FlagAnyWrap, 0), WideTy, 0);
+  const SCEV *LHSB = (this->*Extension)(LHS, WideTy, 0);
+  const SCEV *RHSB = (this->*Extension)(RHS, WideTy, 0);
+  const SCEV *B = (this->*Operation)(LHSB, RHSB, SCEV::FlagAnyWrap, 0);
+
+  // Return the equality predicate.
+  return getEqualPredicate(A, B);
+}
+
 std::optional<SCEV::NoWrapFlags>
 ScalarEvolution::getStrengthenedNoWrapFlagsFromBinOp(
     const OverflowingBinaryOperator *OBO) {