[Mlir-commits] [mlir] [mlir][linalg] Handle existing destination-passing-style ops in `transform.structured.rewrite_in_destination_passing_style` (PR #205034)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: jpwang (jjppp)

<details>
<summary>Changes</summary>

`transform.structured.rewrite_in_destination_passing_style` may be applied to an operation that is already in destination-passing style, e.g. `linalg.add`. In this case, the operation does not need to be rewritten, but the current `TypeSwitch` does not handle `DestinationStyleOpInterface` and falls through to the unreachable case.

Such operations can be handled by returning them unchanged. This makes the transform accept already destination-style operations and avoids the crash.

An regression test for applying `rewrite_in_destination_passing_style` is added to `linalg.add`.

Fixes #<!-- -->204099 


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


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+1) 
- (modified) mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index f44693096b26b..55332447bf168 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -2842,6 +2842,7 @@ transform::RewriteInDestinationPassingStyleOp::applyToOne(
   rewriter.setInsertionPoint(target);
   FailureOr<Operation *> maybeResult =
       TypeSwitch<Operation *, FailureOr<Operation *>>(target)
+          .Case<DestinationStyleOpInterface>([](auto op) { return op; })
           .Case<tensor::FromElementsOp, tensor::GenerateOp, tensor::PadOp>(
               [&rewriter](auto op) {
                 return rewriteInDestinationPassingStyle(rewriter, op);
diff --git a/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir b/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir
index 9b74b5f5cf8d5..b222d575d255e 100644
--- a/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-rewrite-in-destination-passing-style.mlir
@@ -252,3 +252,23 @@ module attributes {transform.with_named_sequence} {
       transform.yield
   }
 }
+
+// -----
+
+// CHECK-LABEL: func @already_destination_passing_style(
+//  CHECK-SAME:   %[[ARG0:.*]]: tensor<134217728xf32>, %[[ARG1:.*]]: tensor<134217728xf32>) -> tensor<134217728xf32>
+//       CHECK:   %[[RESULT:.*]] = linalg.add ins(%[[ARG0]], %[[ARG1]] : tensor<134217728xf32>, tensor<134217728xf32>)
+//  CHECK-SAME:   outs(%[[ARG0]] : tensor<134217728xf32>) -> tensor<134217728xf32>
+//       CHECK:   return %[[RESULT]] : tensor<134217728xf32>
+func.func @already_destination_passing_style(%arg0: tensor<134217728xf32>, %arg1: tensor<134217728xf32>) -> tensor<134217728xf32> {
+  %0 = linalg.add ins(%arg0, %arg1 : tensor<134217728xf32>, tensor<134217728xf32>) outs(%arg0 : tensor<134217728xf32>) -> tensor<134217728xf32>
+  return %0 : tensor<134217728xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.add"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+    %1 = transform.structured.rewrite_in_destination_passing_style %0 : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
\ No newline at end of file

``````````

</details>


https://github.com/llvm/llvm-project/pull/205034