[Mlir-commits] [mlir] [mlir][tosa] Support variadic batch dimensions in MATMUL operator (PR #232371)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 23 10:54:55 PDT 2026
https://github.com/karmeh01 updated https://github.com/llvm/llvm-project/pull/232371
>From 64ea9605a78f69141b63ad4cc5a4afc2bdfc7bf3 Mon Sep 17 00:00:00 2001
From: Kareem Mehanna <kareemhisham.mehanna at arm.com>
Date: Mon, 12 Oct 2026 12:06:27 +0100
Subject: [PATCH] [mlir][tosa] Support variadic batch dimensions in MATMUL
operator
Normalize operator ranks and batch dimensions before lowering.
Use matmul for rank-two inputs, and collapse batch
dimensions to rank-three and use batch matmul for higher rank inputs.
Preserve zero-point handling for quantized inputs.
Add conversion tests covering dynamic shapes, broadcasting, quantized
operations, and unsupported inputs. Update existing lowering checks.
---
.../Conversion/TosaToLinalg/TosaToLinalg.cpp | 46 +-
.../TosaToLinalg/TosaToLinalgNamed.cpp | 150 +++---
.../TosaToLinalg/TosaToLinalgUtils.h | 7 +-
.../tosa-to-linalg-named-invalid.mlir | 36 ++
.../TosaToLinalg/tosa-to-linalg-named.mlir | 433 +++++++++++++++++-
5 files changed, 572 insertions(+), 100 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index c4151146fde68e..9092ddb282a620 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -1451,9 +1451,9 @@ static LogicalResult reduceMatchAndRewriteHelper(OpTy op, uint64_t axis,
}
// Insert missing leading batch dimensions, taking their sizes from `other`.
-static Value broadcastMatMulTLeadingDimensions(PatternRewriter &rewriter,
- Location loc, Value operand,
- Value other) {
+static Value broadcastMatMulLeadingDimensions(PatternRewriter &rewriter,
+ Location loc, Value operand,
+ Value other) {
auto type = cast<RankedTensorType>(operand.getType());
int64_t missingDims =
cast<RankedTensorType>(other.getType()).getRank() - type.getRank();
@@ -1475,13 +1475,12 @@ static Value broadcastMatMulTLeadingDimensions(PatternRewriter &rewriter,
}
// Expand an existing singleton batch dimension to `size`, preserving the rank.
-// Missing dimensions are handled by broadcastMatMulTLeadingDimensions.
+// Missing dimensions are handled by broadcastMatMulLeadingDimensions.
// linalg.broadcast inserts dimensions, so first remove the singleton dimension
// with a rank-reducing slice, then reinsert it at the target size.
-static Value broadcastMatMulTSingletonDimension(PatternRewriter &rewriter,
- Location loc, Value operand,
- int64_t dim,
- OpFoldResult size) {
+static Value broadcastMatMulSingletonDimension(PatternRewriter &rewriter,
+ Location loc, Value operand,
+ int64_t dim, OpFoldResult size) {
auto type = cast<RankedTensorType>(operand.getType());
SmallVector<int64_t> sliceShape(type.getShape());
sliceShape.erase(sliceShape.begin() + dim);
@@ -1503,13 +1502,13 @@ static Value broadcastMatMulTSingletonDimension(PatternRewriter &rewriter,
// Align the ranks, then broadcast each shared batch dimension independently.
static std::pair<Value, Value>
-normalizeMatMulTBatchDimensions(PatternRewriter &rewriter, Location loc,
- Value a, Value b, int64_t batchRank) {
+normalizeMatMulBatchDimensions(PatternRewriter &rewriter, Location loc, Value a,
+ Value b, int64_t batchRank) {
int64_t aRank = cast<RankedTensorType>(a.getType()).getRank();
int64_t bRank = cast<RankedTensorType>(b.getType()).getRank();
int64_t firstSharedDim = std::max(aRank, bRank) - std::min(aRank, bRank);
- a = broadcastMatMulTLeadingDimensions(rewriter, loc, a, b);
- b = broadcastMatMulTLeadingDimensions(rewriter, loc, b, a);
+ a = broadcastMatMulLeadingDimensions(rewriter, loc, a, b);
+ b = broadcastMatMulLeadingDimensions(rewriter, loc, b, a);
IndexPool indexPool;
// The inserted leading dimensions already match their partner's sizes.
@@ -1526,11 +1525,11 @@ normalizeMatMulTBatchDimensions(PatternRewriter &rewriter, Location loc,
auto [size, masterOperand] =
computeTargetSize(rewriter, loc, indexPool, ValueRange{a, b}, dim);
if (aSize == 1) {
- a = broadcastMatMulTSingletonDimension(rewriter, loc, a, dim, size);
+ a = broadcastMatMulSingletonDimension(rewriter, loc, a, dim, size);
continue;
}
if (bSize == 1) {
- b = broadcastMatMulTSingletonDimension(rewriter, loc, b, dim, size);
+ b = broadcastMatMulSingletonDimension(rewriter, loc, b, dim, size);
continue;
}
a = broadcastDynamicDimension(rewriter, loc, indexPool, a, dim, size,
@@ -1541,14 +1540,15 @@ normalizeMatMulTBatchDimensions(PatternRewriter &rewriter, Location loc,
return {a, b};
}
-MatMulTOperands
-mlir::tosa::prepareMatMulTOperands(PatternRewriter &rewriter, Location loc,
- Value a, Value b,
- RankedTensorType outputType) {
+MatMulTOperands mlir::tosa::prepareMatMulOperands(PatternRewriter &rewriter,
+ Location loc, Value a,
+ Value b,
+ RankedTensorType outputType,
+ bool transposeB) {
const int64_t outputRank = outputType.getRank();
const int64_t batchRank = outputRank - 2;
auto normalizedOperands =
- normalizeMatMulTBatchDimensions(rewriter, loc, a, b, batchRank);
+ normalizeMatMulBatchDimensions(rewriter, loc, a, b, batchRank);
a = normalizedOperands.first;
b = normalizedOperands.second;
@@ -1557,7 +1557,8 @@ mlir::tosa::prepareMatMulTOperands(PatternRewriter &rewriter, Location loc,
for (int64_t batchDim = 0; batchDim < batchRank; ++batchDim)
outputShape.push_back(tensor::getMixedSize(rewriter, loc, a, batchDim));
outputShape.push_back(tensor::getMixedSize(rewriter, loc, a, outputRank - 2));
- outputShape.push_back(tensor::getMixedSize(rewriter, loc, b, outputRank - 2));
+ outputShape.push_back(
+ tensor::getMixedSize(rewriter, loc, b, outputRank - 1 - transposeB));
SmallVector<Value> dynamicDims;
for (auto [dim, size] : llvm::enumerate(outputShape))
if (outputType.isDynamicDim(dim))
@@ -1611,8 +1612,9 @@ class QuantizedMatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
const int64_t outputRank = outputType.getRank();
const int64_t batchRank = outputRank - 2;
- auto [a, b, zeroTensor] = prepareMatMulTOperands(
- rewriter, loc, adaptor.getA(), adaptor.getB(), outputType);
+ auto [a, b, zeroTensor] =
+ prepareMatMulOperands(rewriter, loc, adaptor.getA(), adaptor.getB(),
+ outputType, /*transposeB=*/true);
Type outputElementType = outputType.getElementType();
int64_t loopRank = outputRank + 1;
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index a7f7414bf56ed3..f646c3a80f06f2 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -605,6 +605,16 @@ class DepthwiseConvConverter
}
};
+static SmallVector<ReassociationIndices>
+getMatMulBatchCollapseReassociation(int64_t rank) {
+ SmallVector<ReassociationIndices> reassociation(3);
+ for (int64_t dim = 0; dim < rank - 2; ++dim)
+ reassociation[0].push_back(dim);
+ reassociation[1].push_back(rank - 2);
+ reassociation[2].push_back(rank - 1);
+ return reassociation;
+}
+
class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
public:
using OpConversionPattern<tosa::MatMulOp>::OpConversionPattern;
@@ -612,85 +622,110 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
matchAndRewrite(tosa::MatMulOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const final {
Location loc = op.getLoc();
-
- auto outputTy = cast<ShapedType>(op.getType());
- auto outputElementTy = outputTy.getElementType();
-
- SmallVector<Value> dynDims;
- dynDims.resize(cast<ShapedType>(op->getResult(0).getType()).getRank());
-
- if (!outputTy.hasRank() || outputTy.isDynamicDim(0)) {
- dynDims[0] = tensor::DimOp::create(rewriter, loc, op->getOperand(0), 0);
- }
-
- if (!outputTy.hasRank() || outputTy.isDynamicDim(1)) {
- dynDims[1] = tensor::DimOp::create(rewriter, loc, op->getOperand(0), 1);
- }
-
- if (!outputTy.hasRank() || outputTy.isDynamicDim(2)) {
- dynDims[2] = tensor::DimOp::create(rewriter, loc, op->getOperand(1), 2);
- }
-
- SmallVector<Value> filteredDims = condenseValues(dynDims);
-
- auto zeroAttr = rewriter.getZeroAttr(outputElementTy);
- Value zero = arith::ConstantOp::create(rewriter, loc, zeroAttr);
- auto emptyTensor =
- tensor::EmptyOp::create(rewriter, loc, outputTy.getShape(),
- outputTy.getElementType(), filteredDims);
- Value zeroTensor = linalg::FillOp::create(rewriter, loc, ValueRange{zero},
- ValueRange{emptyTensor})
- .result();
+ auto aType = dyn_cast<RankedTensorType>(adaptor.getA().getType());
+ auto bType = dyn_cast<RankedTensorType>(adaptor.getB().getType());
+ auto outputType = dyn_cast<RankedTensorType>(op.getType());
+ if (!aType || !bType || !outputType)
+ return rewriter.notifyMatchFailure(op, "only supports ranked tensors");
FailureOr<int64_t> maybeAZp = op.getAZeroPoint();
FailureOr<int64_t> maybeBZp = op.getBZeroPoint();
if (failed(maybeAZp))
return rewriter.notifyMatchFailure(
- op, "input a zero point cannot be statically determined");
+ op, "input A zero point cannot be statically determined");
if (failed(maybeBZp))
return rewriter.notifyMatchFailure(
- op, "input b zero point cannot be statically determined");
+ op, "input B zero point cannot be statically determined");
const int64_t aZpVal = *maybeAZp;
const int64_t bZpVal = *maybeBZp;
if (op.verifyAZeroPoint(aZpVal).failed())
return rewriter.notifyMatchFailure(
- op, "input a zero point must be zero for non-int8 integer types");
+ op, "input A zero point must be zero for non-int8 integer types");
if (op.verifyBZeroPoint(bZpVal).failed())
return rewriter.notifyMatchFailure(
- op, "input b zero point must be zero for non-int8 integer types");
+ op, "input B zero point must be zero for non-int8 integer types");
+
+ const int64_t outputRank = outputType.getRank();
+ auto [a, b, zeroTensor] =
+ prepareMatMulOperands(rewriter, loc, adaptor.getA(), adaptor.getB(),
+ outputType, /*transposeB=*/false);
+
+ Value aZp, bZp;
+ const bool hasNonzeroZeroPoint = aZpVal != 0 || bZpVal != 0;
+ if (hasNonzeroZeroPoint) {
+ aZp = arith::ConstantOp::create(rewriter, loc,
+ rewriter.getI32IntegerAttr(aZpVal));
+ bZp = arith::ConstantOp::create(rewriter, loc,
+ rewriter.getI32IntegerAttr(bZpVal));
+ }
- if (aZpVal == 0 && bZpVal == 0) {
- rewriter.replaceOpWithNewOp<linalg::BatchMatmulOp>(
- op, TypeRange{op.getType()},
- ValueRange{adaptor.getA(), adaptor.getB()}, ValueRange{zeroTensor});
+ if (outputRank == 2) {
+ if (!hasNonzeroZeroPoint) {
+ rewriter.replaceOpWithNewOp<linalg::MatmulOp>(op, TypeRange{outputType},
+ ValueRange{a, b},
+ ValueRange{zeroTensor});
+ } else {
+ rewriter.replaceOpWithNewOp<linalg::QuantizedMatmulOp>(
+ op, TypeRange{outputType}, ValueRange{a, b, aZp, bZp},
+ ValueRange{zeroTensor});
+ }
return success();
}
- auto aZp = arith::ConstantOp::create(rewriter, loc,
- rewriter.getI32IntegerAttr(aZpVal));
- auto bZp = arith::ConstantOp::create(rewriter, loc,
- rewriter.getI32IntegerAttr(bZpVal));
- rewriter.replaceOpWithNewOp<linalg::QuantizedBatchMatmulOp>(
- op, TypeRange{op.getType()},
- ValueRange{adaptor.getA(), adaptor.getB(), aZp, bZp}, zeroTensor);
+ if (outputRank == 3) {
+ if (!hasNonzeroZeroPoint) {
+ rewriter.replaceOpWithNewOp<linalg::BatchMatmulOp>(
+ op, TypeRange{outputType}, ValueRange{a, b},
+ ValueRange{zeroTensor});
+ } else {
+ rewriter.replaceOpWithNewOp<linalg::QuantizedBatchMatmulOp>(
+ op, TypeRange{outputType}, ValueRange{a, b, aZp, bZp},
+ ValueRange{zeroTensor});
+ }
+ return success();
+ }
+
+ SmallVector<ReassociationIndices> reassociation =
+ getMatMulBatchCollapseReassociation(outputRank);
+ auto collapse = [&](Value value) {
+ auto type = cast<RankedTensorType>(value.getType());
+ auto collapsedType =
+ tensor::CollapseShapeOp::inferCollapsedType(type, reassociation);
+ return tensor::CollapseShapeOp::create(rewriter, loc, collapsedType,
+ value, reassociation);
+ };
+
+ Value collapsedA = collapse(a);
+ Value collapsedB = collapse(b);
+ Value collapsedInit = collapse(zeroTensor);
+ auto collapsedOutputType = cast<RankedTensorType>(collapsedInit.getType());
+ Value collapsedResult;
+ if (!hasNonzeroZeroPoint) {
+ collapsedResult =
+ linalg::BatchMatmulOp::create(
+ rewriter, loc, TypeRange{collapsedOutputType},
+ ValueRange{collapsedA, collapsedB}, ValueRange{collapsedInit})
+ .getResult(0);
+ } else {
+ collapsedResult = linalg::QuantizedBatchMatmulOp::create(
+ rewriter, loc, TypeRange{collapsedOutputType},
+ ValueRange{collapsedA, collapsedB, aZp, bZp},
+ ValueRange{collapsedInit})
+ .getResult(0);
+ }
+
+ Value result = tensor::ExpandShapeOp::create(
+ rewriter, loc, outputType, collapsedResult, reassociation,
+ tensor::getMixedSizes(rewriter, loc, zeroTensor));
+ rewriter.replaceOp(op, result);
return success();
}
};
-static SmallVector<ReassociationIndices>
-getMatMulTBatchCollapseReassociation(int64_t rank) {
- SmallVector<ReassociationIndices> reassociation(3);
- for (int64_t dim = 0; dim < rank - 2; ++dim)
- reassociation[0].push_back(dim);
- reassociation[1].push_back(rank - 2);
- reassociation[2].push_back(rank - 1);
- return reassociation;
-}
-
class MatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
public:
using OpConversionPattern<tosa::MatMulTOp>::OpConversionPattern;
@@ -721,8 +756,9 @@ class MatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
op, "nonzero zero points require the generic MATMUL_T lowering");
const int64_t outputRank = outputType.getRank();
- auto [a, b, zeroTensor] = prepareMatMulTOperands(
- rewriter, loc, adaptor.getA(), adaptor.getB(), outputType);
+ auto [a, b, zeroTensor] =
+ prepareMatMulOperands(rewriter, loc, adaptor.getA(), adaptor.getB(),
+ outputType, /*transposeB=*/true);
if (outputRank == 2) {
rewriter.replaceOpWithNewOp<linalg::MatmulTransposeBOp>(
@@ -737,7 +773,7 @@ class MatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
}
SmallVector<ReassociationIndices> reassociation =
- getMatMulTBatchCollapseReassociation(outputRank);
+ getMatMulBatchCollapseReassociation(outputRank);
auto collapse = [&](Value value) {
auto type = cast<RankedTensorType>(value.getType());
auto collapsedType =
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h
index 2bc56c043a57b2..e40b2d0a64216f 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h
@@ -25,9 +25,10 @@ struct MatMulTOperands {
/// Broadcast the batch dimensions of the ranked operands and create a
/// zero-filled result tensor for either MATMUL_T lowering pattern.
-MatMulTOperands prepareMatMulTOperands(PatternRewriter &rewriter, Location loc,
- Value a, Value b,
- RankedTensorType outputType);
+MatMulTOperands prepareMatMulOperands(PatternRewriter &rewriter, Location loc,
+ Value a, Value b,
+ RankedTensorType outputType,
+ bool transposeB);
} // namespace tosa
} // namespace mlir
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
index d9de3f223f5d0a..66be7f2603d624 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named-invalid.mlir
@@ -25,3 +25,39 @@ func.func @avg_pool2d_adaptive_not_canonicalized(
(tensor<1x4x4x1xf32>, tensor<1xf32>, tensor<1xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x2x2x1xf32>
return %0 : tensor<1x2x2x1xf32>
}
+
+// -----
+
+// Unranked operands must fail conversion cleanly before accessing their shape.
+func.func @matmul_unranked_a(
+ %a: tensor<*xf32>, %b: tensor<1x3x6xf32>) -> tensor<1x5x6xf32> {
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{failed to legalize operation 'tosa.matmul'}}
+ %0 = tosa.matmul %a, %b, %a_zp, %b_zp : (tensor<*xf32>, tensor<1x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x5x6xf32>
+ return %0 : tensor<1x5x6xf32>
+}
+
+// -----
+
+// A's zero point must be statically known; keep B's zero point constant.
+func.func @matmul_nonconstant_a_zero_point(
+ %a: tensor<1x5x3xi8>, %b: tensor<1x3x6xi8>,
+ %a_zp: tensor<1xi8>) -> tensor<1x5x6xi32> {
+ %b_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ // expected-error at +1 {{failed to legalize operation 'tosa.matmul'}}
+ %0 = tosa.matmul %a, %b, %a_zp, %b_zp : (tensor<1x5x3xi8>, tensor<1x3x6xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x5x6xi32>
+ return %0 : tensor<1x5x6xi32>
+}
+
+// -----
+
+// B's zero point must be statically known after A's zero point is accepted.
+func.func @matmul_nonconstant_b_zero_point(
+ %a: tensor<1x5x3xi8>, %b: tensor<1x3x6xi8>,
+ %b_zp: tensor<1xi8>) -> tensor<1x5x6xi32> {
+ %a_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ // expected-error at +1 {{failed to legalize operation 'tosa.matmul'}}
+ %0 = tosa.matmul %a, %b, %a_zp, %b_zp : (tensor<1x5x3xi8>, tensor<1x3x6xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x5x6xi32>
+ return %0 : tensor<1x5x6xi32>
+}
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index 4a052337335a7f..22caf9d409eacc 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -4,15 +4,15 @@
// RUN: mlir-opt --verify-each --split-input-file -pass-pipeline="builtin.module(func.func(tosa-to-linalg-named{allow-non-finites=true}))" %s -verify-diagnostics -o -| FileCheck %s --check-prefixes=CHECK,NONFINITE
// CHECK-LABEL: @matmul
-func.func @matmul(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x3x6xf32>) -> (tensor<1x5x6xf32>) {
+func.func @matmul(%arg0: tensor<2x5x3xf32>, %arg1: tensor<2x3x6xf32>) -> (tensor<2x5x6xf32>) {
// CHECK: [[C0:%.+]] = arith.constant 0
// CHECK: [[INIT:%.+]] = tensor.empty()
- // CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : f32) outs([[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32>
- // CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x5x3xf32>, tensor<1x3x6xf32>) outs([[FILLED]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32>
+ // CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : f32) outs([[INIT]] : tensor<2x5x6xf32>) -> tensor<2x5x6xf32>
+ // CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<2x5x3xf32>, tensor<2x3x6xf32>) outs([[FILLED]] : tensor<2x5x6xf32>) -> tensor<2x5x6xf32>
%a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
%b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
- %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<1x5x3xf32>, tensor<1x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x5x6xf32>
- return %0 : tensor<1x5x6xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xf32>, tensor<2x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x5x6xf32>
+ return %0 : tensor<2x5x6xf32>
}
// -----
@@ -34,14 +34,29 @@ func.func @matmul_quantized(%arg0: tensor<1x5x3xi8>, %arg1: tensor<1x3x6xi8>) ->
// -----
+// Normalize both dynamic batch dimensions before allocating the output.
// CHECK-LABEL: @matmul_dyn_batch
func.func @matmul_dyn_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<?x3x6xf32>) -> (tensor<?x5x6xf32>) {
- // CHECK: %[[C0:.+]] = arith.constant 0
- // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
- // CHECK: %[[C0_0:.+]] = arith.constant 0
- // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM]])
- // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0_0]] : f32) outs(%[[INIT]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32>
- // CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<?x5x3xf32>, tensor<?x3x6xf32>) outs(%[[FILLED]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32>
+ // CHECK: %[[A_BATCH:.+]] = tensor.dim %arg0, %{{.+}}
+ // CHECK: %[[B_BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[BATCH:.+]] = arith.maxui %[[A_BATCH]], %[[B_BATCH]] : index
+ // CHECK: %[[NORMALIZED_A:.+]] = scf.if %{{.*}} -> (tensor<?x5x3xf32>) {
+ // CHECK: %[[A_INIT:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x5x3xf32>
+ // CHECK: %[[A_BROADCAST:.+]] = linalg.generic {{.*}} ins(%arg0 : tensor<?x5x3xf32>) outs(%[[A_INIT]] : tensor<?x5x3xf32>)
+ // CHECK: scf.yield %[[A_BROADCAST]] : tensor<?x5x3xf32>
+ // CHECK: scf.yield %arg0 : tensor<?x5x3xf32>
+ // CHECK: %[[NORMALIZED_B:.+]] = scf.if %{{.*}} -> (tensor<?x3x6xf32>) {
+ // CHECK: %[[B_INIT:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x3x6xf32>
+ // CHECK: %[[B_BROADCAST:.+]] = linalg.generic {{.*}} ins(%arg1 : tensor<?x3x6xf32>) outs(%[[B_INIT]] : tensor<?x3x6xf32>)
+ // CHECK: scf.yield %[[B_BROADCAST]] : tensor<?x3x6xf32>
+ // CHECK: scf.yield %arg1 : tensor<?x3x6xf32>
+ // CHECK: %[[OUT_INDEX:.+]] = arith.constant 0 : index
+ // CHECK: %[[OUT_BATCH:.+]] = tensor.dim %[[NORMALIZED_A]], %[[OUT_INDEX]]
+ // CHECK: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f32
+ // CHECK: %[[INIT:.+]] = tensor.empty(%[[OUT_BATCH]]) : tensor<?x5x6xf32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : f32) outs(%[[INIT]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32>
+ // CHECK: %[[RESULT:.+]] = linalg.batch_matmul ins(%[[NORMALIZED_A]], %[[NORMALIZED_B]] : tensor<?x5x3xf32>, tensor<?x3x6xf32>) outs(%[[FILLED]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32>
+ // CHECK: return %[[RESULT]] : tensor<?x5x6xf32>
%a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
%b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
%0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<?x5x3xf32>, tensor<?x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x5x6xf32>
@@ -50,8 +65,8 @@ func.func @matmul_dyn_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<?x3x6xf32>)
// -----
-// CHECK-LABEL: @matmul_dyn_independent_dim
-func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x3x?xf32>) -> (tensor<1x5x?xf32>) {
+// CHECK-LABEL: @matmul_dynamic_columns
+func.func @matmul_dynamic_columns(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x3x?xf32>) -> (tensor<1x5x?xf32>) {
// CHECK: %[[C2:.+]] = arith.constant 2
// CHECK: %[[DIM:.+]] = tensor.dim %arg1, %[[C2]]
// CHECK: %[[C0:.+]] = arith.constant 0
@@ -66,8 +81,9 @@ func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x
// -----
-// CHECK-LABEL: @matmul_dyn_independent_dim
-func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x?xf32>, %arg1: tensor<1x?x6xf32>) -> (tensor<1x5x6xf32>) {
+// Dynamic reduction size should not make the output shape dynamic
+// CHECK-LABEL: @matmul_dynamic_reduction
+func.func @matmul_dynamic_reduction(%arg0: tensor<1x5x?xf32>, %arg1: tensor<1x?x6xf32>) -> (tensor<1x5x6xf32>) {
// CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32>
@@ -80,12 +96,12 @@ func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x?xf32>, %arg1: tensor<1x
// -----
+// Preserve the declared dynamic output type even when its extent is known.
// CHECK-LABEL: @matmul_dyn_output
func.func @matmul_dyn_output(%arg0: tensor<1x1x8xf32>, %arg1: tensor<1x8x1xf32>) -> tensor<?x1x1xf32> {
- // CHECK: %[[C0:.+]] = arith.constant 0 : index
- // CHECK: %[[DIM0:.+]] = tensor.dim %arg0, %[[C0]] : tensor<1x1x8xf32>
+ // CHECK: %[[ONE:.+]] = arith.constant 1 : index
// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32
- // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM0]]) : tensor<?x1x1xf32>
+ // CHECK: %[[INIT:.+]] = tensor.empty(%[[ONE]]) : tensor<?x1x1xf32>
// CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<?x1x1xf32>) -> tensor<?x1x1xf32>
// CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x1x8xf32>, tensor<1x8x1xf32>) outs(%[[FILLED]] : tensor<?x1x1xf32>) -> tensor<?x1x1xf32>
%a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
@@ -96,6 +112,387 @@ func.func @matmul_dyn_output(%arg0: tensor<1x1x8xf32>, %arg1: tensor<1x8x1xf32>)
// -----
+// CHECK-LABEL: @matmul_broadcast
+func.func @matmul_broadcast(%arg0: tensor<4x5x3xf32>, %arg1: tensor<1x3x6xf32>) -> tensor<4x5x6xf32> {
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 3, 6] [1, 1, 1] : tensor<1x3x6xf32> to tensor<3x6xf32>
+ // CHECK: %[[BCAST_INIT:.+]] = tensor.empty() : tensor<4x3x6xf32>
+ // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<3x6xf32>) outs(%[[BCAST_INIT]] : tensor<4x3x6xf32>) dimensions = [0]
+ // CHECK: %[[C0:.+]] = arith.constant 0
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%arg0, %[[BCAST]] : tensor<4x5x3xf32>, tensor<4x3x6xf32>) outs(%[[FILLED]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<4x5x3xf32>, tensor<1x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x5x6xf32>
+ return %0 : tensor<4x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_broadcast_both_rank4
+func.func @matmul_broadcast_both_rank4(%arg0: tensor<2x1x5x3xf32>, %arg1: tensor<1x3x3x6xf32>) -> tensor<2x3x5x6xf32> {
+ // CHECK: %[[B_SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0, 0] [1, 3, 3, 6] [1, 1, 1, 1] : tensor<1x3x3x6xf32> to tensor<3x3x6xf32>
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%[[B_SLICE]] : tensor<3x3x6xf32>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_SLICE:.+]] = tensor.extract_slice %arg0[0, 0, 0, 0] [2, 1, 5, 3] [1, 1, 1, 1] : tensor<2x1x5x3xf32> to tensor<2x5x3xf32>
+ // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%[[A_SLICE]] : tensor<2x5x3xf32>) {{.*}} dimensions = [1]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x5x3xf32> into tensor<6x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x3x6xf32> into tensor<6x3x6xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<6x5x3xf32>, tensor<6x3x6xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [2, 3, 5, 6] : tensor<6x5x6xf32> into tensor<2x3x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<2x1x5x3xf32>, tensor<1x3x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x5x6xf32>
+ return %0 : tensor<2x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_a_missing_batch_dim
+func.func @matmul_a_missing_batch_dim(%arg0: tensor<3x5x3xf32>, %arg1: tensor<4x3x3x6xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%arg0 : tensor<3x5x3xf32>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<4x3x3x6xf32> into tensor<12x3x6xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x3x6xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [4, 3, 5, 6] : tensor<12x5x6xf32> into tensor<4x3x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<3x5x3xf32>, tensor<4x3x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_b_missing_batch_dim
+func.func @matmul_b_missing_batch_dim(%arg0: tensor<4x3x5x3xf32>, %arg1: tensor<3x3x6xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%arg1 : tensor<3x3x6xf32>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x3x6xf32> into tensor<12x3x6xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x3x6xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [4, 3, 5, 6] : tensor<12x5x6xf32> into tensor<4x3x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<4x3x5x3xf32>, tensor<3x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_rank2_a
+func.func @matmul_rank2_a(%arg0: tensor<5x3xf32>, %arg1: tensor<4x3x3x6xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%arg0 : tensor<5x3xf32>) {{.*}} dimensions = [0, 1]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<4x3x3x6xf32> into tensor<12x3x6xf32>
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x3x6xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<5x3xf32>, tensor<4x3x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_rank2_b
+func.func @matmul_rank2_b(%arg0: tensor<4x3x5x3xf32>, %arg1: tensor<3x6xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%arg1 : tensor<3x6xf32>) {{.*}} dimensions = [0, 1]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x3x6xf32> into tensor<12x3x6xf32>
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x3x6xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<4x3x5x3xf32>, tensor<3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_rank2
+func.func @matmul_rank2(%arg0: tensor<5x3xf32>, %arg1: tensor<3x6xf32>) -> tensor<5x6xf32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<5x6xf32>) -> tensor<5x6xf32>
+ // CHECK: linalg.matmul
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<5x3xf32>, tensor<3x6xf32>) outs(%[[FILLED]] : tensor<5x6xf32>) -> tensor<5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<5x3xf32>, tensor<3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5x6xf32>
+ return %0 : tensor<5x6xf32>
+}
+
+// -----
+
+// A dynamic batch dimension may need broadcasting to a known static size.
+// CHECK-LABEL: @matmul_dynamic_static_batch
+func.func @matmul_dynamic_static_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<4x3x6xf32>) -> tensor<4x5x6xf32> {
+ // CHECK: %[[A:.+]] = scf.if %{{.+}} -> (tensor<?x5x3xf32>) {
+ // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x5x3xf32>
+ // CHECK: %[[BCAST:.+]] = linalg.generic {{.*}} ins(%arg0 : tensor<?x5x3xf32>) outs(%[[EMPTY]] : tensor<4x5x3xf32>)
+ // CHECK: %[[CAST:.+]] = tensor.cast %[[BCAST]] : tensor<4x5x3xf32> to tensor<?x5x3xf32>
+ // CHECK: scf.yield %[[CAST]] : tensor<?x5x3xf32>
+ // CHECK: } else {
+ // CHECK: scf.yield %arg0 : tensor<?x5x3xf32>
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A]], %arg1 : tensor<?x5x3xf32>, tensor<4x3x6xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<?x5x3xf32>, tensor<4x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x5x6xf32>
+ return %0 : tensor<4x5x6xf32>
+}
+
+// -----
+
+// The dynamic operand already determines the target size when its partner is 1.
+// CHECK-LABEL: @matmul_dynamic_singleton_batch
+func.func @matmul_dynamic_singleton_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<1x3x6xf32>) -> tensor<?x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %{{.+}}
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 3, 6] [1, 1, 1] : tensor<1x3x6xf32> to tensor<3x6xf32>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x3x6xf32>
+ // CHECK: %[[B:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<3x6xf32>) outs(%[[EMPTY]] : tensor<?x3x6xf32>) dimensions = [0]
+ // CHECK-NOT: scf.if
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%arg0, %[[B]] : tensor<?x5x3xf32>, tensor<?x3x6xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<?x5x3xf32>, tensor<1x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x5x6xf32>
+ return %0 : tensor<?x5x6xf32>
+}
+
+// -----
+
+// The dynamic B operand determines the target size when A's batch size is 1.
+// CHECK-LABEL: @matmul_singleton_dynamic_batch
+func.func @matmul_singleton_dynamic_batch(%arg0: tensor<1x5x3xf32>, %arg1: tensor<?x3x6xf32>) -> tensor<?x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg0[0, 0, 0] [1, 5, 3] [1, 1, 1] : tensor<1x5x3xf32> to tensor<5x3xf32>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x5x3xf32>
+ // CHECK: %[[A:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<5x3xf32>) outs(%[[EMPTY]] : tensor<?x5x3xf32>) dimensions = [0]
+ // CHECK-NOT: scf.if
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A]], %arg1 : tensor<?x5x3xf32>, tensor<?x3x6xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<1x5x3xf32>, tensor<?x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x5x6xf32>
+ return %0 : tensor<?x5x6xf32>
+}
+
+// -----
+
+// Missing leading dimensions broadcast directly to the partner's runtime size.
+// CHECK-LABEL: @matmul_missing_dynamic_batch
+func.func @matmul_missing_dynamic_batch(%arg0: tensor<3x5x3xf32>, %arg1: tensor<?x3x3x6xf32>) -> tensor<?x3x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x3x5x3xf32>
+ // CHECK: %[[A:.+]] = linalg.broadcast ins(%arg0 : tensor<3x5x3xf32>) outs(%[[EMPTY]] : tensor<?x3x5x3xf32>) dimensions = [0]
+ // CHECK-NOT: scf.if
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A]] {{\[}}[0, 1], [2], [3]] : tensor<?x3x5x3xf32> into tensor<?x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<?x3x3x6xf32> into tensor<?x3x6xf32>
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<?x5x3xf32>, tensor<?x3x6xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<3x5x3xf32>, tensor<?x3x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x3x5x6xf32>
+ return %0 : tensor<?x3x5x6xf32>
+}
+
+// -----
+
+// Resolve dynamic batch dimensions one at a time while preserving matrix sizes.
+// CHECK-LABEL: @matmul_multiple_dynamic_batches
+func.func @matmul_multiple_dynamic_batches(%arg0: tensor<?x?x?x3xf32>, %arg1: tensor<?x?x3x6xf32>) -> tensor<?x?x?x6xf32> {
+ // CHECK: arith.maxui
+ // CHECK: %[[A0:.+]] = scf.if %{{.+}} -> (tensor<?x?x?x3xf32>)
+ // CHECK: %[[B0:.+]] = scf.if %{{.+}} -> (tensor<?x?x3x6xf32>)
+ // CHECK: arith.maxui
+ // CHECK: %[[A1:.+]] = scf.if %{{.+}} -> (tensor<?x?x?x3xf32>)
+ // CHECK: linalg.generic {{.*}} ins(%[[A0]] : tensor<?x?x?x3xf32>)
+ // CHECK: %[[B1:.+]] = scf.if %{{.+}} -> (tensor<?x?x3x6xf32>)
+ // CHECK: linalg.generic {{.*}} ins(%[[B0]] : tensor<?x?x3x6xf32>)
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A1]] {{\[}}[0, 1], [2], [3]] : tensor<?x?x?x3xf32> into tensor<?x?x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B1]] {{\[}}[0, 1], [2], [3]] : tensor<?x?x3x6xf32> into tensor<?x3x6xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<?x?x3xf32>, tensor<?x3x6xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [%{{.+}}, %{{.+}}, %{{.+}}, 6] : tensor<?x?x6xf32> into tensor<?x?x?x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<?x?x?x3xf32>, tensor<?x?x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x?x?x6xf32>
+ return %0 : tensor<?x?x?x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_i8_zero_point
+func.func @matmul_i8_zero_point(%arg0: tensor<2x5x3xi8>, %arg1: tensor<2x3x6xi8>) -> tensor<2x5x6xi32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : i32) outs(%[[INIT]] : tensor<2x5x6xi32>) -> tensor<2x5x6xi32>
+ // CHECK: linalg.batch_matmul
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xi8>, tensor<2x3x6xi8>) outs(%[[FILLED]] : tensor<2x5x6xi32>) -> tensor<2x5x6xi32>
+ %a_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xi8>, tensor<2x3x6xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2x5x6xi32>
+ return %0 : tensor<2x5x6xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_zero_point_rank2
+func.func @matmul_zero_point_rank2(%arg0: tensor<5x3xi8>, %arg1: tensor<3x6xi8>) -> tensor<5x6xi32> {
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<5x6xi32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : i32) outs(%[[INIT]] : tensor<5x6xi32>) -> tensor<5x6xi32>
+ // CHECK: %[[AZP:.+]] = arith.constant 1 : i32
+ // CHECK: %[[BZP:.+]] = arith.constant 2 : i32
+ // CHECK: %[[RESULT:.+]] = linalg.quantized_matmul ins(%arg0, %arg1, %[[AZP]], %[[BZP]] : tensor<5x3xi8>, tensor<3x6xi8>, i32, i32) outs(%[[FILLED]] : tensor<5x6xi32>) -> tensor<5x6xi32>
+ // CHECK: return %[[RESULT]] : tensor<5x6xi32>
+ %a_zp = "tosa.const"() <{values = dense<1> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<2> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<5x3xi8>, tensor<3x6xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<5x6xi32>
+ return %0 : tensor<5x6xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_zero_point
+func.func @matmul_zero_point(%arg0: tensor<4x5x3xi8>, %arg1: tensor<1x3x6xi8>) -> tensor<4x5x6xi32> {
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 3, 6] [1, 1, 1] : tensor<1x3x6xi8> to tensor<3x6xi8>
+ // CHECK: %[[B:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<3x6xi8>) {{.*}} dimensions = [0]
+ // CHECK: %[[ZERO:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<4x5x6xi32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : i32) outs(%[[INIT]] : tensor<4x5x6xi32>) -> tensor<4x5x6xi32>
+ // CHECK: %[[AZP:.+]] = arith.constant 1 : i32
+ // CHECK: %[[BZP:.+]] = arith.constant 2 : i32
+ // CHECK: %[[RESULT:.+]] = linalg.quantized_batch_matmul ins(%arg0, %[[B]], %[[AZP]], %[[BZP]] : tensor<4x5x3xi8>, tensor<4x3x6xi8>, i32, i32) outs(%[[FILLED]] : tensor<4x5x6xi32>) -> tensor<4x5x6xi32>
+ // CHECK: return %[[RESULT]] : tensor<4x5x6xi32>
+ %a_zp = "tosa.const"() <{values = dense<1> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<2> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<4x5x3xi8>, tensor<1x3x6xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<4x5x6xi32>
+ return %0 : tensor<4x5x6xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_zero_point_rank4
+func.func @matmul_zero_point_rank4(%arg0: tensor<2x1x5x3xi8>, %arg1: tensor<1x3x3x6xi8>) -> tensor<2x3x5x6xi32> {
+// CHECK: %[[B_SLICE:.+]] = tensor.extract_slice %arg1
+ // CHECK: %[[B:.+]] = linalg.broadcast ins(%[[B_SLICE]] : tensor<3x3x6xi8>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_SLICE:.+]] = tensor.extract_slice %arg0
+ // CHECK: %[[A:.+]] = linalg.broadcast ins(%[[A_SLICE]] : tensor<2x5x3xi8>) {{.*}} dimensions = [1]
+ // CHECK: %[[ZERO:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x3x5x6xi32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : i32) outs(%[[INIT]] : tensor<2x3x5x6xi32>) -> tensor<2x3x5x6xi32>
+ // CHECK: %[[AZP:.+]] = arith.constant 1 : i32
+ // CHECK: %[[BZP:.+]] = arith.constant 2 : i32
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x5x3xi8> into tensor<6x5x3xi8>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x3x6xi8> into tensor<6x3x6xi8>
+ // CHECK: %[[INIT_COLLAPSED:.+]] = tensor.collapse_shape %[[FILLED]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x5x6xi32> into tensor<6x5x6xi32>
+ // CHECK: %[[RESULT:.+]] = linalg.quantized_batch_matmul ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]], %[[AZP]], %[[BZP]] : tensor<6x5x3xi8>, tensor<6x3x6xi8>, i32, i32) outs(%[[INIT_COLLAPSED]] : tensor<6x5x6xi32>) -> tensor<6x5x6xi32>
+ // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1], [2], [3]] output_shape [2, 3, 5, 6] : tensor<6x5x6xi32> into tensor<2x3x5x6xi32>
+ // CHECK: return %[[EXPANDED]] : tensor<2x3x5x6xi32>
+ %a_zp = "tosa.const"() <{values = dense<1> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<2> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<2x1x5x3xi8>, tensor<1x3x3x6xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2x3x5x6xi32>
+ return %0 : tensor<2x3x5x6xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_b_zero_point
+func.func @matmul_b_zero_point(%arg0: tensor<5x3xi8>, %arg1: tensor<3x6xi8>) -> tensor<5x6xi32> {
+// CHECK: %[[ZERO:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<5x6xi32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : i32) outs(%[[INIT]] : tensor<5x6xi32>) -> tensor<5x6xi32>
+ // CHECK: %[[AZP:.+]] = arith.constant 0 : i32
+ // CHECK: %[[BZP:.+]] = arith.constant 2 : i32
+ // CHECK: %[[RESULT:.+]] = linalg.quantized_matmul ins(%arg0, %arg1, %[[AZP]], %[[BZP]] : tensor<5x3xi8>, tensor<3x6xi8>, i32, i32) outs(%[[FILLED]] : tensor<5x6xi32>) -> tensor<5x6xi32>
+ // CHECK: return %[[RESULT]] : tensor<5x6xi32>
+ %a_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<2> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<5x3xi8>, tensor<3x6xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<5x6xi32>
+ return %0 : tensor<5x6xi32>
+}
+
+// -----
+
+// Derive dynamic M from A and dynamic N from the last dimension of B.
+// CHECK-LABEL: @matmul_rank2_dynamic_matrix
+func.func @matmul_rank2_dynamic_matrix(%arg0: tensor<?x3xf32>, %arg1: tensor<3x?xf32>) -> tensor<?x?xf32> {
+// CHECK: %[[M:.+]] = tensor.dim %arg0, %{{.+}} : tensor<?x3xf32>
+ // CHECK: %[[N:.+]] = tensor.dim %arg1, %{{.+}} : tensor<3x?xf32>
+ // CHECK: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f32
+ // CHECK: %[[INIT:.+]] = tensor.empty(%[[M]], %[[N]]) : tensor<?x?xf32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : f32) outs(%[[INIT]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+ // CHECK: %[[RESULT:.+]] = linalg.matmul ins(%arg0, %arg1 : tensor<?x3xf32>, tensor<3x?xf32>) outs(%[[FILLED]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+ // CHECK: return %[[RESULT]] : tensor<?x?xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<?x3xf32>, tensor<3x?xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x?xf32>
+ return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+// Preserve dynamic batch and column sizes through the quantized collapse/expand path.
+// CHECK-LABEL: @matmul_quantized_dynamic_rank4
+func.func @matmul_quantized_dynamic_rank4(%arg0: tensor<2x?x5x3xi8>, %arg1: tensor<1x?x3x?xi8>) -> tensor<2x?x5x?xi32> {
+// CHECK: %[[B_SLICE:.+]] = tensor.extract_slice %arg1
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%[[B_SLICE]] : tensor<?x3x?xi8>) {{.*}} dimensions = [0]
+ // CHECK: %[[BATCH:.+]] = arith.maxui %{{.+}}, %{{.+}} : index
+ // CHECK: %[[A:.+]] = scf.if %{{.+}} -> (tensor<2x?x5x3xi8>) {
+ // CHECK: tensor.empty(%[[BATCH]]) : tensor<2x?x5x3xi8>
+ // CHECK: %[[B:.+]] = scf.if %{{.+}} -> (tensor<2x?x3x?xi8>) {
+ // CHECK: linalg.generic {{.*}} ins(%[[B_BCAST]] : tensor<2x?x3x?xi8>)
+ // CHECK: %[[OUT_BATCH:.+]] = tensor.dim %[[A]], %{{.+}} : tensor<2x?x5x3xi8>
+ // CHECK: %[[N:.+]] = tensor.dim %[[B]], %{{.+}} : tensor<2x?x3x?xi8>
+ // CHECK: %[[ZERO:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty(%[[OUT_BATCH]], %[[N]]) : tensor<2x?x5x?xi32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : i32) outs(%[[INIT]] : tensor<2x?x5x?xi32>) -> tensor<2x?x5x?xi32>
+ // CHECK: %[[AZP:.+]] = arith.constant -3 : i32
+ // CHECK: %[[BZP:.+]] = arith.constant 0 : i32
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A]] {{\[}}[0, 1], [2], [3]] : tensor<2x?x5x3xi8> into tensor<?x5x3xi8>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B]] {{\[}}[0, 1], [2], [3]] : tensor<2x?x3x?xi8> into tensor<?x3x?xi8>
+ // CHECK: %[[INIT_COLLAPSED:.+]] = tensor.collapse_shape %[[FILLED]] {{\[}}[0, 1], [2], [3]] : tensor<2x?x5x?xi32> into tensor<?x5x?xi32>
+ // CHECK: %[[RESULT:.+]] = linalg.quantized_batch_matmul ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]], %[[AZP]], %[[BZP]] : tensor<?x5x3xi8>, tensor<?x3x?xi8>, i32, i32) outs(%[[INIT_COLLAPSED]] : tensor<?x5x?xi32>) -> tensor<?x5x?xi32>
+ // CHECK: %[[EXPAND_BATCH:.+]] = tensor.dim %[[FILLED]], %{{.+}} : tensor<2x?x5x?xi32>
+ // CHECK: %[[EXPAND_N:.+]] = tensor.dim %[[FILLED]], %{{.+}} : tensor<2x?x5x?xi32>
+ // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1], [2], [3]] output_shape [2, %[[EXPAND_BATCH]], 5, %[[EXPAND_N]]] : tensor<?x5x?xi32> into tensor<2x?x5x?xi32>
+ // CHECK: return %[[EXPANDED]] : tensor<2x?x5x?xi32>
+ %a_zp = "tosa.const"() <{values = dense<-3> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<2x?x5x3xi8>, tensor<1x?x3x?xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2x?x5x?xi32>
+ return %0 : tensor<2x?x5x?xi32>
+}
+
+// -----
+
+// Keep the declared dynamic result type when the input batch sizes are static.
+// CHECK-LABEL: @matmul_static_inputs_dynamic_rank4_output
+func.func @matmul_static_inputs_dynamic_rank4_output(%arg0: tensor<2x3x5x3xf32>, %arg1: tensor<2x3x3x6xf32>) -> tensor<?x?x5x6xf32> {
+// CHECK: %[[BATCH0:.+]] = arith.constant 2 : index
+ // CHECK: %[[BATCH1:.+]] = arith.constant 3 : index
+ // CHECK: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f32
+ // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH0]], %[[BATCH1]]) : tensor<?x?x5x6xf32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[ZERO]] : f32) outs(%[[INIT]] : tensor<?x?x5x6xf32>) -> tensor<?x?x5x6xf32>
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3]] : tensor<2x3x5x3xf32> into tensor<6x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<2x3x3x6xf32> into tensor<6x3x6xf32>
+ // CHECK: %[[INIT_COLLAPSED:.+]] = tensor.collapse_shape %[[FILLED]] {{\[}}[0, 1], [2], [3]] : tensor<?x?x5x6xf32> into tensor<?x5x6xf32>
+ // CHECK: %[[RESULT:.+]] = linalg.batch_matmul ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<6x5x3xf32>, tensor<6x3x6xf32>) outs(%[[INIT_COLLAPSED]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32>
+ // CHECK: %[[EXPAND_BATCH0:.+]] = tensor.dim %[[FILLED]], %{{.+}} : tensor<?x?x5x6xf32>
+ // CHECK: %[[EXPAND_BATCH1:.+]] = tensor.dim %[[FILLED]], %{{.+}} : tensor<?x?x5x6xf32>
+ // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1], [2], [3]] output_shape [%[[EXPAND_BATCH0]], %[[EXPAND_BATCH1]], 5, 6] : tensor<?x5x6xf32> into tensor<?x?x5x6xf32>
+ // CHECK: return %[[EXPANDED]] : tensor<?x?x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %a_zp, %b_zp : (tensor<2x3x5x3xf32>, tensor<2x3x3x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x?x5x6xf32>
+ return %0 : tensor<?x?x5x6xf32>
+}
+
+// -----
+
// CHECK-LABEL: @matmul_t
func.func @matmul_t(%arg0: tensor<2x5x3xf32>, %arg1: tensor<2x6x3xf32>) -> tensor<2x5x6xf32> {
// CHECK: %[[C0:.+]] = arith.constant 0
More information about the Mlir-commits
mailing list