[MLIR] Fix memref get constant bound size and shape method

Fix FlatAffineConstraints::getConstantBoundOnDimSize to ensure that
returned bounds on dim size are always non-negative regardless of the
constraints on that dimension. Add an assertion at the user.

Differential Revision: https://reviews.llvm.org/D105171

mlir/include/mlir/Analysis/AffineStructures.h

@@ -491,18 +491,19 @@ class FlatAffineConstraints {
   /// Returns the smallest known constant bound for the extent of the specified
   /// identifier (pos^th), i.e., the smallest known constant that is greater
   /// than or equal to 'exclusive upper bound' - 'lower bound' of the
-  /// identifier. Returns None if it's not a constant. This method employs
-  /// trivial (low complexity / cost) checks and detection. Symbolic identifiers
-  /// are treated specially, i.e., it looks for constant differences between
-  /// affine expressions involving only the symbolic identifiers. `lb` and
-  /// `ub` (along with the `boundFloorDivisor`) are set to represent the lower
-  /// and upper bound associated with the constant difference: `lb`, `ub` have
-  /// the coefficients, and boundFloorDivisor, their divisor. `minLbPos` and
-  /// `minUbPos` if non-null are set to the position of the constant lower bound
-  /// and upper bound respectively (to the same if they are from an equality).
-  /// Ex: if the lower bound is [(s0 + s2 - 1) floordiv 32] for a system with
-  /// three symbolic identifiers, *lb = [1, 0, 1], lbDivisor = 32. See comments
-  /// at function definition for examples.
+  /// identifier. This constant bound is guaranteed to be non-negative. Returns
+  /// None if it's not a constant. This method employs trivial (low complexity /
+  /// cost) checks and detection. Symbolic identifiers are treated specially,
+  /// i.e., it looks for constant differences between affine expressions
+  /// involving only the symbolic identifiers. `lb` and `ub` (along with the
+  /// `boundFloorDivisor`) are set to represent the lower and upper bound
+  /// associated with the constant difference: `lb`, `ub` have the coefficients,
+  /// and boundFloorDivisor, their divisor. `minLbPos` and `minUbPos` if
+  /// non-null are set to the position of the constant lower bound and upper
+  /// bound respectively (to the same if they are from an equality). Ex: if the
+  /// lower bound is [(s0 + s2 - 1) floordiv 32] for a system with three
+  /// symbolic identifiers, *lb = [1, 0, 1], lbDivisor = 32. See comments at
+  /// function definition for examples.
   Optional<int64_t> getConstantBoundOnDimSize(
       unsigned pos, SmallVectorImpl<int64_t> *lb = nullptr,
       int64_t *boundFloorDivisor = nullptr,

mlir/include/mlir/Analysis/Utils.h

@@ -279,10 +279,11 @@ struct MemRefRegion {
   /// otherwise. Note that the symbols of the region are treated specially,
   /// i.e., the returned bounding constant holds for *any given* value of the
   /// symbol identifiers. The 'shape' vector is set to the corresponding
-  /// dimension-wise bounds major to minor. We use int64_t instead of uint64_t
-  /// since index types can be at most int64_t. `lbs` are set to the lower
-  /// bounds for each of the rank dimensions, and lbDivisors contains the
-  /// corresponding denominators for floorDivs.
+  /// dimension-wise bounds major to minor. The number of elements and all the
+  /// dimension-wise bounds are guaranteed to be non-negative. We use int64_t
+  /// instead of uint64_t since index types can be at most int64_t. `lbs` are
+  /// set to the lower bounds for each of the rank dimensions, and lbDivisors
+  /// contains the corresponding denominators for floorDivs.
   Optional<int64_t> getConstantBoundingSizeAndShape(
       SmallVectorImpl<int64_t> *shape = nullptr,
       std::vector<SmallVector<int64_t, 4>> *lbs = nullptr,

mlir/lib/Analysis/AffineStructures.cpp

@@ -2285,14 +2285,15 @@ void FlatAffineConstraints::constantFoldIdRange(unsigned pos, unsigned num) {
   }
 }
 
-/// Returns the extent (upper bound - lower bound) of the specified
-/// identifier if it is found to be a constant; returns None if it's not a
-/// constant. This methods treats symbolic identifiers specially, i.e.,
-/// it looks for constant differences between affine expressions involving
-/// only the symbolic identifiers. See comments at function definition for
-/// example. 'lb', if provided, is set to the lower bound associated with the
-/// constant difference. Note that 'lb' is purely symbolic and thus will contain
-/// the coefficients of the symbolic identifiers and the constant coefficient.
+/// Returns a non-negative constant bound on the extent (upper bound - lower
+/// bound) of the specified identifier if it is found to be a constant; returns
+/// None if it's not a constant. This methods treats symbolic identifiers
+/// specially, i.e., it looks for constant differences between affine
+/// expressions involving only the symbolic identifiers. See comments at
+/// function definition for example. 'lb', if provided, is set to the lower
+/// bound associated with the constant difference. Note that 'lb' is purely
+/// symbolic and thus will contain the coefficients of the symbolic identifiers
+/// and the constant coefficient.
 //  Egs: 0 <= i <= 15, return 16.
 //       s0 + 2 <= i <= s0 + 17, returns 16. (s0 has to be a symbol)
 //       s0 + s1 + 16 <= d0 <= s0 + s1 + 31, returns 16.
@@ -2384,6 +2385,8 @@ Optional<int64_t> FlatAffineConstraints::getConstantBoundOnDimSize(
       int64_t diff = ceilDiv(atIneq(ubPos, getNumCols() - 1) +
                                  atIneq(lbPos, getNumCols() - 1) + 1,
                              atIneq(lbPos, pos));
+      // This bound is non-negative by definition.
+      diff = std::max<int64_t>(diff, 0);
       if (minDiff == None || diff < minDiff) {
         minDiff = diff;
         minLbPosition = lbPos;

mlir/lib/Analysis/Utils.cpp

@@ -374,6 +374,7 @@ Optional<int64_t> MemRefRegion::getConstantBoundingSizeAndShape(
         cstWithShapeBounds.getConstantBoundOnDimSize(d, &lb, &lbDivisor);
     if (diff.hasValue()) {
       diffConstant = diff.getValue();
+      assert(diffConstant >= 0 && "Dim size bound can't be negative");
       assert(lbDivisor > 0);
     } else {
       // If no constant bound is found, then it can always be bound by the

mlir/test/Dialect/Affine/affine-data-copy.mlir

@@ -273,12 +273,15 @@ func @max_lower_bound(%M: memref<2048x516xf64>, %i : index, %j : index) {
 
 // -----
 
-// CHECK-LABEL: func @empty_loop
-func @empty_loop(%arg0: memref<1024x1024xf64>) {
-  // Empty loop - so no copy generation happens.
+// CHECK-LABEL: func @empty_loops
+func @empty_loops(%arg0: memref<1024x1024xf64>) {
+  // Empty loops - so no copy generation happens.
   affine.for %i = 0 to 0 {
     affine.load %arg0[0, %i] : memref<1024x1024xf64>
   }
+  affine.for %i = 0 to -16 {
+    affine.load %arg0[0, %i] : memref<1024x1024xf64>
+  }
   return
   // CHECK-NOT:    memref.alloc
   // CHECK:        return