[mlir][Linalg] Add return type filter to the transform dialect

[nicolasvasilache]

Summary

This allows matching ops by additionally providing an idiomatic spec for a unique return type.

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

let description = [{
     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.
         Multiple names can be specified. Matched ops must have one of specified
         names.
       - attribute: the matched op must have all specified attributes (with their
         specified values).
+      - filter_result_type: the matched op must return exactly this one type.
       
     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<DictionaryAttr>:$op_attrs);
+                       OptionalAttr<DictionaryAttr>:$op_attrs,
+                       OptionalAttr<TypeAttr>:$filter_result_type);
   // TODO: variadic results when needed.
   let results = (outs PDL_Operation:$results);
 
   let assemblyFormat = [{
     (`ops` `{` $ops^ `}`)?
     (`interface` `{` $interface^ `}`)?
     (`attributes` $op_attrs^)?
+    (`filter_result_type` `=` $filter_result_type^)?
     `in` $target attr-dict
   }];
 }
 
 def MultiTileSizesOp : Op<Transform_Dialect, "structured.multitile_sizes",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{

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

transform::MatchOp::apply(transform::TransformResults &results,
   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 (getOps().has_value() && !strs.contains(op->getName().getStringRef()))
-      return WalkResult::advance();
+      return;
 
     // Interfaces cannot be matched by name, just by ID.
     // So we specifically encode the interfaces we care about for this op.
     if (getInterface().has_value()) {
       auto iface = getInterface().value();
       if (iface == transform::MatchInterfaceEnum::LinalgOp &&
           !isa<linalg::LinalgOp>(op))
-        return WalkResult::advance();
+        return;
       if (iface == transform::MatchInterfaceEnum::TilingInterface &&
           isa<TilingInterface>(op))
-        return WalkResult::advance();
+        return;
     }
 
     // Check if all specified attributes match.
     if (getOpAttrs().has_value()) {
       DictionaryAttr opAttrs = getOpAttrs().value();
       for (NamedAttribute attr : opAttrs) {
         if (attr.getName() == getInterfaceAttrName() ||
             attr.getName() == getOpsAttrName())
           continue;
         if (!op->hasAttr(attr.getName()))
-          return WalkResult::advance();
+          return;
         if (op->getAttr(attr.getName()) != attr.getValue())
-          return WalkResult::advance();
+          return;
       }
     }
 
+    if (getFilterResultType().has_value()) {
+      Type t = getFilterResultType().value();
+      if (op->getNumResults() != 1 || op->getResultTypes().front() != t)
+        return;
+    }
+
     // All constraints are satisfied.
     res.push_back(op);
-    return WalkResult::advance();
+    return;
   };
 
   payloadOps.front()->walk(matchFun);
   results.set(getResult().cast<OpResult>(), res);
   return DiagnosedSilenceableFailure(success());
 }
 
 //===---------------------------------------------------------------------===//

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

^bb1(%arg1: !pdl.operation):
     %match_attr = transform.structured.match ops{["arith.constant"]} attributes{my_attr} in %arg1
     transform.test_print_remark_at_operand %match_attr, "matched attr name"
     transform.test_consume_operand %match_attr
   }
 }
 
 // -----
 
+func.func @by_type() {
+  %0 = arith.constant 0: i32
+  // expected-remark @below {{matched op name}}
+  %1 = arith.constant 1.0 : f32
+  return
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %match_name = transform.structured.match
+      ops{["arith.constant"]} filter_result_type = f32 in %arg1
+    transform.test_print_remark_at_operand %match_name, "matched op name"
+    transform.test_consume_operand %match_name
+  }
+}
+
+// -----
+
 #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 #map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
 func.func @match_complex_attribute(%arg0: tensor<12x128x32xf32>)
     -> tensor<128x12x32xf32> {
   %0 = linalg.init_tensor [128, 12, 32] : tensor<128x12x32xf32>
   // expected-remark @below {{matched complex attr}}
   %1 = linalg.generic {indexing_maps = [#map0, #map1],
                        iterator_types = ["parallel", "parallel", "parallel"]}