[mlir][Transform] Add a new navigation op to retrieve the producer of an operand

[nicolasvasilache]

Summary

Given an opOperand uniquely determined by the operation %op and the operand number num,
the transform.get_producer_of_operand %op[num] returns the handle to the unique operation
that produced the SSA value used as opOperand.

The transform fails if the operand is a block argument.

[chelini]

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

177

corresponds

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

let description = [{
     on the further transformation applied to the handle produced here.
   }];
 
   let arguments = (ins PDL_Operation:$target);
   let results = (outs PDL_Operation:$parent);
   let assemblyFormat = "$target attr-dict";
 }
 
+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 correspond 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
+    a block argument), the transform silently fails.
+    The return handle points to only the subset of successfully produced
+    computational operations, which can be empty.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       I64Attr:$operand_number);
+  let results = (outs PDL_Operation:$parent);
+  let assemblyFormat = "$target `[` $operand_number `]` attr-dict";
+}
+
 def MergeHandlesOp : TransformDialectOp<"merge_handles",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Merges handles into one pointing to the union of payload ops";
   let description = [{
     Creates a new Transform IR handle value that points to the same Payload IR
     operations as the operand handles. The Payload IR operations are listed
     in the same order as they are in the operand handles, grouped by operand

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

for (Operation *target : state.getPayloadOps(getTarget())) {
     }
     parents.insert(parent);
   }
   results.set(getResult().cast<OpResult>(), parents.getArrayRef());
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
+// GetProducerOfOperand
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::GetProducerOfOperand::apply(transform::TransformResults &results,
+                                       transform::TransformState &state) {
+  int64_t operandNumber = getOperandNumber();
+  SmallVector<Operation *> producers;
+  for (Operation *target : state.getPayloadOps(getTarget())) {
+    Operation *producer =
+        target->getNumOperands() <= operandNumber
+            ? nullptr
+            : target->getOperand(operandNumber).getDefiningOp();
+    if (!producer) {
+      DiagnosedSilenceableFailure diag =
+          emitSilenceableError()
+          << "could not find a producer for operand number: " << operandNumber
+          << " of " << *target;
+      diag.attachNote(target->getLoc()) << "target op";
+      results.set(getResult().cast<OpResult>(),
+                  SmallVector<mlir::Operation *>{});
+      return diag;
+    }
+    producers.push_back(producer);
+  }
+  results.set(getResult().cast<OpResult>(), producers);
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // MergeHandlesOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::MergeHandlesOp::apply(transform::TransformResults &results,
                                  transform::TransformState &state) {
   SmallVector<Operation *> operations;
   for (Value operand : getHandles())

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

^bb2(%arg2: !pdl.operation):
       transform.yield %g : !pdl.operation
     }
 
     // expected-remark @below {{3}}
     transform.test_print_number_of_associated_payload_ir_ops %results
     transform.test_print_remark_at_operand %results, "transform applied"
   }
 }
+
+// -----
+
+func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
+  // expected-remark @below {{found muli}}
+  %0 = arith.muli %arg0, %arg1 : index  
+  arith.addi %0, %arg1 : index  
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %addi = transform.structured.match ops{["arith.addi"]} in %arg1
+  %muli = get_producer_of_operand %addi[0]
+  transform.test_print_remark_at_operand %muli, "found muli"
+}
+
+// -----
+
+func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
+  // expected-note @below {{target op}}
+  %0 = arith.muli %arg0, %arg1 : index  
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %muli = transform.structured.match ops{["arith.muli"]} in %arg1
+  // expected-error @below {{could not find a producer for operand number: 0 of}}
+  %bbarg = get_producer_of_operand %muli[0]
+
+}
+