[mlir][affine][transform] Simplify affine.min/max ops with given constraints

[springerm]

Summary

This transform op uses mlir::simplifyConstrainedMinMaxOp to simplify affine.min and affine.max ops based on a given constraints.

Depends On: https://reviews.llvm.org/D141141

[nicolasvasilache]

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

52

this lacks a verifier that args are properly sized and invalid tests.

56

Why do you need this?

Is it practically too annoying to have with { (%1, %2, %3) confined to (0, 0, 0) - (16, 32, 64) } ?

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

56

Why not applyToOne ?

64

llvm::zip will drop stuff, I think there is a recently added asserting version, try zip_equal

85

Seems like this should be a helper in FlatAffineValueConstraints ?

100

I think the typed form of the op gives you that in a uniform way ?

120

Have you considered the handle invalidation/staleness implications of using this, in particular in larger transforms?

This looks fishy to me ..

132

I'd rather have a slightly less nice syntax and drop all this C++.

If you really insist on nicer syntax, I'd want to see something more general and reusable like the subview custom printer/parser slightly adapted to this use case.

[springerm]

address comments

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

56

New syntax: with [%0] within [0] and [32]. I wanted parentheses but couldn't figure out how to print the DenseI64ArrayAttr without []. Same for the -, it is not a valid literal, so changed to and.

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

56

Doesn't work with multiple PDL_Operation operands (bounded_values is the problem).

85

I could add it in this change, but there is a minor inconsistency in the FlatAffineValueConstraints API that should be addressed first. Upper bounds are sometimes inclusive and sometimes exclusive. This needs to be fixed in FlatAffineValueConstraints and IntegerPolyhedron (and all users of the API).

Instead of adding another function that deals with UB, can we land this as is and add the helper together with the cleanup that fixes the API?

100

I checked with Alex and this requires adding a new transform handle type + relatively complex TableGen because we need to support !transform.op<"affine.min OR affine.max">. I could add it but not worth it IMO unless we really want to go properly typed everywhere.

120

Yes, I discussed this with Alex.

This op consumes (invalidates) the target handle and does not return anything. Also there is a check there is no overlap among target ops and the constrained ops, so we cannot come into a situation where a constrained op is erased while the transform is applied.

Note: This applyOpPatternsAndFold does not return LogicalResult, but a bool that indicates whether any IR was rewritten or not. The (void) cast does *not* discard error state here.

Updated the comment: Removed the We do not compose chained affine.min / affine.max ops part, which made no sense. affine.min and affine.max ops cannot be composed. (Ops of the same type could be but that is not good enough, we still need fixpoint iteration.)

[springerm]

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

120

Actually, strict should be set to true here. But then the IR does not simplify as expected. Need to take a closer look at the GreedyPatternRewriter to see what's going on...

[nicolasvasilache]

Thanks!
Please do get to the bottom of the rewriter stuff before landing.

[springerm]

In D140997#4031629, @nicolasvasilache wrote:

Thanks!
Please do get to the bottom of the rewriter stuff before landing.

Turns out this is a bug in the MultiOpPatternRewriteDriver. Fix in D141141.

[springerm]

