[mlir][transform] Add transform.get_operand op #78397

qedawkins · 2024-01-17T05:35:40Z

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.

llvmbot · 2024-01-17T05:36:10Z

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Quinn Dawkins (qedawkins)

Changes

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.

Full diff: https://github.com/llvm/llvm-project/pull/78397.diff

3 Files Affected:

(modified) mlir/include/mlir/Dialect/Transform/IR/TransformOps.td (+26-5)
(modified) mlir/lib/Dialect/Transform/IR/TransformOps.cpp (+36-1)
(modified) mlir/test/Dialect/Transform/test-interpreter.mlir (+73)

diff --git a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
index fe2c28f45aea04..6637d81dab5e2a 100644
--- a/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
+++ b/mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
@@ -725,22 +725,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.
+    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)";
 }
 
diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index b80fc09751d2aa..56baae9b5fadf2 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/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";
diff --git a/mlir/test/Dialect/Transform/test-interpreter.mlir b/mlir/test/Dialect/Transform/test-interpreter.mlir
index 96f2122e976df5..b89b52e2f403d5 100644
--- a/mlir/test/Dialect/Transform/test-interpreter.mlir
+++ b/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

ftynse

Looks good in principle.

An overall design comment: we have a similar mechanism in Linalg matcher operations, e.g. https://mlir.llvm.org/docs/Dialects/Transform/#transformmatchstructuredinit-transformmatchstructuredinitop. It has more advanced addressing modes, e.g., associating a list of operands with the result. Not necessarily arguing to adopt all of that here immediately or at all, but I'd recommend using an explicit token/attribute to indicate that all operands are requested instead of omitting the attribute.

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

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

qedawkins · 2024-01-17T16:44:22Z

Looks good in principle.

An overall design comment: we have a similar mechanism in Linalg matcher operations, e.g. https://mlir.llvm.org/docs/Dialects/Transform/#transformmatchstructuredinit-transformmatchstructuredinitop. It has more advanced addressing modes, e.g., associating a list of operands with the result. Not necessarily arguing to adopt all of that here immediately or at all, but I'd recommend using an explicit token/attribute to indicate that all operands are requested instead of omitting the attribute.

Good suggestion, I like the linalg matcher approach more so I just imported it. I'll leave this open for a day or so to give a chance to review if you want to check that I got the code shuffling right.

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.

…spec

qedawkins requested review from ftynse, matthias-springer and nicolasvasilache January 17, 2024 05:35

llvmbot added the mlir label Jan 17, 2024

qedawkins mentioned this pull request Jan 17, 2024

[mlir][transform] Add an op for replacing values with function calls #78398

Merged

ftynse approved these changes Jan 17, 2024

View reviewed changes

qedawkins requested a review from dcaballe as a code owner January 17, 2024 16:41

llvmbot added the mlir:linalg label Jan 17, 2024

qedawkins added 2 commits January 17, 2024 11:46

Address comments and switch to mirroring the linalg.match positional …

59fae09

…spec

ftynse approved these changes Jan 17, 2024

View reviewed changes

Drop unused function and rebase

9cd797f

qedawkins force-pushed the get_operand branch from f3ec382 to 9cd797f Compare January 17, 2024 16:59

Fix python test

c57bbef

llvmbot added the mlir:python MLIR Python bindings label Jan 18, 2024

qedawkins merged commit 5caab8b into llvm:main Jan 18, 2024

qedawkins deleted the get_operand branch January 18, 2024 14:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][transform] Add transform.get_operand op #78397

[mlir][transform] Add transform.get_operand op #78397

Uh oh!

qedawkins commented Jan 17, 2024

Uh oh!

llvmbot commented Jan 17, 2024 •

edited

Loading

Uh oh!

ftynse left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

qedawkins commented Jan 17, 2024

Uh oh!

Uh oh!

[mlir][transform] Add transform.get_operand op #78397

[mlir][transform] Add transform.get_operand op #78397

Uh oh!

Conversation

qedawkins commented Jan 17, 2024

Uh oh!

llvmbot commented Jan 17, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

qedawkins commented Jan 17, 2024

Uh oh!

Uh oh!

llvmbot commented Jan 17, 2024 •

edited

Loading