[mlir] Add transform.foreach_match

Add a new transform op combinator that implements an "if-then-else"
style of mechanism for applying transformations. Its main purpose is to
serve as a higher-level driver when applying multiple transform scripts
to potentially overlapping pieces of the payload IR. This is similar to
how the various rewrite drivers operate in C++, but at a higher level
and with more declarative expressions. This is not intended to replace
existing pattern-based rewrites, but to to drive more complex
transformations that are exposed in the transform dialect and are too
complex to be expressed as simple declarative rewrites.

Reviewed By: springerm

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

Author: ftynse

mlir/include/mlir/Dialect/Transform/IR/CMakeLists.txt

@@ -28,6 +28,10 @@ add_mlir_doc(TransformOps TransformOps Dialects/ -gen-op-doc -dialect=transform)
 add_mlir_interface(TransformInterfaces)
 add_mlir_doc(TransformInterfaces TransformOpInterfaces Dialects/ -gen-op-interface-docs)
 
+add_mlir_interface(MatchInterfaces)
+add_dependencies(MLIRMatchInterfacesIncGen MLIRTransformInterfacesIncGen)
+add_mlir_doc(TransformInterfaces MatchOpInterfaces Dialects/ -gen-op-interface-docs)
+
 set(LLVM_TARGET_DEFINITIONS TransformInterfaces.td)
 mlir_tablegen(TransformTypeInterfaces.h.inc -gen-type-interface-decls)
 mlir_tablegen(TransformTypeInterfaces.cpp.inc -gen-type-interface-defs)

mlir/include/mlir/Dialect/Transform/IR/MatchInterfaces.h

@@ -0,0 +1,73 @@
+//===- MatchInterfaces.h - Transform Dialect Interfaces ---------*- 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_TRANSFORM_IR_MATCHINTERFACES_H
+#define MLIR_DIALECT_TRANSFORM_IR_MATCHINTERFACES_H
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+namespace transform {
+class MatchOpInterface;
+
+template <typename OpTy>
+class SingleOpMatcherOpTrait
+    : public OpTrait::TraitBase<OpTy, SingleOpMatcherOpTrait> {
+  template <typename T>
+  using has_get_operand_handle =
+      decltype(std::declval<T &>().getOperandHandle());
+  template <typename T>
+  using has_match_operation = decltype(std::declval<T &>().matchOperation(
+      std::declval<Operation *>(), std::declval<TransformResults &>(),
+      std::declval<TransformState &>()));
+
+public:
+  static LogicalResult verifyTrait(Operation *op) {
+    static_assert(llvm::is_detected<has_get_operand_handle, OpTy>::value,
+                  "SingleOpMatcherOpTrait expects operation type to have the "
+                  "getOperandHandle() method");
+    static_assert(llvm::is_detected<has_match_operation, OpTy>::value,
+                  "SingleOpMatcherOpTrait expected operation type to have the "
+                  "matchOperation(Operation *, TransformResults &, "
+                  "TransformState &) method");
+
+    // This must be a dynamic assert because interface registration is dynamic.
+    assert(isa<MatchOpInterface>(op) &&
+           "SingleOpMatchOpTrait is only available on operations with "
+           "MatchOpInterface");
+    Value operandHandle = cast<OpTy>(op).getOperandHandle();
+    if (!operandHandle.getType().isa<TransformHandleTypeInterface>()) {
+      return op->emitError() << "SingleOpMatchOpTrait requires the op handle "
+                                "to be of TransformHandleTypeInterface";
+    }
+
+    return success();
+  }
+
+  DiagnosedSilenceableFailure apply(TransformResults &results,
+                                    TransformState &state) {
+    Value operandHandle = cast<OpTy>(this->getOperation()).getOperandHandle();
+    ArrayRef<Operation *> payload = state.getPayloadOps(operandHandle);
+    if (payload.size() != 1) {
+      return emitDefiniteFailure(this->getOperation()->getLoc())
+             << "SingleOpMatchOpTrait requires the operand handle to point to "
+                "a single payload op";
+    }
+
+    return cast<OpTy>(this->getOperation())
+        .matchOperation(payload[0], results, state);
+  }
+};
+
+} // namespace transform
+} // namespace mlir
+
+#include "mlir/Dialect/Transform/IR/MatchInterfaces.h.inc"
+
+#endif // MLIR_DIALECT_TRANSFORM_IR_MATCHINTERFACES_H

mlir/include/mlir/Dialect/Transform/IR/MatchInterfaces.td

@@ -0,0 +1,26 @@
+//===- MatchInterfaces.td - Transform dialect interfaces ---*- 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
+//
+//===----------------------------------------------------------------------===//
+
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+
+def MatchOpInterface
+    : OpInterface<"MatchOpInterface", [TransformOpInterface]> {
+  let cppNamespace = "::mlir::transform";
+}
+
+def SingleOpMatcher : NativeOpTrait<"SingleOpMatcherOpTrait"> {
+  let cppNamespace = "::mlir::transform";
+
+  string extraDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure matchOperation(
+          ::mlir::Operation *current,
+          ::mlir::transform::TransformResults &results,
+          ::mlir::transform::TransformState &state);
+  }];
+}

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h

