[Mlir-commits] [mlir] 8a8c655 - [mlir][SCF] Fix off-by-one bug in affine analysis

[Matthias Springer]

Author: Matthias Springer
Date: 2021-11-25T12:37:02+09:00
New Revision: 8a8c655fe7c37665471d06b3e8cdff74f6b8ed3e

URL: https://github.com/llvm/llvm-project/commit/8a8c655fe7c37665471d06b3e8cdff74f6b8ed3e
DIFF: https://github.com/llvm/llvm-project/commit/8a8c655fe7c37665471d06b3e8cdff74f6b8ed3e.diff

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

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.

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 5065f2f92b2d9..c54f634dcac04 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -174,6 +174,15 @@ static LogicalResult alignAndAddBound(FlatAffineValueConstraints &constraints,
   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.
@@ -228,18 +237,23 @@ canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap 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`.)