[Mlir-commits] [mlir] 0c6a0c4 - [mlir][tosa] Add support for matmul_t (#203894)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 23 14:02:28 PDT 2026
Author: Luke Hutton
Date: 2026-06-23T22:02:23+01:00
New Revision: 0c6a0c4ccf234c5d9969294b454cb615277ede14
URL: https://github.com/llvm/llvm-project/commit/0c6a0c4ccf234c5d9969294b454cb615277ede14
DIFF: https://github.com/llvm/llvm-project/commit/0c6a0c4ccf234c5d9969294b454cb615277ede14.diff
LOG: [mlir][tosa] Add support for matmul_t (#203894)
Adds support for matmul_t added to the specification in
https://github.com/arm/tosa-specification/pull/50.
This includes:
- Operator definition
- Verification logic for the operator
- Output shape inference for the operator
- Validation checks to ensure compliance with the TOSA specification
including profile compliance and level checks.
- Support for matmul_t to matmul in the downgrade pass.
Note that MXFP support will be added in a follow-up commit.
Added:
Modified:
mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h
mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
mlir/test/Dialect/Tosa/dynamic_extension.mlir
mlir/test/Dialect/Tosa/invalid_extension.mlir
mlir/test/Dialect/Tosa/level_check.mlir
mlir/test/Dialect/Tosa/ops.mlir
mlir/test/Dialect/Tosa/profile_all_unsupported.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-1p1-valid.mlir
mlir/test/Dialect/Tosa/verifier.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 7db696cc935ea..45d56416fdba0 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
@@ -60,6 +60,13 @@ profileComplianceMap = {
{{{fp16T, fp16T, fp16T, fp16T, fp16T}, SpecificationVersion::V_1_0},
{{fp16T, fp16T, fp16T, fp16T, fp32T}, SpecificationVersion::V_1_0},
{{fp32T, fp32T, fp32T, fp32T, fp32T}, SpecificationVersion::V_1_0}}}}},
+ {"tosa.matmul_t",
+ {{{Profile::pro_int},
+ {{{i8T, i8T, i32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Profile::pro_fp},
+ {{{fp16T, fp16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, fp16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, fp32T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}}}},
{"tosa.max_pool2d",
{{{Profile::pro_int}, {{{i8T, i8T}, SpecificationVersion::V_1_0}}},
{{Profile::pro_fp},
@@ -687,6 +694,54 @@ extensionComplianceMap = {
allOf},
{{Extension::bf16},
{{{bf16T, bf16T, bf16T, bf16T, fp32T}, SpecificationVersion::V_1_0}}}}},
+ {"tosa.matmul_t",
+ {{{Extension::int16},
+ {{{i16T, i16T, i48T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e4m3},
+ {{{fp8e4m3T, fp8e4m3T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e4m3T, fp16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, fp8e4m3T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e4m3T, fp8e4m3T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e4m3T, fp16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e4m3T, fp32T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, fp8e4m3T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, fp8e4m3T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::fp8e4m3, Extension::fp8e5m2},
+ {{{fp8e4m3T, fp8e5m2T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e5m2T, fp8e4m3T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e4m3T, fp8e5m2T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e5m2T, fp8e4m3T, fp32T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::bf16, Extension::fp8e4m3},
+ {{{fp8e4m3T, bf16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, fp8e4m3T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e4m3T, bf16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, fp8e4m3T, fp32T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::fp8e5m2},
+ {{{fp8e5m2T, fp8e5m2T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e5m2T, fp16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, fp8e5m2T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e5m2T, fp8e5m2T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e5m2T, fp16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e5m2T, fp32T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, fp8e5m2T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, fp8e5m2T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}},
+ {{Extension::bf16, Extension::fp8e5m2},
+ {{{fp8e5m2T, bf16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, fp8e5m2T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp8e5m2T, bf16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, fp8e5m2T, fp32T}, SpecificationVersion::V_1_1_DRAFT}},
+ allOf},
+ {{Extension::bf16},
+ {{{fp16T, bf16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, fp16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, bf16T, fp16T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp16T, bf16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, fp16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, fp32T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{fp32T, bf16T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
+ {{bf16T, bf16T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}}}},
{"tosa.matmul_t_block_scaled",
{{{Extension::mxfp},
{{{fp4e2m1T, fp8ue8m0T, fp4e2m1T, fp8ue8m0T, fp32T},
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
index 591073e9985ae..3c7d870a44879 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td
@@ -166,6 +166,12 @@ def Tosa_MatMulOpQuantInfoBuilder : OpBuilder<
a, b);
}]>;
+def Tosa_MatMulTOpQuantInfoBuilder
+ : OpBuilder<(ins "Type":$outputType, "Value":$a, "Value":$b), [{
+ buildMatMulTOpWithQuantInfo($_builder, $_state, outputType,
+ a, b);
+ }]>;
+
// Both the tosa.avg_pool2d and unary ops use the same
// UnaryOpQuantizationAttr but the avg_pool operator has its own builder as it
// has additional parameters not part of the unary ops.
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h
index e0626368175ee..2d96fb6891139 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h
@@ -135,6 +135,12 @@ RankedTensorType getVariableType(VariableOp variableOp);
// Returns the bitwidth of a TOSA tensor element type
unsigned getBitWidth(Type type);
+// Returns the storage element type for a given type
+Type getStorageElementTypeOrSelf(Type type);
+
+// Returns the storage element type for a given value
+Type getStorageElementTypeOrSelf(Value value);
+
} // namespace tosa
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index a333505082b7e..c01f32239a59f 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -454,6 +454,44 @@ def Tosa_MatMulOp : Tosa_InferShapedTypeOp<"matmul", [NoMemoryEffect]> {
"operands attr-dict `:` functional-type(operands, results)";
}
+//===----------------------------------------------------------------------===//
+// Operator: matmul_t
+//===----------------------------------------------------------------------===//
+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.
+ }];
+
+ let arguments = (ins Tosa_Tensor3D:$a, Tosa_Tensor3D:$b,
+ Tosa_ScalarIntOrFloatTensor:$a_zp, Tosa_ScalarIntOrFloatTensor:$b_zp);
+
+ let results = (outs Tosa_Tensor3D:$output);
+
+ list<Availability> availability =
+ [Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
+ Extension<[Tosa_EXT_INT16, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2,
+ Tosa_EXT_BF16]>,
+ ];
+
+ let extraClassDeclaration = [{
+ FailureOr<int64_t> getAZeroPoint();
+ FailureOr<int64_t> getBZeroPoint();
+ LogicalResult verifyAZeroPoint(int64_t zp);
+ LogicalResult verifyBZeroPoint(int64_t zp);
+ }];
+
+ let builders = [Tosa_MatMulTOpQuantInfoBuilder];
+ let hasVerifier = 1;
+
+ let assemblyFormat =
+ "operands attr-dict `:` functional-type(operands, results)";
+}
+
//===----------------------------------------------------------------------===//
// Operator: matmul_t_block_scaled
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index 3c54d7448d021..c4705c8080069 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -581,14 +581,14 @@ static std::optional<int64_t> idivCheck(const int64_t lhs, const int64_t rhs) {
return lhs / rhs;
}
-static Type getStorageElementTypeOrSelf(Type type) {
+Type mlir::tosa::getStorageElementTypeOrSelf(Type type) {
auto srcType = getElementTypeOrSelf(type);
if (auto quantType = llvm::dyn_cast<mlir::quant::QuantizedType>(srcType))
srcType = getStorageElementTypeFromQuantized(quantType);
return srcType;
}
-static Type getStorageElementTypeOrSelf(Value value) {
+Type mlir::tosa::getStorageElementTypeOrSelf(Value value) {
return getStorageElementTypeOrSelf(value.getType());
}
@@ -1447,18 +1447,14 @@ buildTransConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,
result.addTypes(finalOutputType);
}
-/// The tosa.matmul op is also intended to be generated where a fully_connected
-/// op must be constructed where the weight is not a constant. In this case,
-/// the fully_connected op must be expressed using matmul.
-/// TODO: Add link to the leglization document explaining this.
-static void buildMatMulOpWithQuantInfo(OpBuilder &builder,
- OperationState &result, Type outputType,
- Value a, Value b) {
- auto zps = createZPsAsConst(builder, a, b);
+static void buildMatMulLikeOpWithQuantInfo(OpBuilder &builder,
+ OperationState &result,
+ Type outputType, Value a, Value b) {
+ const std::pair<Value, Value> zps = createZPsAsConst(builder, a, b);
result.addOperands({a, b, zps.first, zps.second});
Type finalOutputType{outputType};
- if (auto quantAttr = buildMatMulOpQuantizationAttr(builder, a, b)) {
+ if (buildMatMulOpQuantizationAttr(builder, a, b)) {
auto eType = getStorageElementTypeOrSelf(a.getType());
auto inputBits = eType.getIntOrFloatBitWidth();
@@ -1476,6 +1472,18 @@ static void buildMatMulOpWithQuantInfo(OpBuilder &builder,
result.addTypes(finalOutputType);
}
+static void buildMatMulOpWithQuantInfo(OpBuilder &builder,
+ OperationState &result, Type outputType,
+ Value a, Value b) {
+ buildMatMulLikeOpWithQuantInfo(builder, result, outputType, a, b);
+}
+
+static void buildMatMulTOpWithQuantInfo(OpBuilder &builder,
+ OperationState &result, Type outputType,
+ Value a, Value b) {
+ buildMatMulLikeOpWithQuantInfo(builder, result, outputType, a, b);
+}
+
/// Both the tosa.avg_pool2d and unary ops use the same
/// UnaryOpQuantizationAttr but avg_pool operator has its own builder as it
/// has additional parameters not part of the unary ops.
@@ -2052,12 +2060,9 @@ LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
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();
-
+template <typename T>
+static LogicalResult verifyMatMulQuantizedOperandsType(T op, Type aElementType,
+ Type bElementType) {
const auto aQuantizedEType =
llvm::dyn_cast<quant::UniformQuantizedType>(aElementType);
const auto bQuantizedEType =
@@ -2065,33 +2070,52 @@ LogicalResult MatMulOp::verify() {
if (aQuantizedEType || bQuantizedEType) {
if (!aQuantizedEType || !bQuantizedEType) {
- return emitOpError("expect operands to be both quantized or both not "
- "quantized, got ")
+ return op.emitOpError("expect operands to be both quantized or both not "
+ "quantized, got ")
<< aElementType << " and " << bElementType;
}
// both a and b have quantized element types
auto aQuantWidth = aQuantizedEType.getStorageTypeIntegralWidth();
auto bQuantWidth = bQuantizedEType.getStorageTypeIntegralWidth();
if (aQuantWidth != bQuantWidth) {
- return emitOpError("expect quantized operands to have same widths, got ")
+ return op.emitOpError("expect quantized operands to have same widths, "
+ "got ")
<< aQuantWidth << " and " << bQuantWidth;
}
}
- // check a_zp and b_zp
- auto aEType = getStorageElementTypeOrSelf(aElementType);
- auto aZpEType = getStorageElementTypeOrSelf(getAZp().getType());
- if (aEType != aZpEType)
- return emitOpError("expect input a and a_zp have the same "
- "element type, got ")
- << aEType << " and " << aZpEType;
+ return success();
+}
- const Type bEType = getStorageElementTypeOrSelf(bElementType);
- const Type bZpEType = getStorageElementTypeOrSelf(getBZp().getType());
- if (bEType != bZpEType)
- return emitOpError("expect input b and b_zp have the same "
- "element type, got ")
- << bEType << " and " << bZpEType;
+template <typename T>
+static LogicalResult verifyMatMulZeroPointType(T op, Value input, Value zp,
+ StringRef inputName,
+ StringRef zpName) {
+ const Type inputStorageElementType = getStorageElementTypeOrSelf(input);
+ const Type zpElementType = getStorageElementTypeOrSelf(zp);
+
+ if (inputStorageElementType != zpElementType)
+ return op.emitOpError("expect input ")
+ << inputName << " and " << zpName
+ << " have the same element type, got " << inputStorageElementType
+ << " and " << zpElementType;
+
+ 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();
+
+ if (failed(
+ verifyMatMulQuantizedOperandsType(*this, aElementType, bElementType)))
+ return failure();
+
+ if (failed(verifyMatMulZeroPointType(*this, getA(), getAZp(), "a", "a_zp")) ||
+ failed(verifyMatMulZeroPointType(*this, getB(), getBZp(), "b", "b_zp")))
+ return failure();
FailureOr<int64_t> maybeAZp = getAZeroPoint();
if (succeeded(maybeAZp) && verifyAZeroPoint(*maybeAZp).failed())
@@ -2125,15 +2149,97 @@ LogicalResult MatMulOp::verify() {
const SmallVector<int64_t, 3> expectedOutputShape = {N, H, W};
const auto outputType = cast<ShapedType>(getResult().getType());
if (outputType.hasRank() &&
- failed(
- verifyCompatibleShape(outputType.getShape(), expectedOutputShape))) {
- InFlightDiagnostic opError = emitOpError("expected output shape ");
- printShapeToDiagnostic(opError, outputType.getShape());
- opError << " to be compatible with expected output shape ";
- printShapeToDiagnostic(opError, expectedOutputShape);
- return opError;
+ failed(verifyOutputShapeCompatibleWithExpected(getOperation(), outputType,
+ expectedOutputShape)))
+ return failure();
+
+ return success();
+}
+
+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();
+}
+
+LogicalResult MatMulTOp::verify() {
+ const ShapeAdaptor aShape(getA().getType());
+ const ShapeAdaptor bShape(getB().getType());
+ const Type aElementType = aShape.getElementType();
+ const Type bElementType = bShape.getElementType();
+
+ if (failed(
+ verifyMatMulQuantizedOperandsType(*this, aElementType, bElementType)))
+ return failure();
+
+ if (failed(verifyMatMulZeroPointType(*this, getA(), getAZp(), "a", "a_zp")) ||
+ failed(verifyMatMulZeroPointType(*this, getB(), getBZp(), "b", "b_zp")))
+ return failure();
+
+ FailureOr<int64_t> maybeAZp = getAZeroPoint();
+ if (succeeded(maybeAZp) && verifyAZeroPoint(*maybeAZp).failed())
+ return failure();
+
+ FailureOr<int64_t> maybeBZp = getBZeroPoint();
+ 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();
}
@@ -3195,6 +3301,8 @@ ZERO_POINT_HELPER(AvgPool2dAdaptiveOp, Input, true)
ZERO_POINT_HELPER(AvgPool2dAdaptiveOp, Output, true)
ZERO_POINT_HELPER(MatMulOp, A, true)
ZERO_POINT_HELPER(MatMulOp, B, true)
+ZERO_POINT_HELPER(MatMulTOp, A, true)
+ZERO_POINT_HELPER(MatMulTOp, B, true)
ZERO_POINT_HELPER(NegateOp, Input1, true)
ZERO_POINT_HELPER(NegateOp, Output, true)
ZERO_POINT_HELPER(RescaleOp, Input, !getInputUnsigned())
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
index cfd2dd9d29650..163850f741166 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
@@ -16,6 +16,8 @@
#include "mlir/Dialect/Tosa/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -140,6 +142,83 @@ class BoolScatterRewrite : public OpRewritePattern<tosa::ScatterOp> {
}
};
+static LogicalResult isMatMulTTypeCompatibleForDowngrade(tosa::MatMulTOp op) {
+ const Type aElementType = getStorageElementTypeOrSelf(op.getA().getType());
+ const Type bElementType = getStorageElementTypeOrSelf(op.getB().getType());
+ const Type outputElementType =
+ getStorageElementTypeOrSelf(op.getOutput().getType());
+
+ if (aElementType != bElementType)
+ return failure();
+
+ if ((aElementType.isF16() && outputElementType.isF16()) ||
+ (aElementType.isF16() && outputElementType.isF32()) ||
+ (aElementType.isF32() && outputElementType.isF32()) ||
+ (aElementType.isBF16() && outputElementType.isF32()) ||
+ (aElementType.isInteger(8) && outputElementType.isInteger(32)) ||
+ (aElementType.isInteger(16) && outputElementType.isInteger(48)) ||
+ (isa<Float8E5M2Type>(aElementType) && outputElementType.isF16()) ||
+ (isa<Float8E4M3FNType>(aElementType) && outputElementType.isF16()))
+ return success();
+
+ return failure();
+}
+
+class MatMulTRewrite : public OpRewritePattern<tosa::MatMulTOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::MatMulTOp op,
+ PatternRewriter &rewriter) const override {
+ if (failed(isMatMulTTypeCompatibleForDowngrade(op)))
+ return rewriter.notifyMatchFailure(
+ op, "expected 1.0-compatible matmul_t element types");
+
+ const Type aType = op.getA().getType();
+ const Type bType = op.getB().getType();
+ const ShapeAdaptor aShape(aType);
+ const ShapeAdaptor bShape(bType);
+ if (!aShape.hasRank() || !bShape.hasRank())
+ return rewriter.notifyMatchFailure(op, "expected ranked A and B tensors");
+
+ const int64_t dSize = bShape.getDimSize(0);
+ const int64_t nSize = aShape.getDimSize(0);
+
+ // To convert broadcasting behaviour to TOSA 1.0, we're required to tile the
+ // input. TOSA 1.0 does not support shape expressions, so the batch size
+ // must be known at compile time.
+ if (ShapedType::isDynamic(dSize) ||
+ (dSize == 1 && ShapedType::isDynamic(nSize)))
+ return rewriter.notifyMatchFailure(
+ op, "expected known batch size for broadcast");
+
+ const int64_t wSize = bShape.getDimSize(1);
+ const int64_t cSize = bShape.getDimSize(2);
+ const Location loc = op.getLoc();
+ const RankedTensorType transposedBType =
+ cast<RankedTensorType>(bType).clone({dSize, cSize, wSize});
+ auto transpose =
+ tosa::TransposeOp::create(rewriter, loc, transposedBType, op.getB(),
+ rewriter.getDenseI32ArrayAttr({0, 2, 1}));
+ Value matMulB = transpose.getOutput();
+
+ // Matmul does not support broadcasting, so tile b if required
+ if (dSize == 1 && nSize != 1) {
+ const RankedTensorType tiledBType =
+ cast<RankedTensorType>(bType).clone({nSize, cSize, wSize});
+ const Value multiples = getTosaConstShape(rewriter, loc, {nSize, 1, 1});
+ auto tile =
+ tosa::TileOp::create(rewriter, loc, tiledBType, matMulB, multiples);
+ matMulB = tile.getOutput();
+ }
+
+ auto matmul = tosa::MatMulOp::create(rewriter, loc, op.getType(), op.getA(),
+ matMulB, op.getAZp(), op.getBZp());
+ rewriter.replaceOp(op, matmul.getOutput());
+ return success();
+ }
+};
+
struct TosaDowngrade1p1To1p0Pass
: public tosa::impl::TosaDowngrade1p1To1p0PassBase<
TosaDowngrade1p1To1p0Pass> {
@@ -150,8 +229,8 @@ struct TosaDowngrade1p1To1p0Pass
func::FuncOp func = getOperation();
RewritePatternSet patterns(&context);
- patterns.add<BoolFp32CastRewrite, BoolGatherRewrite, BoolScatterRewrite>(
- &context);
+ patterns.add<BoolFp32CastRewrite, BoolGatherRewrite, BoolScatterRewrite,
+ MatMulTRewrite>(&context);
FrozenRewritePatternSet frozenPatterns(std::move(patterns));
if (failed(applyPatternsGreedily(func, frozenPatterns)))
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
index 1bf92ff562a84..0b4983bea4bc6 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
@@ -256,6 +256,14 @@ LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::MatMulOp op) {
return success();
}
+template <>
+LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::MatMulTOp op) {
+ addValue(op.getA());
+ addValue(op.getB());
+ addValue(op.getOutput());
+ return success();
+}
+
template <>
LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::VariableOp op) {
addType(op.getType());
@@ -315,6 +323,7 @@ LogicalResult ProfileInfoDepot::populatationDispatch(Operation *op) {
POPULATE_PROFILE_INFO_CUSTOM(Select)
POPULATE_PROFILE_INFO_CUSTOM(Rescale)
POPULATE_PROFILE_INFO_CUSTOM(MatMul)
+ POPULATE_PROFILE_INFO_CUSTOM(MatMulT)
POPULATE_PROFILE_INFO_CUSTOM(Variable)
POPULATE_PROFILE_INFO_CUSTOM(VariableWrite)
POPULATE_PROFILE_INFO_CUSTOM(Dim)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 4a29e5327cc55..5d4a2affbd7cd 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -107,7 +107,8 @@ static LogicalResult checkConstantOperandConvOps(Operation *op,
static LogicalResult checkConstantOperandMatMul(Operation *op,
const TargetEnv &env) {
- if (!env.allows(Extension::dynamic) && isa<tosa::MatMulOp>(op)) {
+ if (!env.allows(Extension::dynamic) &&
+ isa<tosa::MatMulOp, tosa::MatMulTOp>(op)) {
// Check 'A_zp' and 'B_zp'
return checkConstantOperands(op, {2, 3});
}
@@ -838,6 +839,7 @@ LogicalResult TosaValidation::levelCheckRanksAndSizes(Operation *op) {
CHECK_SIZES(TransposeConv2D);
CHECK_SIZES(FFT2d);
CHECK_SIZES(MatMul);
+ CHECK_SIZES(MatMulT);
CHECK_SIZES(MatmulTBlockScaled);
CHECK_SIZES(MaxPool2d);
CHECK_SIZES(MaxPool2dAdaptive);
diff --git a/mlir/test/Dialect/Tosa/dynamic_extension.mlir b/mlir/test/Dialect/Tosa/dynamic_extension.mlir
index 5f5ab795459f7..bb27f4bbbd8af 100644
--- a/mlir/test/Dialect/Tosa/dynamic_extension.mlir
+++ b/mlir/test/Dialect/Tosa/dynamic_extension.mlir
@@ -73,6 +73,22 @@ func.func @test_matmul_non_const_zps(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1
// -----
+func.func @test_matmul_t_non_const_a_zp(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>, %a_zp: tensor<1xf32>) -> tensor<1x14x28xf32> {
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_non_const_b_zp(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>, %b_zp: tensor<1xf32>) -> tensor<1x14x28xf32> {
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
func.func @test_negate_non_const_zps(%arg0: tensor<1xf32>, %input1_zp: tensor<1xf32>, %output_zp: tensor<1xf32>) -> tensor<1xf32> {
%0 = tosa.negate %arg0, %input1_zp, %output_zp {} : (tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
return %0 : tensor<1xf32>
diff --git a/mlir/test/Dialect/Tosa/invalid_extension.mlir b/mlir/test/Dialect/Tosa/invalid_extension.mlir
index cd82e314e3b6e..c69ea83f5ef36 100644
--- a/mlir/test/Dialect/Tosa/invalid_extension.mlir
+++ b/mlir/test/Dialect/Tosa/invalid_extension.mlir
@@ -518,6 +518,23 @@ func.func @test_matmul_non_const_b_zp(%arg0: tensor<1x14x19xf32>, %arg1: tensor<
// -----
+func.func @test_matmul_t_non_const_a_zp(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>, %a_zp: tensor<1xf32>, %b_zp: tensor<1xf32>) -> tensor<1x14x28xf32> {
+ // expected-error at +1 {{'tosa.matmul_t' op expected compile time resolvable constant, but got variable value for operand #2}}
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_non_const_b_zp(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>, %b_zp: tensor<1xf32>) -> tensor<1x14x28xf32> {
+ %a_zp = "tosa.const"() {values = dense<0.0> : tensor<1xf32> } : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul_t' op expected compile time resolvable constant, but got variable value for operand #3}}
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
func.func @test_mul_non_const(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x1x3xi8>, %shift: tensor<1xi8>) -> tensor<13x21x3xi32> {
// expected-error at +1 {{'tosa.mul' op expected compile time resolvable constant, but got variable value for operand #2}}
%0 = tosa.mul %arg0, %arg1, %shift : (tensor<13x21x3xi8>, tensor<13x1x3xi8>, tensor<1xi8>) -> tensor<13x21x3xi32>
diff --git a/mlir/test/Dialect/Tosa/level_check.mlir b/mlir/test/Dialect/Tosa/level_check.mlir
index 85b5d002434bc..8f91d2f40d37d 100644
--- a/mlir/test/Dialect/Tosa/level_check.mlir
+++ b/mlir/test/Dialect/Tosa/level_check.mlir
@@ -1278,6 +1278,15 @@ func.func @test_matmul_tensor_size_invalid(%arg0: tensor<23178x20000x19xf32>, %a
// -----
+func.func @test_matmul_t_tensor_size_invalid(%arg0: tensor<23178x20000x19xf32>, %arg1: tensor<23178x28x19xf32>) -> tensor<23178x20000x28xf32> {
+ %zero = "tosa.const"() {values = dense<0.0> : tensor<1xf32>} : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul_t' op failed level check: operand tensor size (in bytes) <= (1 << MAX_LOG2_SIZE - 1)}}
+ %0 = tosa.matmul_t %arg0, %arg1, %zero, %zero : (tensor<23178x20000x19xf32>, tensor<23178x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<23178x20000x28xf32>
+ return %0 : tensor<23178x20000x28xf32>
+}
+
+// -----
+
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/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index 50f28df37a169..262b8aa548506 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -232,6 +232,87 @@ func.func @test_matmul(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x19x28xf32>) -
return %0 : tensor<1x14x28xf32>
}
+// -----
+// CHECK-LABEL: test_matmul_t
+func.func @test_matmul_t(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>) -> tensor<1x14x28xf32> {
+%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<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_i8
+func.func @test_matmul_t_i8(%arg0: tensor<2x14x19xi8>, %arg1: tensor<2x28x19xi8>) -> tensor<2x14x28xi32> {
+%azp0 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+%bzp0 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x14x19xi8>, tensor<2x28x19xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2x14x28xi32>
+ return %0 : tensor<2x14x28xi32>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_quantized_i8
+func.func @test_matmul_t_quantized_i8(%arg0: tensor<2x14x19x!quant.uniform<i8:f32, 0.125>>, %arg1: tensor<2x28x19x!quant.uniform<i8:f32, 0.25>>) -> tensor<2x14x28xi32> {
+%azp0 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+%bzp0 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x14x19x!quant.uniform<i8:f32, 0.125>>, tensor<2x28x19x!quant.uniform<i8:f32, 0.25>>, tensor<1xi8>, tensor<1xi8>) -> tensor<2x14x28xi32>
+ return %0 : tensor<2x14x28xi32>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_f16_acc_f32
+func.func @test_matmul_t_f16_acc_f32(%arg0: tensor<2x14x19xf16>, %arg1: tensor<2x28x19xf16>) -> tensor<2x14x28xf32> {
+%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
+%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x14x19xf16>, tensor<2x28x19xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x14x28xf32>
+ return %0 : tensor<2x14x28xf32>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_bf16
+func.func @test_matmul_t_bf16(%arg0: tensor<2x14x19xbf16>, %arg1: tensor<2x28x19xbf16>) -> tensor<2x14x28xf32> {
+%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xbf16>}> : () -> tensor<1xbf16>
+%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xbf16>}> : () -> tensor<1xbf16>
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x14x19xbf16>, tensor<2x28x19xbf16>, tensor<1xbf16>, tensor<1xbf16>) -> tensor<2x14x28xf32>
+ return %0 : tensor<2x14x28xf32>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_fp8_mixed
+func.func @test_matmul_t_fp8_mixed(%arg0: tensor<2x14x19xf8E4M3FN>, %arg1: tensor<2x28x19xf8E5M2>) -> tensor<2x14x28xf16> {
+%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E4M3FN>}> : () -> tensor<1xf8E4M3FN>
+%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E5M2>}> : () -> tensor<1xf8E5M2>
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x14x19xf8E4M3FN>, tensor<2x28x19xf8E5M2>, tensor<1xf8E4M3FN>, tensor<1xf8E5M2>) -> tensor<2x14x28xf16>
+ return %0 : tensor<2x14x28xf16>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_dynamic_dims
+func.func @test_matmul_t_dynamic_dims(%arg0: tensor<?x14x19xf32>, %arg1: tensor<?x28x19xf32>) -> tensor<?x14x28xf32> {
+%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<?x14x19xf32>, tensor<?x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x14x28xf32>
+ return %0 : tensor<?x14x28xf32>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_static_broadcast
+func.func @test_matmul_t_static_broadcast(%arg0: tensor<4x14x19xf32>, %arg1: tensor<1x28x19xf32>) -> tensor<4x14x28xf32> {
+%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<4x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+ return %0 : tensor<4x14x28xf32>
+}
+
+// -----
+// CHECK-LABEL: test_matmul_t_unranked
+func.func @test_matmul_t_unranked(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
+%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<*xf32>, tensor<*xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<*xf32>
+ return %0 : tensor<*xf32>
+}
+
// -----
// CHECK-LABEL: max_pool2d_f32
func.func @test_max_pool2d_f32(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
diff --git a/mlir/test/Dialect/Tosa/profile_all_unsupported.mlir b/mlir/test/Dialect/Tosa/profile_all_unsupported.mlir
index f9011766c885d..caa547a0830e7 100644
--- a/mlir/test/Dialect/Tosa/profile_all_unsupported.mlir
+++ b/mlir/test/Dialect/Tosa/profile_all_unsupported.mlir
@@ -88,6 +88,13 @@ func.func @test_matmul(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x19x28xf32>, %
return %0 : tensor<1x14x28xf32>
}
+// -----
+func.func @test_matmul_t(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>, %arg2: tensor<1xf32>) -> tensor<1x14x28xf32> {
+ // expected-error at +1 {{'tosa.matmul_t' op illegal: requires specification version compatible with 1.1.draft (got 1.0) and requires any of [pro_fp] profiles/extensions to be specified in the target environment}}
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg2: (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
// -----
func.func @test_sigmoid(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
// expected-error at +1 {{'tosa.sigmoid' op illegal: requires any of [pro_fp] profiles/extensions to be specified in the target environment}}
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 5427d9119b7af..378e5eb7b0147 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
@@ -116,3 +116,37 @@ func.func @test_preserve_scatter_i8_i32(%arg0: tensor<13x52x3xi8>, %arg1: tensor
%0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x3xi8>, tensor<13x26xi32>, tensor<13x26x3xi8>) -> tensor<13x52x3xi8>
return %0 : tensor<13x52x3xi8>
}
+
+// -----
+
+// CHECK-LABEL: @test_matmul_t_static_batch
+// CHECK: %[[TRANSPOSE:.+]] = tosa.transpose %arg1 {perms = array<i32: 0, 2, 1>} : (tensor<4x28x19xf32>) -> tensor<4x19x28xf32>
+// CHECK: %[[MATMUL:.+]] = tosa.matmul %arg0, %[[TRANSPOSE]], %arg2, %arg3 : (tensor<4x14x19xf32>, tensor<4x19x28xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+// CHECK: return %[[MATMUL]]
+func.func @test_matmul_t_static_batch(%arg0: tensor<4x14x19xf32>, %arg1: tensor<4x28x19xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<4x14x28xf32> {
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<4x14x19xf32>, tensor<4x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+ return %0 : tensor<4x14x28xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_matmul_t_static_broadcast
+// CHECK: %[[MULTIPLES:.+]] = tosa.const_shape {values = dense<[4, 1, 1]> : tensor<3xindex>} : () -> !tosa.shape<3>
+// CHECK: %[[TRANSPOSE:.+]] = tosa.transpose %arg1 {perms = array<i32: 0, 2, 1>} : (tensor<1x28x19xf32>) -> tensor<1x19x28xf32>
+// CHECK: %[[TILE:.+]] = tosa.tile %[[TRANSPOSE]], %[[MULTIPLES]] : (tensor<1x19x28xf32>, !tosa.shape<3>) -> tensor<4x19x28xf32>
+// CHECK: %[[MATMUL:.+]] = tosa.matmul %arg0, %[[TILE]], %arg2, %arg3 : (tensor<4x14x19xf32>, tensor<4x19x28xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+// CHECK: return %[[MATMUL]]
+func.func @test_matmul_t_static_broadcast(%arg0: tensor<4x14x19xf32>, %arg1: tensor<1x28x19xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<4x14x28xf32> {
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<4x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+ return %0 : tensor<4x14x28xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_matmul_t_dynamic_broadcast
+// CHECK: %[[MATMUL_T:.+]] = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<4x14x19xf32>, tensor<?x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+// CHECK: return %[[MATMUL_T]]
+func.func @test_preserve_matmul_t_dynamic_broadcast(%arg0: tensor<4x14x19xf32>, %arg1: tensor<?x28x19xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<4x14x28xf32> {
+ %0 = tosa.matmul_t %arg0, %arg1, %arg2, %arg3 : (tensor<4x14x19xf32>, tensor<?x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+ return %0 : tensor<4x14x28xf32>
+}
diff --git a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
index cfe4ea4782037..d0d1906524f83 100644
--- a/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir
@@ -365,6 +365,72 @@ func.func @test_unranked_zero_points_matmul(%arg0: tensor<1x2x3xf32>, %arg1: ten
// -----
+// CHECK-LABEL: @test_static_matmul_t
+func.func @test_static_matmul_t(%arg0 : tensor<2x3x4xi32>, %arg1 : tensor<2x5x4xi32>) -> () {
+ // CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<2x3x4xi32>, tensor<2x5x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x3x5xi32>
+ %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<2x3x4xi32>, tensor<2x5x4xi32>, 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>
+ %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<?x?x?xi32>, tensor<2x5x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_dynamic_rhs_matmul_t
+func.func @test_dynamic_rhs_matmul_t(%arg0 : tensor<2x3x4xi32>, %arg1 : tensor<?x?x?xi32>) -> () {
+ // CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<2x3x4xi32>, tensor<?x?x?xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<2x3x?xi32>
+ %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<2x3x4xi32>, tensor<?x?x?xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_broadcast_matmul_t
+func.func @test_broadcast_matmul_t(%arg0 : tensor<4x3x4xi32>, %arg1 : tensor<1x5x4xi32>) -> () {
+ // CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<4x3x4xi32>, tensor<1x5x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<4x3x5xi32>
+ %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<4x3x4xi32>, tensor<1x5x4xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_dynamic_broadcast_matmul_t
+func.func @test_dynamic_broadcast_matmul_t(%arg0 : tensor<?x?x?xi32>, %arg1 : tensor<1x?x?xi32>) -> () {
+ // CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<?x?x?xi32>, tensor<1x?x?xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ %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<?x?x?xi32>, tensor<1x?x?xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<*xi32>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @test_unranked_matmul_t
+func.func @test_unranked_matmul_t(%arg0 : tensor<*xi32>, %arg1 : tensor<*xi32>) -> () {
+ // CHECK: tosa.matmul_t %arg0, %arg1, %0, %1 : (tensor<*xi32>, tensor<*xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ %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<*xi32>, tensor<*xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<?x?x?xi32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: @test_accepts_unranked_scalar_tensor
func.func @test_accepts_unranked_scalar_tensor(%arg0: tensor<1x2x2xf32>, %arg1: tensor<1xf32>) -> tensor<*xf32> {
// CHECK-DAG: %[[SHAPE:.*]] = tosa.const_shape {values = dense<[0, 0, 0, 1, 0, 1]> : tensor<6xindex>} : () -> !tosa.shape<6>
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 3dc10eff15a62..cbff768ec1ecb 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir
@@ -22,6 +22,16 @@ func.func @test_matmul_fp8_input_fp32_acc_type(%arg0: tensor<1x14x19xf8E4M3FN>,
// -----
+func.func @test_matmul_t_fp32(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>) -> tensor<1x14x28xf32> {
+ %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 illegal: requires specification version compatible with 1.1.draft (got 1.0) to be specified in the target environment}}
+ %0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
func.func @test_conv2d_fp8_acc32(%arg0: tensor<1x4x4x4xf8E5M2>, %arg1: tensor<8x1x1x4xf8E5M2>, %arg2: tensor<8xf32>) -> tensor<1x4x4x8xf32> {
%input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E5M2>}> : () -> tensor<1xf8E5M2>
%weight_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E5M2>}> : () -> tensor<1xf8E5M2>
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 584dc4e9aa9ed..f4b5ab5dd11fd 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -22,6 +22,16 @@ func.func @test_matmul_fp8_input_fp32_acc_type(%arg0: tensor<1x14x19xf8E4M3FN>,
// -----
+// CHECK-LABEL: test_matmul_t_fp8_input_fp32_acc_type
+func.func @test_matmul_t_fp8_input_fp32_acc_type(%arg0: tensor<1x14x19xf8E4M3FN>, %arg1: tensor<1x28x19xf8E4M3FN>) -> tensor<1x14x28xf32> {
+ %azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E4M3FN>}> : () -> tensor<1xf8E4M3FN>
+ %bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E4M3FN>}> : () -> tensor<1xf8E4M3FN>
+ %0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<1x14x19xf8E4M3FN>, tensor<1x28x19xf8E4M3FN>, tensor<1xf8E4M3FN>, tensor<1xf8E4M3FN>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
// CHECK-LABEL: test_conv2d_fp8_acc32
func.func @test_conv2d_fp8_acc32(%arg0: tensor<1x4x4x4xf8E5M2>, %arg1: tensor<8x1x1x4xf8E5M2>, %arg2: tensor<8xf32>) -> tensor<1x4x4x8xf32> {
%input_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E5M2>}> : () -> tensor<1xf8E5M2>
diff --git a/mlir/test/Dialect/Tosa/verifier.mlir b/mlir/test/Dialect/Tosa/verifier.mlir
index 9d97df2f43c6d..0eb679ca7ced3 100644
--- a/mlir/test/Dialect/Tosa/verifier.mlir
+++ b/mlir/test/Dialect/Tosa/verifier.mlir
@@ -1314,7 +1314,7 @@ func.func @test_matmul_output_channel_mismatch(%arg0: tensor<2x3x4xf32>, %arg1:
func.func @test_matmul_output_shape_mismatch(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x4x6xf32>) -> tensor<2x5x6xf32> {
%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 2, 5, 6 to be compatible with expected output shape 2, 3, 6}}
+ // expected-error at +1 {{'tosa.matmul' op expected output shape 2, 5, 6 to be compatible with inferred shape 2, 3, 6}}
%0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4xf32>, tensor<2x4x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x5x6xf32>
return %0 : tensor<2x5x6xf32>
}
@@ -1345,7 +1345,7 @@ func.func @test_matmul_dynamic_channel_mismatch(%arg0: tensor<?x3x4xf32>, %arg1:
func.func @test_matmul_dynamic_output_shape_mismatch(%arg0: tensor<?x3x4xf32>, %arg1: tensor<2x4x6xf32>) -> tensor<5x3x6xf32> {
%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, 6 to be compatible with expected output shape 2, 3, 6}}
+ // expected-error at +1 {{'tosa.matmul' op expected output shape 5, 3, 6 to be compatible with inferred shape 2, 3, 6}}
%0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<?x3x4xf32>, tensor<2x4x6xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5x3x6xf32>
return %0 : tensor<5x3x6xf32>
}
@@ -1356,7 +1356,7 @@ 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>
- // expected-error at +1 {{'tosa.matmul' op expected output shape 2, 5, ? to be compatible with expected output shape 2, 3, ?}}
+ // expected-error at +1 {{'tosa.matmul' op expected output shape 2, 5, ? to be compatible with inferred shape 2, 3, ?}}
%0 = tosa.matmul %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4xf32>, tensor<*xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x5x?xf32>
return %0 : tensor<2x5x?xf32>
}
@@ -1384,6 +1384,26 @@ func.func @test_matmul_quantized_width_mismatch(%arg0: tensor<2x3x4x!quant.unifo
// -----
+func.func @test_matmul_t_quantized_mixed_operands(%arg0: tensor<2x3x4x!quant.uniform<i8:f32, 0.125>>, %arg1: tensor<2x6x4xf32>) -> tensor<2x3x6xi32> {
+ %azp0 = "tosa.const"() {values = dense<0> : tensor<1xi8>} : () -> tensor<1xi8>
+ %bzp0 = "tosa.const"() {values = dense<0.0> : tensor<1xf32>} : () -> tensor<1xf32>
+ // expected-error at +1 {{'tosa.matmul_t' op expect operands to be both quantized or both not quantized, got '!quant.uniform<i8:f32, 1.250000e-01>' and 'f32'}}
+ %0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4x!quant.uniform<i8:f32, 0.125>>, tensor<2x6x4xf32>, tensor<1xi8>, tensor<1xf32>) -> tensor<2x3x6xi32>
+ return %0 : tensor<2x3x6xi32>
+}
+
+// -----
+
+func.func @test_matmul_t_quantized_width_mismatch(%arg0: tensor<2x3x4x!quant.uniform<i8:f32, 0.125>>, %arg1: tensor<2x6x4x!quant.uniform<i16:f32, 0.125>>) -> tensor<2x3x6xi32> {
+ %azp0 = "tosa.const"() {values = dense<0> : tensor<1xi8>} : () -> tensor<1xi8>
+ %bzp0 = "tosa.const"() {values = dense<0> : tensor<1xi16>} : () -> tensor<1xi16>
+ // expected-error at +1 {{'tosa.matmul_t' op expect quantized operands to have same widths, got 8 and 16}}
+ %0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<2x3x4x!quant.uniform<i8:f32, 0.125>>, tensor<2x6x4x!quant.uniform<i16:f32, 0.125>>, tensor<1xi8>, tensor<1xi16>) -> tensor<2x3x6xi32>
+ return %0 : tensor<2x3x6xi32>
+}
+
+// -----
+
// CHECK-LABEL: test_matmul_a_zp_same_element_type
func.func @test_matmul_a_zp_same_element_type(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x19x28xf32>) -> tensor<1x14x28xf32> {
%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
@@ -1428,6 +1448,56 @@ func.func @test_matmul_b_zp_non_zero(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1
// -----
+func.func @test_matmul_t_channel_mismatch(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x20xf32>) -> tensor<1x14x28xf32> {
+%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 expected channels of b to match size 19, got 20}}
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<1x14x19xf32>, tensor<1x28x20xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
+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}}
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<4x14x19xf32>, tensor<2x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x28xf32>
+ return %0 : tensor<4x14x28xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_output_shape_mismatch(%arg0: tensor<4x14x19xf32>, %arg1: tensor<1x28x19xf32>) -> tensor<4x14x29xf32> {
+%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 expected output shape 4, 14, 29 to be compatible with inferred shape 4, 14, 28}}
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<4x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x14x29xf32>
+ return %0 : tensor<4x14x29xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_a_zp_same_element_type(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>) -> tensor<1x14x28xf32> {
+%azp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
+%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+// expected-error at +1 {{'tosa.matmul_t' op expect input a and a_zp have the same element type, got 'f32' and 'f16'}}
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf16>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
+func.func @test_matmul_t_a_zp_non_zero(%arg0: tensor<1x14x19xf32>, %arg1: tensor<1x28x19xf32>) -> tensor<1x14x28xf32> {
+%azp0 = "tosa.const"() <{values = dense<1.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+%bzp0 = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+// expected-error at +1 {{'tosa.matmul_t' op a zero point must be zero for non-int8 integer types}}
+%0 = tosa.matmul_t %arg0, %arg1, %azp0, %bzp0 : (tensor<1x14x19xf32>, tensor<1x28x19xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x14x28xf32>
+ return %0 : tensor<1x14x28xf32>
+}
+
+// -----
+
func.func @test_matmul_t_block_scaled_data_mismatch(%arg0: tensor<4x8x32xf8E4M3FN>, %arg1: tensor<4x8x1xf8E8M0FNU>, %arg2: tensor<4x16x32xf8E5M2>, %arg3: tensor<4x16x1xf8E8M0FNU>) -> tensor<4x8x16xf32> {
// expected-error at +1 {{'tosa.matmul_t_block_scaled' op expect A_data and B_data to have same element type, got 'f8E4M3FN' and 'f8E5M2'}}
%0 = tosa.matmul_t_block_scaled %arg0, %arg1, %arg2, %arg3 {block_size = #tosa.block_size<BLOCK_SIZE_32> : i32} : (tensor<4x8x32xf8E4M3FN>, tensor<4x8x1xf8E8M0FNU>, tensor<4x16x32xf8E5M2>, tensor<4x16x1xf8E8M0FNU>) -> tensor<4x8x16xf32>
More information about the Mlir-commits
mailing list