@@ -17,9 +17,11 @@
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
+
 namespace transform {
 
 class TransformOpInterface;
+class TransformResults;
 
 /// Options controlling the application of transform operations by the
 /// TransformState.
@@ -400,6 +402,11 @@ class TransformState {
     return it->second;
   }
 
+  /// Updates the state to include the associations between op results and the
+  /// provided result of applying a transform op.
+  LogicalResult updateStateFromResults(const TransformResults &results,
+                                       ResultRange opResults);
+
   /// Sets the payload IR ops associated with the given transform IR value
   /// (handle). A payload op may be associated multiple handles as long as
   /// at most one of them gets consumed by further transformations.
@@ -690,6 +697,11 @@ LogicalResult verifyTransformOpInterface(Operation *op);
 void prepareValueMappings(
     SmallVectorImpl<SmallVector<transform::MappedValue>> &mappings,
     ValueRange values, const transform::TransformState &state);
+
+/// Populates `results` with payload associations that match exactly those of
+/// the operands to `block`'s terminator.
+void forwardTerminatorOperands(Block *block, transform::TransformState &state,
+                               transform::TransformResults &results);
 } // namespace detail
 
 /// This trait is supposed to be attached to Transform dialect operations that

mlir/include/mlir/Dialect/Transform/IR/TransformOps.h

@@ -10,6 +10,7 @@
 #define MLIR_DIALECT_TRANSFORM_IR_TRANSFORMOPS_H
 
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Transform/IR/MatchInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/FunctionInterfaces.h"

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

@@ -17,6 +17,7 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
+include "mlir/Dialect/Transform/IR/MatchInterfaces.td"
 include "mlir/Dialect/Transform/IR/TransformAttrs.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
@@ -116,6 +117,69 @@ def CastOp : TransformDialectOp<"cast",
   }];
 }
 
