[Mlir-commits] [mlir] [MLIR][Vector] Fix FlattenElementwiseOp on broadcasted linalgs (PR #207005)

[Hugo Trachino]

https://github.com/nujaa created https://github.com/llvm/llvm-project/pull/207005

Applying FlattenElementwiseLinalgOp on Broadcasted linalg would break. This transformation is not valid for those usecase, so I added a condition to exit gracefully instead.

>From d453196fa77a41fa2c54d13a98da731353c10b0f Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Wed, 1 Jul 2026 21:08:28 +0800
Subject: [PATCH] [MLIR][Vector] Fix FlattenElementwiseOp on broadcasted
 linalgs

---
 .../TransformOps/LinalgTransformOps.cpp       |  8 +++
 .../Dialect/Linalg/flatten-elementwise.mlir   | 53 +++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index f44693096b26b..d39c1cb0e4a20 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4297,6 +4297,14 @@ DiagnosedSilenceableFailure transform::FlattenElementwiseLinalgOp::applyToOne(
     return DiagnosedSilenceableFailure::success();
   }
 
+  // Bail out on linalgs with broadcasting semantics
+  if (!llvm::all_of(target.getIndexingMapsArray(), [](AffineMap m) {
+        return m.isProjectedPermutation(/*allowZeroInResults=*/false);
+      })) {
+    results.push_back(target);
+    return DiagnosedSilenceableFailure::success();
+  }
+
   // Attempt to flatten all dims to one.
   ReassociationIndices reassociation(target.getNumLoops());
   std::iota(reassociation.begin(), reassociation.end(), 0);
diff --git a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
index 9fe50a521d2d8..2ca08215a166a 100644
--- a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
@@ -118,3 +118,56 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+// CHECK-LABEL: func.func @generic
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9_]*]]: tensor<1x2x1xi32>
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9_]*]]: tensor<32x2x2xi32>
+// CHECK-NEXT:         %[[RESULT:.*]] = linalg.generic
+// CHECK-NEXT:         return %[[RESULT]]
+#map1 = affine_map<(d0, d1, d2) -> (0, d1, 0)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+func.func @generic(%arg0: tensor<1x2x1xi32>, %arg1: tensor<32x2x2xi32>) -> tensor<32x2x2xi32> {
+    %13 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0 : tensor<1x2x1xi32>) outs(%arg1 : tensor<32x2x2xi32>) {
+      ^bb0(%in: i32, %out: i32):
+        linalg.yield %in : i32
+    } -> tensor<32x2x2xi32>
+    return %13 : tensor<32x2x2xi32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match interface{LinalgOp} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %flattened = transform.structured.flatten_elementwise %0
+      : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// CHECK-LABEL: func.func @generic
+// CHECK-SAME:     %[[ARG0:[a-zA-Z0-9_]*]]: tensor<2x1xi32>
+// CHECK-SAME:     %[[ARG1:[a-zA-Z0-9_]*]]: tensor<32x2x2xi32>
+// CHECK-NEXT:         %[[RESULT:.*]] = linalg.generic
+// CHECK-NEXT:         return %[[RESULT]]
+#map1 = affine_map<(d0, d1, d2) -> (d1, 0)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+func.func @generic(%arg0: tensor<2x1xi32>, %arg1: tensor<32x2x2xi32>) -> tensor<32x2x2xi32> {
+    %13 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x1xi32>) outs(%arg1 : tensor<32x2x2xi32>) {
+      ^bb0(%in: i32, %out: i32):
+        linalg.yield %in : i32
+    } -> tensor<32x2x2xi32>
+    return %13 : tensor<32x2x2xi32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match interface{LinalgOp} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %flattened = transform.structured.flatten_elementwise %0
+      : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+}
\ No newline at end of file