[mlir][transform] Add transform.get_operand op

Similar to `transform.get_result`, except it returns a handle to the
operand indicated by `operand_number`, or all operands if no index is
given.

Additionally updates `get_result` to make the `result_number`
optional. This makes the use case of wanting to get all of the
results of an operation easier by no longer requiring the user to
reconstruct the list of results one-by-one.

Author: qedawkins

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

@@ -728,23 +728,43 @@ def GetProducerOfOperand : TransformDialectOp<"get_producer_of_operand",
                        "functional-type(operands, results)";
 }
 
+def GetOperandOp : TransformDialectOp<"get_operand",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     NavigationTransformOpTrait, MatchOpInterface, MemoryEffectsOpInterface]> {
+  let summary = "Get a handle to the operand(s) of the targeted op";
+  let description = [{
+    The handle defined by this Transform op corresponds to the Operands of the
+    given `target` operation. Optionally `operand_number` can be specified to
+    select a specific operand.
+    
+    This transform fails silently if the targeted operation does not have enough
+    operands. It reads the target handle and produces the result handle.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                       OptionalAttr<I64Attr>:$operand_number);
+  let results = (outs TransformValueHandleTypeInterface:$result);
+  let assemblyFormat = "$target (`[` $operand_number^ `]`)? attr-dict `:` "
+                       "functional-type(operands, results)";
+}
+
 def GetResultOp : TransformDialectOp<"get_result",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      NavigationTransformOpTrait, MemoryEffectsOpInterface]> {
-  let summary = "Get handle to the a result of the targeted op";
+  let summary = "Get a handle to the result(s) of the targeted op";
   let description = [{
-    The handle defined by this Transform op corresponds to the OpResult with
-    `result_number` that is defined by the given `target` operation.
-
-    This transform produces a silenceable failure if the targeted operation
-    does not have enough results. It reads the target handle and produces the
-    result handle.
+    The handle defined by this Transform op correspond to the OpResults of the
+    given `target` operation. Optionally `result_number` can be specified to
+    select a specific result.
+    
+    This transform fails silently if the targeted operation does not have enough
+    results. It reads the target handle and produces the result handle.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
-                       I64Attr:$result_number);
+                       OptionalAttr<I64Attr>:$result_number);
   let results = (outs TransformValueHandleTypeInterface:$result);
-  let assemblyFormat = "$target `[` $result_number `]` attr-dict `:` "
+  let assemblyFormat = "$target (`[` $result_number^ `]`)? attr-dict `:` "
                        "functional-type(operands, results)";
 }
 

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

@@ -1464,6 +1464,35 @@ transform::GetProducerOfOperand::apply(transform::TransformRewriter &rewriter,
   return DiagnosedSilenceableFailure::success();
 }
 
+//===----------------------------------------------------------------------===//
+// GetOperandOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::GetOperandOp::apply(transform::TransformRewriter &rewriter,
+                               transform::TransformResults &results,
+                               transform::TransformState &state) {
+  std::optional<int64_t> maybeOperandNumber = getOperandNumber();
+  SmallVector<Value> operands;
+  for (Operation *target : state.getPayloadOps(getTarget())) {
+    if (!maybeOperandNumber) {
+      for (Value operand : target->getOperands())
+        operands.push_back(operand);
+      continue;
+    }
+    int64_t operandNumber = *maybeOperandNumber;
+    if (operandNumber >= target->getNumOperands()) {
+      DiagnosedSilenceableFailure diag =
+          emitSilenceableError() << "targeted op does not have enough operands";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+    operands.push_back(target->getOperand(operandNumber));
+  }
+  results.setValues(llvm::cast<OpResult>(getResult()), operands);
+  return DiagnosedSilenceableFailure::success();
+}
+
 //===----------------------------------------------------------------------===//
 // GetResultOp
 //===----------------------------------------------------------------------===//
@@ -1472,9 +1501,15 @@ DiagnosedSilenceableFailure
 transform::GetResultOp::apply(transform::TransformRewriter &rewriter,
                               transform::TransformResults &results,
                               transform::TransformState &state) {
-  int64_t resultNumber = getResultNumber();
+  std::optional<int64_t> maybeResultNumber = getResultNumber();
   SmallVector<Value> opResults;
   for (Operation *target : state.getPayloadOps(getTarget())) {
+    if (!maybeResultNumber) {
+      for (Value result : target->getResults())
+        opResults.push_back(result);
+      continue;
+    }
+    int64_t resultNumber = *maybeResultNumber;
     if (resultNumber >= target->getNumResults()) {
       DiagnosedSilenceableFailure diag =
           emitSilenceableError() << "targeted op does not have enough results";

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

@@ -1483,6 +1483,60 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// expected-remark @below {{addi operand}}
+// expected-note @below {{value handle points to a block argument #0}}
+func.func @get_operand_of_op(%arg0: index, %arg1: index) -> index {
+  %r = arith.addi %arg0, %arg1 : index
+  return %r : index
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op) {
+    %addi = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %operand = transform.get_operand %addi[0] : (!transform.any_op) -> !transform.any_value
+    transform.debug.emit_remark_at %operand, "addi operand" : !transform.any_value
+    transform.yield
+  }
+}
+
+// -----
+
+func.func @get_out_of_bounds_operand_of_op(%arg0: index, %arg1: index) -> index {
+  // expected-note @below {{target op}}
+  %r = arith.addi %arg0, %arg1 : index
+  return %r : index
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%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 operands}}
+    %operand = transform.get_operand %addi[2] : (!transform.any_op) -> !transform.any_value
+    transform.debug.emit_remark_at %operand, "addi operand" : !transform.any_value
+    transform.yield
+  }
+}
+
+// -----
+
+func.func @get_multiple_operands_of_op(%arg0: index, %arg1: index) -> index {
+  %r = arith.addi %arg0, %arg1 : index
+  return %r : index
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op) {
+    %addui = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %operands = transform.get_operand %addui : (!transform.any_op) -> !transform.any_value
+    %p = transform.num_associations %operands : (!transform.any_value) -> !transform.param<i64>
+    // expected-remark @below {{2}}
+    transform.debug.emit_param_as_remark %p : !transform.param<i64>
+    transform.yield
+  }
+}
+
+// -----
+
 func.func @get_result_of_op(%arg0: index, %arg1: index) -> index {
   // expected-remark @below {{addi result}}
   // expected-note @below {{value handle points to an op result #0}}
@@ -1537,6 +1591,25 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+func.func @get_multiple_result_of_op(%arg0: index, %arg1: index) -> (index, i1) {
+  // expected-remark @below {{matched bool}}
+  %r, %b = arith.addui_extended %arg0, %arg1 : index, i1
+  return %r, %b : index, i1
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op) {
+    %addui = transform.structured.match ops{["arith.addui_extended"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %results = transform.get_result %addui : (!transform.any_op) -> !transform.any_value
+    %adds = transform.get_defining_op %results : (!transform.any_value) -> !transform.any_op
+    %_, %add_again = transform.split_handle %adds : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.debug.emit_remark_at %add_again, "matched bool" : !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
 // expected-note @below {{target value}}
 func.func @get_result_of_op_bbarg(%arg0: index, %arg1: index) -> index {
   %r = arith.addi %arg0, %arg1 : index