+def ForeachMatchOp : TransformDialectOp<"foreach_match", [
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+    DeclareOpInterfaceMethods<SymbolUserOpInterface>,
+    DeclareOpInterfaceMethods<TransformOpInterface>]> {
+  let summary = "Applies named sequences when a named matcher succeeds";
+  let description = [{
+    Given a pair of co-indexed lists of transform dialect symbols (such as
+    `transform.named_sequence`), walks the payload IR associated with the root
+    handle and interprets the symbols as matcher/action pairs by applying the
+    body of the corresponding symbol definition. The symbol from the first list
+    is the matcher part: if it results in a silenceable error, the error is
+    silenced and the next matcher is attempted. Definite failures from any
+    matcher stop the application immediately and are propagated unconditionally.
+    If none of the matchers succeeds, the next payload operation in walk order
+    (post-order at the moment of writing, double check `Operation::walk`) is
+    matched. If a matcher succeeds, the co-indexed action symbol is applied and
+    the following matchers are not applied to the same payload operation. If the
+    action succeeds, the next payload operation in walk order is matched. If it
+    fails, both silenceable and definite errors are propagated as the result of
+    this op.
+
+    The matcher symbol must take one operand of a type that implements the same
+    transform dialect interface as the `root` operand (a check is performed at
+    application time to see if the associated payload satisfies the constraints
+    of the actual type). It must not consume the operand as multiple matchers
+    may be applied. The matcher may produce any number of results. The action
+    symbol paired with the matcher must take the same number of arguments as the
+    matcher has results, and these arguments must implement the same transform
+    dialect interfaces, but not necessarily have the exact same type (again, a
+    check is performed at application time to see if the associated payload
+    satisfies the constraints of actual types on both sides). The action symbol
+    may not have results. The actions are expected to only modify payload
+    operations nested in the `root` payload operations associated with the
+    operand of this transform operation.
+
+    This operation consumes the operand and produces a new handle associated
+    with the same payload. This is necessary to trigger invalidation of handles
+    to any of the payload operations nested in the payload operations associated
+    with the operand, as those are likely to be modified by actions. Note that
+    the root payload operation associated with the operand are not matched.
+
+    The operation succeeds if none of the matchers produced a definite failure
+    during application and if all of the applied actions produced success. Note
+    that it also succeeds if all the matchers failed on all payload operations,
+    i.e. failure to apply is not an error. The operation produces a silenceable
+    failure if any applied action produced a silenceable failure. In this case,
+    the resulting handle is associated with an empty payload. The operation
+    produces a definite failure if any of the applied matchers or actions
+    produced a definite failure.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$root,
+                       SymbolRefArrayAttr:$matchers,
+                       SymbolRefArrayAttr:$actions);
+  let results = (outs TransformHandleTypeInterface:$updated);
+
+  let assemblyFormat =
+      "`in` $root custom<ForeachMatchSymbols>($matchers, $actions) "
+      "attr-dict `:` functional-type($root, $updated)";
+
+  let hasVerifier = 1;
+}
+
 def ForeachOp : TransformDialectOp<"foreach",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
@@ -270,6 +334,7 @@ def GetResultOp : TransformDialectOp<"get_result",
 
 def IncludeOp : TransformDialectOp<"include",
     [CallOpInterface,
+     MatchOpInterface,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<SymbolUserOpInterface>,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {

mlir/lib/Dialect/Transform/IR/CMakeLists.txt

@@ -1,10 +1,12 @@
 add_mlir_dialect_library(MLIRTransformDialect
+  MatchInterfaces.cpp
   TransformDialect.cpp
   TransformInterfaces.cpp
   TransformOps.cpp
   TransformTypes.cpp
 
   DEPENDS
+  MLIRMatchInterfacesIncGen
   MLIRTransformDialectIncGen
   MLIRTransformInterfacesIncGen
 

mlir/lib/Dialect/Transform/IR/MatchInterfaces.cpp

@@ -0,0 +1,17 @@
+//===- MatchInterfaces.cpp - Transform Dialect Interfaces -----------------===//
+//
+// 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/Transform/IR/MatchInterfaces.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Generated interface implementation.
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Transform/IR/MatchInterfaces.cpp.inc"

mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp

@@ -920,40 +920,44 @@ transform::TransformState::applyTransform(TransformOpInterface transform) {
   }
 #endif // LLVM_ENABLE_ABI_BREAKING_CHECKS
 
-  for (OpResult result : transform->getResults()) {
-    assert(result.getDefiningOp() == transform.getOperation() &&
-           "payload IR association for a value other than the result of the "
-           "current transform op");
+  if (failed(updateStateFromResults(results, transform->getResults())))
+    return DiagnosedSilenceableFailure::definiteFailure();
+
+  printOnFailureRAII.release();
+  DEBUG_WITH_TYPE(DEBUG_PRINT_AFTER_ALL, {
+    DBGS() << "Top-level payload:\n";
+    getTopLevel()->print(llvm::dbgs());
+  });
+  return result;
+}
+
+LogicalResult transform::TransformState::updateStateFromResults(
+    const TransformResults &results, ResultRange opResults) {
+  for (OpResult result : opResults) {
     if (result.getType().isa<TransformParamTypeInterface>()) {
       assert(results.isParam(result.getResultNumber()) &&
              "expected parameters for the parameter-typed result");
       if (failed(
               setParams(result, results.getParams(result.getResultNumber())))) {
-        return DiagnosedSilenceableFailure::definiteFailure();
+        return failure();
       }
     } else if (result.getType().isa<TransformValueHandleTypeInterface>()) {
       assert(results.isValue(result.getResultNumber()) &&
              "expected values for value-type-result");
       if (failed(setPayloadValues(
               result, results.getValues(result.getResultNumber())))) {
-        return DiagnosedSilenceableFailure::definiteFailure();
+        return failure();
       }
     } else {
       assert(!results.isParam(result.getResultNumber()) &&
              "expected payload ops for the non-parameter typed result");
       if (failed(
               setPayloadOps(result, results.get(result.getResultNumber())))) {
-        return DiagnosedSilenceableFailure::definiteFailure();
+        return failure();
       }
     }
   }
-
-  printOnFailureRAII.release();
-  DEBUG_WITH_TYPE(DEBUG_PRINT_AFTER_ALL, {
-    DBGS() << "Top-level payload:\n";
-    getTopLevel()->print(llvm::dbgs());
-  });
-  return result;
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
@@ -1193,7 +1197,7 @@ void transform::detail::setApplyToOneResults(
 }
 
 //===----------------------------------------------------------------------===//
-// Utilities for PossibleTopLevelTransformOpTrait.
+// Utilities for implementing transform ops with regions.
 //===----------------------------------------------------------------------===//
 
 void transform::detail::prepareValueMappings(
@@ -1213,6 +1217,29 @@ void transform::detail::prepareValueMappings(
   }
 }
 
+void transform::detail::forwardTerminatorOperands(
+    Block *block, transform::TransformState &state,
+    transform::TransformResults &results) {
+  for (auto &&[terminatorOperand, result] :
+       llvm::zip(block->getTerminator()->getOperands(),
+                 block->getParentOp()->getOpResults())) {
+    if (result.getType().isa<transform::TransformHandleTypeInterface>()) {
+      results.set(result, state.getPayloadOps(terminatorOperand));
+    } else if (result.getType()
+                   .isa<transform::TransformValueHandleTypeInterface>()) {
+      results.setValues(result, state.getPayloadValues(terminatorOperand));
+    } else {
+      assert(result.getType().isa<transform::TransformParamTypeInterface>() &&
+             "unhandled transform type interface");
+      results.setParams(result, state.getParams(terminatorOperand));
+    }
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// Utilities for PossibleTopLevelTransformOpTrait.
+//===----------------------------------------------------------------------===//
+
 LogicalResult transform::detail::mapPossibleTopLevelTransformOpBlockArguments(
     TransformState &state, Operation *op, Region &region) {
   SmallVector<Operation *> targets;