[mlir][transform] Add transform.get_defining_op op

[springerm]

Summary

This op is the inverse of transform.get_result.

Depends On: D144397

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

def GetConsumersOfResult : TransformDialectOp<"get_consumers_of_result",
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        I64Attr:$result_number);
   let results = (outs TransformHandleTypeInterface:$consumers);
   let assemblyFormat = "$target `[` $result_number `]` attr-dict `:` "
                        "functional-type(operands, results)";
 }
 
+def GetDefiningOp : TransformDialectOp<"get_defining_op",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     NavigationTransformOpTrait, MemoryEffectsOpInterface]> {
+  let summary = "Get handle to the defining op of a value";
+  let description = [{
+    The handle defined by this Transform op corresponds to the defining op of
+    the targeted value.
+
+    This transform fails silently if the targeted value is a block argument.
+  }];
+
+  let arguments = (ins TransformValueHandleTypeInterface:$target);
+  let results = (outs TransformHandleTypeInterface:$result);
+  let assemblyFormat = "$target attr-dict `:` "
+                       "functional-type(operands, results)";
+}
+
 def GetProducerOfOperand : TransformDialectOp<"get_producer_of_operand",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      NavigationTransformOpTrait, MemoryEffectsOpInterface]> {
   let summary = "Get handle to the producer of this operation's operand number";
   let description = [{
     The handle defined by this Transform op corresponds to operation that
     produces the SSA value defined by the `target` and `operand_number`
     arguments. If the origin of the SSA value is not an operations (i.e. it is

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

transform::GetConsumersOfResult::apply(transform::TransformResults &results,
   if (target->getNumResults() <= resultNumber)
     return emitDefiniteFailure() << "result number overflow";
   results.set(getResult().cast<OpResult>(),
               llvm::to_vector(target->getResult(resultNumber).getUsers()));
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
+// GetDefiningOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::GetDefiningOp::apply(transform::TransformResults &results,
+                              transform::TransformState &state) {
+  SmallVector<Operation *> definingOps;
+  for (Value v : state.getPayloadValues(getTarget())) {
+    if (v.isa<BlockArgument>()) {
+      DiagnosedSilenceableFailure diag =
+          emitSilenceableError() << "cannot get defining op of block argument";
+      diag.attachNote(v.getLoc()) << "target value";
+      return diag;
+    }
+    definingOps.push_back(v.getDefiningOp());
+  }
+  results.set(getResult().cast<OpResult>(), definingOps);
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // GetProducerOfOperand
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::GetProducerOfOperand::apply(transform::TransformResults &results,
                                        transform::TransformState &state) {
   int64_t operandNumber = getOperandNumber();
   SmallVector<Operation *> producers;

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

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %addi = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{targeted op does not have enough results}}
   %result = transform.get_result %addi[1] : (!transform.any_op) -> !transform.any_value
   transform.test_print_remark_at_operand_value %result, "addi result" : !transform.any_value
 }
+
+// -----
+
+func.func @get_result_of_op(%arg0: index, %arg1: index) -> index {
+  // expected-remark @below {{matched}}
+  %r = arith.addi %arg0, %arg1 : index
+  return %r : index
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %addi = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %result = transform.get_result %addi[0] : (!transform.any_op) -> !transform.any_value
+  %op = transform.get_defining_op %result : (!transform.any_value) -> !transform.any_op
+  transform.test_print_remark_at_operand %op, "matched" : !transform.any_op
+}
+
+// -----
+
+// expected-note @below {{target value}}
+func.func @get_result_of_op_bbarg(%arg0: index, %arg1: index) -> index {
+  %r = arith.addi %arg0, %arg1 : index
+  return %r : index
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %addi = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %bbarg = test_produce_value_handle_to_argument_of_parent_block %addi, 0 : (!transform.any_op) -> !transform.any_value
+  // expected-error @below {{cannot get defining op of block argument}}
+  %op = transform.get_defining_op %bbarg : (!transform.any_value) -> !transform.any_op
+  transform.test_print_remark_at_operand %op, "matched" : !transform.any_op
+}