[mlir][SCF] Further simplify affine maps during `for-loop-canonicalization`

[springerm]

Summary

• Implement FlatAffineConstraints::getConstantBound(EQ).
• Inject a simpler constraint for loops that have at most 1 iteration.
• Taking into account constant EQ bounds of FlatAffineConstraint dims/symbols during canonicalization of the resulting affine map in canonicalizeMinMaxOp.

[springerm]

update

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

public:
   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 bound for the pos^th identifier if there is one;
   /// None otherwise.
-  // TODO: Support EQ bounds.
   Optional<int64_t> getConstantBound(BoundType type, unsigned pos) const;
 
   /// Removes constraints that are independent of (i.e., do not have a
   /// coefficient) identifiers in the range [pos, pos + num).
   void removeIndependentConstraints(unsigned pos, unsigned num);
 
   /// Returns true if the set can be trivially detected as being
   /// hyper-rectangular on the specified contiguous set of identifiers.

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

if (isLower) {
         minOrMaxConst = boundConst;
     }
   }
   return minOrMaxConst;
 }
 
 Optional<int64_t> IntegerPolyhedron::getConstantBound(BoundType type,
                                                       unsigned pos) const {
-  assert(type != BoundType::EQ && "EQ not implemented");
-  IntegerPolyhedron tmpCst(*this);
   if (type == BoundType::LB)
-    return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
-  return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
+    return IntegerPolyhedron(*this)
+        .computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
+  if (type == BoundType::UB)
+    return IntegerPolyhedron(*this)
+        .computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
+
+  assert(type == BoundType::EQ && "expected EQ");
+  Optional<int64_t> lb =
+      IntegerPolyhedron(*this)
+          .computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
+  Optional<int64_t> ub =
+      IntegerPolyhedron(*this)
+          .computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
+  return (lb && ub && *lb == *ub) ? Optional<int64_t>(*ub) : None;
 }
 
 // A simple (naive and conservative) check for hyper-rectangularity.
 bool IntegerPolyhedron::isHyperRectangular(unsigned pos, unsigned num) const {
   assert(pos < getNumCols() - 1);
   // Check for two non-zero coefficients in the range [pos, pos + sum).
   for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
     unsigned sum = 0;

mlir/lib/Dialect/SCF/Utils/AffineCanonicalizationUtils.cpp

for (unsigned i = resultDimStart; i < resultDimStart + numResults; ++i) {
     if (!newConstr.isEmpty())
       return failure();
   }
 
   // Lower and upper bound of `op` are equal. Replace `minOp` with its bound.
   AffineMap newMap = alignedBoundMap;
   SmallVector<Value> newOperands;
   unpackOptionalValues(constraints.getMaybeDimAndSymbolValues(), newOperands);
+  // If dims/symbols have known constant values, use those in order to simplify
+  // the affine map further.
+  for (int64_t i = 0, e = constraints.getNumIds(); i < e; ++i) {
+    // Skip unused operands and operands that are already constants.
+    if (!newOperands[i] || getConstantIntValue(newOperands[i]))
+      continue;
+    if (auto bound = constraints.getConstantBound(FlatAffineConstraints::EQ, i))
+      newOperands[i] =
+          rewriter.create<arith::ConstantIndexOp>(op->getLoc(), *bound);
+  }
   mlir::canonicalizeMapAndOperands(&newMap, &newOperands);
   rewriter.setInsertionPoint(op);
   rewriter.replaceOpWithNewOp<AffineApplyOp>(op, newMap, newOperands);
   return success();
 }
 
 static LogicalResult
 addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv,
addLoopRangeConstraints(FlatAffineValueConstraints &constraints, Value iv,
   // If loop lower/upper bounds are constant: Add EQ constraint.
   Optional<int64_t> lbInt = getConstantIntValue(lb);
   Optional<int64_t> ubInt = getConstantIntValue(ub);
   if (lbInt)
     constraints.addBound(FlatAffineConstraints::EQ, dimLb, *lbInt);
   if (ubInt)
     constraints.addBound(FlatAffineConstraints::EQ, dimUb, *ubInt);
 
-  // iv >= lb (equiv.: iv - lb >= 0)
+  // Lower bound: iv >= lb (equiv.: iv - lb >= 0)
   SmallVector<int64_t> ineqLb(constraints.getNumCols(), 0);
   ineqLb[dimIv] = 1;
   ineqLb[dimLb] = -1;
   constraints.addInequality(ineqLb);
 
+  // Upper bound
+  AffineExpr ivUb;
+  if (lbInt && ubInt && (*lbInt + *stepInt >= *ubInt)) {
+    // The loop has at most one iteration.
+    // iv < lb + 1
+    // TODO: Try to derive this constraint by simplifying the expression in
+    // the else-branch.
+    ivUb = rewriter.getAffineDimExpr(dimLb) + 1;
+  } else {
+    // The loop may have more than one iteration.
     // iv < lb + step * ((ub - lb - 1) floorDiv step) + 1
     AffineExpr exprLb = lbInt ? rewriter.getAffineConstantExpr(*lbInt)
                               : rewriter.getAffineDimExpr(dimLb);
     AffineExpr exprUb = ubInt ? rewriter.getAffineConstantExpr(*ubInt)
                               : rewriter.getAffineDimExpr(dimUb);
-  AffineExpr ivUb =
-      exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
+    ivUb = exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
+  }
   auto map = AffineMap::get(
       /*dimCount=*/constraints.getNumDimIds(),
       /*symbolCount=*/constraints.getNumSymbolIds(), /*result=*/ivUb);
 
   return constraints.addBound(FlatAffineConstraints::UB, dimIv, map);
 }
 
 /// Canonicalize min/max operations in the context of for loops with a known

mlir/test/Dialect/SCF/for-loop-canonicalization.mlir

func @tensor_dim_of_loop_result_no_canonicalize(%t : tensor<?x?xf32>,
   %c10 = arith.constant 10 : index
   %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %u)
       -> (tensor<?x?xf32>, tensor<?x?xf32>) {
     scf.yield %arg0, %u : tensor<?x?xf32>, tensor<?x?xf32>
   }
   %dim = tensor.dim %1, %c0 : tensor<?x?xf32>
   return %dim : index
 }
+
+// -----
+
+// CHECK-LABEL: func @one_trip_scf_for_canonicalize_min
+//       CHECK:   %[[C4:.*]] = arith.constant 4 : i64
+//       CHECK:   scf.for
+//       CHECK:     memref.store %[[C4]], %{{.*}}[] : memref<i64>
+func @one_trip_scf_for_canonicalize_min(%A : memref<i64>) {
+  %c0 = arith.constant 0 : index
+  %c2 = arith.constant 2 : index
+  %c4 = arith.constant 4 : index
+
+  scf.for %i = %c0 to %c4 step %c4 {
+    %1 = affine.min affine_map<(d0, d1)[] -> (4, d1 - d0)> (%i, %c4)
+    %2 = arith.index_cast %1: index to i64
+    memref.store %2, %A[]: memref<i64>
+  }
+  return
+}