[mlir][linalg] Add attribute matcher to structured.match transform op

[springerm]

Summary

This is useful for building small test cases and will be utilized in a subsequent commit that adds a fusion example.

Depends On D130244

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

def MatchOp : Op<Transform_Dialect, "structured.match",
     [MemoryEffectsOpInterface,
      NavigationTransformOpTrait,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let description = [{
     Match op with the specified constraints, within the target op.
 
     The following constraints are supported:
       - interface: an optional MatchInterfaceEnum specifying an enum
         representation for an interface to target.
-      - ops: an optional StrArrayAttr specifying the concrete name of an op. 
+      - ops: an optional StrArrayAttr specifying the concrete name of an op.
+        Multiple names can be specified. Matched ops must have one of specified
+        names.
+      - attribute: an optional Str specifying the name of an attribute that
+        matched ops must have.
       
-    Note: either `ops` or `interface` must be specified.
+    Note: Only ops that satisfy all specified constraints are matched.
 
     TODO: Extend with regions to allow a limited form of constraints.
 
     #### Return modes
     
     This op traverses the ops nested under `target` and returns the handles to
     all the operations that match the requirements.
 
     This op fails if the target is not a handle to exactly one operation. 
     Otherwise it succeeds.
   
     This operation does not consume the target handle and produces new handles:
     it is a navigation op.
   }];
 
   let arguments = (ins PDL_Operation:$target,
                        OptionalAttr<StrArrayAttr>:$ops,
-                       OptionalAttr<MatchInterfaceEnum>:$interface);
+                       OptionalAttr<MatchInterfaceEnum>:$interface,
+                       OptionalAttr<StrAttr>:$attribute);
   // TODO: variadic results when needed.
   let results = (outs PDL_Operation:$results);
 
-  let hasCustomAssemblyFormat = 1;
-  let hasVerifier = 1;
+  let assemblyFormat = [{
+    (`ops` `{` $ops^ `}`)? 
+    (`interface` `{` $interface^ `}`)? 
+    (`attribute` `{` $attribute^ `}`)? 
+    `in` $target attr-dict
+  }];
 }
 
 def MultiTileSizesOp : Op<Transform_Dialect, "structured.multitile_sizes",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Emits the IR computing the tile sizes `s1` and `s2` such that:
 

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

LogicalResult transform::InterchangeOp::verify() {
   }
   return success();
 }
 
 //===---------------------------------------------------------------------===//
 // MatchOp
 //===---------------------------------------------------------------------===//
 
-LogicalResult transform::MatchOp::verify() {
-  bool opXorIface = getOps().hasValue() ^ getInterface().hasValue();
-  if (!opXorIface)
-    return this->emitOpError(
-        "requires a either a match_op or a match_interface attribute (but not "
-        "both)");
-  return success();
-}
-
 DiagnosedSilenceableFailure
 transform::MatchOp::apply(transform::TransformResults &results,
                           transform::TransformState &state) {
   llvm::StringSet<> strs;
   if (getOps().hasValue())
     strs.insert(getOps()->getAsValueRange<StringAttr>().begin(),
                 getOps()->getAsValueRange<StringAttr>().end());
 
   ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
   if (payloadOps.size() != 1)
     return DiagnosedSilenceableFailure(
         this->emitOpError("requires exactly one target handle"));
 
   SmallVector<Operation *> res;
-
   auto matchFun = [&](Operation *op) {
-    if (strs.contains(op->getName().getStringRef()))
-      res.push_back(op);
+    if (getOps().hasValue() && !strs.contains(op->getName().getStringRef()))
+      return WalkResult::advance();
+
     // Interfaces cannot be matched by name, just by ID.
     // So we specifically encode the interfaces we care about for this op.
     if (getInterface().hasValue()) {
       auto iface = getInterface().getValue();
       if (iface == transform::MatchInterfaceEnum::LinalgOp &&
-          isa<linalg::LinalgOp>(op))
-        res.push_back(op);
+          !isa<linalg::LinalgOp>(op))
+        return WalkResult::advance();
       if (iface == transform::MatchInterfaceEnum::TilingInterface &&
           isa<TilingInterface>(op))
-        res.push_back(op);
+        return WalkResult::advance();
     }
+
+    if (getAttribute().hasValue() && !op->hasAttr(getAttribute().getValue()))
+      return WalkResult::advance();
+
+    // All constraints are satisfied.
+    res.push_back(op);
+    return WalkResult::advance();
   };
 
   payloadOps.front()->walk(matchFun);
   results.set(getResult().cast<OpResult>(), res);
   return DiagnosedSilenceableFailure(success());
 }
 
