[Mlir-commits] [mlir] [MLIR][Linalg] Recompute linalg.broadcast dimensions when flattening (PR #213641)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 3 04:33:26 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Chibuoyim (Wilson) Ogbonna (bruteforceboy)
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/213641.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+40)
- (modified) mlir/test/Dialect/Linalg/flatten-elementwise.mlir (+45)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index db46de75abd1a..5282514f9c5dc 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1809,12 +1809,49 @@ GenericOp cloneToCollapsedOp<GenericOp>(RewriterBase &rewriter,
return collapsedOp;
}
+/// Collapse a `BroadcastOp`, recomputing its `dimensions` for the collapsed
+/// iteration space. Returns null if the collapse is not expressible as a
+/// broadcast.
+template <>
+BroadcastOp
+cloneToCollapsedOp<BroadcastOp>(RewriterBase &rewriter, BroadcastOp origOp,
+ const CollapsingInfo &collapsingInfo) {
+ ArrayRef<int64_t> broadcastDims = origOp.getDimensions();
+ SmallVector<int64_t> newDimensions;
+ for (auto [collapsedDim, foldedDims] :
+ llvm::enumerate(collapsingInfo.getCollapsedOpToOrigOpMapping())) {
+ size_t numBroadcast = llvm::count_if(foldedDims, [&](int64_t d) {
+ return llvm::is_contained(broadcastDims, d);
+ });
+ // A collapsed dimension is a broadcast dimension iff all the dimensions it
+ // folds are; a mix of the two cannot be represented as a broadcast.
+ if (numBroadcast != 0 && numBroadcast != foldedDims.size())
+ return nullptr;
+ if (numBroadcast != 0)
+ newDimensions.push_back(collapsedDim);
+ }
+
+ SmallVector<Value> inputOperands, outputOperands;
+ SmallVector<Type> resultTypes;
+ collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
+ outputOperands, resultTypes);
+
+ 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())) {
+ BroadcastOp collapsedOp =
+ cloneToCollapsedOp(rewriter, broadcastOp, collapsingInfo);
+ return collapsedOp ? cast<LinalgOp>(collapsedOp.getOperation())
+ : LinalgOp();
+ }
return cloneToCollapsedOp(rewriter, op, collapsingInfo);
}
@@ -1871,6 +1908,9 @@ FailureOr<CollapseResult> mlir::linalg::collapseOpIterationDims(
}
LinalgOp collapsedOp = createCollapsedOp(op, collapsingInfo, rewriter);
+ if (!collapsedOp)
+ return rewriter.notifyMatchFailure(
+ op, "failed to create collapsed op for the specified dimensions");
Location loc = op->getLoc();
SmallVector<OpFoldResult> loopBound =
diff --git a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
index ca06062f61840..451563d25ca1c 100644
--- a/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
+++ b/mlir/test/Dialect/Linalg/flatten-elementwise.mlir
@@ -71,6 +71,51 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-LABEL: func.func @broadcast_rank0_named_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_rank0_named_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_identity_named_tensor(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<4x8xf32>,
+// CHECK-SAME: %[[ARG1:.*]]: tensor<4x8xf32>
+// CHECK-NEXT: %[[IN:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]]
+// CHECK-NEXT: %[[OUT:.*]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1]]
+// CHECK-NEXT: %[[FLATTENED_RESULT:.*]] = linalg.broadcast ins(%[[IN]] : tensor<32xf32>) outs(%[[OUT]] : tensor<32xf32>) dimensions = []
+// CHECK: %[[RESULT:.*]] = tensor.expand_shape %[[FLATTENED_RESULT]] {{\[}}[0, 1]] output_shape [4, 8] : tensor<32xf32> into tensor<4x8xf32>
+func.func @broadcast_identity_named_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>
``````````
</details>
https://github.com/llvm/llvm-project/pull/213641
More information about the Mlir-commits
mailing list