[mlir][transform] Add transform.select op

[springerm]

Summary

This transform op can be used to select all payload ops with a given name from a handle.

[ftynse]

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

863–864

Longer-term, we should use a region with MatchOps as a conjunction/disjunction of constraints.

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

def ReplicateOp : TransformDialectOp<"replicate",
 
   let arguments = (ins TransformHandleTypeInterface:$pattern,
                        Variadic<Transform_AnyHandleOrParamType>:$handles);
   let results = (outs Variadic<Transform_AnyHandleOrParamType>:$replicated);
   let assemblyFormat = "`num` `(` $pattern `)` $handles attr-dict `:` "
                        "type($pattern) `,` type($handles)";
 }
 
+def SelectOp : TransformDialectOp<"select",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     NavigationTransformOpTrait, MemoryEffectsOpInterface]> {
+  let summary = "Select payload ops by name";
+  let description = [{
+    The handle defined by this Transform op corresponds to all operations among
+    `target` that have the specified properties. Currently the following
+    properties are supported:
+
+    - `op_name`: The op must have the specified name.
+
+    The result payload ops are in the same relative order as the targeted ops.
+    This transform op reads the `target` handle and produces the `result`
+    handle. It reads the payload, but does not modify it.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                       StrAttr:$op_name);
+  let results = (outs TransformHandleTypeInterface:$result);
+  let assemblyFormat = [{
+    $op_name `in` $target attr-dict `:` functional-type(operands, results)
+  }];
+}
+
 def SequenceOp : TransformDialectOp<"sequence",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getSuccessorEntryOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,
      MatchOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      OpAsmOpInterface, PossibleTopLevelTransformOpTrait,

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

 }
 
 LogicalResult transform::NamedSequenceOp::verify() {
   // Actual verification happens in a separate function for reusability.
   return verifyNamedSequenceOp(*this, /*emitWarnings=*/true).checkAndReport();
 }
 
 //===----------------------------------------------------------------------===//
+// SelectOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::SelectOp::apply(transform::TransformRewriter &rewriter,
+                           transform::TransformResults &results,
+                           transform::TransformState &state) {
+  SmallVector<Operation *> result;
+  auto payloadOps = state.getPayloadOps(getTarget());
+  for (Operation *op : payloadOps) {
+    if (op->getName().getStringRef() == getOpName())
+      result.push_back(op);
+  }
+  results.set(cast<OpResult>(getResult()), result);
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // SplitHandleOp
 //===----------------------------------------------------------------------===//
 
 void transform::SplitHandleOp::build(OpBuilder &builder, OperationState &result,
                                      Value target, int64_t numResultHandles) {
   result.addOperands(target);
   result.addTypes(SmallVector<Type>(numResultHandles, target.getType()));
 }

mlir/test/Dialect/Transform/test-interpreter.mlir

 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
   transform.test_produce_invalid_ir %0 : !transform.any_op
   // expected-error @below{{failed to verify payload op}}
   transform.verify %0 : !transform.any_op
 }
+
+// -----
+
+func.func @select() {
+  // expected-remark @below{{found foo}}
+  "test.foo"() : () -> ()
+  // expected-remark @below{{found bar}}
+  "test.bar"() : () -> ()
+  // expected-remark @below{{found foo}}
+  "test.foo"() : () -> ()
+  func.return
+}
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !transform.any_op):
+  // Match all ops inside the function (including the function itself).
+  %func_op = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+  %0 = transform.structured.match in %func_op : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below{{5}}
+  test_print_number_of_associated_payload_ir_ops %0 : !transform.any_op
+
+  // Select "test.foo".
+  %foo = transform.select "test.foo" in %0 : (!transform.any_op) -> !transform.any_op
+  test_print_remark_at_operand %foo, "found foo" : !transform.any_op
+
+  // Select "test.bar".
+  %bar = transform.select "test.bar" in %0 : (!transform.any_op) -> !transform.any_op
+  test_print_remark_at_operand %bar, "found bar" : !transform.any_op
+}