[mlir][SCF] Fix off-by-one bug in affine analysis

[springerm]

Summary

This change is NFC. There were two issues when passing/reading upper bounds into/from FlatAffineConstraints that negate each other, so the bug was not apparent. However, it made debugging harder because some constraints in the FlatAffineConstraints were off by one when dumping all constraints.

[nicolasvasilache]

LGTM, accept conditioned on writing an extra unittest to lock in the new behavior without things cancelling out.

[springerm]

I tried adding a unit test, but the effect of this change is localized to canonicalizeMinMaxOp and not visible outside of the function. Basically, this was not a problem of FlatAffineConstraints itself, but just incorrect API usage. To test this change, I would have to split canonicalizeMinMaxOp into separate functions, so that I can call them separately from a unit test. Splitting the function in such a way does not make sense conceptually though and would pollute the loop specialization API with functions that should never be called by users. I'll land this commit as is.

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

static LogicalResult alignAndAddBound(FlatAffineValueConstraints &constraints,
 
   AffineMap alignedMap =
       alignAffineMapWithValues(map, operands, dims, syms, &newSyms);
   for (unsigned i = syms.size(); i < newSyms.size(); ++i)
     constraints.appendSymbolId(newSyms[i]);
   return constraints.addBound(type, pos, alignedMap);
 }
 
+/// Add `val` to each result of `map`.
+static AffineMap addConstToResults(AffineMap map, int64_t val) {
+  SmallVector<AffineExpr> newResults;
+  for (AffineExpr r : map.getResults())
+    newResults.push_back(r + val);
+  return AffineMap::get(map.getNumDims(), map.getNumSymbols(), newResults,
+                        map.getContext());
+}
+
 /// This function tries to canonicalize min/max operations by proving that their
 /// value is bounded by the same lower and upper bound. In that case, the
 /// operation can be folded away.
 ///
 /// Bounds are computed by FlatAffineValueConstraints. Invariants required for
 /// finding/proving bounds should be supplied via `constraints`.
 ///
 /// 1. Add dimensions for `op` and `opBound` (lower or upper bound of `op`).
canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map,
   unsigned dimOp = constraints.appendDimId();      // `op`
   unsigned dimOpBound = constraints.appendDimId(); // `op` lower/upper bound
   unsigned resultDimStart = constraints.appendDimId(/*num=*/numResults);
 
   // Add an inequality for each result expr_i of map:
   // isMin: op <= expr_i, !isMin: op >= expr_i
   auto boundType =
       isMin ? FlatAffineConstraints::UB : FlatAffineConstraints::LB;
-  if (failed(alignAndAddBound(constraints, boundType, dimOp, map, operands)))
+  // Upper bounds are exclusive, so add 1. (`affine.min` ops are inclusive.)
+  AffineMap mapLbUb = isMin ? addConstToResults(map, 1) : map;
+  if (failed(
+          alignAndAddBound(constraints, boundType, dimOp, mapLbUb, operands)))
     return failure();
 
   // Try to compute a lower/upper bound for op, expressed in terms of the other
   // `dims` and extra symbols.
   SmallVector<AffineMap> opLb(1), opUb(1);
   constraints.getSliceBounds(dimOp, 1, rewriter.getContext(), &opLb, &opUb);
-  AffineMap boundMap = isMin ? opUb[0] : opLb[0];
+  AffineMap sliceBound = isMin ? opUb[0] : opLb[0];
   // TODO: `getSliceBounds` may return multiple bounds at the moment. This is
   // a TODO of `getSliceBounds` and not handled here.
-  if (!boundMap || boundMap.getNumResults() != 1)
+  if (!sliceBound || sliceBound.getNumResults() != 1)
     return failure(); // No or multiple bounds found.
+  // Recover the inclusive UB in the case of an `affine.min`.
+  AffineMap boundMap = isMin ? addConstToResults(sliceBound, -1) : sliceBound;
 
   // Add an equality: Set dimOpBound to computed bound.
   // Add back dimension for op. (Was removed by `getSliceBounds`.)
   AffineMap alignedBoundMap = boundMap.shiftDims(/*shift=*/1, /*offset=*/dimOp);
   if (failed(constraints.addBound(FlatAffineConstraints::EQ, dimOpBound,
                                   alignedBoundMap)))
     return failure();