[Mlir-commits] [mlir] 455f8ba - [mlir][tosa] Support variadic batch dimensions in matmul ops (#226179)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 2 02:12:12 PDT 2026
Author: Deeptanshu Sekhri
Date: 2026-10-02T10:12:04+01:00
New Revision: 455f8ba10a9f1d2e6cede13b8c08bdfcd2de6fc5
URL: https://github.com/llvm/llvm-project/commit/455f8ba10a9f1d2e6cede13b8c08bdfcd2de6fc5
DIFF: https://github.com/llvm/llvm-project/commit/455f8ba10a9f1d2e6cede13b8c08bdfcd2de6fc5.diff
LOG: [mlir][tosa] Support variadic batch dimensions in matmul ops (#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.
---------
Signed-off-by: Deeptanshu Sekhri <deeptanshu.sekhri at arm.com>
Added:
Modified:
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
mlir/test/Dialect/Tosa/error_if_check.mlir
mlir/test/Dialect/Tosa/level_check.mlir
mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
mlir/test/Dialect/Tosa/tosa-validation-version-1p0-valid.mlir
mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
mlir/test/Dialect/Tosa/verifier.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 01ef662cc4f71..5f34856859bfa 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -416,21 +416,31 @@ def Tosa_MatMulOp : Tosa_InferShapedTypeOp<"matmul", [NoMemoryEffect]> {
let summary = "Matrix multiplication operator.";
let description = [{
- Performs two dimensional matrix multiplications.
+ Performs matrix multiplications over the two innermost dimensions while
+ treating any leading dimensions as broadcastable batch dimensions. `A` has
+ shape `[A0, ..., A(n-1), H, C]`. `B` has shape
+ `[B0, ..., B(m-1), C, W]`. `[A0, ..., A(n-1)]` and
+ `[B0, ..., B(m-1)]` must be compatible as leading batch dimensions.
+
+ Right-align these two dimension lists and treat any missing leading
+ dimensions as size 1, producing common batch dimensions
+ `[D0, ..., D(k-1)]` where `k = max(n, m)` and each `Di` is the larger of
+ the two aligned dimension sizes. `output` has shape
+ `[D0, ..., D(k-1), H, W]`.
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 +464,28 @@ 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 matrix multiplications over the two innermost dimensions while
+ implicitly transposing the last two dimensions of `B`. Any leading
+ dimensions are broadcastable batch dimensions. `A` has shape
+ `[A0, ..., A(n-1), H, C]`. `B` has shape
+ `[B0, ..., B(m-1), W, C]` before swapping those last two dimensions to
+ `[B0, ..., B(m-1), C, W]`. `[A0, ..., A(n-1)]` and
+ `[B0, ..., B(m-1)]` must be compatible as leading batch dimensions.
+
+ Right-align these two dimension lists and treat any missing leading
+ dimensions as size 1, producing common batch dimensions
+ `[D0, ..., D(k-1)]` where `k = max(n, m)` and each `Di` is the larger of
+ the two aligned dimension sizes. `output` has shape
+ `[D0, ..., D(k-1), H, W]`.
+
+ This operation is not pure. Undefined behaviour may occur if the accumulated
+ result overflows.
}];
- 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 b5701853aaf68..4d6727ed76c56 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 feac95c5a021a..4cdd0b097e7b6 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
diff erent 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 1d3d7a344341e..2a3c7aedc0311 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 9a37e57af0d2a..cb82e87a5cfd0 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -174,6 +174,34 @@ 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 checkSpecificationVersionConstraint(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());
+ auto getBatchDimOrDynamic = [](RankedTensorType type) {
+ return type ? type.getDimSize(0) : ShapedType::kDynamic;
+ };
+ if ((!aType || aType.getRank() == 3) && (!bType || bType.getRank() == 3) &&
+ (!outputType || outputType.getRank() == 3) &&
+ succeeded(verifyCompatibleDims({getBatchDimOrDynamic(aType),
+ getBatchDimOrDynamic(bType),
+ getBatchDimOrDynamic(outputType)})))
+ 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 +887,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 +1668,10 @@ void TosaValidation::runOnOperation() {
failed(profileComp.checkExtension(op, targetEnv)))
return signalPassFailure();
+ if (strictOpSpecAlignment &&
+ failed(checkSpecificationVersionConstraint(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 1e5eedd07e621..b87f124f64ac2 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 f97b2fe283152..5c2514057ea2d 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 402f8801f35c5..85bdd40d82ee9 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 38488160d8a45..1c64c97708a6d 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_
diff erent_rank_broadcast_matmul_t
+func.func @test_rank6_
diff erent_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 681025deaebb8..3f020494d7c22 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 562f1c87e14e5..21d956450824e 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 db25c3b1e685b..c84496018dfc6 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 5c37e4eda1da4..79c5bb1af1c80 100644
--- a/mlir/test/Dialect/Tosa/verifier.mlir
+++ b/mlir/test/Dialect/Tosa/verifier.mlir
@@ -1342,10 +1342,64 @@ 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_channel_mismatch_rank4(%arg0: tensor<2x3x4x7xf32>, %arg1: tensor<2x3x8x6xf32>) -> tensor<2x3x4x6xf32> {
+ %zero = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul' op expected channels of b to match size 7, got 8}}
+ %0 = tosa.matmul %arg0, %arg1, %zero, %zero : (tensor<2x3x4x7xf32>, tensor<2x3x8x6xf32>, 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 +1430,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 +1457,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 +1467,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 +1594,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