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

[bondhugula]

Summary

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.

[bondhugula]

Drop blank line. Update comment.

[vinayaka-polymage]

It seems correct to me, that the constant bound is non-negative always, on principle. Was this bound being negative creating any problem in some cases ? If yes, it might be useful to mention that in comments.

[bondhugula]

In D105171#2858192, @vinayaka-polymage wrote:

It seems correct to me, that the constant bound is non-negative always, on principle. Was this bound being negative creating any problem in some cases ? If yes, it might be useful to mention that in comments.

Yes, for any user of that method, a negative bound would be problematic and wrong by definition. For eg. MemRefRegion::getConstantBoundingSizeAndShape is in turn used by memref region compute, which is in turn used by many things including affineDataCopyGenerate. Since this is a lower-level API, we'd like to avoid referring to its users - since the latter change. By design/definition this bound has to be non-negative.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   }
 
   /// Clears this list of constraints and copies other into it.
   void clearAndCopyFrom(const FlatAffineConstraints &other);
 
   /// 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,
       SmallVectorImpl<int64_t> *ub = nullptr, unsigned *minLbPos = nullptr,
       unsigned *minUbPos = nullptr) const;
 
   /// Returns the constant lower bound for the pos^th identifier if there is
   /// one; None otherwise.

mlir/include/mlir/Analysis/Utils.h

struct MemRefRegion {
   bool isWrite() const { return write; }
   void setWrite(bool flag) { write = flag; }
 
   /// Returns a constant upper bound on the number of elements in this region if
   /// bounded by a known constant (always possible for static shapes), None
   /// 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,
       SmallVectorImpl<int64_t> *lbDivisors = nullptr) const;
 
   /// Gets the lower and upper bound map for the dimensional identifier at
   /// `pos`.
   void getLowerAndUpperBound(unsigned pos, AffineMap &lbMap,

mlir/lib/Analysis/AffineStructures.cpp

 
 void FlatAffineConstraints::constantFoldIdRange(unsigned pos, unsigned num) {
   for (unsigned s = pos, t = pos, e = pos + num; s < e; s++) {
     if (failed(constantFoldId(t)))
       t++;
   }
 }
 
-/// 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.
 //       s0 - 7 <= 8*j <= s0 returns 1 with lb = s0, lbDivisor = 8 (since lb =
 //       ceil(s0 - 7 / 8) = floor(s0 / 8)).
 Optional<int64_t> FlatAffineConstraints::getConstantBoundOnDimSize(
     unsigned pos, SmallVectorImpl<int64_t> *lb, int64_t *boundFloorDivisor,
     SmallVectorImpl<int64_t> *ub, unsigned *minLbPos,
for (auto lbPos : lbIndices) {
         if (atIneq(ubPos, j) != -atIneq(lbPos, j)) {
           break;
         }
       if (j < getNumCols() - 1)
         continue;
       int64_t diff = ceilDiv(atIneq(ubPos, getNumCols() - 1) +
                                  atIneq(lbPos, getNumCols() - 1) + 1,
                              atIneq(lbPos, pos));
+      // This bound can't be negative by definition.
+      diff = std::max<int64_t>(diff, 0);
       if (minDiff == None || diff < minDiff) {
         minDiff = diff;
         minLbPosition = lbPos;
         minUbPosition = ubPos;
       }
     }
   }
   if (lb && minDiff.hasValue()) {

mlir/lib/Analysis/Utils.cpp

Optional<int64_t> MemRefRegion::getConstantBoundingSizeAndShape(
   int64_t diffConstant;
   int64_t lbDivisor;
   for (unsigned d = 0; d < rank; d++) {
     SmallVector<int64_t, 4> lb;
     Optional<int64_t> diff =
         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
       // memref's dim size if the latter has a constant size along this dim.
       auto dimSize = memRefType.getDimSize(d);
       if (dimSize == -1)
         return None;
       diffConstant = dimSize;

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

 // CHECK-NEXT: }
 // CHECK-NEXT: affine.for %[[ii_:.*]] = 0 to 2048 {
 // CHECK-NEXT:   affine.for %[[jj_:.*]] = max #[[$LB]]()[%{{.*}}, %{{.*}}] to min #[[$UB]]()[%{{.*}}, %{{.*}}] {
 // CHECK-NEXT:     affine.load %[[BUF]][%[[ii_]], %[[jj_]] - symbol(%[[j]]) * 6] : memref<2048x6xf64>
 // CHECK-NEXT:    }
 // CHECK-NEXT: }
 // CHECK-NEXT: memref.dealloc %[[BUF]] : memref<2048x6xf64>
 
+
 // -----
 
-// 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
 }