[Mlir-commits] [mlir] [mlir][tosa] Add support for matmul_t (PR #203894)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 15 07:19:48 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Luke Hutton (lhutton1)
<details>
<summary>Changes</summary>
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.
---
Patch is 53.71 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/203894.diff
18 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc (+56)
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOpBase.td (+6)
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.h (+6)
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (+38)
- (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+148-40)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp (+81-2)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp (+9)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp (+3-1)
- (modified) mlir/test/Dialect/Tosa/dynamic_extension.mlir (+16)
- (modified) mlir/test/Dialect/Tosa/invalid_extension.mlir (+17)
- (modified) mlir/test/Dialect/Tosa/level_check.mlir (+9)
- (modified) mlir/test/Dialect/Tosa/ops.mlir (+81)
- (modified) mlir/test/Dialect/Tosa/profile_all_unsupported.mlir (+7)
- (modified) mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir (+34)
- (modified) mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir (+66)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir (+10)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir (+10)
- (modified) mlir/test/Dialect/Tosa/verifier.mlir (+73-3)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc b/mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc
index 50bb9f69c6242..9dd740b4f445f 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},
@@ -660,6 +667,55 @@ 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, bf16T, 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 a99fb2fcae547..aaa8a243d6d40 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -453,6 +453,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 `D` 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 f05399cf6b00b..99942ff3d1837 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());
}
@@ -1444,18 +1444,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();
@@ -1473,6 +1469,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.
@@ -2049,12 +2057,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 =
@@ -2062,33 +2067,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())
@@ -2122,15 +2146,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();
}
@@ -3204,6 +3310,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<Flo...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/203894
More information about the Mlir-commits
mailing list