From 73992bf4e13e40d5c0c33df2f32f2c246ba6bf28 Mon Sep 17 00:00:00 2001 From: Hugo 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, +// CHECK-SAME: %[[ARG1:.*]]: tensor<32x2xi32> +// CHECK-NEXT: %[[FLATTENED:.*]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1]] +// CHECK-NEXT: %[[FLATTENED_RESULT:.*]] = linalg.generic {{.*}} ins(%[[ARG0]] : tensor) 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, %arg1: tensor<32x2xi32>) -> tensor<32x2xi32> { + %0 = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor) 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 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> }