[Mlir-commits] [mlir] Feat/tosa variadic matmul sharding (PR #226180)
Deeptanshu Sekhri
llvmlistbot at llvm.org
Thu Sep 24 07:20:11 PDT 2026
https://github.com/deeptanshusekhri created https://github.com/llvm/llvm-project/pull/226180
None
>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 1/2] [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>
}
>From baf30d4cd3736061402e0c55a832fc0d93c78ec4 Mon Sep 17 00:00:00 2001
From: Deeptanshu Sekhri <deeptanshu.sekhri at arm.com>
Date: Wed, 23 Sep 2026 13:31:03 +0100
Subject: [PATCH 2/2] [mlir][tosa] Support variadic matmul sharding
Build MATMUL and MATMUL_T indexing maps from operand and result ranks,
right align batch dimensions, and keep broadcast singleton dimensions
replicated. Decline propagation when dynamic batch dimensions make the
indexing ambiguous.
Signed-off-by: Deeptanshu Sekhri <deeptanshu.sekhri at arm.com>
---
.../Shard/Interfaces/ShardingInterface.cpp | 29 ++++-
.../Dialect/Tosa/IR/ShardingInterfaceImpl.cpp | 105 +++++++++++++---
.../Shard/sharding-propagation-failed.mlir | 11 ++
.../Dialect/Shard/sharding-propagation.mlir | 117 ++++++++++++++++++
4 files changed, 243 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp b/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp
index cbdc274940d1e3..50d2bebe1853ac 100644
--- a/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp
+++ b/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp
@@ -73,6 +73,9 @@ checkOperandAffineExprRecursively(AffineExpr expr,
seenIds[position] = true;
return success();
}
+ case AffineExprKind::Constant:
+ // A constant indexing expression does not reference a loop dimension.
+ return success();
default:
return failure();
}
@@ -304,7 +307,14 @@ shard::detail::defaultGetShardingOption(Operation *op,
if (!shardAttr)
continue;
- anyShardingInResultsOrOperands = !shardAttr.getSplitAxes().empty();
+ // Record the grid before processing mapped dimensions. Constant-indexed
+ // dimensions represent replication and do not call fillShardingOption.
+ FlatSymbolRefAttr grid = shardAttr.getGridAttr();
+ if (shardingOption.grid && grid && shardingOption.grid != grid)
+ return failure();
+ if (grid)
+ shardingOption.grid = grid;
+ anyShardingInResultsOrOperands |= !shardAttr.getSplitAxes().empty();
AffineMap map = maps[shardingIt.index()];
unsigned numDims = map.getNumDims();
@@ -312,17 +322,26 @@ shard::detail::defaultGetShardingOption(Operation *op,
//
// TODO: Change to process the operands with single loop index first and
// then the operands with multiple loop indices.
- for (auto it : llvm::zip(map.getResults(), shardAttr.getSplitAxes())) {
- AffineExpr expr = std::get<0>(it);
- ArrayRef<GridAxis> axes = std::get<1>(it).asArrayRef();
+ for (auto [dim, exprAndAxes] : llvm::enumerate(
+ llvm::zip(map.getResults(), shardAttr.getSplitAxes()))) {
+ AffineExpr expr = std::get<0>(exprAndAxes);
+ ArrayRef<GridAxis> axes = std::get<1>(exprAndAxes).asArrayRef();
FailureOr<llvm::SmallSet<unsigned, 2>> loopIndices =
checkOperandAffineExpr(expr, numDims);
if (failed(loopIndices))
return op->emitOpError()
<< "operand's affine expression is restricted to const_i * "
"dim_i + const_j + dim_j + ...";
- if (loopIndices->empty())
+ // A constant expression selects the same element for every loop
+ // iteration. It can represent replication, but cannot carry split axes.
+ if (loopIndices->empty()) {
+ if (!axes.empty())
+ return op->emitOpError()
+ << "cannot shard operand " << shardingIt.index()
+ << " dimension " << dim + 1
+ << ": its indexing expression is constant";
continue;
+ }
if (loopIndices->size() == 1) {
unsigned loopIdx = *loopIndices->begin();
visitedLoopIndices.insert(loopIdx);
diff --git a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
index 45994a7ec679f1..b0ca3e355dd9b8 100644
--- a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
@@ -23,13 +23,14 @@ using namespace mlir::shard;
namespace {
-// loop types: [parallel, parallel, parallel, reduction_sum]
-// indexing maps:
-// (d0, d1, d2, d3) -> (d0, d1, d3)
-// (d0, d1, d2, d3) -> (d0, d3, d2)
-// (d0, d1, d2, d3) -> (d0, d1, d2)
-struct MatMulOpSharding
- : public ShardingInterface::ExternalModel<MatMulOpSharding, MatMulOp> {
+// For an output of rank R, use R parallel loops followed by one reduction
+// loop. Right align the input batch dimensions with the output batch loops and
+// map A[..., H, C] to [..., d(R-2), d(R)]. MATMUL maps B[..., C, W] to
+// [..., d(R), d(R-1)], while MATMUL_T maps B[..., W, C] to
+// [..., d(R-1), d(R)].
+template <typename OpType, bool TransposeB>
+struct MatMulSharding : public ShardingInterface::ExternalModel<
+ MatMulSharding<OpType, TransposeB>, OpType> {
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
auto tensorType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
if (!tensorType)
@@ -47,20 +48,95 @@ struct MatMulOpSharding
}
SmallVector<AffineMap> getIndexingMaps(Operation *op) const {
- auto tensorType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
- if (!tensorType)
+ auto aType = dyn_cast<RankedTensorType>(op->getOperand(0).getType());
+ auto bType = dyn_cast<RankedTensorType>(op->getOperand(1).getType());
+ auto outputType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
+ if (!aType || !bType || !outputType || aType.getRank() < 2 ||
+ bType.getRank() < 2 || outputType.getRank() < aType.getRank() ||
+ outputType.getRank() < bType.getRank())
return {};
+
MLIRContext *ctx = op->getContext();
+ const int64_t outputRank = outputType.getRank();
+ const int64_t loopRank = outputRank + 1;
+ const int64_t reductionLoop = outputRank;
+
+ auto getAlignedBatchDim = [&](RankedTensorType operandType,
+ int64_t outputDim) {
+ const int64_t operandDim =
+ outputDim - (outputRank - operandType.getRank());
+ return operandDim < 0 ? 1 : operandType.getDimSize(operandDim);
+ };
+ for (int64_t i = 0; i < outputRank - 2; ++i) {
+ const int64_t aDim = getAlignedBatchDim(aType, i);
+ const int64_t bDim = getAlignedBatchDim(bType, i);
+ // A dynamic dimension may be either one (broadcast) or the corresponding
+ // output extent. These cases require different indexing maps unless the
+ // other operand is known to be a singleton.
+ if ((ShapedType::isDynamic(aDim) && bDim != 1) ||
+ (ShapedType::isDynamic(bDim) && aDim != 1))
+ return {};
+ }
+
+ auto getOperandMap = [&](RankedTensorType operandType, bool isB) {
+ const int64_t operandRank = operandType.getRank();
+ const int64_t operandBatchRank = operandRank - 2;
+ const int64_t batchRankDiff = outputRank - operandRank;
+ SmallVector<AffineExpr> results;
+ results.reserve(operandRank);
+ for (int64_t i = 0; i < operandBatchRank; ++i) {
+ // A statically sized batch dimension of one is broadcast and must
+ // remain replicated rather than inherit the corresponding sharding.
+ if (operandType.getDimSize(i) == 1)
+ results.push_back(getAffineConstantExpr(0, ctx));
+ else
+ results.push_back(getAffineDimExpr(batchRankDiff + i, ctx));
+ }
+ if (isB) {
+ if constexpr (TransposeB) {
+ results.push_back(getAffineDimExpr(outputRank - 1, ctx));
+ results.push_back(getAffineDimExpr(reductionLoop, ctx));
+ } else {
+ results.push_back(getAffineDimExpr(reductionLoop, ctx));
+ results.push_back(getAffineDimExpr(outputRank - 1, ctx));
+ }
+ } else {
+ results.push_back(getAffineDimExpr(outputRank - 2, ctx));
+ results.push_back(getAffineDimExpr(reductionLoop, ctx));
+ }
+ return AffineMap::get(loopRank, 0, results, ctx);
+ };
+
+ SmallVector<unsigned> outputTargets;
+ outputTargets.reserve(outputRank);
+ for (int64_t i = 0; i < outputRank; ++i)
+ outputTargets.push_back(i);
+
SmallVector<AffineMap> maps;
- maps.push_back(AffineMap::getMultiDimMapWithTargets(4, {0, 1, 3}, ctx));
- maps.push_back(AffineMap::getMultiDimMapWithTargets(4, {0, 3, 2}, ctx));
- maps.push_back(AffineMap::get(0, 0, {}, ctx));
- maps.push_back(AffineMap::get(0, 0, {}, ctx));
- maps.push_back(AffineMap::getMultiDimMapWithTargets(4, {0, 1, 2}, ctx));
+ maps.push_back(getOperandMap(aType, /*isB=*/false));
+ maps.push_back(getOperandMap(bType, /*isB=*/true));
+ maps.push_back(AffineMap::get(loopRank, 0, {}, ctx));
+ maps.push_back(AffineMap::get(loopRank, 0, {}, ctx));
+ maps.push_back(
+ AffineMap::getMultiDimMapWithTargets(loopRank, outputTargets, ctx));
return maps;
}
+
+ FailureOr<ShardingOption>
+ getShardingOption(Operation *op, ArrayRef<Sharding> operandShardings,
+ ArrayRef<Sharding> resultShardings) const {
+ // Decline propagation when a dynamic batch dimension makes the operand
+ // indexing ambiguous between replication and the corresponding loop.
+ if (getIndexingMaps(op).empty())
+ return ShardingOption::makeEmpty();
+ return shard::detail::defaultGetShardingOption(op, operandShardings,
+ resultShardings);
+ }
};
+using MatMulOpSharding = MatMulSharding<MatMulOp, /*TransposeB=*/false>;
+using MatMulTOpSharding = MatMulSharding<MatMulTOp, /*TransposeB=*/true>;
+
struct NegateOpSharding
: public ShardingInterface::ExternalModel<NegateOpSharding, NegateOp> {
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
@@ -126,6 +202,7 @@ void mlir::tosa::registerShardingInterfaceExternalModels(
GreaterOp, GreaterEqualOp>(ctx);
MatMulOp::attachInterface<MatMulOpSharding>(*ctx);
+ MatMulTOp::attachInterface<MatMulTOpSharding>(*ctx);
NegateOp::attachInterface<NegateOpSharding>(*ctx);
});
}
diff --git a/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir b/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir
index c80bf754a7bb39..4e8af9417a75a7 100644
--- a/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir
+++ b/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir
@@ -1,6 +1,7 @@
// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation))" %s -verify-diagnostics
shard.grid "private" @grid(shape = 1)
+shard.grid "private" @grid_2(shape = 2)
// expected-error @+1 {{'func.func' op only one block is supported!}}
func.func @multi_block_function(%arg0 : tensor<6x6xi32>) -> tensor<6x6xi32> {
%sharding = shard.sharding @grid split_axes = [[0]] : !shard.sharding
@@ -9,3 +10,13 @@ func.func @multi_block_function(%arg0 : tensor<6x6xi32>) -> tensor<6x6xi32> {
^bb1:
return %sharded : tensor<6x6xi32>
}
+
+// -----
+
+func.func @matmul_shard_broadcast_dimension(%a: tensor<1x4x8xf32>, %b: tensor<2x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x16xf32> {
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %a_sharded = shard.shard %a to %sharding annotate_for_users : tensor<1x4x8xf32>
+ // expected-error @+1 {{'tosa.matmul' op cannot shard operand 0 dimension 1: its indexing expression is constant}}
+ %0 = tosa.matmul %a_sharded, %b, %zp, %zp : (tensor<1x4x8xf32>, tensor<2x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x16xf32>
+ return %0 : tensor<2x4x16xf32>
+}
diff --git a/mlir/test/Dialect/Shard/sharding-propagation.mlir b/mlir/test/Dialect/Shard/sharding-propagation.mlir
index 34aaf0598b3f01..e1953526162e21 100644
--- a/mlir/test/Dialect/Shard/sharding-propagation.mlir
+++ b/mlir/test/Dialect/Shard/sharding-propagation.mlir
@@ -5,6 +5,123 @@ shard.grid @grid_1d(shape = ?)
shard.grid @grid_2d(shape = 2x4)
shard.grid @grid_3d(shape = ?x?x?)
+// CHECK-LABEL: func.func @matmul_broadcast_batches
+// CHECK-SAME: %[[A:.*]]: tensor<2x1x4x8xf32>, %[[B:.*]]: tensor<1x4x8x16xf32>
+func.func @matmul_broadcast_batches(%a: tensor<2x1x4x8xf32>, %b: tensor<1x4x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}0]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}], [1]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<2x1x4x8xf32>, tensor<1x4x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2d split_axes = [[0], [1], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_missing_batch
+// CHECK-SAME: %[[A:.*]]: tensor<4x4x8xf32>, %[[B:.*]]: tensor<2x4x8x16xf32>
+func.func @matmul_missing_batch(%a: tensor<4x4x8xf32>, %b: tensor<2x4x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}0]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<4x4x8xf32>, tensor<2x4x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0], [], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_t_rank2
+// CHECK-SAME: %[[A:.*]]: tensor<4x8xf32>, %[[B:.*]]: tensor<16x8xf32>
+func.func @matmul_t_rank2(%a: tensor<4x8xf32>, %b: tensor<16x8xf32>, %zp: tensor<1xf32>) -> tensor<4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}0]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}1]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul_t %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul_t %a, %b, %zp, %zp : (tensor<4x8xf32>, tensor<16x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x16xf32>
+ %sharding = shard.sharding @grid_2d split_axes = [[0], [1]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<4x16xf32>
+ return %sharded : tensor<4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_t_broadcast_batches
+// CHECK-SAME: %[[A:.*]]: tensor<2x1x4x8xf32>, %[[B:.*]]: tensor<1x4x16x8xf32>
+func.func @matmul_t_broadcast_batches(%a: tensor<2x1x4x8xf32>, %b: tensor<1x4x16x8xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}0]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}], [1]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul_t %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul_t %a, %b, %zp, %zp : (tensor<2x1x4x8xf32>, tensor<1x4x16x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2d split_axes = [[0], [1], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_t_missing_batch
+// CHECK-SAME: %[[A:.*]]: tensor<4x4x8xf32>, %[[B:.*]]: tensor<2x4x16x8xf32>
+func.func @matmul_t_missing_batch(%a: tensor<4x4x8xf32>, %b: tensor<2x4x16x8xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}0]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul_t %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul_t %a, %b, %zp, %zp : (tensor<4x4x8xf32>, tensor<2x4x16x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0], [], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_dynamic_lhs_batch
+// CHECK-NOT: annotate_for_users
+// CHECK: tosa.matmul
+func.func @matmul_dynamic_lhs_batch(%a: tensor<?x4x8xf32>, %b: tensor<2x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x16xf32> {
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<?x4x8xf32>, tensor<2x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x16xf32>
+ return %sharded : tensor<2x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_dynamic_rhs_batch
+// CHECK-NOT: annotate_for_users
+// CHECK: tosa.matmul
+func.func @matmul_dynamic_rhs_batch(%a: tensor<2x4x8xf32>, %b: tensor<?x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x16xf32> {
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<2x4x8xf32>, tensor<?x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x16xf32>
+ return %sharded : tensor<2x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_dynamic_and_singleton_batch
+// CHECK-SAME: %[[A:.*]]: tensor<?x4x8xf32>, %[[B:.*]]: tensor<1x8x16xf32>
+// CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}0]]
+// CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+// CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+// CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+// CHECK: tosa.matmul %[[A_SHARDED]], %[[B_SHARDED]]
+func.func @matmul_dynamic_and_singleton_batch(%a: tensor<?x4x8xf32>, %b: tensor<1x8x16xf32>, %zp: tensor<1xf32>) -> tensor<?x4x16xf32> {
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<?x4x8xf32>, tensor<1x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<?x4x16xf32>
+ return %sharded : tensor<?x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_operand_replication
+// CHECK-SAME: %[[A:.*]]: tensor<1x4x8xf32>, %[[B:.*]]: tensor<1x8x16xf32>
+// CHECK: %[[REPLICATION:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+// CHECK: %[[A_REPLICATED:.*]] = shard.shard %[[A]] to %[[REPLICATION]] annotate_for_users
+// CHECK: %[[B_REPLICATED:.*]] = shard.shard %[[B]] to %[[REPLICATION]] annotate_for_users
+// CHECK: tosa.matmul %[[A_REPLICATED]], %[[B_REPLICATED]]
+func.func @matmul_operand_replication(%a: tensor<1x4x8xf32>, %b: tensor<1x8x16xf32>, %zp: tensor<1xf32>) -> tensor<1x4x16xf32> {
+ %replication = shard.sharding @grid_2 split_axes = [[]] : !shard.sharding
+ %a_replicated = shard.shard %a to %replication annotate_for_users : tensor<1x4x8xf32>
+ %0 = tosa.matmul %a_replicated, %b, %zp, %zp : (tensor<1x4x8xf32>, tensor<1x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x4x16xf32>
+ return %0 : tensor<1x4x16xf32>
+}
+
// CHECK-LABEL: func.func @element_wise_empty_sharding_info
func.func @element_wise_empty_sharding_info(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
// CHECK-NEXT: tosa.sigmoid
More information about the Mlir-commits
mailing list