[Mlir-commits] [mlir] bb81031 - [MLIR][Linalg] Recompute linalg.broadcast dimensions when flattening (#213641)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 13 03:58:44 PDT 2026
Author: Chibuoyim (Wilson) Ogbonna
Date: 2026-08-13T11:58:39+01:00
New Revision: bb810316fb4cb8e1c3ebbee81467de62b24295aa
URL: https://github.com/llvm/llvm-project/commit/bb810316fb4cb8e1c3ebbee81467de62b24295aa
DIFF: https://github.com/llvm/llvm-project/commit/bb810316fb4cb8e1c3ebbee81467de62b24295aa.diff
LOG: [MLIR][Linalg] Recompute linalg.broadcast dimensions when flattening (#213641)
per discussion in
[211203](https://github.com/llvm/llvm-project/pull/211203), @Nujaa found
the following case being rejected currently:
```
func.func @broadcast_rank0_tensor(%arg0: tensor<i32>, %arg1: tensor<32x2xi32>) -> tensor<32x2xi32> {
%0 = linalg.broadcast ins(%arg0 : tensor<i32>) outs(%arg1 : tensor<32x2xi32>) dimensions = [0, 1]
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
}
}
```
with:
```
~/llvm-project/build$ bin/mlir-opt test.mlir --transform-interpreter
test.mlir:7:8: error: 'linalg.broadcast' op input rank plus added dimensions does not match init rank. input rank: 0, dimensions size: 2, init rank: 1
%0 = linalg.broadcast ins(%arg0 : tensor<i32>) outs(%arg1 : tensor<32x2xi32>) dimensions = [0, 1]
```
After the `flatten_elementwise` pass runs, the `linalg.broadcast`
*still* carries its original dimensions attribute `[0, 1]`, leaving us
with something like:
```
linalg.broadcast ins(%0 : tensor<i32>) outs(%1 : tensor<64xi32>) dimensions = [0, 1]
```
which, of course, gets rejected by the `linalg.broadcast` verifier
because `len(ins) + len(dimensions) != len(outs)`, as expected.
Flattening/collapsing needs to also recompute the dimensions attribute
for a broadcast, and this patch does this.
---------
Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
Added:
Modified:
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
mlir/test/Dialect/Linalg/flatten-elementwise.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index e8aa2e5c3b60b..22724e3c31121 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -4464,6 +4464,14 @@ DiagnosedSilenceableFailure transform::FlattenElementwiseLinalgOp::applyToOne(
return DiagnosedSilenceableFailure::success();
}
+ // Only broadcasts with a 0-D input are handled; leave anything else
+ // unchanged.
+ if (auto broadcastOp = dyn_cast<linalg::BroadcastOp>(target.getOperation());
+ broadcastOp && broadcastOp.getInput().getType().getRank() != 0) {
+ 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/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index db46de75abd1a..37386bee2def4 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1809,12 +1809,33 @@ GenericOp cloneToCollapsedOp<GenericOp>(RewriterBase &rewriter,
return collapsedOp;
}
+/// Collapse a `BroadcastOp` with a 0-D input into the single flattened
+/// dimension (`dimensions = [0]`).
+template <>
+BroadcastOp
+cloneToCollapsedOp<BroadcastOp>(RewriterBase &rewriter, BroadcastOp origOp,
+ const CollapsingInfo &collapsingInfo) {
+ assert(origOp.getInput().getType().getRank() == 0 && "expected a 0-D input");
+
+ SmallVector<Value> inputOperands, outputOperands;
+ SmallVector<Type> resultTypes;
+ collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
+ outputOperands, resultTypes);
+
+ SmallVector<int64_t> newDimensions = {0};
+ return BroadcastOp::create(rewriter, origOp.getLoc(), inputOperands[0],
+ outputOperands[0], newDimensions);
+}
+
static LinalgOp createCollapsedOp(LinalgOp op,
const CollapsingInfo &collapsingInfo,
RewriterBase &rewriter) {
if (GenericOp genericOp = dyn_cast<GenericOp>(op.getOperation())) {
return cloneToCollapsedOp(rewriter, genericOp, collapsingInfo);
}
+ if (BroadcastOp broadcastOp = dyn_cast<BroadcastOp>(op.getOperation())) {
+ return cloneToCollapsedOp(rewriter, broadcastOp, collapsingInfo);
+ }
return cloneToCollapsedOp(rewriter, op, collapsingInfo);
}
diff --git a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
index ca06062f61840..0a282b16b8b97 100644
--- a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
@@ -43,7 +43,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// CHECK-LABEL: func.func @broadcast_rank0_tensor(
+// CHECK-LABEL: func.func @broadcast_as_generic_rank0_tensor(
// CHECK-SAME: %[[ARG0:.*]]: tensor<i32>,
// CHECK-SAME: %[[ARG1:.*]]: tensor<32x2xi32>
// CHECK-NEXT: %[[FLATTENED:.*]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1]]
@@ -52,7 +52,7 @@ module attributes {transform.with_named_sequence} {
#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> {
+func.func @broadcast_as_generic_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
@@ -71,6 +71,49 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-LABEL: func.func @broadcast_as_named_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.broadcast ins(%[[ARG0]] : tensor<i32>) outs(%[[FLATTENED]] : tensor<64xi32>) dimensions = [0]
+// CHECK: %[[RESULT:.*]] = tensor.expand_shape %[[FLATTENED_RESULT]] {{\[}}[0, 1]] output_shape [32, 2] : tensor<64xi32> into tensor<32x2xi32>
+func.func @broadcast_as_named_rank0_tensor(%arg0: tensor<i32>, %arg1: tensor<32x2xi32>) -> tensor<32x2xi32> {
+ %0 = linalg.broadcast ins(%arg0 : tensor<i32>) outs(%arg1 : tensor<32x2xi32>) dimensions = [0, 1]
+ 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 @broadcast_as_named_non_rank0_tensor(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<4x8xf32>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<4x8xf32>
+// CHECK-NEXT: %[[RESULT:.*]] = linalg.broadcast ins(%[[ARG0]] : tensor<4x8xf32>) outs(%[[ARG1]] : tensor<4x8xf32>) dimensions = []
+// CHECK-NEXT: return %[[RESULT]] : tensor<4x8xf32>
+func.func @broadcast_as_named_non_rank0_tensor(%arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %0 = linalg.broadcast ins(%arg0 : tensor<4x8xf32>) outs(%arg1 : tensor<4x8xf32>) dimensions = []
+ return %0 : tensor<4x8xf32>
+}
+
+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>
More information about the Mlir-commits
mailing list