[Mlir-commits] [mlir] [MLIR][Linalg] Fix FlattenElementwiseOp for linalg.broadcast (PR #210738)

[Hugo Trachino]

https://github.com/nujaa updated https://github.com/llvm/llvm-project/pull/210738

>From 73992bf4e13e40d5c0c33df2f32f2c246ba6bf28 Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Mon, 20 Jul 2026 23:07:09 +0800
Subject: [PATCH 1/2] [MLIR][Linalg] Fix FlattenElementwiseOp for
 linalg.broadcast

---
 .../TransformOps/LinalgTransformOps.cpp       |  2 +-
 .../Dialect/Linalg/flatten-elementwise.mlir   | 28 +++++++++++++++++++
 .../Dialect/Linalg/flatten-unsupported.mlir   | 12 ++------
 3 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 42ce13b6f214c..01957c91497c8 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4311,7 +4311,7 @@ DiagnosedSilenceableFailure transform::FlattenElementwiseLinalgOp::applyToOne(
            << "only elementwise flattening is supported";
 
   if (!llvm::all_of(target.getIndexingMapsArray(), [](AffineMap m) {
-        return m.isProjectedPermutation(/*allowZeroInResults=*/false);
+        return m.isPermutation() || m.getNumResults() == 0;
       })) {
     results.push_back(target);
     return mlir::emitSilenceableFailure(target->getLoc())
diff --git a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
index f9f8f9e2b8b44..ca06062f61840 100644
--- a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
@@ -43,6 +43,34 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK-LABEL: func.func @broadcast_rank0_tensor(
+// CHECK-SAME:                         %[[ARG0:.*]]: tensor<i32>,
+// CHECK-SAME:                         %[[ARG1:.*]]: tensor<32x2xi32>
+// CHECK-NEXT:    %[[FLATTENED:.*]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1]]
+// CHECK-NEXT:    %[[FLATTENED_RESULT:.*]] = linalg.generic {{.*}} ins(%[[ARG0]] : tensor<i32>) outs(%[[FLATTENED]] : tensor<64xi32>)
+// CHECK:         %[[RESULT:.*]] = tensor.expand_shape %[[FLATTENED_RESULT]] {{\[}}[0, 1]] output_shape [32, 2] : tensor<64xi32> into tensor<32x2xi32>
+#map0 = affine_map<(d0, d1) -> ()>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+
+func.func @broadcast_rank0_tensor(%arg0: tensor<i32>, %arg1: tensor<32x2xi32>) -> tensor<32x2xi32> {
+  %0 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<i32>) outs(%arg1 : tensor<32x2xi32>) {
+    ^bb0(%in: i32, %out: i32):
+      linalg.yield %in : i32
+  } -> tensor<32x2xi32>
+  return %0 : tensor<32x2xi32>
+}
+
+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 @map_memref(
 // CHECK-SAME:                 %[[ARG0:[a-zA-Z0-9_]*]]: memref<32x7xf32>
 // CHECK-SAME:                 %[[ARG1:[a-zA-Z0-9_]*]]: memref<32x7xf32>
diff --git a/mlir/test/Dialect/Linalg/flatten-unsupported.mlir b/mlir/test/Dialect/Linalg/flatten-unsupported.mlir
index f0b8b8b153bbb..7ecafcb6fdd9e 100644
--- a/mlir/test/Dialect/Linalg/flatten-unsupported.mlir
+++ b/mlir/test/Dialect/Linalg/flatten-unsupported.mlir
@@ -56,16 +56,10 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-#map1 = affine_map<(d0, d1, d2) -> (d1, 0)>
-#map2 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-
-func.func @unsupported_rank_expanding_broadcasting_elementwise(%arg0: tensor<2x1xi32>, %arg1: tensor<32x2x2xi32>) -> tensor<32x2x2xi32> {
+func.func @unsupported_dim_expanding_broadcast(%arg0: tensor<64xi16>, %arg1: tensor<32x64xi16>) -> tensor<32x64xi16> {
   // expected-error @below {{operators with broadcasting semantics are not supported}}
-  %0 = 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 %0 : tensor<32x2x2xi32>
+  %broadcasted = linalg.broadcast ins(%arg0 : tensor<64xi16>) outs(%arg1 : tensor<32x64xi16>) dimensions = [0] 
+  return %broadcasted : tensor<32x64xi16>
 }
 
 module attributes {transform.with_named_sequence} {

>From 851168bcdf1be59fc0bc1cab6075f7ef526d64fd Mon Sep 17 00:00:00 2001
From: Hugo <hugo.trachino at huawei.com>
Date: Tue, 21 Jul 2026 21:59:35 +0800
Subject: [PATCH 2/2] fixup! [MLIR][Linalg] Fix FlattenElementwiseOp for
 linalg.broadcast

---
 mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp | 2 +-
 mlir/test/Dialect/Linalg/flatten-unsupported.mlir           | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 01957c91497c8..8c5bd50bd6f8b 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4315,7 +4315,7 @@ DiagnosedSilenceableFailure transform::FlattenElementwiseLinalgOp::applyToOne(
       })) {
     results.push_back(target);
     return mlir::emitSilenceableFailure(target->getLoc())
-           << "operators with broadcasting semantics are not supported";
+           << "broadcasting of non scalar operands is not supported";
   }
 
   // If rank <= 1, do nothing
diff --git a/mlir/test/Dialect/Linalg/flatten-unsupported.mlir b/mlir/test/Dialect/Linalg/flatten-unsupported.mlir
index 7ecafcb6fdd9e..492211a5e5141 100644
--- a/mlir/test/Dialect/Linalg/flatten-unsupported.mlir
+++ b/mlir/test/Dialect/Linalg/flatten-unsupported.mlir
@@ -37,7 +37,7 @@ module attributes {transform.with_named_sequence} {
 #map2 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 
 func.func @unsupported_broadcasting_elementwise(%arg0: tensor<1x2x1xi32>, %arg1: tensor<32x2x2xi32>) -> tensor<32x2x2xi32> {
-  // expected-error @below {{operators with broadcasting semantics are not supported}}
+  // expected-error @below {{broadcasting of non scalar operands is not supported}}
   %0 = 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
@@ -57,7 +57,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 func.func @unsupported_dim_expanding_broadcast(%arg0: tensor<64xi16>, %arg1: tensor<32x64xi16>) -> tensor<32x64xi16> {
-  // expected-error @below {{operators with broadcasting semantics are not supported}}
+  // expected-error @below {{broadcasting of non scalar operands is not supported}}
   %broadcasted = linalg.broadcast ins(%arg0 : tensor<64xi16>) outs(%arg1 : tensor<32x64xi16>) dimensions = [0] 
   return %broadcasted : tensor<32x64xi16>
 }