[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]

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

`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 


>From 91ce3e7783e1f9f69b429826469ca8d458a54e5c Mon Sep 17 00:00:00 2001
From: jpwang <jpwang at smail.nju.edu.cn>
Date: Mon, 22 Jun 2026 12:36:50 +0800
Subject: [PATCH] [mlir] Ignore Op that is already in destination-passing style
 during RewriteInDestinationPassingStyleOp

---
 .../TransformOps/LinalgTransformOps.cpp       |  1 +
 ...-rewrite-in-destination-passing-style.mlir | 20 +++++++++++++++++++
 2 files changed, 21 insertions(+)

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