address reviewer comments

Author: fabianmcg

mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.td

@@ -65,38 +65,33 @@ def SimplifyBoundedAffineOpsOp
 
 def SimplifyMinMaxAffineOpsOp :
   Op<Transform_Dialect, "affine.simplify_min_max_affine_ops", [
-    TransformOpInterface,
-    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-    TransformEachOpTrait
+    DeclareOpInterfaceMethods<TransformOpInterface>,
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
   ]> {
   let description = [{
-    Simplify all the affine.min / affine.max ops being targeted or nested in the
-    target operation, using the `mlir::affine::simplifyAffineMinMaxOps`
-    transform.
+    Simplify the targeted `affine.min` / `affine.max` ops using the
+    `mlir::affine::simplifyAffineMinMaxOps` transform.
 
     Example:
     ```
-    %0 = transform.structured.match ops{["gpu.launch", "affine.max"]} in %arg1
+    %0 = transform.structured.match ops{["affine.max"]} in %arg1
     transform.affine.simplify_min_max_affine_ops %0 : !transform.any_op
     ```
 
     #### Return modes
 
     This transform consumes the target handle and does not produce any results.
-    This transforms never produces errors.
+    This transforms definitely fails if any of the targeted operations is not an
+    `affine.min` or `affine.max` operation, or if the canonicalization patterns
+    failed to converge.
+    This transform silently fails if none of the operations were simplified.
+    Otherwise, it succeeds.
   }];
   let arguments = (ins TransformHandleTypeInterface:$target);
   let results = (outs);
   let assemblyFormat = [{
       $target attr-dict `:` type($target)
   }];
-  let extraClassDeclaration = [{
-    ::mlir::DiagnosedSilenceableFailure applyToOne(
-        ::mlir::transform::TransformRewriter &rewriter,
-        ::mlir::Operation *target,
-        ::mlir::transform::ApplyToEachResultList &results,
-        ::mlir::transform::TransformState &state);
-  }];
 }
 
 #endif // Affine_TRANSFORM_OPS

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

@@ -129,23 +129,37 @@ OpFoldResult materializeComputedBound(
     OpBuilder &b, Location loc, AffineMap boundMap,
     ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands);
 
-/// Tries to simplify all affine min or max operations under `topOp`. The
-/// transform works by finding disjoint sets of affine result expressions
-/// bounded by a common affine expression on the min/max operation. It populates
-/// `modifiedOps` with all the operations modified by the transform.
+/// This transform tries to simplify the affine min operation `op`, by finding a
+/// common lower bound for a set of expressions in the affine map results. It
+/// returns whether the transform updated `op`'s affine map.
 ///
 /// In concrete terms, given an operation like:
 /// `affine.min affine_map<(d0)[s0, s1] -> (d0, s1, s0, 128)>(%i)[%s0, %s1]`
-/// If `d0 < 128` and `128 < s1 < s0`, the transform will update the op to:
+/// If `d0 < 128` and `128 < s1 < s0`, the transform will update `op` to:
 /// `affine.min affine_map<(d0)[s0, s1] -> (d0, 128)>(%i)[%s0, %s1]`.
-void simplifyAffineMinMaxOps(RewriterBase &rewriter, Operation *topOp,
-                             SmallVectorImpl<Operation *> &modifiedOps);
-/// Applies `simplifyAffineMinMaxOps` to a single operation and returns whether
-/// the operation was modified.
 bool simplifyAffineMinOp(RewriterBase &rewriter, AffineMinOp op);
-/// Applies `simplifyAffineMinMaxOps` to a single operation and returns whether
-/// the operation was modified.
+
+/// This transform tries to simplify the affine max operation `op`, by finding a
+/// common upper bound for a set of expressions in the affine map results. It
+/// returns whether the transform updated `op`'s affine map.
+///
+/// In concrete terms, given an operation like:
+/// `affine.max affine_map<(d0)[s0, s1] -> (d0, s1, s0, 128)>(%i)[%s0, %s1]`
+/// If `d0 > 128` and `s0 > s1 > 128`, the transform will update `op` to:
+/// `affine.max affine_map<(d0)[s0, s1] -> (d0, s0)>(%i)[%s0, %s1]`.
 bool simplifyAffineMaxOp(RewriterBase &rewriter, AffineMaxOp op);