-ParseResult transform::MatchOp::parse(OpAsmParser &parser,
-                                      OperationState &result) {
-  // Parse 'match_op' or 'interface' clause.
-  if (succeeded(parser.parseOptionalKeyword("ops"))) {
-    ArrayAttr opsAttr;
-    if (parser.parseLBrace() ||
-        parser.parseCustomAttributeWithFallback(
-            opsAttr, parser.getBuilder().getType<NoneType>(), "ops",
-            result.attributes) ||
-        parser.parseRBrace())
-      return failure();
-  } else if (succeeded(parser.parseOptionalKeyword("interface"))) {
-    if (parser.parseLBrace())
-      return failure();
-    StringRef attrStr;
-    auto loc = parser.getCurrentLocation();
-    if (parser.parseKeyword(&attrStr))
-      return failure();
-    auto interfaceEnum = transform::symbolizeMatchInterfaceEnum(attrStr);
-    if (!interfaceEnum)
-      return parser.emitError(loc, "invalid ")
-             << "match_interface attribute specification: \"" << attrStr << '"';
-    transform::MatchInterfaceEnumAttr match_interfaceAttr =
-        transform::MatchInterfaceEnumAttr::get(parser.getBuilder().getContext(),
-                                               interfaceEnum.value());
-    result.addAttribute("interface", match_interfaceAttr);
-    if (parser.parseRBrace())
-      return failure();
-  } else {
-    auto loc = parser.getCurrentLocation();
-    return parser.emitError(loc, "expected ops or interface");
-  }
-
-  OpAsmParser::UnresolvedOperand targetRawOperands[1];
-  ArrayRef<OpAsmParser::UnresolvedOperand> targetOperands(targetRawOperands);
-  if (parser.parseKeyword("in") || parser.parseOperand(targetRawOperands[0]) ||
-      parser.parseOptionalAttrDict(result.attributes))
-    return failure();
-  Type pdlOpType = parser.getBuilder().getType<pdl::OperationType>();
-  result.addTypes(pdlOpType);
-  if (parser.resolveOperands(targetOperands, pdlOpType, result.operands))
-    return failure();
-  return success();
-}
-
-void transform::MatchOp::print(OpAsmPrinter &p) {
-  if ((*this)->getAttr("ops")) {
-    p << " ops{";
-    p.printAttributeWithoutType(getOpsAttr());
-    p << "}";
-  }
-  if ((*this)->getAttr("interface")) {
-    p << " interface{" << stringifyMatchInterfaceEnum(*getInterface()) << "}";
-  }
-  p << " in " << getTarget();
-  p.printOptionalAttrDict((*this)->getAttrs(),
-                          /*elidedAttrs=*/{"ops", "interface"});
-}
-
 //===---------------------------------------------------------------------===//
 // MultiTileSizesOp
 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
     LinalgOp target, SmallVector<Operation *> &results, TransformState &state) {
   OpBuilder builder(target.getContext());
   builder.setInsertionPoint(target);

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

+// RUN: mlir-opt %s --test-transform-dialect-interpreter -allow-unregistered-dialect --split-input-file --verify-diagnostics
+
+func.func @bar() {
+  // expected-remark @below {{matched op name}}
+  // expected-remark @below {{matched attr name}}
+  %0 = arith.constant {my_attr} 0: i32
+  // expected-remark @below {{matched op name}}
+  %1 = arith.constant 1 : i32
+  return
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %match_name = transform.structured.match ops{["arith.constant"]} in %arg1
+    transform.test_print_remark_at_operand %match_name, "matched op name"
+    transform.test_consume_operand %match_name
+
+    %match_attr = transform.structured.match ops{["arith.constant"]} attribute{"my_attr"} in %arg1
+    transform.test_print_remark_at_operand %match_attr, "matched attr name"
+    transform.test_consume_operand %match_attr
+  }
+}