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

This reverts commit ee1bf18672.

It breaks IREE lowering. Reverting the commit for now while we
investigate what's going on.

mlir/include/mlir/Analysis/AffineStructures.h

@@ -385,6 +385,7 @@ public:
 
   /// 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;
 
   /// Gets the lower and upper bound of the `offset` + `pos`th identifier

mlir/lib/Analysis/AffineStructures.cpp

@@ -2836,22 +2836,11 @@ FlatAffineConstraints::computeConstantLowerOrUpperBound(unsigned pos) {
 
 Optional<int64_t> FlatAffineConstraints::getConstantBound(BoundType type,
                                                           unsigned pos) const {
+  assert(type != BoundType::EQ && "EQ not implemented");
   FlatAffineConstraints tmpCst(*this);
   if (type == BoundType::LB)
-    return FlatAffineConstraints(*this)
-        .computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
-  if (type == BoundType::UB)
-    return FlatAffineConstraints(*this)
-        .computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
-
-  assert(type == BoundType::EQ && "expected EQ");
-  Optional<int64_t> lb =
-      FlatAffineConstraints(*this)
-          .computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
-  Optional<int64_t> ub =
-      FlatAffineConstraints(*this)
-          .computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
-  return (lb && ub && *lb == *ub) ? Optional<int64_t>(*ub) : None;
+    return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
+  return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
 }
 
 // A simple (naive and conservative) check for hyper-rectangularity.

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

@@ -305,16 +305,6 @@ canonicalizeMinMaxOp(RewriterBase &rewriter, Operation *op, AffineMap map,
   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; i < constraints.getNumDimAndSymbolIds(); ++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);
@@ -467,30 +457,19 @@ mlir::scf::canonicalizeMinMaxOpInLoop(RewriterBase &rewriter, Operation *op,
     if (ubInt)
       constraints.addBound(FlatAffineConstraints::EQ, dimUb, *ubInt);
 
-    // Lower bound: iv >= lb (equiv.: iv - lb >= 0)
+    // 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);
-      ivUb =
-          exprLb + 1 + (*stepInt * ((exprUb - exprLb - 1).floorDiv(*stepInt)));
-    }
+    // 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)));
     auto map = AffineMap::get(
         /*dimCount=*/constraints.getNumDimIds(),
         /*symbolCount=*/constraints.getNumSymbolIds(), /*result=*/ivUb);

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

@@ -348,22 +348,3 @@ func @tensor_dim_of_loop_result_no_canonicalize(%t : 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
-}