[Mlir-commits] [mlir] [mlir][tosa] Lower MATMUL_T to Linalg (PR #222661)
Jack Frankland
llvmlistbot at llvm.org
Tue Sep 22 05:24:09 PDT 2026
https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/222661
>From b5b86fc73c6c702d35381d9299af04ed4012d2fd Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Tue, 22 Sep 2026 10:31:05 +0100
Subject: [PATCH 1/2] [mlir][arith] Support equal-width float conversions
Use arith.convertf in convertScalarToDtype when converting between
floating-point types that have the same bit width but different
semantics.
Add focused unit coverage for f16-to-bf16 and bf16-to-f16 conversions.
Assisted-by: Codex
---
mlir/lib/Dialect/Arith/Utils/Utils.cpp | 5 +-
.../Dialect/Arith/ArithUtilsTest.cpp | 50 +++++++++++++++++++
mlir/unittests/Dialect/Arith/CMakeLists.txt | 7 +++
mlir/unittests/Dialect/CMakeLists.txt | 1 +
4 files changed, 61 insertions(+), 2 deletions(-)
create mode 100644 mlir/unittests/Dialect/Arith/ArithUtilsTest.cpp
create mode 100644 mlir/unittests/Dialect/Arith/CMakeLists.txt
diff --git a/mlir/lib/Dialect/Arith/Utils/Utils.cpp b/mlir/lib/Dialect/Arith/Utils/Utils.cpp
index 990fb2cd2bde0..8cbddd920c91f 100644
--- a/mlir/lib/Dialect/Arith/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Arith/Utils/Utils.cpp
@@ -169,7 +169,6 @@ static Value convertScalarToIntDtype(ImplicitLocOpBuilder &b, Value operand,
static Value convertScalarToFpDtype(ImplicitLocOpBuilder &b, Value operand,
FloatType toType, bool isUnsigned) {
// If operand is integer, cast directly to the float type.
- // Note that it is unclear how to cast from BF16<->FP16.
if (isa<IntegerType>(operand.getType())) {
if (isUnsigned)
return arith::UIToFPOp::create(b, toType, operand);
@@ -181,7 +180,9 @@ static Value convertScalarToFpDtype(ImplicitLocOpBuilder &b, Value operand,
arith::FastMathFlagsAttr{});
if (toType.getWidth() < fromFpTy.getWidth())
return arith::TruncFOp::create(b, toType, operand);
- return operand;
+
+ return arith::ConvertFOp::create(b, toType, operand,
+ arith::ConvertFOp::Properties{}, {});
}
return {};
diff --git a/mlir/unittests/Dialect/Arith/ArithUtilsTest.cpp b/mlir/unittests/Dialect/Arith/ArithUtilsTest.cpp
new file mode 100644
index 0000000000000..174b1e41e054e
--- /dev/null
+++ b/mlir/unittests/Dialect/Arith/ArithUtilsTest.cpp
@@ -0,0 +1,50 @@
+//===- ArithUtilsTest.cpp - Arith utils unit tests ------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/Builders.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+
+class ArithUtilsTest : public ::testing::Test {
+protected:
+ ArithUtilsTest() : builder(&context), loc(builder.getUnknownLoc()) {
+ context.loadDialect<arith::ArithDialect>();
+ builder.setInsertionPointToStart(&block);
+ }
+
+ void expectConvertFOp(Type inputType, Type resultType) {
+ Value input = block.addArgument(inputType, loc);
+ Value result = convertScalarToDtype(builder, loc, input, resultType,
+ /*isUnsignedCast=*/false);
+
+ EXPECT_EQ(result.getType(), resultType);
+ auto convertOp = result.getDefiningOp<arith::ConvertFOp>();
+ ASSERT_TRUE(convertOp);
+ EXPECT_EQ(convertOp.getIn(), input);
+ }
+
+ MLIRContext context;
+ OpBuilder builder;
+ Location loc;
+ Block block;
+};
+
+TEST_F(ArithUtilsTest, ConvertF16ToBF16) {
+ expectConvertFOp(builder.getF16Type(), builder.getBF16Type());
+}
+
+TEST_F(ArithUtilsTest, ConvertBF16ToF16) {
+ expectConvertFOp(builder.getBF16Type(), builder.getF16Type());
+}
+
+} // namespace
diff --git a/mlir/unittests/Dialect/Arith/CMakeLists.txt b/mlir/unittests/Dialect/Arith/CMakeLists.txt
new file mode 100644
index 0000000000000..3f0ba6d19378d
--- /dev/null
+++ b/mlir/unittests/Dialect/Arith/CMakeLists.txt
@@ -0,0 +1,7 @@
+add_mlir_unittest(MLIRArithTests
+ ArithUtilsTest.cpp
+)
+mlir_target_link_libraries(MLIRArithTests
+ PRIVATE
+ MLIRArithUtils
+)
diff --git a/mlir/unittests/Dialect/CMakeLists.txt b/mlir/unittests/Dialect/CMakeLists.txt
index 269eccb1f93c3..5fd835189ccc7 100644
--- a/mlir/unittests/Dialect/CMakeLists.txt
+++ b/mlir/unittests/Dialect/CMakeLists.txt
@@ -7,6 +7,7 @@ mlir_target_link_libraries(MLIRDialectTests
MLIRDialect)
add_subdirectory(AMDGPU)
+add_subdirectory(Arith)
add_subdirectory(ArmSME)
add_subdirectory(Index)
add_subdirectory(Linalg)
>From 8b37ada663765608075770ad219b08b8bd8ae083 Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Thu, 10 Sep 2026 14:09:26 +0100
Subject: [PATCH 2/2] [mlir][tosa] Lower MATMUL_T to Linalg
Broadcast a singleton RHS batch explicitly before lowering
floating-point operations to linalg.batch_matmul_transpose_b. Lower
quantized operations with linalg.generic so zero-point adjustment is
represented in the reduction.
Assisted-by: Codex
---
.../Conversion/TosaToLinalg/CMakeLists.txt | 1 +
.../TosaToLinalg/TosaToLinalgNamed.cpp | 181 ++++++++++++++++++
.../TosaToLinalg/TosaToLinalgNamedPass.cpp | 1 +
.../TosaToLinalg/tosa-to-linalg-named.mlir | 105 ++++++++++
4 files changed, 288 insertions(+)
diff --git a/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt b/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt
index 5123c2a7cf916..98d8637686ba2 100644
--- a/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt
+++ b/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt
@@ -20,6 +20,7 @@ add_mlir_conversion_library(MLIRTosaToLinalg
MLIRLinalgUtils
MLIRMathDialect
MLIRPass
+ MLIRSCFDialect
MLIRSupport
MLIRTensorDialect
MLIRTosaDialect
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 4a11a927090bf..b7f916ddc874d 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -14,6 +14,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
@@ -679,6 +680,185 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
}
};
+// Broadcast the batch dimension of B from one to the batch dimension of A.
+// linalg.broadcast adds dimensions rather than expanding dimensions of size
+// one, so first rank-reduce B from 1xWxC to WxC.
+static Value broadcastMatMulTRhs(PatternRewriter &rewriter, Location loc,
+ Value a, Value b) {
+ auto bType = cast<RankedTensorType>(b.getType());
+ auto rankReducedBType =
+ RankedTensorType::get(bType.getShape().drop_front(),
+ bType.getElementType(), bType.getEncoding());
+
+ SmallVector<OpFoldResult> offsets(3, rewriter.getIndexAttr(0));
+ SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
+ tensor::getMixedSize(rewriter, loc, b, 1),
+ tensor::getMixedSize(rewriter, loc, b, 2)};
+ SmallVector<OpFoldResult> strides(3, rewriter.getIndexAttr(1));
+ Value rankReducedB = tensor::ExtractSliceOp::create(
+ rewriter, loc, rankReducedBType, b, offsets, sizes, strides);
+
+ SmallVector<OpFoldResult> broadcastShape = {
+ tensor::getMixedSize(rewriter, loc, a, 0),
+ tensor::getMixedSize(rewriter, loc, b, 1),
+ tensor::getMixedSize(rewriter, loc, b, 2)};
+ Value broadcastInit =
+ tensor::EmptyOp::create(rewriter, loc, broadcastShape,
+ bType.getElementType(), bType.getEncoding());
+ return linalg::BroadcastOp::create(rewriter, loc, rankReducedB, broadcastInit,
+ ArrayRef<int64_t>{0})
+ .getResult()
+ .front();
+}
+
+// Materialize MATMUL_T's optional B batch broadcast. When the B batch is
+// dynamic, select the broadcasted value at runtime if its size is one.
+static Value normalizeMatMulTRhsBatch(PatternRewriter &rewriter, Location loc,
+ Value a, Value b) {
+ auto aType = cast<RankedTensorType>(a.getType());
+ auto bType = cast<RankedTensorType>(b.getType());
+ const int64_t aBatch = aType.getDimSize(0);
+ const int64_t bBatch = bType.getDimSize(0);
+
+ if (bBatch == 1 && aBatch != 1)
+ return broadcastMatMulTRhs(rewriter, loc, a, b);
+
+ if (ShapedType::isStatic(bBatch) || aBatch == 1)
+ return b;
+
+ Value bBatchSize = tensor::DimOp::create(rewriter, loc, b, 0);
+ Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ Value requiresBroadcast = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, bBatchSize, one);
+
+ auto thenBuilder = [&](OpBuilder &builder, Location nestedLoc) {
+ Value broadcastedB = broadcastMatMulTRhs(rewriter, nestedLoc, a, b);
+ Value castB =
+ builder.createOrFold<tensor::CastOp>(nestedLoc, bType, broadcastedB);
+ scf::YieldOp::create(builder, nestedLoc, castB);
+ };
+ auto elseBuilder = [&](OpBuilder &builder, Location nestedLoc) {
+ scf::YieldOp::create(builder, nestedLoc, b);
+ };
+ return scf::IfOp::create(rewriter, loc, requiresBroadcast, thenBuilder,
+ elseBuilder)
+ .getResult(0);
+}
+
+class MatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
+public:
+ using OpConversionPattern<tosa::MatMulTOp>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(tosa::MatMulTOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const final {
+ Location loc = op.getLoc();
+ auto aType = dyn_cast<RankedTensorType>(adaptor.getA().getType());
+ auto bType = dyn_cast<RankedTensorType>(adaptor.getB().getType());
+ auto outputType = dyn_cast<RankedTensorType>(op.getType());
+ if (!aType || !bType || !outputType)
+ return rewriter.notifyMatchFailure(op, "only supports ranked tensors");
+
+ if (aType.getRank() != 3 || bType.getRank() != 3)
+ return rewriter.notifyMatchFailure(
+ op, "only supports rank-3 A and B tensors");
+
+ const int64_t aBatch = aType.getDimSize(0);
+ const int64_t bBatch = bType.getDimSize(0);
+ if (aBatch == 1 && ShapedType::isStatic(bBatch) && bBatch != 1)
+ return rewriter.notifyMatchFailure(
+ op, "does not support broadcasting the A matrix");
+
+ FailureOr<int64_t> maybeAZp = op.getAZeroPoint();
+ FailureOr<int64_t> maybeBZp = op.getBZeroPoint();
+ if (failed(maybeAZp))
+ return rewriter.notifyMatchFailure(
+ op, "input a zero point cannot be statically determined");
+ if (failed(maybeBZp))
+ return rewriter.notifyMatchFailure(
+ op, "input b zero point cannot be statically determined");
+
+ const int64_t aZpVal = *maybeAZp;
+ const int64_t bZpVal = *maybeBZp;
+ if (op.verifyAZeroPoint(aZpVal).failed())
+ return rewriter.notifyMatchFailure(
+ op, "input a zero point must be zero for non-int8 integer types");
+ if (op.verifyBZeroPoint(bZpVal).failed())
+ return rewriter.notifyMatchFailure(
+ op, "input b zero point must be zero for non-int8 integer types");
+
+ Value b =
+ normalizeMatMulTRhsBatch(rewriter, loc, adaptor.getA(), adaptor.getB());
+
+ SmallVector<Value> dynamicDims(3);
+ if (outputType.isDynamicDim(0))
+ dynamicDims[0] = tensor::DimOp::create(rewriter, loc, adaptor.getA(), 0);
+ if (outputType.isDynamicDim(1))
+ dynamicDims[1] = tensor::DimOp::create(rewriter, loc, adaptor.getA(), 1);
+ if (outputType.isDynamicDim(2))
+ dynamicDims[2] = tensor::DimOp::create(rewriter, loc, adaptor.getB(), 1);
+
+ Value zero = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getZeroAttr(outputType.getElementType()));
+ Value emptyTensor = tensor::EmptyOp::create(
+ rewriter, loc, outputType.getShape(), outputType.getElementType(),
+ condenseValues(dynamicDims));
+ Value zeroTensor = linalg::FillOp::create(rewriter, loc, ValueRange{zero},
+ ValueRange{emptyTensor})
+ .result();
+
+ Type outputElementType = outputType.getElementType();
+ if (!isa<IntegerType>(outputElementType)) {
+ rewriter.replaceOpWithNewOp<linalg::BatchMatmulTransposeBOp>(
+ op, TypeRange{outputType}, ValueRange{adaptor.getA(), b},
+ ValueRange{zeroTensor});
+ return success();
+ }
+
+ AffineExpr batch, height, width, channels;
+ bindDims(rewriter.getContext(), batch, height, width, channels);
+ SmallVector<AffineMap> indexingMaps = {
+ AffineMap::get(4, 0, {batch, height, channels}, rewriter.getContext()),
+ AffineMap::get(4, 0, {batch, width, channels}, rewriter.getContext()),
+ AffineMap::get(4, 0, {batch, height, width}, rewriter.getContext())};
+ SmallVector<utils::IteratorType> iteratorTypes = {
+ utils::IteratorType::parallel, utils::IteratorType::parallel,
+ utils::IteratorType::parallel, utils::IteratorType::reduction};
+
+ auto generic = linalg::GenericOp::create(
+ rewriter, loc, outputType, ValueRange{adaptor.getA(), b}, zeroTensor,
+ indexingMaps, iteratorTypes,
+ [=](OpBuilder &builder, Location nestedLoc, ValueRange args) {
+ auto extendToOutputType = [&](Value value) {
+ if (value.getType() == outputElementType)
+ return value;
+ return arith::ExtSIOp::create(builder, nestedLoc, outputElementType,
+ value)
+ .getResult();
+ };
+
+ Value aValue = extendToOutputType(args[0]);
+ Value bValue = extendToOutputType(args[1]);
+ Value aZp = arith::ConstantOp::create(
+ builder, nestedLoc,
+ builder.getIntegerAttr(outputElementType, aZpVal));
+ Value bZp = arith::ConstantOp::create(
+ builder, nestedLoc,
+ builder.getIntegerAttr(outputElementType, bZpVal));
+ Value adjustedA =
+ arith::SubIOp::create(builder, nestedLoc, aValue, aZp);
+ Value adjustedB =
+ arith::SubIOp::create(builder, nestedLoc, bValue, bZp);
+ Value product =
+ arith::MulIOp::create(builder, nestedLoc, adjustedA, adjustedB);
+ Value sum =
+ arith::AddIOp::create(builder, nestedLoc, args[2], product);
+ linalg::YieldOp::create(builder, nestedLoc, sum);
+ });
+ rewriter.replaceOp(op, generic.getResult(0));
+ return success();
+ }
+};
+
class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
public:
using OpConversionPattern::OpConversionPattern;
@@ -1134,6 +1314,7 @@ void mlir::tosa::populateTosaToLinalgNamedConversionPatterns(
ConvConverter<tosa::Conv3DOp, linalg::Conv3DNdhwcDhwcfOp, linalg::Conv3DNdhwcDhwcfQOp>,
DepthwiseConvConverter,
MatMulConverter,
+ MatMulTConverter,
AvgPool2dConverter,
TransposeConverter
>(patterns->getContext());
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
index df066a5cac261..db84f9ae6cfe1 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -60,6 +60,7 @@ struct TosaToLinalgNamed
target.addIllegalOp<tosa::MaxPool2dOp>();
target.addIllegalOp<tosa::AvgPool2dOp>();
target.addIllegalOp<tosa::MatMulOp>();
+ target.addIllegalOp<tosa::MatMulTOp>();
target.addIllegalOp<tosa::TransposeOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index 801144c01829a..f339caf56d5cd 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -95,6 +95,111 @@ func.func @matmul_dyn_output(%arg0: tensor<1x1x8xf32>, %arg1: tensor<1x8x1xf32>)
// -----
+// CHECK-LABEL: @matmul_t
+func.func @matmul_t(%arg0: tensor<2x5x3xf32>, %arg1: tensor<2x6x3xf32>) -> tensor<2x5x6xf32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<2x5x6xf32>) -> tensor<2x5x6xf32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xf32>, tensor<2x6x3xf32>) outs(%[[FILLED]] : tensor<2x5x6xf32>) -> tensor<2x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xf32>, tensor<2x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x5x6xf32>
+ return %0 : tensor<2x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_f16_fp8e4m3_to_bf16
+func.func @matmul_t_f16_fp8e4m3_to_bf16(%arg0: tensor<2x5x3xf16>, %arg1: tensor<2x6x3xf8E4M3FN>) -> tensor<2x5x6xbf16> {
+ // CHECK: %[[FILLED:.+]] = linalg.fill {{.*}} -> tensor<2x5x6xbf16>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xf16>, tensor<2x6x3xf8E4M3FN>) outs(%[[FILLED]] : tensor<2x5x6xbf16>) -> tensor<2x5x6xbf16>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E4M3FN>}> : () -> tensor<1xf8E4M3FN>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xf16>, tensor<2x6x3xf8E4M3FN>, tensor<1xf16>, tensor<1xf8E4M3FN>) -> tensor<2x5x6xbf16>
+ return %0 : tensor<2x5x6xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_fp8e5m2_f16_to_bf16
+func.func @matmul_t_fp8e5m2_f16_to_bf16(%arg0: tensor<2x5x3xf8E5M2>, %arg1: tensor<2x6x3xf16>) -> tensor<2x5x6xbf16> {
+ // CHECK: %[[FILLED:.+]] = linalg.fill {{.*}} -> tensor<2x5x6xbf16>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xf8E5M2>, tensor<2x6x3xf16>) outs(%[[FILLED]] : tensor<2x5x6xbf16>) -> tensor<2x5x6xbf16>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E5M2>}> : () -> tensor<1xf8E5M2>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xf8E5M2>, tensor<2x6x3xf16>, tensor<1xf8E5M2>, tensor<1xf16>) -> tensor<2x5x6xbf16>
+ return %0 : tensor<2x5x6xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_broadcast
+func.func @matmul_t_broadcast(%arg0: tensor<4x5x3xf32>, %arg1: tensor<1x6x3xf32>) -> tensor<4x5x6xf32> {
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3] [1, 1, 1] : tensor<1x6x3xf32> to tensor<6x3xf32>
+ // CHECK: %[[BCAST_INIT:.+]] = tensor.empty() : tensor<4x6x3xf32>
+ // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xf32>) outs(%[[BCAST_INIT]] : tensor<4x6x3xf32>) dimensions = [0]
+ // CHECK: %[[C0:.+]] = arith.constant 0
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %[[BCAST]] : tensor<4x5x3xf32>, tensor<4x6x3xf32>) outs(%[[FILLED]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<4x5x3xf32>, tensor<1x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x5x6xf32>
+ return %0 : tensor<4x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_dynamic_batch
+func.func @matmul_t_dynamic_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<?x6x3xf32>) -> tensor<?x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[ONE:.+]] = arith.constant 1 : index
+ // CHECK: %[[NEEDS_BCAST:.+]] = arith.cmpi eq, %[[BATCH]], %[[ONE]]
+ // CHECK: %[[B:.+]] = scf.if %[[NEEDS_BCAST]] -> (tensor<?x6x3xf32>) {
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3]
+ // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xf32>)
+ // CHECK: scf.yield %[[BCAST]] : tensor<?x6x3xf32>
+ // CHECK: } else {
+ // CHECK: scf.yield %arg1 : tensor<?x6x3xf32>
+ // CHECK: }
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %[[B]] : tensor<?x5x3xf32>, tensor<?x6x3xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<?x5x3xf32>, tensor<?x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x5x6xf32>
+ return %0 : tensor<?x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_zero_point
+func.func @matmul_t_zero_point(%arg0: tensor<4x5x3xi8>, %arg1: tensor<1x6x3xi8>) -> tensor<4x5x6xi32> {
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3] [1, 1, 1] : tensor<1x6x3xi8> to tensor<6x3xi8>
+ // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xi8>)
+ // CHECK: %[[RESULT:.+]] = linalg.generic
+ // CHECK-SAME: ins(%arg0, %[[BCAST]] : tensor<4x5x3xi8>, tensor<4x6x3xi8>)
+ // CHECK: ^bb0(%[[A:.+]]: i8, %[[BVAL:.+]]: i8, %[[ACC:.+]]: i32):
+ // CHECK: %[[EXT_A:.+]] = arith.extsi %[[A]] : i8 to i32
+ // CHECK: %[[EXT_B:.+]] = arith.extsi %[[BVAL]] : i8 to i32
+ // CHECK: %[[AZP:.+]] = arith.constant 1 : i32
+ // CHECK: %[[BZP:.+]] = arith.constant 2 : i32
+ // CHECK: %[[ADJUSTED_A:.+]] = arith.subi %[[EXT_A]], %[[AZP]] : i32
+ // CHECK: %[[ADJUSTED_B:.+]] = arith.subi %[[EXT_B]], %[[BZP]] : i32
+ // CHECK: %[[PRODUCT:.+]] = arith.muli %[[ADJUSTED_A]], %[[ADJUSTED_B]] : i32
+ // CHECK: %[[SUM:.+]] = arith.addi %[[ACC]], %[[PRODUCT]] : i32
+ // CHECK: linalg.yield %[[SUM]] : i32
+ %a_zp = "tosa.const"() <{values = dense<1> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<2> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<4x5x3xi8>, tensor<1x6x3xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<4x5x6xi32>
+ return %0 : tensor<4x5x6xi32>
+}
+
+// -----
+
// CHECK-LABEL: @max_pool
func.func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () {
// CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38
More information about the Mlir-commits
mailing list