+
+/// This transform applies `simplifyAffineMinOp` and `simplifyAffineMaxOp` to
+/// all the `affine.min` or `affine.max` operations in `ops`. After
+/// simplification, it invokes the `affine.min/max` canonicalization patterns on
+/// `ops`.
+///
+/// This transform returns failure if the greedy pattern rewriter failed to
+/// converge during canonicalization, otherwise it returns success. If provided,
+/// `modified` is set to `true` if the IR was modified in any way.
+LogicalResult simplifyAffineMinMaxOps(RewriterBase &rewriter,
+                                      ArrayRef<Operation *> ops,
+                                      bool *modified = nullptr);
 } // namespace affine
 } // namespace mlir
 

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

@@ -113,7 +113,7 @@ SimplifyBoundedAffineOpsOp::apply(transform::TransformRewriter &rewriter,
     }
     if (boundedOps.contains(target)) {
       auto diag = emitDefiniteFailure()
-                  << "target op result must not be constrainted";
+                  << "target op result must not be constrained";
       diag.attachNote(target->getLoc()) << "target/constrained op";
       return diag;
     }
@@ -152,12 +152,30 @@ void SimplifyBoundedAffineOpsOp::getEffects(
 //===----------------------------------------------------------------------===//
 // SimplifyMinMaxAffineOpsOp
 //===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure SimplifyMinMaxAffineOpsOp::applyToOne(
-    TransformRewriter &rewriter, Operation *target,
-    ApplyToEachResultList &results, TransformState &state) {
-  SmallVector<Operation *> modifiedOps;
-  simplifyAffineMinMaxOps(rewriter, target, modifiedOps);
+DiagnosedSilenceableFailure
+SimplifyMinMaxAffineOpsOp::apply(transform::TransformRewriter &rewriter,
+                                 TransformResults &results,
+                                 TransformState &state) {
+  SmallVector<Operation *> targets;
+  for (Operation *target : state.getPayloadOps(getTarget())) {
+    if (!isa<AffineMinOp, AffineMaxOp>(target)) {
+      auto diag = emitDefiniteFailure()
+                  << "target must be affine.min or affine.max";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+    targets.push_back(target);
+  }
+  bool modified = false;
+  if (failed(mlir::affine::simplifyAffineMinMaxOps(rewriter, targets,
+                                                   &modified))) {
+    return emitDefiniteFailure()
+           << "affine.min/max simplification did not converge";
+  }
+  if (!modified) {
+    return emitSilenceableError()
+           << "the transform failed to simplify any of the target operations";
+  }
   return DiagnosedSilenceableFailure::success();
 }
 

mlir/lib/Dialect/Affine/Transforms/SimplifyAffineMinMax.cpp

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/Affine/Transforms/Transforms.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/IntEqClasses.h"
 #include "llvm/Support/Debug.h"
 
@@ -69,6 +70,10 @@ static bool simplifyAffineMinMaxOp(RewriterBase &rewriter, AffineOp affineOp) {
     // Check against the other variables.
     for (size_t j = i + 1; j < variables.size(); ++j) {
       unsigned jEqClass = boundedClasses.findLeader(j);
+      // Skip if the class is the same.
+      if (jEqClass == eqClass)
+        continue;
+
       // Get the bound of the equivalence class or itself.
       Variable *nv = bounds.lookup_or(jEqClass, &variables[j]);
 
@@ -93,7 +98,7 @@ static bool simplifyAffineMinMaxOp(RewriterBase &rewriter, AffineOp affineOp) {
       // Join the equivalent classes and update the bound if necessary.
       LLVM_DEBUG({
         DBGS() << "-- merging classes: #" << i << ", #" << jEqClass
-               << ", is lhs <= rhs: " << *cmpResult << "`\n";
+               << ", is cmp(lhs, rhs): " << *cmpResult << "`\n";
       });
       if (*cmpResult) {
         boundedClasses.join(eqClass, jEqClass);
@@ -137,17 +142,33 @@ bool mlir::affine::simplifyAffineMaxOp(RewriterBase &rewriter, AffineMaxOp op) {
   return simplifyAffineMinMaxOp(rewriter, op);
 }
 
-void mlir::affine::simplifyAffineMinMaxOps(
-    RewriterBase &rewriter, Operation *topOp,
-    SmallVectorImpl<Operation *> &modifiedOps) {
-  assert(topOp && "null-op");
-  topOp->walk([&](Operation *op) {
-    if (auto affineOp = dyn_cast<AffineMinOp>(op)) {
-      if (simplifyAffineMinMaxOp(rewriter, affineOp))
-        modifiedOps.push_back(op);
-    } else if (auto affineOp = dyn_cast<AffineMaxOp>(op)) {
-      if (simplifyAffineMinMaxOp(rewriter, affineOp))
-        modifiedOps.push_back(op);
-    }
-  });
+LogicalResult mlir::affine::simplifyAffineMinMaxOps(RewriterBase &rewriter,
+                                                    ArrayRef<Operation *> ops,
+                                                    bool *modified) {
+  bool changed = false;
+  for (Operation *op : ops) {
+    if (auto minOp = dyn_cast<AffineMinOp>(op))
+      changed = simplifyAffineMinOp(rewriter, minOp) || changed;
+    else if (auto maxOp = cast<AffineMaxOp>(op))
+      changed = simplifyAffineMaxOp(rewriter, maxOp) || changed;
+  }
+  RewritePatternSet patterns(rewriter.getContext());
+  AffineMaxOp::getCanonicalizationPatterns(patterns, rewriter.getContext());
+  AffineMinOp::getCanonicalizationPatterns(patterns, rewriter.getContext());
+  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+  if (modified)
+    *modified = changed;
+  // Canonicalize to a fixpoint.
+  if (failed(applyOpPatternsGreedily(
+          ops, frozenPatterns,
+          GreedyRewriteConfig()
+              .setListener(
+                  static_cast<RewriterBase::Listener *>(rewriter.getListener()))
+              .setStrictness(GreedyRewriteStrictness::ExistingAndNewOps),
+          &changed))) {
+    return failure();
+  }
+  if (modified)
+    *modified = changed;
+  return success();
 }

mlir/test/Dialect/Affine/transform-op-simplify-min-max-ops.mlir

@@ -61,12 +61,8 @@ func.func @overlapping_constraints() -> (index, index) {
 
 module attributes {transform.with_named_sequence} {
   transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
-    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    %0 = transform.structured.match ops{["affine.min", "affine.max"]} in %arg0 : (!transform.any_op) -> !transform.any_op
     transform.affine.simplify_min_max_affine_ops %0 : !transform.any_op
-    %1 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
-    transform.apply_patterns to %1 {
-      transform.apply_patterns.canonicalization
-    } {apply_cse} : !transform.any_op
     transform.yield 
   }
 }

mlir/test/Dialect/Linalg/transform-op-pad.mlir

@@ -464,8 +464,8 @@ module attributes {transform.with_named_sequence} {
 // CHECK: %[[LHS:.*]] = tensor.pad
 // CHECK: %[[RHS:.*]] = tensor.pad
 // CHECK: scf.for
-// CHECK: tensor.extract_slice %[[LHS]][0, %{{.*}}] [%{{.*}}, 32]
-// CHECK: tensor.extract_slice %[[RHS]][0, %{{.*}}] [%{{.*}}, 32]
+// CHECK-DAG: tensor.extract_slice %[[LHS]][0, %{{.*}}] [%{{.*}}, 32]
+// CHECK-DAG: tensor.extract_slice %[[RHS]][0, %{{.*}}] [%{{.*}}, 32]
 func.func @dyn_pad_tiling(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
   %0 = linalg.matmul_transpose_b ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
@@ -480,9 +480,9 @@ module attributes {transform.with_named_sequence} {
     transform.apply_patterns to %2 {
       transform.apply_patterns.canonicalization
     } {apply_cse} : !transform.any_op
-    transform.affine.simplify_min_max_affine_ops %2 : !transform.any_op
-    %3 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
-    transform.apply_patterns to %3 {
+    %3 = transform.structured.match ops{["affine.min", "affine.max"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    transform.affine.simplify_min_max_affine_ops %3 : !transform.any_op
+    transform.apply_patterns to %2 {
       transform.apply_patterns.canonicalization
     } {apply_cse} : !transform.any_op
     transform.yield