[Mlir-commits] [mlir] [mlir][tosa] Support variadic batch dimensions in matmul ops (PR #226179)
Deeptanshu Sekhri
llvmlistbot at llvm.org
Thu Sep 24 07:19:55 PDT 2026
https://github.com/deeptanshusekhri created https://github.com/llvm/llvm-project/pull/226179
Allow tosa.matmul and tosa.matmul_t operands and results to have rank two or greater. Infer and verify right aligned broadcasting across their leading batch dimensions.
Preserve TOSA 1.0 compatibility through target validation and restrict the 1.1-to-1.0 downgrade to supported rank 3 cases. Extend level, shape inference, verifier, validation, and downgrade coverage.
>From fc1bc85e29044aeed9d8cb777f6f8d664b734556 Mon Sep 17 00:00:00 2001
From: Deeptanshu Sekhri <deeptanshu.sekhri at arm.com>
Date: Tue, 22 Sep 2026 10:34:00 +0100
Subject: [PATCH] [mlir][tosa] Support variadic batch dimensions in matmul ops
Allow tosa.matmul and tosa.matmul_t operands and results to have rank
two or greater. Infer and verify right aligned broadcasting across their
leading batch dimensions.
Preserve TOSA 1.0 compatibility through target validation and restrict
the 1.1-to-1.0 downgrade to supported rank 3 cases. Extend level,
shape inference, verifier, validation, and downgrade coverage.
Signed-off-by: Deeptanshu Sekhri <deeptanshu.sekhri at arm.com>
---
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 32 ++-
.../mlir/Dialect/Tosa/IR/TosaTypesBase.td | 8 +
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp | 253 ++++++++++--------
.../Tosa/Transforms/TosaDowngrade1p1To1p0.cpp | 11 +
.../Tosa/Transforms/TosaValidation.cpp | 41 ++-
mlir/test/Dialect/Tosa/error_if_check.mlir | 20 ++
mlir/test/Dialect/Tosa/level_check.mlir | 36 +++
.../Tosa/tosa-downgrade-1-1-to-1-0.mlir | 44 +++
mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir | 88 ++++++
.../tosa-validation-version-1p0-invalid.mlir | 30 +++
.../tosa-validation-version-1p0-valid.mlir | 10 +
.../tosa-validation-version-1p1-valid.mlir | 40 +++
mlir/test/Dialect/Tosa/verifier.mlir | 81 +++++-
13 files changed, 566 insertions(+), 128 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 7cbe8e0b6c70e7..e55727ffdd2dd0 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -416,21 +416,28 @@ def Tosa_MatMulOp : Tosa_InferShapedTypeOp<"matmul", [NoMemoryEffect]> {
let summary = "Matrix multiplication operator.";
let description = [{
- Performs two dimensional matrix multiplications.
+ Performs batched matrix multiplication of `A` with shape `[..., H, C]` and
+ `B` with shape `[..., C, W]`. Their leading batch dimensions are aligned
+ from the right. Missing leading dimensions are treated as size 1, and each
+ aligned pair must either match or include a dimension of size 1. The output
+ shape is the resulting broadcast batch shape followed by `[H, W]`.
+
+ Both inputs must have rank 2 or greater. Rank 1 inputs are not accepted
+ directly and must first be reshaped to rank 2.
This operation is not pure. Undefined behaviour may occur if the accumulated
result overflows.
}];
let arguments = (ins
- Tosa_Tensor3D:$a,
- Tosa_Tensor3D:$b,
+ Tosa_TensorAtLeast2D:$a,
+ Tosa_TensorAtLeast2D:$b,
Tosa_ScalarIntOrFloatTensor:$a_zp,
Tosa_ScalarIntOrFloatTensor:$b_zp
);
let results = (outs
- Tosa_Tensor3D:$output
+ Tosa_TensorAtLeast2D:$output
);
let extraClassDeclaration = [{
@@ -454,16 +461,21 @@ def Tosa_MatMulTOp : Tosa_InferShapedTypeOp<"matmul_t", [NoMemoryEffect]> {
let summary = "Matrix multiplication with transposed B operator.";
let description = [{
- Performs two dimensional matrix multiplications. `A` matrix is of shape
- `N x H x C`. `B` matrix is of shape `D x W x C`. This is effectively a
- matrix multiply of `A` by the transposed `B` matrix. If the batched
- dimension of input `B` is of size 1, the `B` matrix is broadcast.
+ Performs batched matrix multiplication of `A` with shape `[..., H, C]` and
+ `B` with shape `[..., W, C]`, implicitly transposing only the two innermost
+ dimensions of `B`. Their leading batch dimensions are aligned from the
+ right. Missing leading dimensions are treated as size 1, and each aligned
+ pair must either match or include a dimension of size 1. The output shape
+ is the resulting broadcast batch shape followed by `[H, W]`.
+
+ Both inputs must have rank 2 or greater. Rank 1 inputs are not accepted
+ directly and must first be reshaped to rank 2.
}];
- let arguments = (ins Tosa_Tensor3D:$a, Tosa_Tensor3D:$b,
+ let arguments = (ins Tosa_TensorAtLeast2D:$a, Tosa_TensorAtLeast2D:$b,
Tosa_ScalarIntOrFloatTensor:$a_zp, Tosa_ScalarIntOrFloatTensor:$b_zp);
- let results = (outs Tosa_Tensor3D:$output);
+ let results = (outs Tosa_TensorAtLeast2D:$output);
let extraClassDeclaration = [{
FailureOr<int64_t> getAZeroPoint();
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
index b5701853aaf683..4d6727ed76c564 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
@@ -171,6 +171,10 @@ def AtLeastRankOne : And<[
IsRankedTensorTypePred,
CPred<"::llvm::cast<::mlir::RankedTensorType>($_self).getRank() >= 1">]>;
+def AtLeastRankTwo : And<[
+ IsRankedTensorTypePred,
+ CPred<"::llvm::cast<::mlir::RankedTensorType>($_self).getRank() >= 2">]>;
+
def IsValidBlockScaledTensorType
: CPred<"::mlir::succeeded(::mlir::tosa::"
"verifyBlockScaledTensorType($_self))">;
@@ -257,6 +261,10 @@ def Tosa_TensorAtLeast1D
: TosaTensorOf<[Tosa_AnyNumber], [AtLeastRankOne],
"tosa-conformant tensor of at least rank 1">;
+def Tosa_TensorAtLeast2D
+ : TosaTensorOf<[Tosa_AnyNumber], [AtLeastRankTwo],
+ "tosa-conformant tensor of at least rank 2">;
+
def Tosa_MXFPDataTensor3D
: TosaTensorRankOf<[Tosa_MXFPValue], [3]>;
def Tosa_MXFPScaleTensor3D
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index feac95c5a021a5..4cdd0b097e7b6a 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -1899,32 +1899,86 @@ bool tosa::EqualOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
return succeeded(verifyCompatibleShape(l[0], r[0]));
}
-LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
- MLIRContext *context, ::std::optional<Location> location,
- MatMulOp::Adaptor adaptor,
- SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
- ShapeAdaptor lhsShape(adaptor.getA().getType());
- ShapeAdaptor rhsShape(adaptor.getB().getType());
-
- // All shapes are dynamic.
- SmallVector<int64_t> outShape;
- outShape.resize(3, ShapedType::kDynamic);
+// MATMUL batch shapes are right aligned and may have different ranks. Missing
+// leading dimensions are treated as one, while an unranked input remains
+// unknown because it may contain any number of batch dimensions.
+static int64_t getMatMulBatchDim(const ShapeAdaptor &shape, int64_t outputRank,
+ int64_t axis) {
+ if (!shape.hasRank())
+ return ShapedType::kDynamic;
+ const int64_t inputAxis = axis - (outputRank - shape.getRank());
+ return inputAxis < 0 ? 1 : shape.getDimSize(inputAxis);
+}
+
+static FailureOr<SmallVector<int64_t>>
+resolveMatMulOutputShape(const ShapeAdaptor &aShape, const ShapeAdaptor &bShape,
+ int64_t outputRank, bool transposeB) {
+ if (outputRank < 2 ||
+ (aShape.hasRank() &&
+ (aShape.getRank() < 2 || aShape.getRank() > outputRank)) ||
+ (bShape.hasRank() &&
+ (bShape.getRank() < 2 || bShape.getRank() > outputRank)))
+ return failure();
- if (lhsShape.hasRank()) {
- outShape[0] = lhsShape.getDimSize(0);
- outShape[1] = lhsShape.getDimSize(1);
+ SmallVector<int64_t> outputShape(outputRank, ShapedType::kDynamic);
+ for (int64_t axis = 0; axis < outputRank - 2; ++axis) {
+ const int64_t aDim = getMatMulBatchDim(aShape, outputRank, axis);
+ const int64_t bDim = getMatMulBatchDim(bShape, outputRank, axis);
+ FailureOr<int64_t> resolvedDim = resolveBroadcastDim(aDim, bDim);
+ if (failed(resolvedDim))
+ return failure();
+ outputShape[axis] = *resolvedDim;
}
- if (rhsShape.hasRank()) {
- outShape[0] = outShape[0] == ShapedType::kDynamic ? rhsShape.getDimSize(0)
- : outShape[0];
- outShape[2] = rhsShape.getDimSize(2);
+ if (aShape.hasRank())
+ outputShape[outputRank - 2] = aShape.getDimSize(aShape.getRank() - 2);
+ if (bShape.hasRank())
+ outputShape[outputRank - 1] =
+ bShape.getDimSize(bShape.getRank() - (transposeB ? 2 : 1));
+ return outputShape;
+}
+
+static LogicalResult inferMatMulReturnTypeComponents(
+ const ShapeAdaptor &aShape, const ShapeAdaptor &bShape, bool transposeB,
+ SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+ if ((aShape.hasRank() && aShape.getRank() < 2) ||
+ (bShape.hasRank() && bShape.getRank() < 2))
+ return failure();
+
+ // An unranked input can have an arbitrary batch prefix, so the output rank
+ // cannot yet be inferred.
+ if (!aShape.hasRank() || !bShape.hasRank()) {
+ inferredReturnShapes.emplace_back();
+ return success();
}
- inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
+ const int64_t aChannels = aShape.getDimSize(aShape.getRank() - 1);
+ const int64_t bChannels =
+ bShape.getDimSize(bShape.getRank() - (transposeB ? 1 : 2));
+ if (ShapedType::isStatic(aChannels) && ShapedType::isStatic(bChannels) &&
+ aChannels != bChannels)
+ return failure();
+
+ const int64_t outputRank = std::max(aShape.getRank(), bShape.getRank());
+ FailureOr<SmallVector<int64_t>> outputShape =
+ resolveMatMulOutputShape(aShape, bShape, outputRank, transposeB);
+ if (failed(outputShape))
+ return failure();
+
+ inferredReturnShapes.emplace_back(*outputShape);
return success();
}
+LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
+ MLIRContext *context, ::std::optional<Location> location,
+ MatMulOp::Adaptor adaptor,
+ SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+ return inferMatMulReturnTypeComponents(ShapeAdaptor(adaptor.getA().getType()),
+ ShapeAdaptor(adaptor.getB().getType()),
+ /*transposeB=*/false,
+ inferredReturnShapes);
+}
+
template <typename T>
static LogicalResult verifyMatMulQuantizedOperandsType(T op, Type aElementType,
Type bElementType) {
@@ -1978,11 +2032,68 @@ static LogicalResult verifyMatMulZeroPointType(T op, Value input, Value zp,
return diag;
}
+static SmallVector<int64_t> getMatMulBatchShape(const ShapeAdaptor &shape) {
+ SmallVector<int64_t> batchShape;
+ if (!shape.hasRank() || shape.getRank() < 2)
+ return batchShape;
+ batchShape.reserve(shape.getRank() - 2);
+ for (int64_t i = 0, e = shape.getRank() - 2; i < e; ++i)
+ batchShape.push_back(shape.getDimSize(i));
+ return batchShape;
+}
+
+template <typename T>
+static LogicalResult verifyMatMulShapes(T op, bool transposeB) {
+ const ShapeAdaptor aShape(op.getA().getType());
+ const ShapeAdaptor bShape(op.getB().getType());
+ const auto outputType = cast<ShapedType>(op.getResult().getType());
+
+ int64_t channels = aShape.hasRank() ? aShape.getDimSize(aShape.getRank() - 1)
+ : ShapedType::kDynamic;
+ if (bShape.hasRank() &&
+ failed(tryUpdateDimOrFailure(
+ op, channels,
+ bShape.getDimSize(bShape.getRank() - (transposeB ? 1 : 2)), "b",
+ "channels")))
+ return failure();
+
+ const int64_t minimumOutputRank =
+ std::max(aShape.hasRank() ? aShape.getRank() : 2,
+ bShape.hasRank() ? bShape.getRank() : 2);
+ if (outputType.hasRank() && outputType.getRank() < minimumOutputRank)
+ return op.emitOpError("expected output rank of at least ")
+ << minimumOutputRank << ", got " << outputType.getRank();
+
+ const bool bothInputsRanked = aShape.hasRank() && bShape.hasRank();
+ if (!bothInputsRanked && !outputType.hasRank())
+ return success();
+
+ // When an input is unranked, use the declared output rank to validate every
+ // result dimension constrained by the ranked input.
+ const int64_t expectedOutputRank =
+ bothInputsRanked ? minimumOutputRank : outputType.getRank();
+ FailureOr<SmallVector<int64_t>> expectedOutputShape =
+ resolveMatMulOutputShape(aShape, bShape, expectedOutputRank, transposeB);
+ if (failed(expectedOutputShape)) {
+ InFlightDiagnostic diag = op.emitOpError(
+ "expected batch dimensions of a and b to be broadcast compatible, "
+ "got a=[");
+ printShapeToDiagnostic(diag, getMatMulBatchShape(aShape));
+ diag << "] and b=[";
+ printShapeToDiagnostic(diag, getMatMulBatchShape(bShape));
+ diag << "]";
+ return diag;
+ }
+
+ if (outputType.hasRank())
+ return verifyOutputShapeCompatibleWithExpected(
+ op.getOperation(), outputType, *expectedOutputShape);
+ return success();
+}
+
LogicalResult MatMulOp::verify() {
- const ShapeAdaptor aShape(getA().getType());
- const ShapeAdaptor bShape(getB().getType());
- const Type aElementType = aShape.getElementType();
- const Type bElementType = bShape.getElementType();
+ const Type aElementType = getElementTypeOrSelf(getA());
+ const Type bElementType = getElementTypeOrSelf(getB());
if (failed(
verifyMatMulQuantizedOperandsType(*this, aElementType, bElementType)))
@@ -2000,67 +2111,22 @@ LogicalResult MatMulOp::verify() {
if (succeeded(maybeBZp) && verifyBZeroPoint(*maybeBZp).failed())
return failure();
- // Verify input/output shapes
- int64_t N = ShapedType::kDynamic;
- int64_t H = ShapedType::kDynamic;
- int64_t W = ShapedType::kDynamic;
- int64_t C = ShapedType::kDynamic;
-
- if (aShape.hasRank()) {
- N = aShape.getDimSize(0);
- H = aShape.getDimSize(1);
- C = aShape.getDimSize(2);
- }
-
- if (bShape.hasRank()) {
- if (failed(tryUpdateDimOrFailure(*this, N, bShape.getDimSize(0), "b",
- "batch")) ||
- failed(tryUpdateDimOrFailure(*this, C, bShape.getDimSize(1), "b",
- "channels")))
- return failure();
- W = bShape.getDimSize(2);
- }
-
- const SmallVector<int64_t, 3> expectedOutputShape = {N, H, W};
- const auto outputType = cast<ShapedType>(getResult().getType());
- if (outputType.hasRank() &&
- failed(verifyOutputShapeCompatibleWithExpected(getOperation(), outputType,
- expectedOutputShape)))
- return failure();
-
- return success();
+ return verifyMatMulShapes(*this, /*transposeB=*/false);
}
LogicalResult tosa::MatMulTOp::inferReturnTypeComponents(
MLIRContext *context, ::std::optional<Location> location,
MatMulTOp::Adaptor adaptor,
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
- const ShapeAdaptor lhsShape(adaptor.getA().getType());
- const ShapeAdaptor rhsShape(adaptor.getB().getType());
-
- SmallVector<int64_t, 3> outShape(3, ShapedType::kDynamic);
-
- if (lhsShape.hasRank()) {
- outShape[0] = lhsShape.getDimSize(0);
- outShape[1] = lhsShape.getDimSize(1);
- }
-
- if (rhsShape.hasRank()) {
- const int64_t bBatchSize = rhsShape.getDimSize(0);
- if (bBatchSize != 1 && ShapedType::isDynamic(outShape[0]))
- outShape[0] = bBatchSize;
- outShape[2] = rhsShape.getDimSize(1);
- }
-
- inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
- return success();
+ return inferMatMulReturnTypeComponents(ShapeAdaptor(adaptor.getA().getType()),
+ ShapeAdaptor(adaptor.getB().getType()),
+ /*transposeB=*/true,
+ inferredReturnShapes);
}
LogicalResult MatMulTOp::verify() {
- const ShapeAdaptor aShape(getA().getType());
- const ShapeAdaptor bShape(getB().getType());
- const Type aElementType = aShape.getElementType();
- const Type bElementType = bShape.getElementType();
+ const Type aElementType = getElementTypeOrSelf(getA());
+ const Type bElementType = getElementTypeOrSelf(getB());
if (failed(
verifyMatMulQuantizedOperandsType(*this, aElementType, bElementType)))
@@ -2078,44 +2144,7 @@ LogicalResult MatMulTOp::verify() {
if (succeeded(maybeBZp) && verifyBZeroPoint(*maybeBZp).failed())
return failure();
- // Verify input/output shapes
- int64_t N = ShapedType::kDynamic;
- int64_t D = ShapedType::kDynamic;
- int64_t H = ShapedType::kDynamic;
- int64_t W = ShapedType::kDynamic;
- int64_t C = ShapedType::kDynamic;
-
- if (aShape.hasRank()) {
- N = aShape.getDimSize(0);
- H = aShape.getDimSize(1);
- C = aShape.getDimSize(2);
- }
-
- if (bShape.hasRank()) {
- D = bShape.getDimSize(0);
- W = bShape.getDimSize(1);
- if (failed(tryUpdateDimOrFailure(*this, C, bShape.getDimSize(2), "b",
- "channels")))
- return failure();
- }
-
- // Verify B batch size is broadcast compatible with A.
- if (ShapedType::isStatic(N) && ShapedType::isStatic(D) && N != D && D != 1)
- return emitOpError("expect B matrix batch size to be broadcast compatible "
- "with A, got D=")
- << D << " vs N=" << N;
-
- if (ShapedType::isDynamic(N) && ShapedType::isStatic(D) && D != 1)
- N = D;
-
- const SmallVector<int64_t, 3> expectedOutputShape = {N, H, W};
- const auto outputType = cast<ShapedType>(getResult().getType());
- if (outputType.hasRank() &&
- failed(verifyOutputShapeCompatibleWithExpected(getOperation(), outputType,
- expectedOutputShape)))
- return failure();
-
- return success();
+ return verifyMatMulShapes(*this, /*transposeB=*/true);
}
LogicalResult tosa::MatmulTBlockScaledOp::inferReturnTypeComponents(
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
index 1d3d7a344341e9..2a3c7aedc0311f 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
@@ -188,6 +188,11 @@ class MatMulTRewrite : public OpRewritePattern<tosa::MatMulTOp> {
if (!aShape.hasRank() || !bShape.hasRank())
return rewriter.notifyMatchFailure(op, "expected ranked A and B tensors");
+ // TOSA 1.0 MATMUL requires rank 3 inputs. MATMUL_T operations with other
+ // ranks require separate rank and batch broadcast legalization.
+ if (aShape.getRank() != 3 || bShape.getRank() != 3)
+ return rewriter.notifyMatchFailure(op, "expected rank 3 A and B tensors");
+
const int64_t dSize = bShape.getDimSize(0);
const int64_t nSize = aShape.getDimSize(0);
@@ -199,6 +204,12 @@ class MatMulTRewrite : public OpRewritePattern<tosa::MatMulTOp> {
return rewriter.notifyMatchFailure(
op, "expected known batch size for broadcast");
+ // TOSA 1.0 MATMUL does not broadcast batch dimensions. The lowering below
+ // handles broadcasting B by tiling it, but does not implement tiling A.
+ if (dSize != 1 && (nSize == 1 || ShapedType::isDynamic(nSize)))
+ return rewriter.notifyMatchFailure(
+ op, "cannot downgrade A batch broadcasting");
+
const int64_t wSize = bShape.getDimSize(1);
const int64_t cSize = bShape.getDimSize(2);
const Location loc = op.getLoc();
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 9a37e57af0d2a9..6cd840128cae0e 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -174,6 +174,39 @@ static LogicalResult checkConstantOperandSilceShape(Operation *op,
return success();
}
+// MATMUL's data type availability predates variadic batch support, so the
+// generated availability checks cannot distinguish its 1.0 and 1.1 shapes.
+static LogicalResult
+checkMatMulShapeVersionCompatibility(Operation *op, const TargetEnv &env) {
+ auto matmul = dyn_cast<tosa::MatMulOp>(op);
+ if (!matmul ||
+ env.getSpecVersion().isBackwardsCompatibleWith(
+ TosaSpecificationVersion(SpecificationVersion::V_1_1_DRAFT)))
+ return success();
+
+ auto aType = dyn_cast<RankedTensorType>(matmul.getA().getType());
+ auto bType = dyn_cast<RankedTensorType>(matmul.getB().getType());
+ auto outputType = dyn_cast<RankedTensorType>(matmul.getOutput().getType());
+ const bool hasNonRankThreeTensor = (aType && aType.getRank() != 3) ||
+ (bType && bType.getRank() != 3) ||
+ (outputType && outputType.getRank() != 3);
+ auto knownBatchSizesDiffer = [](RankedTensorType lhs, RankedTensorType rhs) {
+ return lhs && rhs && lhs.getRank() == 3 && rhs.getRank() == 3 &&
+ !lhs.isDynamicDim(0) && !rhs.isDynamicDim(0) &&
+ lhs.getDimSize(0) != rhs.getDimSize(0);
+ };
+ const bool hasIncompatibleKnownBatchSizes =
+ knownBatchSizesDiffer(aType, bType) ||
+ knownBatchSizesDiffer(aType, outputType) ||
+ knownBatchSizesDiffer(bType, outputType);
+ if (!hasNonRankThreeTensor && !hasIncompatibleKnownBatchSizes)
+ return success();
+
+ return op->emitOpError(
+ "MATMUL ranks other than 3 or batch broadcasting require TOSA "
+ "specification version 1.1.draft");
+}
+
//===----------------------------------------------------------------------===//
// TOSA Validation Pass.
//===----------------------------------------------------------------------===//
@@ -859,8 +892,8 @@ LogicalResult TosaValidation::levelCheckRanksAndSizes(Operation *op) {
CHECK_SIZES(DepthwiseConv2D);
CHECK_SIZES(TransposeConv2D);
CHECK_SIZES(FFT2d);
- CHECK_SIZES(MatMul);
- CHECK_SIZES(MatMulT);
+ CHECK_RANKS_AND_SIZES(MatMul);
+ CHECK_RANKS_AND_SIZES(MatMulT);
CHECK_SIZES(MatmulTBlockScaled);
CHECK_SIZES(MaxPool2d);
CHECK_SIZES(MaxPool2dAdaptive);
@@ -1640,6 +1673,10 @@ void TosaValidation::runOnOperation() {
failed(profileComp.checkExtension(op, targetEnv)))
return signalPassFailure();
+ if (strictOpSpecAlignment &&
+ failed(checkMatMulShapeVersionCompatibility(op, targetEnv)))
+ return signalPassFailure();
+
if (!allowInvalidOpDatatypeCombinations &&
failed(profileComp.checkInvalid(op)))
return signalPassFailure();
diff --git a/mlir/test/Dialect/Tosa/error_if_check.mlir b/mlir/test/Dialect/Tosa/error_if_check.mlir
index 1e5eedd07e6215..b87f124f64ac27 100644
--- a/mlir/test/Dialect/Tosa/error_if_check.mlir
+++ b/mlir/test/Dialect/Tosa/error_if_check.mlir
@@ -336,3 +336,23 @@ func.func @test_while_loop_body_not_isolated_from_above(%arg0: tensor<i32>, %arg
}) : (tensor<i32>) -> (tensor<i32>)
return
}
+
+// -----
+
+func.func @test_matmul_rank3_lhs_broadcast_dynamic_rhs(%arg0: tensor<1x4x7xf32>, %arg1: tensor<?x7x5xf32>) -> tensor<2x4x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op MATMUL ranks other than 3 or batch broadcasting require TOSA specification version 1.1.draft}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<1x4x7xf32>, tensor<?x7x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x5xf32>
+ return %0 : tensor<2x4x5xf32>
+}
+
+// -----
+
+func.func @test_matmul_rank3_rhs_broadcast_dynamic_lhs(%arg0: tensor<?x4x7xf32>, %arg1: tensor<1x7x5xf32>) -> tensor<2x4x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op MATMUL ranks other than 3 or batch broadcasting require TOSA specification version 1.1.draft}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<?x4x7xf32>, tensor<1x7x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x5xf32>
+ return %0 : tensor<2x4x5xf32>
+}
diff --git a/mlir/test/Dialect/Tosa/level_check.mlir b/mlir/test/Dialect/Tosa/level_check.mlir
index f97b2fe2831523..5c2514057ea2d5 100644
--- a/mlir/test/Dialect/Tosa/level_check.mlir
+++ b/mlir/test/Dialect/Tosa/level_check.mlir
@@ -1269,6 +1269,24 @@ func.func @test_rfft2d_tensor_size_invalid(%arg0: tensor<536870912x8x16xf32>) ->
// -----
+func.func @test_matmul_rank_invalid(%arg0: tensor<1x1x1x1x1x3x4xf32>, %arg1: tensor<1x1x1x1x1x4x5xf32>) -> tensor<1x1x1x1x1x3x5xf32> {
+ %zero = "tosa.const"() {values = dense<0.0> : tensor<1xf32>} : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op failed level check: operand rank(shape) <= MAX_RANK}}
+ %0 = tosa.matmul %arg0, %arg1, %zero, %zero : (tensor<1x1x1x1x1x3x4xf32>, tensor<1x1x1x1x1x4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x1x1x1x1x3x5xf32>
+ return %0 : tensor<1x1x1x1x1x3x5xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_rank_invalid(%arg0: tensor<1x1x1x1x1x3x4xf32>, %arg1: tensor<1x1x1x1x1x5x4xf32>) -> tensor<1x1x1x1x1x3x5xf32> {
+ %zero = "tosa.const"() {values = dense<0.0> : tensor<1xf32>} : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul_t' op failed level check: operand rank(shape) <= MAX_RANK}}
+ %0 = tosa.matmul_t %arg0, %arg1, %zero, %zero : (tensor<1x1x1x1x1x3x4xf32>, tensor<1x1x1x1x1x5x4xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x1x1x1x1x3x5xf32>
+ return %0 : tensor<1x1x1x1x1x3x5xf32>
+}
+
+// -----
+
func.func @test_matmul_tensor_size_invalid(%arg0: tensor<23178x20000x19xf32>, %arg1: tensor<23178x19x28xf32>) -> tensor<23178x20000x28xf32> {
%zero = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
// expected-error at +1 {{'tosa.matmul' op failed level check: operand tensor size (in bytes) <= (1 << MAX_LOG2_SIZE - 1)}}
@@ -1287,6 +1305,24 @@ func.func @test_matmul_t_tensor_size_invalid(%arg0: tensor<23178x20000x19xf32>,
// -----
+func.func @test_matmul_broadcast_output_size_invalid(%arg0: tensor<32768x1x1x1xf32>, %arg1: tensor<1x32768x1x1xf32>) {
+ %zero = "tosa.const"() {values = dense<0.0> : tensor<1xf32>} : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op failed level check: result tensor size (in bytes) <= (1 << MAX_LOG2_SIZE - 1)}}
+ %0 = tosa.matmul %arg0, %arg1, %zero, %zero : (tensor<32768x1x1x1xf32>, tensor<1x32768x1x1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<32768x32768x1x1xf32>
+ return
+}
+
+// -----
+
+func.func @test_matmul_t_broadcast_output_size_invalid(%arg0: tensor<32768x1x1x1xf32>, %arg1: tensor<1x32768x1x1xf32>) {
+ %zero = "tosa.const"() {values = dense<0.0> : tensor<1xf32>} : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul_t' op failed level check: result tensor size (in bytes) <= (1 << MAX_LOG2_SIZE - 1)}}
+ %0 = tosa.matmul_t %arg0, %arg1, %zero, %zero : (tensor<32768x1x1x1xf32>, tensor<1x32768x1x1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<32768x32768x1x1xf32>
+ return
+}
+
+// -----
+
func.func @test_gather_tensor_size_invalid(%arg0: tensor<536870912x21x3xf32>, %arg1: tensor<536870912x26xi32>) -> tensor<536870912x26x3xf32> {
// expected-error at +1 {{'tosa.gather' op failed level check: operand tensor size (in bytes) <= (1 << MAX_LOG2_SIZE - 1)}}
%0 = tosa.gather %arg0, %arg1 : (tensor<536870912x21x3xf32>, tensor<536870912x26xi32>) -> tensor<536870912x26x3xf32>
diff --git a/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir b/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
index 402f8801f35c59..85bdd40d82ee95 100644
--- a/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
@@ -160,3 +160,47 @@ func.func @test_preserve_matmul_t_block_scaled(%arg0: tensor<1x14x32x!tosa.block
%0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<1x14x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>, tensor<1x28x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f8E4M3FN>>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
return %0 : tensor<1x14x28xf32>
}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_matmul_t_rank2
+// CHECK-NOT: tosa.transpose
+// CHECK: %[[MATMUL_T:.+]] = tosa.matmul_t %arg0, %arg1, %arg2, %arg3
+// CHECK: return %[[MATMUL_T]]
+func.func @test_preserve_matmul_t_rank2(%arg0: tensor<4x7xf32>, %arg1: tensor<5x7xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<4x5xf32> {
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<4x7xf32>, tensor<5x7xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x5xf32>
+ return %0 : tensor<4x5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_matmul_t_rank4
+// CHECK-NOT: tosa.transpose
+// CHECK: %[[MATMUL_T:.+]] = tosa.matmul_t %arg0, %arg1, %arg2, %arg3
+// CHECK: return %[[MATMUL_T]]
+func.func @test_preserve_matmul_t_rank4(%arg0: tensor<2x1x4x7xf32>, %arg1: tensor<1x3x5x7xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x3x4x5xf32> {
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<2x1x4x7xf32>, tensor<1x3x5x7xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x4x5xf32>
+ return %0 : tensor<2x3x4x5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_matmul_t_broadcast_a
+// CHECK-NOT: tosa.transpose
+// CHECK: %[[MATMUL_T:.+]] = tosa.matmul_t %arg0, %arg1, %arg2, %arg3
+// CHECK: return %[[MATMUL_T]]
+func.func @test_preserve_matmul_t_broadcast_a(%arg0: tensor<1x4x7xf32>, %arg1: tensor<3x5x7xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<3x4x5xf32> {
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<1x4x7xf32>, tensor<3x5x7xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3x4x5xf32>
+ return %0 : tensor<3x4x5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_matmul_t_dynamic_broadcast_a
+// CHECK-NOT: tosa.transpose
+// CHECK: %[[MATMUL_T:.+]] = tosa.matmul_t %arg0, %arg1, %arg2, %arg3
+// CHECK: return %[[MATMUL_T]]
+func.func @test_preserve_matmul_t_dynamic_broadcast_a(%arg0: tensor<?x4x7xf32>, %arg1: tensor<3x5x7xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<3x4x5xf32> {
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<?x4x7xf32>, tensor<3x5x7xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3x4x5xf32>
+ return %0 : tensor<3x4x5xf32>
+}
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index 3cd1b964a8a1c4..4ded720650cf08 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -318,6 +318,61 @@ func.func @test_static_matmul(%arg0 : tensor<2x3x4xi32>, %arg1 : tensor<2x4x5xi3
// -----
+// CHECK-LABEL: @test_rank2_matmul
+func.func @test_rank2_matmul(%arg0 : tensor<3x4xi32>, %arg1 : tensor<4x5xi32>) -> () {
+ // CHECK: tosa.matmul %arg0, %arg1, %0, %1 : (tensor<3x4xi32>, tensor<4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3x5xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul %arg0, %arg1, %0, %1 : (tensor<3x4xi32>, tensor<4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_rank6_broadcast_matmul
+func.func @test_rank6_broadcast_matmul(%arg0 : tensor<2x1x3x1x4x7xi32>, %arg1 : tensor<1x5x1x6x7x5xi32>) -> () {
+ // CHECK: tosa.matmul %arg0, %arg1, %0, %1 : (tensor<2x1x3x1x4x7xi32>, tensor<1x5x1x6x7x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x5x3x6x4x5xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul %arg0, %arg1, %0, %1 : (tensor<2x1x3x1x4x7xi32>, tensor<1x5x1x6x7x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_dynamic_batch_broadcast_matmul
+func.func @test_dynamic_batch_broadcast_matmul(%arg0 : tensor<?x3x?x4xi32>, %arg1 : tensor<2x1x4x5xi32>) -> () {
+ // CHECK: tosa.matmul %arg0, %arg1, %0, %1 : (tensor<?x3x?x4xi32>, tensor<2x1x4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x3x?x5xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul %arg0, %arg1, %0, %1 : (tensor<?x3x?x4xi32>, tensor<2x1x4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_missing_lhs_batch_matmul
+func.func @test_missing_lhs_batch_matmul(%arg0 : tensor<3x4x7xi32>, %arg1 : tensor<4x3x7x5xi32>) -> () {
+ // CHECK: tosa.matmul %arg0, %arg1, %0, %1 : (tensor<3x4x7xi32>, tensor<4x3x7x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4x3x4x5xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul %arg0, %arg1, %0, %1 : (tensor<3x4x7xi32>, tensor<4x3x7x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_unit_lhs_dynamic_batch_matmul
+func.func @test_unit_lhs_dynamic_batch_matmul(%arg0 : tensor<1x4x7xi32>, %arg1 : tensor<?x7x5xi32>) -> () {
+ // CHECK: tosa.matmul %arg0, %arg1, %0, %1 : (tensor<1x4x7xi32>, tensor<?x7x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?x4x5xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul %arg0, %arg1, %0, %1 : (tensor<1x4x7xi32>, tensor<?x7x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @test_dynamic_lhs_matmul
func.func @test_dynamic_lhs_matmul(%arg0 : tensor<?x?x?xi32>, %arg1 : tensor<2x4x5xi32>) -> () {
// CHECK: tosa.matmul %arg0, %arg1, %0, %1 : (tensor<?x?x?xi32>, tensor<2x4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x?x5xi32>
@@ -354,6 +409,17 @@ func.func @test_dynamic_mixed_matmul(%arg0 : tensor<?x3x?xi32>, %arg1 : tensor<?
// -----
+// CHECK-LABEL: @test_unranked_input_matmul
+func.func @test_unranked_input_matmul(%arg0 : tensor<*xi32>, %arg1 : tensor<2x4x5xi32>) -> () {
+ // CHECK: tosa.matmul %arg0, %arg1, %0, %1 : (tensor<*xi32>, tensor<2x4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul %arg0, %arg1, %0, %1 : (tensor<*xi32>, tensor<2x4x5xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @test_unranked_zero_points_matmul
func.func @test_unranked_zero_points_matmul(%arg0: tensor<1x2x3xf32>, %arg1: tensor<1x3x4xf32>, %zero_point: tensor<1xf32>) -> tensor<1x2x4xf32> {
// CHECK: %[[ZP:.*]] = tosa.cast %arg2 input_unsigned(false) : (tensor<1xf32>) -> tensor<1xf32>
@@ -376,6 +442,28 @@ func.func @test_static_matmul_t(%arg0 : tensor<2x3x4xi32>, %arg1 : tensor<2x5x4x
// -----
+// CHECK-LABEL: @test_rank2_matmul_t
+func.func @test_rank2_matmul_t(%arg0 : tensor<3x4xi32>, %arg1 : tensor<5x4xi32>) -> () {
+ // CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<3x4xi32>, tensor<5x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3x5xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<3x4xi32>, tensor<5x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_rank6_different_rank_broadcast_matmul_t
+func.func @test_rank6_different_rank_broadcast_matmul_t(%arg0 : tensor<2x1x3x1x4x7xi32>, %arg1 : tensor<5x3x1x6x7xi32>) -> () {
+ // CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<2x1x3x1x4x7xi32>, tensor<5x3x1x6x7xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x5x3x1x4x6xi32>
+ %0 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %1 = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %2 = tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<2x1x3x1x4x7xi32>, tensor<5x3x1x6x7xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @test_dynamic_lhs_matmul_t
func.func @test_dynamic_lhs_matmul_t(%arg0 : tensor<?x?x?xi32>, %arg1 : tensor<2x5x4xi32>) -> () {
// CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<?x?x?xi32>, tensor<2x5x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x?x5xi32>
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
index 5236c568a55805..f132b036e4aeeb 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -2,6 +2,36 @@
// -----
+func.func @test_matmul_rank2(%arg0: tensor<3x4xf32>, %arg1: tensor<4x5xf32>) -> tensor<3x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op MATMUL ranks other than 3 or batch broadcasting require TOSA specification version 1.1.draft}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<3x4xf32>, tensor<4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3x5xf32>
+ return %0 : tensor<3x5xf32>
+}
+
+// -----
+
+func.func @test_matmul_rank4(%arg0: tensor<2x3x3x4xf32>, %arg1: tensor<2x3x4x6xf32>) -> tensor<2x3x3x6xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op MATMUL ranks other than 3 or batch broadcasting require TOSA specification version 1.1.draft}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x3x4xf32>, tensor<2x3x4x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x3x6xf32>
+ return %0 : tensor<2x3x3x6xf32>
+}
+
+// -----
+
+func.func @test_matmul_rank3_batch_broadcast(%arg0: tensor<2x3x4xf32>, %arg1: tensor<1x4x5xf32>) -> tensor<2x3x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op MATMUL ranks other than 3 or batch broadcasting require TOSA specification version 1.1.draft}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4xf32>, tensor<1x4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x5xf32>
+ return %0 : tensor<2x3x5xf32>
+}
+
+// -----
+
func.func @test_matmul_fp8_mixed_precision_operands(%arg0: tensor<1x14x19xf8E4M3FN>, %arg1: tensor<1x19x28xf8E5M2>) -> tensor<1x14x28xf16> {
%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E4M3FN>}> : () -> tensor<1xf8E4M3FN>
%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E5M2>}> : () -> tensor<1xf8E5M2>
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-valid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-valid.mlir
index 562f1c87e14e5a..21d956450824e5 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-valid.mlir
@@ -1,5 +1,15 @@
// RUN: mlir-opt %s -split-input-file -verify-diagnostics -tosa-attach-target="specification_version=1.0 profiles=pro_int,pro_fp" -tosa-validate="strict-op-spec-alignment" | FileCheck %s
+// CHECK-LABEL: test_matmul_rank3
+func.func @test_matmul_rank3(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x4x5xf32>) -> tensor<2x3x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4xf32>, tensor<2x4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x5xf32>
+ return %0 : tensor<2x3x5xf32>
+}
+
+// -----
+
// CHECK-LABEL: test_gather_i8_i32
func.func @test_gather_i8_i32(%input: tensor<13x27x3xi8>, %indices: tensor<13x26xi32>) -> tensor<13x26x3xi8> {
%gather = tosa.gather %input, %indices : (tensor<13x27x3xi8>, tensor<13x26xi32>) -> tensor<13x26x3xi8>
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
index 5adecaf0ea6156..3e7ac0879deb16 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -2,6 +2,46 @@
// -----
+// CHECK-LABEL: test_matmul_rank2
+func.func @test_matmul_rank2(%arg0: tensor<3x4xf32>, %arg1: tensor<4x5xf32>) -> tensor<3x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<3x4xf32>, tensor<4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3x5xf32>
+ return %0 : tensor<3x5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_matmul_rank4
+func.func @test_matmul_rank4(%arg0: tensor<2x3x3x4xf32>, %arg1: tensor<2x3x4x6xf32>) -> tensor<2x3x3x6xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x3x4xf32>, tensor<2x3x4x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x3x6xf32>
+ return %0 : tensor<2x3x3x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_matmul_rank3_batch_broadcast
+func.func @test_matmul_rank3_batch_broadcast(%arg0: tensor<2x3x4xf32>, %arg1: tensor<1x4x5xf32>) -> tensor<2x3x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4xf32>, tensor<1x4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x5xf32>
+ return %0 : tensor<2x3x5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_matmul_t_rank4
+func.func @test_matmul_t_rank4(%arg0: tensor<2x1x3x4xf32>, %arg1: tensor<1x5x6x4xf32>) -> tensor<2x5x3x6xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x1x3x4xf32>, tensor<1x5x6x4xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x5x3x6xf32>
+ return %0 : tensor<2x5x3x6xf32>
+}
+
+// -----
+
// CHECK-LABEL: test_matmul_fp8_mixed_precision_operands
func.func @test_matmul_fp8_mixed_precision_operands(%arg0: tensor<1x14x19xf8E4M3FN>, %arg1: tensor<1x19x28xf8E5M2>) -> tensor<1x14x28xf16> {
%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E4M3FN>}> : () -> tensor<1xf8E4M3FN>
diff --git a/mlir/test/Dialect/Tosa/verifier.mlir b/mlir/test/Dialect/Tosa/verifier.mlir
index 5229ccffce78ea..b8b619a966a482 100644
--- a/mlir/test/Dialect/Tosa/verifier.mlir
+++ b/mlir/test/Dialect/Tosa/verifier.mlir
@@ -1342,10 +1342,55 @@ func.func @scatter_invalid_K_W(%arg0 : tensor<2x4x5xi32>, %arg1 : tensor<2x6xi32
// -----
+func.func @test_matmul_a_rank_too_low(%arg0: tensor<4xf32>, %arg1: tensor<4x5xf32>) -> tensor<1x5xf32> {
+ %zero = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op operand #0 must be tosa-conformant tensor of at least rank 2 of number values, but got 'tensor<4xf32>'}}
+ %0 = tosa.matmul %arg0, %arg1, %zero, %zero : (tensor<4xf32>, tensor<4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x5xf32>
+ return %0 : tensor<1x5xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_b_rank_too_low(%arg0: tensor<3x4xf32>, %arg1: tensor<4xf32>) -> tensor<3x1xf32> {
+ %zero = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul_t' op operand #1 must be tosa-conformant tensor of at least rank 2 of number values, but got 'tensor<4xf32>'}}
+ %0 = tosa.matmul_t %arg0, %arg1, %zero, %zero : (tensor<3x4xf32>, tensor<4xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3x1xf32>
+ return %0 : tensor<3x1xf32>
+}
+
+// -----
+
+func.func @test_matmul_output_rank_too_low(%arg0: tensor<3x4xf32>, %arg1: tensor<4x5xf32>) -> tensor<5xf32> {
+ %zero = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op result #0 must be tosa-conformant tensor of at least rank 2 of number values, but got 'tensor<5xf32>'}}
+ %0 = tosa.matmul %arg0, %arg1, %zero, %zero : (tensor<3x4xf32>, tensor<4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5xf32>
+ return %0 : tensor<5xf32>
+}
+
+// -----
+
+func.func @test_matmul_batch_mismatch_rank4(%arg0: tensor<2x3x4x7xf32>, %arg1: tensor<5x4x7x6xf32>) -> tensor<2x3x4x6xf32> {
+ %zero = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op expected batch dimensions of a and b to be broadcast compatible, got a=[2, 3] and b=[5, 4]}}
+ %0 = tosa.matmul %arg0, %arg1, %zero, %zero : (tensor<2x3x4x7xf32>, tensor<5x4x7x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x4x6xf32>
+ return %0 : tensor<2x3x4x6xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_output_shape_mismatch_rank5(%arg0: tensor<2x1x3x4x7xf32>, %arg1: tensor<5x3x6x7xf32>) -> tensor<2x5x3x4x8xf32> {
+ %zero = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul_t' op expected output shape 2, 5, 3, 4, 8 to be compatible with inferred shape 2, 5, 3, 4, 6}}
+ %0 = tosa.matmul_t %arg0, %arg1, %zero, %zero : (tensor<2x1x3x4x7xf32>, tensor<5x3x6x7xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x5x3x4x8xf32>
+ return %0 : tensor<2x5x3x4x8xf32>
+}
+
+// -----
+
func.func @test_matmul_output_batch_mismatch(%arg0: tensor<2x3x4xf32>, %arg1: tensor<5x4x6xf32>) -> tensor<2x3x6xf32> {
%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
- // expected-error at +1 {{'tosa.matmul' op expected batch of b to match size 2, got 5}}
+ // expected-error at +1 {{'tosa.matmul' op expected batch dimensions of a and b to be broadcast compatible, got a=[2] and b=[5]}}
%0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4xf32>, tensor<5x4x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x6xf32>
return %0 : tensor<2x3x6xf32>
}
@@ -1376,7 +1421,7 @@ func.func @test_matmul_output_shape_mismatch(%arg0: tensor<2x3x4xf32>, %arg1: te
func.func @test_matmul_dynamic_batch_mismatch(%arg0: tensor<2x?x4xf32>, %arg1: tensor<5x4x6xf32>) -> tensor<2x?x6xf32> {
%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
- // expected-error at +1 {{'tosa.matmul' op expected batch of b to match size 2, got 5}}
+ // expected-error at +1 {{'tosa.matmul' op expected batch dimensions of a and b to be broadcast compatible, got a=[2] and b=[5]}}
%0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x?x4xf32>, tensor<5x4x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x?x6xf32>
return %0 : tensor<2x?x6xf32>
}
@@ -1403,7 +1448,6 @@ func.func @test_matmul_dynamic_output_shape_mismatch(%arg0: tensor<?x3x4xf32>, %
// -----
-
func.func @test_matmul_unranked_b_output_shape_mismatch(%arg0: tensor<2x3x4xf32>, %arg1: tensor<*xf32>) -> tensor<2x5x?xf32> {
%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
@@ -1414,6 +1458,35 @@ func.func @test_matmul_unranked_b_output_shape_mismatch(%arg0: tensor<2x3x4xf32>
// -----
+func.func @test_matmul_unranked_b_output_batch_mismatch(%arg0: tensor<2x3x4x7xf32>, %arg1: tensor<*xf32>) -> tensor<5x3x4x?xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op expected output shape 5, 3, 4, ? to be compatible with inferred shape 2, 3, 4, ?}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4x7xf32>, tensor<*xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5x3x4x?xf32>
+ return %0 : tensor<5x3x4x?xf32>
+}
+
+// -----
+
+func.func @test_matmul_unranked_b_output_rank_too_low(%arg0: tensor<2x3x4x7xf32>, %arg1: tensor<*xf32>) -> tensor<3x4x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op expected output rank of at least 4, got 3}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4x7xf32>, tensor<*xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<3x4x5xf32>
+ return %0 : tensor<3x4x5xf32>
+}
+
+// -----
+
+func.func @test_matmul_output_rank_too_high(%arg0: tensor<3x4xf32>, %arg1: tensor<4x5xf32>) -> tensor<1x3x5xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op expected output shape 1, 3, 5 to be compatible with inferred shape 3, 5}}
+ %0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<3x4xf32>, tensor<4x5xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x3x5xf32>
+ return %0 : tensor<1x3x5xf32>
+}
+
+// -----
func.func @test_matmul_quantized_mixed_operands(%arg0: tensor<2x3x4x!quant.uniform<i8:f32, 0.125>>, %arg1: tensor<2x4x6xf32>) -> tensor<2x3x6xi32> {
%azp0 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
@@ -1512,7 +1585,7 @@ func.func @test_matmul_t_channel_mismatch(%arg0: tensor<1x14x19xf32>, %arg1: ten
func.func @test_matmul_t_batch_mismatch(%arg0: tensor<4x14x19xf32>, %arg1: tensor<2x28x19xf32>) -> tensor<4x14x28xf32> {
%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
-// expected-error at +1 {{'tosa.matmul_t' op expect B matrix batch size to be broadcast compatible with A, got D=2 vs N=4}}
+// expected-error at +1 {{'tosa.matmul_t' op expected batch dimensions of a and b to be broadcast compatible, got a=[4] and b=[2]}}
%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<4x14x19xf32>, tensor<2x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
return %0 : tensor<4x14x28xf32>
}
More information about the Mlir-commits
mailing list