update

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

 //===- AffineTransformOps.h - Affine transformation ops ---------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
 #define MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
 
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
 class AffineForOp;
 namespace func {
 class FuncOp;

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

 //===- AffineTransformOps.td - Affine transformation ops ---*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef AFFINE_TRANSFORM_OPS
 #define AFFINE_TRANSFORM_OPS
 
+include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
 def Transform_AffineForOp : Transform_ConcreteOpType<"affine.for">;
 
+def SimplifyBoundedAffineOpsOp
+    : Op<Transform_Dialect, "affine.simplify_bounded_affine_ops",
+         [DeclareOpInterfaceMethods<TransformOpInterface>,
+          DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let description = [{
+    Simplify the targeted affine.min / affine.max ops given the supplied
+    lower and upper bounds for values that may be used as target op operands.
+
+    Example:
+    ```
+    %0 = transform.structured.match ops{["affine.min", "affine.max"]} in %arg1
+    %1 = transform.structured.match ops{["gpu.lane_id"]} in %arg1
+    // Multiple bounds can be specified.
+    transform.affine.simplify_bounded_affine_ops %0 with {%1 in (0, 32)}
+    ```
+
+    Bounded op handles (`%1`) must be mapped to ops that have a single result of
+    index type. The sets of target ops and bounded ops must not overlap.
+
+    #### Return modes
+
+    Target ops must be affine.min or affine.max ops. This transform consumes the
+    target handle and does not produce any handle. It reads the bounded op
+    handles.
+
+    TODO: Support affine.apply targets.
+    TODO: Allow mixed PDL_Operation/int64_t for lower_bounds and upper_bounds.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       Variadic<PDL_Operation>:$bounded_values,
+                       DenseI64ArrayAttr:$lower_bounds,
+                       DenseI64ArrayAttr:$upper_bounds);
+  let results = (outs);
+  let hasCustomAssemblyFormat = 1;
+}
+
 #endif // Affine_TRANSFORM_OPS

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

 //=== AffineTransformOps.cpp - Implementation of Affine transformation ops ===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
+#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
+#include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/PDL/IR/PDL.h"
+#include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformUtils.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
+using namespace mlir::transform;
+
+//===----------------------------------------------------------------------===//
+// SimplifyBoundedAffineOpsOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Simplify affine.min / affine.max ops with the given constraints. They are
+/// either rewritten to affine.apply or left unchanged.
+template <typename OpTy>
+struct SimplifyAffineMinMaxOp : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+  SimplifyAffineMinMaxOp(MLIRContext *ctx,
+                         const FlatAffineValueConstraints &constraints,
+                         PatternBenefit benefit = 1)
+      : OpRewritePattern<OpTy>(ctx, benefit), constraints(constraints) {}
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    FailureOr<AffineValueMap> simplified =
+        simplifyConstrainedMinMaxOp(op, constraints);
+    if (failed(simplified))
+      return failure();
+    rewriter.replaceOpWithNewOp<AffineApplyOp>(op, simplified->getAffineMap(),
+                                               simplified->getOperands());
+    return success();
+  }
+
+  const FlatAffineValueConstraints &constraints;
+};
+} // namespace
+
+DiagnosedSilenceableFailure
+SimplifyBoundedAffineOpsOp::apply(TransformResults &results,
+                                  TransformState &state) {
+  // Get constraints for bounded values.
+  SmallVector<int64_t> lbs;
+  SmallVector<int64_t> ubs;
+  SmallVector<Value> boundedValues;
+  DenseSet<Operation *> boundedOps;
+  for (const auto &it :
+       llvm::zip(getBoundedValues(), getLowerBounds(), getUpperBounds())) {
+    Value handle = std::get<0>(it);
+    ArrayRef<Operation *> boundedValueOps = state.getPayloadOps(handle);
+    for (Operation *op : boundedValueOps) {
+      if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
+        auto diag =
+            emitDefiniteFailure()
+            << "expected bounded value handle to point to one or multiple "
+               "single-result index-typed ops";
+        diag.attachNote(op->getLoc()) << "multiple/non-index result";
+        return diag;
+      }
+      boundedValues.push_back(op->getResult(0));
+      boundedOps.insert(op);
+      lbs.push_back(std::get<1>(it));
+      ubs.push_back(std::get<2>(it));
+    }
+  }
+
+  // Build constraint set.
+  FlatAffineValueConstraints cstr;
+  for (const auto &it : llvm::zip(boundedValues, lbs, ubs)) {
+    unsigned pos;
+    if (!cstr.findVar(std::get<0>(it), &pos))
+      pos = cstr.appendSymbolVar(std::get<0>(it));
+    cstr.addBound(FlatAffineValueConstraints::BoundType::LB, pos,
+                  std::get<1>(it));
+    // Note: addBound bounds are inclusive, but specified UB is exclusive.
+    cstr.addBound(FlatAffineValueConstraints::BoundType::UB, pos,
+                  std::get<2>(it) - 1);
+  }
+
+  // Transform all targets.
+  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
+  for (Operation *target : targets) {
+    if (!isa<AffineMinOp, AffineMaxOp>(target)) {
+      auto diag = emitDefiniteFailure()
+                  << "target must be affine.min or affine.max";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+    if (boundedOps.contains(target)) {
+      auto diag = emitDefiniteFailure()
+                  << "target op result must not be constrainted";
+      diag.attachNote(target->getLoc()) << "target/constrained op";
+      return diag;
+    }
+  }
+  SmallVector<Operation *> transformed;
+  RewritePatternSet patterns(getContext());
+  AffineMinOp::getCanonicalizationPatterns(patterns, getContext());
+  patterns.insert<SimplifyAffineMinMaxOp<AffineMinOp>,
+                  SimplifyAffineMinMaxOp<AffineMaxOp>>(getContext(), cstr);
+  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+  // We do not compose chained affine.min / affine.max ops before simplifying
+  // them. Instead, we apply the simplification pattern to a fixpoint.
+  (void)applyOpPatternsAndFold(targets, frozenPatterns, /*strict=*/false);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void SimplifyBoundedAffineOpsOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  consumesHandle(getTarget(), effects);
+  for (Value v : getBoundedValues())
+    onlyReadsHandle(v, effects);
+  modifiesPayload(effects);
+}
+
+ParseResult SimplifyBoundedAffineOpsOp::parse(OpAsmParser &parser,
+                                              OperationState &result) {
+  OpAsmParser::UnresolvedOperand target;
+  auto pdlOperationType =
+      pdl::OperationType::get(parser.getBuilder().getContext());
+  if (parser.parseOperand(target) ||
+      parser.resolveOperand(target, pdlOperationType, result.operands))
+    return failure();
+
+  SmallVector<int64_t> lowerBounds, upperBounds;
+  if (parser.parseOptionalKeyword("with").succeeded()) {
+    if (parser.parseLBrace())
+      return failure();
+    while (true) {
+      // Parse $bounded_value => (lb, ub).
+      OpAsmParser::UnresolvedOperand boundedValue;
+      int64_t lb, ub;
+      if (parser.parseOperand(boundedValue) ||
+          parser.resolveOperand(boundedValue, pdlOperationType,
+                                result.operands) ||
+          parser.parseKeyword("in") || parser.parseLParen() ||
+          parser.parseInteger(lb) || parser.parseComma() ||
+          parser.parseInteger(ub) || parser.parseRParen())
+        return failure();
+      lowerBounds.push_back(lb);
+      upperBounds.push_back(ub);
+      if (parser.parseOptionalRBrace().succeeded())
+        break;
+      if (parser.parseComma())
+        return failure();
+    }
+  }
+  result.addAttribute(
+      SimplifyBoundedAffineOpsOp::getLowerBoundsAttrName(result.name)
+          .getValue(),
+      parser.getBuilder().getDenseI64ArrayAttr(lowerBounds));
+  result.addAttribute(
+      SimplifyBoundedAffineOpsOp::getUpperBoundsAttrName(result.name)
+          .getValue(),
+      parser.getBuilder().getDenseI64ArrayAttr(upperBounds));
+
+  if (failed(parser.parseOptionalAttrDict(result.attributes)))
+    return failure();
+
+  return success();
+}
+
+void SimplifyBoundedAffineOpsOp::print(OpAsmPrinter &p) {
+  p << " " << getTarget() << " ";
+  if (!getBoundedValues().empty()) {
+    p << "with {";
+    for (int i = 0; i < getBoundedValues().size(); ++i) {
+      p << getBoundedValues()[i] << " in (" << getLowerBounds()[i] << ", "
+        << getUpperBounds()[i] << ")";
+      if (i < getBoundedValues().size() - 1)
+        p << ", ";
+    }
+    p << "} ";
+  }
+  p.printOptionalAttrDict(getOperation()->getAttrs(),
+                          {getLowerBoundsAttrName(), getUpperBoundsAttrName()});
+}
 
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 class AffineTransformDialectExtension
     : public transform::TransformDialectExtension<

mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineTransformOps
   AffineTransformOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine/TransformOps
 
   DEPENDS
   MLIRAffineTransformOpsIncGen
 
   LINK_LIBS PUBLIC
+  MLIRAffineAnalysis
   MLIRAffineDialect
   MLIRFuncDialect
   MLIRIR
   MLIRPDLDialect
   MLIRAffineDialect
   MLIRAffineTransforms
   MLIRAffineUtils
   MLIRTransformDialect
   MLIRVectorDialect
 )

mlir/test/Dialect/Affine/transform-op-simplify-bounded-affine-ops.mlir

+// RUN: mlir-opt  %s -allow-unregistered-dialect --test-transform-dialect-interpreter --split-input-file | FileCheck %s
+
+//     CHECK: func @simplify_min_max()
+// CHECK-DAG:   %[[c50:.*]] = arith.constant 50 : index
+// CHECK-DAG:   %[[c100:.*]] = arith.constant 100 : index
+//     CHECK:   return %[[c50]], %[[c100]]
+func.func @simplify_min_max() -> (index, index) {
+  %0 = "test.some_op"() : () -> (index)
+  %1 = affine.min affine_map<()[s0] -> (50, 100 - s0)>()[%0]
+  %2 = affine.max affine_map<()[s0] -> (100, 80 + s0)>()[%0]
+  return %1, %2 : index, index
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["affine.min", "affine.max"]} in %arg1
+  %1 = transform.structured.match ops{["test.some_op"]} in %arg1
+  transform.affine.simplify_bounded_affine_ops %0 with {%1 in (0, 20)}
+}
+
+// -----
+
+// CHECK: func @simplify_min_sequence()
+// CHECK:   %[[c1:.*]] = arith.constant 1 : index
+// CHECK:   return %[[c1]]
+func.func @simplify_min_sequence() -> index {
+  %1 = "test.workgroup_id"() : () -> (index)
+  %2 = affine.min affine_map<()[s0] -> (s0 * -32 + 1023, 32)>()[%1]
+  %3 = "test.thread_id"() : () -> (index)
+  %4 = affine.min affine_map<()[s0, s1] -> (s0 - s1 * (s0 ceildiv 32), s0 ceildiv 32)>()[%2, %3]
+  return %4 : index
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["affine.min"]} in %arg1
+  %1 = transform.structured.match ops{["test.workgroup_id"]} in %arg1
+  %2 = transform.structured.match ops{["test.thread_id"]} in %arg1
+  transform.affine.simplify_bounded_affine_ops %0 with {%1 in (0, 31), %2 in (0, 31)}
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

 )
 
 cc_library(
     name = "AffineTransformOps",
     srcs = glob(["lib/Dialect/Affine/TransformOps/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/Affine/TransformOps/*.h"]),
     includes = ["include"],
     deps = [
+        ":AffineAnalysis",
         ":AffineDialect",
         ":AffineTransformOpsIncGen",
         ":AffineTransforms",
         ":AffineUtils",
         ":FuncDialect",
         ":IR",
+        ":PDLDialect",
         ":TransformDialect",
+        ":Transforms",
         ":VectorDialect",
     ],
 )
 
 ##---------------------------------------------------------------------------##
 # AMDGPU dialect.
 ##---------------------------------------------------------------------------##