[GlobalOpt] Add pass to convert broadcast batch_matmul to matmul (#24670)
When a `linalg.batch_matmul` multiplies a batched input by a weight that
is the same for every batch, every batch slice runs the same matmul
against the same weight. This pass rewrites that case into a single
`linalg.matmul`.
For an input of shape `[B, M, K]` and a weight `[K, N]` broadcast to
`[B, K, N]`, the pass:
1. collapses the input to `[B*M, K]` using tensor.collapse_shape
2. multiplies it by the `[K, N]` weight using linalg.matmul
3. expands the result back to `[B, M, N]` using tensor.expand_shape
This removes the broadcast and turns the batched contraction into a
plain
matmul, which downstream passes can handle better.
---------
Signed-off-by: Roberto Laudani <laudani@roofline.ai>
Co-authored-by: ziereis <ziereis@roofline.ai>
diff --git a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
index 03dd920..5b5ac17 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
+++ b/compiler/src/iree/compiler/GlobalOptimization/BUILD.bazel
@@ -50,6 +50,7 @@
name = "GlobalOptimization",
srcs = [
"CleanupNumericNarrowing.cpp",
+ "ConvertBatchMatmulToMatmul.cpp",
"ConvertConv2DToImg2Col.cpp",
"ConvertStridedContractionToContraction.cpp",
"DataLayoutPropagation.cpp",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
index 050b3c9..c3b994d 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
+++ b/compiler/src/iree/compiler/GlobalOptimization/CMakeLists.txt
@@ -41,6 +41,7 @@
"Utils.h"
SRCS
"CleanupNumericNarrowing.cpp"
+ "ConvertBatchMatmulToMatmul.cpp"
"ConvertConv2DToImg2Col.cpp"
"ConvertStridedContractionToContraction.cpp"
"DataLayoutPropagation.cpp"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/ConvertBatchMatmulToMatmul.cpp b/compiler/src/iree/compiler/GlobalOptimization/ConvertBatchMatmulToMatmul.cpp
new file mode 100644
index 0000000..648edf1
--- /dev/null
+++ b/compiler/src/iree/compiler/GlobalOptimization/ConvertBatchMatmulToMatmul.cpp
@@ -0,0 +1,247 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/GlobalOptimization/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::GlobalOptimization {
+
+#define GEN_PASS_DEF_CONVERTBATCHMATMULTOMATMULPASS
+#include "iree/compiler/GlobalOptimization/Passes.h.inc"
+
+namespace {
+
+/// If `op` broadcasts a weight along the batch dimension (dim 0) only, returns
+/// the original pre-broadcast weight; otherwise returns failure. Handles both
+/// linalg.broadcast and its generic form (after generalization).
+static FailureOr<Value> getOriginalBroadcastWeight(Operation *op) {
+ if (auto broadcastOp = dyn_cast<linalg::BroadcastOp>(op)) {
+ ArrayRef<int64_t> dims = broadcastOp.getDimensions();
+ if (dims.size() != 1 || dims[0] != 0) {
+ return failure();
+ }
+ return broadcastOp.getInput();
+ }
+
+ auto genericOp = dyn_cast<linalg::GenericOp>(op);
+ if (!genericOp) {
+ return failure();
+ }
+
+ // Check generic is a broadcast: 1 input, 1 output, all parallel iterators
+ if (genericOp.getNumDpsInputs() != 1 || genericOp.getNumDpsInits() != 1) {
+ return failure();
+ }
+
+ auto iterTypes = genericOp.getIteratorTypesArray();
+ if (!llvm::all_of(iterTypes, [](utils::IteratorType t) {
+ return t == utils::IteratorType::parallel;
+ })) {
+ return failure();
+ }
+
+ SmallVector<AffineMap> maps = genericOp.getIndexingMapsArray();
+ AffineMap inputMap = maps[0];
+ AffineMap outputMap = maps[1];
+
+ unsigned numDims = outputMap.getNumDims();
+ if (!outputMap.isIdentity()) {
+ return failure();
+ }
+
+ if (!inputMap.isProjectedPermutation()) {
+ return failure();
+ }
+
+ // Check that the input map is exactly (d0, d1, ..., dN) -> (d1, d2, ..., dN)
+ // i.e., only dimension 0 is broadcast (missing in the result affine expr) and
+ // the remaining dimensions are in order (not permuted). A permuted map like
+ // (d0,d1,d2) -> (d2,d1) represents a fused broadcast+transpose, which this
+ // pattern cannot handle.
+ if (inputMap.getNumResults() != numDims - 1) {
+ return failure();
+ }
+ for (unsigned i = 0; i < inputMap.getNumResults(); ++i) {
+ auto dimExpr = dyn_cast<AffineDimExpr>(inputMap.getResult(i));
+ if (!dimExpr || dimExpr.getPosition() != i + 1) {
+ return failure();
+ }
+ }
+
+ // Check body simply yields the input argument
+ Block *body = genericOp.getBody();
+ if (body->getNumArguments() != 2) {
+ return failure();
+ }
+ auto yieldOp = dyn_cast<linalg::YieldOp>(body->getTerminator());
+ if (!yieldOp || yieldOp.getNumOperands() != 1) {
+ return failure();
+ }
+ if (yieldOp.getOperand(0) != body->getArgument(0)) {
+ return failure();
+ }
+
+ return genericOp.getDpsInputs()[0];
+}
+
+static int64_t computeCollapsedDim(int64_t dim0, int64_t dim1) {
+ if (dim0 == ShapedType::kDynamic || dim1 == ShapedType::kDynamic) {
+ return ShapedType::kDynamic;
+ }
+ return dim0 * dim1;
+}
+
+static Value createCollapseShape(OpBuilder &builder, Location loc, Value input,
+ RankedTensorType inputType) {
+ SmallVector<ReassociationIndices> reassoc = {{0, 1}, {2}};
+ int64_t collapsedDim0 =
+ computeCollapsedDim(inputType.getDimSize(0), inputType.getDimSize(1));
+ auto collapsedType = RankedTensorType::get(
+ {collapsedDim0, inputType.getDimSize(2)}, inputType.getElementType());
+ return tensor::CollapseShapeOp::create(builder, loc, collapsedType, input,
+ reassoc);
+}
+
+static Value createExpandShape(OpBuilder &builder, Location loc, Value input,
+ RankedTensorType outputType, Value originalAct,
+ Value originalOut) {
+ SmallVector<ReassociationIndices> reassoc = {{0, 1}, {2}};
+
+ // Build output shape - extract dynamic dimensions from original tensors
+ SmallVector<OpFoldResult> outputShape;
+ for (int64_t i = 0; i < 3; ++i) {
+ if (outputType.isDynamicDim(i)) {
+ // For batch (i=0) and M (i=1), get from original activation
+ // For N (i=2), get from original output
+ Value source = (i < 2) ? originalAct : originalOut;
+ Value idx = arith::ConstantIndexOp::create(builder, loc, i);
+ Value dimVal = tensor::DimOp::create(builder, loc, source, idx);
+ outputShape.push_back(dimVal);
+ } else {
+ outputShape.push_back(builder.getIndexAttr(outputType.getDimSize(i)));
+ }
+ }
+
+ return tensor::ExpandShapeOp::create(builder, loc, outputType, input, reassoc,
+ outputShape);
+}
+
+/// Pattern to convert broadcast + batch_matmul to collapse_shape + matmul +
+/// expand_shape. This handles BatchMatmulOp and BatchMatmulTransposeBOp
+/// (which share the same underlying op but with different indexing maps).
+/// BatchMatmulTransposeAOp is not supported because the collapse would
+/// produce mismatched shapes (LHS collapses batch*K but output collapses
+/// batch*M).
+class ConvertBroadcastBatchMatmulToMatmul
+ : public OpRewritePattern<linalg::BatchMatmulOp> {
+public:
+ using OpRewritePattern<linalg::BatchMatmulOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(linalg::BatchMatmulOp batchMatmul,
+ PatternRewriter &rewriter) const override {
+ // The RHS must be a batch-dim broadcast of a 2-D weight.
+ Value rhs = batchMatmul.getDpsInputOperand(1)->get();
+ Operation *rhsDefOp = rhs.getDefiningOp();
+ if (!rhsDefOp) {
+ return failure();
+ }
+ FailureOr<Value> weight = getOriginalBroadcastWeight(rhsDefOp);
+ if (failed(weight)) {
+ return failure();
+ }
+
+ // Get LHS activation
+ Value act = batchMatmul.getDpsInputOperand(0)->get();
+ auto actType = cast<RankedTensorType>(act.getType());
+
+ Value out = batchMatmul.getDpsInitOperand(0)->get();
+ auto outType = cast<RankedTensorType>(out.getType());
+
+ // Classify the contraction layout by comparing the indexing maps directly.
+ // batch_matmul has four loop dimensions (batch, m, n, k); the canonical
+ // maps are LHS (batch, m, k), RHS (batch, k, n) and OUT (batch, m, n).
+ MLIRContext *ctx = batchMatmul.getContext();
+ AffineExpr bDim, mDim, nDim, kDim;
+ bindDims(ctx, bDim, mDim, nDim, kDim);
+ auto mapOf = [&](ArrayRef<AffineExpr> results) {
+ return AffineMap::get(/*dimCount=*/4, /*symbolCount=*/0, results, ctx);
+ };
+ SmallVector<AffineMap> maps = batchMatmul.getIndexingMapsArray();
+ if (maps.size() != 3) {
+ return failure();
+ }
+ // Only fold when the LHS is [batch, M, K] and the output is [batch, M, N].
+ // A transposed LHS [batch, K, M] cannot be folded: collapsing it merges the
+ // batch and leading dims into [batch*K, M], which no longer matches the
+ // output's [batch*M, N] collapse.
+ if (maps[0] != mapOf({bDim, mDim, kDim}) ||
+ maps[2] != mapOf({bDim, mDim, nDim})) {
+ return failure();
+ }
+ // The RHS may be plain (batch, k, n) or transpose_b (batch, n, k).
+ bool transposeB;
+ if (maps[1] == mapOf({bDim, kDim, nDim})) {
+ transposeB = false;
+ } else if (maps[1] == mapOf({bDim, nDim, kDim})) {
+ transposeB = true;
+ } else {
+ return failure();
+ }
+
+ Location loc = batchMatmul.getLoc();
+
+ // Collapse activation: [batch, dim1, dim2] -> [batch*dim1, dim2]
+ Value collapsedAct = createCollapseShape(rewriter, loc, act, actType);
+
+ // Collapse output init tensor
+ Value collapsedOut = createCollapseShape(rewriter, loc, out, outType);
+
+ // Create matmul - check if this is a transpose variant
+ auto collapsedOutType = cast<RankedTensorType>(collapsedOut.getType());
+ Value matmulResult;
+
+ if (transposeB) {
+ matmulResult = linalg::MatmulTransposeBOp::create(
+ rewriter, loc, collapsedOutType,
+ ValueRange{collapsedAct, *weight}, collapsedOut)
+ .getResult(0);
+ } else {
+ matmulResult = linalg::MatmulOp::create(rewriter, loc, collapsedOutType,
+ ValueRange{collapsedAct, *weight},
+ collapsedOut)
+ .getResult(0);
+ }
+
+ // Expand result back to 3D: [batch*dim1, dim2] -> [batch, dim1, dim2]
+ Value expandedResult =
+ createExpandShape(rewriter, loc, matmulResult, outType, act, out);
+
+ rewriter.replaceOp(batchMatmul, expandedResult);
+ return success();
+ }
+};
+
+struct ConvertBatchMatmulToMatmulPass
+ : public impl::ConvertBatchMatmulToMatmulPassBase<
+ ConvertBatchMatmulToMatmulPass> {
+ void runOnOperation() override {
+ MLIRContext *context = &getContext();
+ RewritePatternSet patterns(context);
+ patterns.add<ConvertBroadcastBatchMatmulToMatmul>(context);
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+ return signalPassFailure();
+ }
+ }
+};
+
+} // namespace
+} // namespace mlir::iree_compiler::GlobalOptimization
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
index 2db2ee1..049053d 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
+++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.cpp
@@ -173,6 +173,12 @@
createFuseDequantizationMatmulPass)
.addPass(IREE::Flow::createCanonicalizePass)
.addPass(mlir::createCSEPass)
+ // Convert broadcast + batch_matmul before transpose propagation so we
+ // can match the pure broadcast patterns before they get fused with
+ // transposes.
+ .addPass(GlobalOptimization::createConvertBatchMatmulToMatmulPass);
+
+ FunctionLikeNest(mainPassManager)
// Propagate transposes immediately before set encoding/data tiling
// because transpose propagation cannot take an opinion on the preferred
// layout of various operations. This simplifies local propagation
diff --git a/compiler/src/iree/compiler/GlobalOptimization/Passes.td b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
index 8b9260e..b05e76d 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/Passes.td
+++ b/compiler/src/iree/compiler/GlobalOptimization/Passes.td
@@ -14,6 +14,16 @@
let summary = "Cleans up any numeric narrowing ops inserted by iree-global-opt-infer-numeric-narrowing.";
}
+def ConvertBatchMatmulToMatmulPass :
+ Pass<"iree-global-opt-convert-batch-matmul-to-matmul", ""> {
+ let summary = "Converts batch_matmul to reshape + matmul + reshape.";
+ let dependentDialects = [
+ "mlir::arith::ArithDialect",
+ "mlir::linalg::LinalgDialect",
+ "mlir::tensor::TensorDialect",
+ ];
+}
+
def ConvertConv2DToImg2ColPass :
Pass<"iree-global-opt-convert-conv2d-to-img2col", ""> {
let summary = "Convert linalg convolution ops to matmul img2col based implementation";
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
index 8d33426..a8481d7 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/BUILD.bazel
@@ -19,6 +19,7 @@
[
"cleanup_numeric_narrowing.mlir",
"conv2d_to_img2col.mlir",
+ "convert_batch_matmul_to_matmul.mlir",
"data_layout_propagation.mlir",
"demote_contraction_inputs.mlir",
"detach_elementwise_from_named_ops.mlir",
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
index e36a746..95ae4f9 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/CMakeLists.txt
@@ -16,6 +16,7 @@
SRCS
"cleanup_numeric_narrowing.mlir"
"conv2d_to_img2col.mlir"
+ "convert_batch_matmul_to_matmul.mlir"
"data_layout_propagation.mlir"
"demote_contraction_inputs.mlir"
"detach_elementwise_from_named_ops.mlir"
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/convert_batch_matmul_to_matmul.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/convert_batch_matmul_to_matmul.mlir
new file mode 100644
index 0000000..1d36e5e
--- /dev/null
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/convert_batch_matmul_to_matmul.mlir
@@ -0,0 +1,290 @@
+// RUN: iree-opt --split-input-file --mlir-print-local-scope -iree-global-opt-convert-batch-matmul-to-matmul %s | FileCheck %s
+
+// Static shapes with named broadcast
+// CHECK-LABEL: util.func public @static_broadcast_batch_matmul
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<4x10x10xf32>, %[[ARG1:[a-zA-Z0-9]+]]: tensor<10x10xf32>
+// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x10x10xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[EMPTY]] : tensor<4x10x10xf32>) -> tensor<4x10x10xf32>
+// CHECK: %[[COLLAPSED_ACT:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2]{{\]}} : tensor<4x10x10xf32> into tensor<40x10xf32>
+// CHECK: %[[COLLAPSED_OUT:.*]] = tensor.collapse_shape %[[FILL]] {{\[}}[0, 1], [2]{{\]}} : tensor<4x10x10xf32> into tensor<40x10xf32>
+// CHECK: %[[MATMUL:.*]] = linalg.matmul ins(%[[COLLAPSED_ACT]], %[[ARG1]] : tensor<40x10xf32>, tensor<10x10xf32>) outs(%[[COLLAPSED_OUT]] : tensor<40x10xf32>) -> tensor<40x10xf32>
+// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2]{{\]}} output_shape [4, 10, 10] : tensor<40x10xf32> into tensor<4x10x10xf32>
+// CHECK: util.return %[[EXPANDED]] : tensor<4x10x10xf32>
+util.func public @static_broadcast_batch_matmul(%act: tensor<4x10x10xf32>, %weight: tensor<10x10xf32>) -> tensor<4x10x10xf32> {
+ %init_broadcast = tensor.empty() : tensor<4x10x10xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<10x10xf32>)
+ outs(%init_broadcast : tensor<4x10x10xf32>) dimensions = [0]
+ %init_out = tensor.empty() : tensor<4x10x10xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<4x10x10xf32>) -> tensor<4x10x10xf32>
+ %result = linalg.batch_matmul ins(%act, %broadcast : tensor<4x10x10xf32>, tensor<4x10x10xf32>)
+ outs(%fill : tensor<4x10x10xf32>) -> tensor<4x10x10xf32>
+ util.return %result : tensor<4x10x10xf32>
+}
+
+// -----
+
+// Dynamic batch dimension (M is static)
+// CHECK-LABEL: util.func public @dynamic_batch_broadcast_batch_matmul
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x10x10xf32>, %[[ARG1:[a-zA-Z0-9]+]]: tensor<10x10xf32>
+// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x10x10xf32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[DIM]]) : tensor<?x10x10xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[EMPTY]] : tensor<?x10x10xf32>) -> tensor<?x10x10xf32>
+// CHECK: %[[COLLAPSED_ACT:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2]{{\]}} : tensor<?x10x10xf32> into tensor<?x10xf32>
+// CHECK: %[[COLLAPSED_OUT:.*]] = tensor.collapse_shape %[[FILL]] {{\[}}[0, 1], [2]{{\]}} : tensor<?x10x10xf32> into tensor<?x10xf32>
+// CHECK: %[[MATMUL:.*]] = linalg.matmul ins(%[[COLLAPSED_ACT]], %[[ARG1]] : tensor<?x10xf32>, tensor<10x10xf32>) outs(%[[COLLAPSED_OUT]] : tensor<?x10xf32>) -> tensor<?x10xf32>
+// CHECK: %[[BATCH:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x10x10xf32>
+// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2]{{\]}} output_shape [%[[BATCH]], 10, 10] : tensor<?x10xf32> into tensor<?x10x10xf32>
+// CHECK: util.return %[[EXPANDED]] : tensor<?x10x10xf32>
+util.func public @dynamic_batch_broadcast_batch_matmul(%act: tensor<?x10x10xf32>, %weight: tensor<10x10xf32>) -> tensor<?x10x10xf32> {
+ %c0 = arith.constant 0 : index
+ %batch = tensor.dim %act, %c0 : tensor<?x10x10xf32>
+ %init_broadcast = tensor.empty(%batch) : tensor<?x10x10xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<10x10xf32>)
+ outs(%init_broadcast : tensor<?x10x10xf32>) dimensions = [0]
+ %init_out = tensor.empty(%batch) : tensor<?x10x10xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<?x10x10xf32>) -> tensor<?x10x10xf32>
+ %result = linalg.batch_matmul ins(%act, %broadcast : tensor<?x10x10xf32>, tensor<?x10x10xf32>)
+ outs(%fill : tensor<?x10x10xf32>) -> tensor<?x10x10xf32>
+ util.return %result : tensor<?x10x10xf32>
+}
+
+// -----
+
+// Generic broadcast form (after generalization)
+// CHECK-LABEL: util.func public @generic_broadcast_batch_matmul
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x10x10xf32>, %[[ARG1:[a-zA-Z0-9]+]]: tensor<10x10xf32>
+// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x10x10xf32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[DIM]]) : tensor<?x10x10xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[EMPTY]] : tensor<?x10x10xf32>) -> tensor<?x10x10xf32>
+// CHECK: %[[COLLAPSED_ACT:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2]{{\]}} : tensor<?x10x10xf32> into tensor<?x10xf32>
+// CHECK: %[[COLLAPSED_OUT:.*]] = tensor.collapse_shape %[[FILL]] {{\[}}[0, 1], [2]{{\]}} : tensor<?x10x10xf32> into tensor<?x10xf32>
+// CHECK: %[[MATMUL:.*]] = linalg.matmul ins(%[[COLLAPSED_ACT]], %[[ARG1]] : tensor<?x10xf32>, tensor<10x10xf32>) outs(%[[COLLAPSED_OUT]] : tensor<?x10xf32>) -> tensor<?x10xf32>
+// CHECK: %[[BATCH:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x10x10xf32>
+// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2]{{\]}} output_shape [%[[BATCH]], 10, 10] : tensor<?x10xf32> into tensor<?x10x10xf32>
+// CHECK: util.return %[[EXPANDED]] : tensor<?x10x10xf32>
+util.func public @generic_broadcast_batch_matmul(%act: tensor<?x10x10xf32>, %weight: tensor<10x10xf32>) -> tensor<?x10x10xf32> {
+ %c0 = arith.constant 0 : index
+ %batch = tensor.dim %act, %c0 : tensor<?x10x10xf32>
+ %init_broadcast = tensor.empty(%batch) : tensor<?x10x10xf32>
+ // Generic form of broadcast (as produced by GeneralizeLinalgNamedOps)
+ %broadcast = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ } ins(%weight : tensor<10x10xf32>) outs(%init_broadcast : tensor<?x10x10xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ linalg.yield %in : f32
+ } -> tensor<?x10x10xf32>
+ %init_out = tensor.empty(%batch) : tensor<?x10x10xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<?x10x10xf32>) -> tensor<?x10x10xf32>
+ %result = linalg.batch_matmul ins(%act, %broadcast : tensor<?x10x10xf32>, tensor<?x10x10xf32>)
+ outs(%fill : tensor<?x10x10xf32>) -> tensor<?x10x10xf32>
+ util.return %result : tensor<?x10x10xf32>
+}
+
+// -----
+
+// Negative test - broadcast on LHS (should NOT transform)
+// CHECK-LABEL: util.func public @broadcast_on_lhs
+// CHECK: linalg.broadcast
+// CHECK: linalg.batch_matmul
+// CHECK-NOT: tensor.collapse_shape
+util.func public @broadcast_on_lhs(%lhs: tensor<10x10xf32>, %rhs: tensor<4x10x10xf32>) -> tensor<4x10x10xf32> {
+ %init_broadcast = tensor.empty() : tensor<4x10x10xf32>
+ %broadcast = linalg.broadcast ins(%lhs : tensor<10x10xf32>)
+ outs(%init_broadcast : tensor<4x10x10xf32>) dimensions = [0]
+ %init_out = tensor.empty() : tensor<4x10x10xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<4x10x10xf32>) -> tensor<4x10x10xf32>
+ %result = linalg.batch_matmul ins(%broadcast, %rhs : tensor<4x10x10xf32>, tensor<4x10x10xf32>)
+ outs(%fill : tensor<4x10x10xf32>) -> tensor<4x10x10xf32>
+ util.return %result : tensor<4x10x10xf32>
+}
+
+// -----
+
+// Both batch and M are dynamic (should transform with explicit output_shape)
+// CHECK-LABEL: util.func public @both_batch_and_m_dynamic
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?x10xf32>, %[[ARG1:[a-zA-Z0-9]+]]: tensor<10x10xf32>
+// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[BATCH:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?x10xf32>
+// CHECK: %[[M:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?x10xf32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[BATCH]], %[[M]]) : tensor<?x?x10xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[EMPTY]] : tensor<?x?x10xf32>) -> tensor<?x?x10xf32>
+// CHECK: %[[COLLAPSED_ACT:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2]{{\]}} : tensor<?x?x10xf32> into tensor<?x10xf32>
+// CHECK: %[[COLLAPSED_OUT:.*]] = tensor.collapse_shape %[[FILL]] {{\[}}[0, 1], [2]{{\]}} : tensor<?x?x10xf32> into tensor<?x10xf32>
+// CHECK: %[[MATMUL:.*]] = linalg.matmul ins(%[[COLLAPSED_ACT]], %[[ARG1]] : tensor<?x10xf32>, tensor<10x10xf32>) outs(%[[COLLAPSED_OUT]] : tensor<?x10xf32>) -> tensor<?x10xf32>
+// CHECK: %[[DIM_BATCH:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?x10xf32>
+// CHECK: %[[DIM_M:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?x10xf32>
+// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2]{{\]}} output_shape [%[[DIM_BATCH]], %[[DIM_M]], 10] : tensor<?x10xf32> into tensor<?x?x10xf32>
+// CHECK: util.return %[[EXPANDED]] : tensor<?x?x10xf32>
+util.func public @both_batch_and_m_dynamic(%act: tensor<?x?x10xf32>, %weight: tensor<10x10xf32>) -> tensor<?x?x10xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %batch = tensor.dim %act, %c0 : tensor<?x?x10xf32>
+ %m = tensor.dim %act, %c1 : tensor<?x?x10xf32>
+ %init_broadcast = tensor.empty(%batch) : tensor<?x10x10xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<10x10xf32>)
+ outs(%init_broadcast : tensor<?x10x10xf32>) dimensions = [0]
+ %init_out = tensor.empty(%batch, %m) : tensor<?x?x10xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<?x?x10xf32>) -> tensor<?x?x10xf32>
+ %result = linalg.batch_matmul ins(%act, %broadcast : tensor<?x?x10xf32>, tensor<?x10x10xf32>)
+ outs(%fill : tensor<?x?x10xf32>) -> tensor<?x?x10xf32>
+ util.return %result : tensor<?x?x10xf32>
+}
+
+// -----
+
+// Negative test - broadcast adds dimension other than batch (should NOT transform)
+// CHECK-LABEL: util.func public @broadcast_non_batch_dim
+// CHECK: linalg.broadcast
+// CHECK: linalg.batch_matmul
+// CHECK-NOT: tensor.collapse_shape
+util.func public @broadcast_non_batch_dim(%act: tensor<4x10x10xf32>, %weight: tensor<4x10xf32>) -> tensor<4x10x10xf32> {
+ %init_broadcast = tensor.empty() : tensor<4x10x10xf32>
+ // Broadcast adds dimension 1, not 0
+ %broadcast = linalg.broadcast ins(%weight : tensor<4x10xf32>)
+ outs(%init_broadcast : tensor<4x10x10xf32>) dimensions = [1]
+ %init_out = tensor.empty() : tensor<4x10x10xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<4x10x10xf32>) -> tensor<4x10x10xf32>
+ %result = linalg.batch_matmul ins(%act, %broadcast : tensor<4x10x10xf32>, tensor<4x10x10xf32>)
+ outs(%fill : tensor<4x10x10xf32>) -> tensor<4x10x10xf32>
+ util.return %result : tensor<4x10x10xf32>
+}
+
+// -----
+
+// Transpose B variant (RHS weight is [N, K] instead of [K, N])
+// batch_matmul with indexing_maps specifying transpose_b: RHS accesses (batch, n, k)
+// Activation: [4, 8, 10] = [batch, M, K]
+// Weight: [6, 10] = [N, K] (transposed)
+// After broadcast: [4, 6, 10] = [batch, N, K]
+// Output: [4, 8, 6] = [batch, M, N]
+// CHECK-LABEL: util.func public @broadcast_batch_matmul_transpose_b
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<4x8x10xf32>, %[[ARG1:[a-zA-Z0-9]+]]: tensor<6x10xf32>
+// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8x6xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[EMPTY]] : tensor<4x8x6xf32>) -> tensor<4x8x6xf32>
+// CHECK: %[[COLLAPSED_ACT:.*]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2]{{\]}} : tensor<4x8x10xf32> into tensor<32x10xf32>
+// CHECK: %[[COLLAPSED_OUT:.*]] = tensor.collapse_shape %[[FILL]] {{\[}}[0, 1], [2]{{\]}} : tensor<4x8x6xf32> into tensor<32x6xf32>
+// CHECK: %[[MATMUL:.*]] = linalg.matmul indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>] ins(%[[COLLAPSED_ACT]], %[[ARG1]] : tensor<32x10xf32>, tensor<6x10xf32>) outs(%[[COLLAPSED_OUT]] : tensor<32x6xf32>) -> tensor<32x6xf32>
+// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2]{{\]}} output_shape [4, 8, 6] : tensor<32x6xf32> into tensor<4x8x6xf32>
+// CHECK: util.return %[[EXPANDED]] : tensor<4x8x6xf32>
+util.func public @broadcast_batch_matmul_transpose_b(%act: tensor<4x8x10xf32>, %weight: tensor<6x10xf32>) -> tensor<4x8x6xf32> {
+ %init_broadcast = tensor.empty() : tensor<4x6x10xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<6x10xf32>)
+ outs(%init_broadcast : tensor<4x6x10xf32>) dimensions = [0]
+ %init_out = tensor.empty() : tensor<4x8x6xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<4x8x6xf32>) -> tensor<4x8x6xf32>
+ // batch_matmul with transpose_b indexing: RHS is (batch, n, k) instead of (batch, k, n)
+ %result = linalg.batch_matmul
+ indexing_maps = [affine_map<(b, m, n, k) -> (b, m, k)>,
+ affine_map<(b, m, n, k) -> (b, n, k)>,
+ affine_map<(b, m, n, k) -> (b, m, n)>]
+ ins(%act, %broadcast : tensor<4x8x10xf32>, tensor<4x6x10xf32>)
+ outs(%fill : tensor<4x8x6xf32>) -> tensor<4x8x6xf32>
+ util.return %result : tensor<4x8x6xf32>
+}
+
+// -----
+
+// Negative test - fused broadcast + transpose in generic (should NOT transform)
+// When transpose propagation folds a transpose into the broadcast generic, the
+// input map becomes (d0,d1,d2) -> (d2,d1) instead of (d0,d1,d2) -> (d1,d2).
+// The pass must not treat this as a pure broadcast.
+// CHECK-LABEL: util.func public @fused_broadcast_transpose_generic
+// CHECK: linalg.generic
+// CHECK: linalg.batch_matmul
+// CHECK-NOT: tensor.collapse_shape
+util.func public @fused_broadcast_transpose_generic(%act: tensor<4x10x8xf32>, %weight: tensor<6x8xf32>) -> tensor<4x10x6xf32> {
+ %init_broadcast = tensor.empty() : tensor<4x8x6xf32>
+ // Generic that broadcasts on dim 0 AND transposes the inner dims
+ %broadcast_transpose = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d2, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ } ins(%weight : tensor<6x8xf32>) outs(%init_broadcast : tensor<4x8x6xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ linalg.yield %in : f32
+ } -> tensor<4x8x6xf32>
+ %init_out = tensor.empty() : tensor<4x10x6xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<4x10x6xf32>) -> tensor<4x10x6xf32>
+ %result = linalg.batch_matmul ins(%act, %broadcast_transpose : tensor<4x10x8xf32>, tensor<4x8x6xf32>)
+ outs(%fill : tensor<4x10x6xf32>) -> tensor<4x10x6xf32>
+ util.return %result : tensor<4x10x6xf32>
+}
+
+// -----
+
+// Transpose A variant (LHS activation is [batch, K, M] instead of [batch, M, K])
+// This should NOT transform because the collapse would produce mismatched shapes:
+// - LHS [batch, K, M] collapses to [batch*K, M]
+// - Out [batch, M, N] collapses to [batch*M, N]
+// These don't match, so transformation is invalid.
+// CHECK-LABEL: util.func public @broadcast_batch_matmul_transpose_a
+// CHECK: linalg.broadcast
+// CHECK: linalg.batch_matmul
+// CHECK-NOT: tensor.collapse_shape
+util.func public @broadcast_batch_matmul_transpose_a(%act: tensor<4x10x8xf32>, %weight: tensor<10x6xf32>) -> tensor<4x8x6xf32> {
+ %init_broadcast = tensor.empty() : tensor<4x10x6xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<10x6xf32>)
+ outs(%init_broadcast : tensor<4x10x6xf32>) dimensions = [0]
+ %init_out = tensor.empty() : tensor<4x8x6xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<4x8x6xf32>) -> tensor<4x8x6xf32>
+ // batch_matmul with transpose_a indexing: LHS is (batch, k, m) instead of (batch, m, k)
+ %result = linalg.batch_matmul
+ indexing_maps = [affine_map<(b, m, n, k) -> (b, k, m)>,
+ affine_map<(b, m, n, k) -> (b, k, n)>,
+ affine_map<(b, m, n, k) -> (b, m, n)>]
+ ins(%act, %broadcast : tensor<4x10x8xf32>, tensor<4x10x6xf32>)
+ outs(%fill : tensor<4x8x6xf32>) -> tensor<4x8x6xf32>
+ util.return %result : tensor<4x8x6xf32>
+}
+
+// -----
+
+// Negative test - combined transpose_a and transpose_b (should NOT transform).
+// The LHS activation is [batch, K, M] and the RHS weight is [N, K]. For the
+// same reason as the transpose_a variant, because the LHS is transposed the
+// collapse would produce mismatched shapes:
+// - LHS [batch, K, M] collapses to [batch*K, M]
+// - Out [batch, M, N] collapses to [batch*M, N]
+// These don't match, so transformation is invalid. The combined form must be
+// rejected even though it is neither a pure transpose_a nor a pure transpose_b.
+// CHECK-LABEL: util.func public @broadcast_batch_matmul_transpose_a_and_b
+// CHECK: linalg.broadcast
+// CHECK: linalg.batch_matmul
+// CHECK-NOT: tensor.collapse_shape
+// CHECK-NOT: linalg.matmul
+util.func public @broadcast_batch_matmul_transpose_a_and_b(%act: tensor<4x10x8xf32>, %weight: tensor<6x10xf32>) -> tensor<4x8x6xf32> {
+ %init_broadcast = tensor.empty() : tensor<4x6x10xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<6x10xf32>)
+ outs(%init_broadcast : tensor<4x6x10xf32>) dimensions = [0]
+ %init_out = tensor.empty() : tensor<4x8x6xf32>
+ %cst = arith.constant 0.0 : f32
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<4x8x6xf32>) -> tensor<4x8x6xf32>
+ // batch_matmul with transpose_a and transpose_b indexing: LHS is (b, k, m) and
+ // RHS is (b, n, k) instead of (b, m, k) and (b, k, n)
+ %result = linalg.batch_matmul
+ indexing_maps = [affine_map<(b, m, n, k) -> (b, k, m)>,
+ affine_map<(b, m, n, k) -> (b, n, k)>,
+ affine_map<(b, m, n, k) -> (b, m, n)>]
+ ins(%act, %broadcast : tensor<4x10x8xf32>, tensor<4x6x10xf32>)
+ outs(%fill : tensor<4x8x6xf32>) -> tensor<4x8x6xf32>
+ util.return %result : tensor<4x8x6xf32>
+}
diff --git a/tests/e2e/regression/BUILD.bazel b/tests/e2e/regression/BUILD.bazel
index e4bb04b..0991c05 100644
--- a/tests/e2e/regression/BUILD.bazel
+++ b/tests/e2e/regression/BUILD.bazel
@@ -68,6 +68,7 @@
iree_check_single_backend_test_suite(
name = "check_regression_llvm-cpu",
srcs = [
+ "broadcast_batch_matmul_to_matmul.mlir",
"layernorm.mlir",
"pack_pad_transpose_1x9_into_2x1x8x4_issue_12546.mlir",
"split_reduction_using_tiling.mlir",
diff --git a/tests/e2e/regression/CMakeLists.txt b/tests/e2e/regression/CMakeLists.txt
index 8ce9ed7..b8a5fc2 100644
--- a/tests/e2e/regression/CMakeLists.txt
+++ b/tests/e2e/regression/CMakeLists.txt
@@ -51,6 +51,7 @@
NAME
check_regression_llvm-cpu
SRCS
+ "broadcast_batch_matmul_to_matmul.mlir"
"i1_inlined_constant.mlir"
"layernorm.mlir"
"linalg_ops.mlir"
diff --git a/tests/e2e/regression/broadcast_batch_matmul_to_matmul.mlir b/tests/e2e/regression/broadcast_batch_matmul_to_matmul.mlir
new file mode 100644
index 0000000..316d6f3
--- /dev/null
+++ b/tests/e2e/regression/broadcast_batch_matmul_to_matmul.mlir
@@ -0,0 +1,60 @@
+// End-to-end correctness test for the batch-matmul-to-matmul pass
+// (--iree-global-opt-convert-batch-matmul-to-matmul), which runs as
+// part of the default global optimization pipeline.
+
+func.func @broadcast_batch_matmul() {
+ %act = util.unfoldable_constant dense<[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+ [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]]> : tensor<2x2x3xf32>
+ %weight = util.unfoldable_constant dense<[[1.0, 0.0, 0.0, 1.0],
+ [0.0, 1.0, 0.0, 1.0],
+ [0.0, 0.0, 1.0, 1.0]]> : tensor<3x4xf32>
+ %cst = arith.constant 0.0 : f32
+ %init_broadcast = tensor.empty() : tensor<2x3x4xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<3x4xf32>)
+ outs(%init_broadcast : tensor<2x3x4xf32>) dimensions = [0]
+ %init_out = tensor.empty() : tensor<2x2x4xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+
+ %opt = linalg.batch_matmul ins(%act, %broadcast : tensor<2x2x3xf32>, tensor<2x3x4xf32>)
+ outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+
+ %broadcast_ref = util.optimization_barrier %broadcast : tensor<2x3x4xf32>
+ %ref = linalg.batch_matmul ins(%act, %broadcast_ref : tensor<2x2x3xf32>, tensor<2x3x4xf32>)
+ outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+
+ check.expect_almost_eq(%opt, %ref) : tensor<2x2x4xf32>
+ return
+}
+
+func.func @broadcast_batch_matmul_transpose_b() {
+ %act = util.unfoldable_constant dense<[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
+ [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]]> : tensor<2x2x3xf32>
+ %weight = util.unfoldable_constant dense<[[1.0, 0.0, 0.0],
+ [0.0, 1.0, 0.0],
+ [0.0, 0.0, 1.0],
+ [1.0, 1.0, 1.0]]> : tensor<4x3xf32>
+ %cst = arith.constant 0.0 : f32
+ %init_broadcast = tensor.empty() : tensor<2x4x3xf32>
+ %broadcast = linalg.broadcast ins(%weight : tensor<4x3xf32>)
+ outs(%init_broadcast : tensor<2x4x3xf32>) dimensions = [0]
+ %init_out = tensor.empty() : tensor<2x2x4xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%init_out : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+
+ %opt = linalg.batch_matmul
+ indexing_maps = [affine_map<(b, m, n, k) -> (b, m, k)>,
+ affine_map<(b, m, n, k) -> (b, n, k)>,
+ affine_map<(b, m, n, k) -> (b, m, n)>]
+ ins(%act, %broadcast : tensor<2x2x3xf32>, tensor<2x4x3xf32>)
+ outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+
+ %broadcast_ref = util.optimization_barrier %broadcast : tensor<2x4x3xf32>
+ %ref = linalg.batch_matmul
+ indexing_maps = [affine_map<(b, m, n, k) -> (b, m, k)>,
+ affine_map<(b, m, n, k) -> (b, n, k)>,
+ affine_map<(b, m, n, k) -> (b, m, n)>]
+ ins(%act, %broadcast_ref : tensor<2x2x3xf32>, tensor<2x4x3xf32>)
+ outs(%fill : tensor<2x2x4xf32>) -> tensor<2x2x4xf32>
+
+ check.expect_almost_eq(%opt, %ref) : tensor<2x2x4xf32>
+ return
+}
