[Mlir-commits] [mlir] [mlir][tosa] Lower MATMUL_T to Linalg (PR #222661)
Jack Frankland
llvmlistbot at llvm.org
Tue Oct 6 03:08:13 PDT 2026
https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/222661
>From 0a62d75f63fec5c44980ba8a28e5d0ff1df1ae2a 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] [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/TosaToLinalg.h | 6 +
.../Conversion/TosaToLinalg/CMakeLists.txt | 1 +
.../Conversion/TosaToLinalg/TosaToLinalg.cpp | 9 +
.../TosaToLinalg/TosaToLinalgNamed.cpp | 267 +++++++++++++++
.../TosaToLinalg/TosaToLinalgNamedPass.cpp | 1 +
.../TosaToLinalg/tosa-to-linalg-named.mlir | 322 ++++++++++++++++++
6 files changed, 606 insertions(+)
diff --git a/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h b/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
index 06914be338aaf..b4b610c8a0db7 100644
--- a/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
+++ b/mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
@@ -46,6 +46,12 @@ void addTosaToLinalgPasses(
/// Currently, this includes only the "tosa-to-linalg-pipeline".
void registerTosaToLinalgPipelines();
+/// Broadcast a dynamically sized dimension of `operand` to `targetSize` when
+/// its runtime size is one. Otherwise, return `operand` unchanged.
+Value broadcastDynamicDimension(PatternRewriter &rewriter, Location loc,
+ Value operand, int64_t dim,
+ OpFoldResult targetSize);
+
/// Populates conversion passes from TOSA dialect to Linalg dialect.
void populateTosaToLinalgConversionPatterns(
const TypeConverter &converter, RewritePatternSet *patterns,
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/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index b2a8360ee0c2b..bc08c59730529 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -949,6 +949,15 @@ static Value broadcastDynamicDimension(PatternRewriter &rewriter, Location loc,
return ifOp.getResult(0);
}
+Value mlir::tosa::broadcastDynamicDimension(PatternRewriter &rewriter,
+ Location loc, Value operand,
+ int64_t dim,
+ OpFoldResult targetSize) {
+ IndexPool indexPool;
+ return ::broadcastDynamicDimension(rewriter, loc, indexPool, operand, dim,
+ targetSize, /*masterOperand=*/nullptr);
+}
+
static Value broadcastDynamicDimensions(PatternRewriter &rewriter, Location loc,
IndexPool &indexPool, Value operand,
ArrayRef<OpFoldResult> targetShape,
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index ae5390b36ddc7..ecd1eae7380d0 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -12,6 +12,7 @@
#include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -679,6 +680,271 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
}
};
+// Insert missing leading batch dimensions, taking their sizes from `other`.
+static Value broadcastMatMulTLeadingDimensions(PatternRewriter &rewriter,
+ Location loc, Value operand,
+ Value other) {
+ auto type = cast<RankedTensorType>(operand.getType());
+ int64_t missingDims =
+ cast<RankedTensorType>(other.getType()).getRank() - type.getRank();
+ if (missingDims <= 0)
+ return operand;
+
+ SmallVector<OpFoldResult> shape;
+ SmallVector<int64_t> dimensions;
+ for (int64_t dim = 0; dim < missingDims; ++dim) {
+ shape.push_back(tensor::getMixedSize(rewriter, loc, other, dim));
+ dimensions.push_back(dim);
+ }
+ llvm::append_range(shape, tensor::getMixedSizes(rewriter, loc, operand));
+ Value empty = tensor::EmptyOp::create(
+ rewriter, loc, shape, type.getElementType(), type.getEncoding());
+ return linalg::BroadcastOp::create(rewriter, loc, operand, empty, dimensions)
+ .getResult()
+ .front();
+}
+
+// Replace a known singleton batch dimension with `size`.
+static Value broadcastMatMulTSingletonDimension(PatternRewriter &rewriter,
+ Location loc, Value operand,
+ int64_t dim,
+ OpFoldResult size) {
+ auto type = cast<RankedTensorType>(operand.getType());
+ SmallVector<int64_t> sliceShape(type.getShape());
+ sliceShape.erase(sliceShape.begin() + dim);
+ auto sliceType = RankedTensorType::get(sliceShape, type.getElementType(),
+ type.getEncoding());
+ Value slice = tensor::createCanonicalRankReducingExtractSliceOp(
+ rewriter, loc, operand, sliceType);
+
+ SmallVector<OpFoldResult> shape =
+ tensor::getMixedSizes(rewriter, loc, operand);
+ shape[dim] = size;
+ Value empty = tensor::EmptyOp::create(
+ rewriter, loc, shape, type.getElementType(), type.getEncoding());
+ return linalg::BroadcastOp::create(rewriter, loc, slice, empty,
+ ArrayRef<int64_t>{dim})
+ .getResult()
+ .front();
+}
+
+// Align the ranks, then broadcast each shared batch dimension independently.
+static std::pair<Value, Value>
+normalizeMatMulTBatchDimensions(PatternRewriter &rewriter, Location loc,
+ Value a, Value b, int64_t batchRank) {
+ int64_t aRank = cast<RankedTensorType>(a.getType()).getRank();
+ int64_t bRank = cast<RankedTensorType>(b.getType()).getRank();
+ int64_t firstSharedDim = std::max(aRank, bRank) - std::min(aRank, bRank);
+ a = broadcastMatMulTLeadingDimensions(rewriter, loc, a, b);
+ b = broadcastMatMulTLeadingDimensions(rewriter, loc, b, a);
+
+ // The inserted leading dimensions already match their partner's sizes.
+ for (int64_t dim = firstSharedDim; dim < batchRank; ++dim) {
+ auto aType = cast<RankedTensorType>(a.getType());
+ auto bType = cast<RankedTensorType>(b.getType());
+ int64_t aSize = aType.getDimSize(dim);
+ int64_t bSize = bType.getDimSize(dim);
+ if (aSize == 1 && bSize == 1)
+ continue;
+ if (aSize == 1) {
+ a = broadcastMatMulTSingletonDimension(
+ rewriter, loc, a, dim, tensor::getMixedSize(rewriter, loc, b, dim));
+ continue;
+ }
+ if (bSize == 1) {
+ b = broadcastMatMulTSingletonDimension(
+ rewriter, loc, b, dim, tensor::getMixedSize(rewriter, loc, a, dim));
+ continue;
+ }
+ if (!aType.isDynamicDim(dim) && !bType.isDynamicDim(dim))
+ continue;
+
+ OpFoldResult size;
+ if (!aType.isDynamicDim(dim)) {
+ size = rewriter.getIndexAttr(aSize);
+ } else if (!bType.isDynamicDim(dim)) {
+ size = rewriter.getIndexAttr(bSize);
+ } else {
+ Value aDim = tensor::DimOp::create(rewriter, loc, a, dim);
+ Value bDim = tensor::DimOp::create(rewriter, loc, b, dim);
+ Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ Value aIsSingleton = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, aDim, one);
+ // If A is a singleton, B determines the size; otherwise A does.
+ size = arith::SelectOp::create(rewriter, loc, aIsSingleton, bDim, aDim)
+ .getResult();
+ }
+ a = broadcastDynamicDimension(rewriter, loc, a, dim, size);
+ b = broadcastDynamicDimension(rewriter, loc, b, dim, size);
+ }
+ return {a, b};
+}
+
+static SmallVector<ReassociationIndices>
+getMatMulTBatchCollapseReassociation(int64_t rank) {
+ SmallVector<ReassociationIndices> reassociation(3);
+ for (int64_t dim = 0; dim < rank - 2; ++dim)
+ reassociation[0].push_back(dim);
+ reassociation[1].push_back(rank - 2);
+ reassociation[2].push_back(rank - 1);
+ return reassociation;
+}
+
+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");
+
+ 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;
+
+ const int64_t outputRank = outputType.getRank();
+ const int64_t batchRank = outputRank - 2;
+ auto [a, b] = normalizeMatMulTBatchDimensions(rewriter, loc, adaptor.getA(),
+ adaptor.getB(), batchRank);
+
+ SmallVector<OpFoldResult> outputShape;
+ outputShape.reserve(outputRank);
+ for (int64_t batchDim = 0; batchDim < batchRank; ++batchDim)
+ outputShape.push_back(tensor::getMixedSize(rewriter, loc, a, batchDim));
+ outputShape.push_back(
+ tensor::getMixedSize(rewriter, loc, a, outputRank - 2));
+ outputShape.push_back(
+ tensor::getMixedSize(rewriter, loc, b, outputRank - 2));
+ SmallVector<Value> dynamicDims;
+ for (auto [dim, size] : llvm::enumerate(outputShape))
+ if (outputType.isDynamicDim(dim))
+ dynamicDims.push_back(
+ getValueOrCreateConstantIndexOp(rewriter, loc, size));
+
+ 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 (aZpVal == 0 && bZpVal == 0) {
+ if (outputRank == 2) {
+ rewriter.replaceOpWithNewOp<linalg::MatmulTransposeBOp>(
+ op, TypeRange{outputType}, ValueRange{a, b},
+ ValueRange{zeroTensor});
+ return success();
+ }
+
+ if (outputRank == 3) {
+ rewriter.replaceOpWithNewOp<linalg::BatchMatmulTransposeBOp>(
+ op, TypeRange{outputType}, ValueRange{a, b},
+ ValueRange{zeroTensor});
+ return success();
+ }
+
+ SmallVector<ReassociationIndices> reassociation =
+ getMatMulTBatchCollapseReassociation(outputRank);
+ auto collapse = [&](Value value) {
+ auto type = cast<RankedTensorType>(value.getType());
+ auto collapsedType =
+ tensor::CollapseShapeOp::inferCollapsedType(type, reassociation);
+ return tensor::CollapseShapeOp::create(rewriter, loc, collapsedType,
+ value, reassociation);
+ };
+ Value collapsedA = collapse(a);
+ Value collapsedB = collapse(b);
+ Value collapsedInit = collapse(zeroTensor);
+ auto collapsedOutputType =
+ cast<RankedTensorType>(collapsedInit.getType());
+ Value collapsedResult =
+ linalg::BatchMatmulTransposeBOp::create(
+ rewriter, loc, TypeRange{collapsedOutputType},
+ ValueRange{collapsedA, collapsedB}, ValueRange{collapsedInit})
+ .getResult(0);
+ Value result = tensor::ExpandShapeOp::create(
+ rewriter, loc, outputType, collapsedResult, reassociation,
+ tensor::getMixedSizes(rewriter, loc, zeroTensor));
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+
+ int64_t loopRank = outputRank + 1;
+ SmallVector<AffineExpr> dimensions =
+ llvm::map_to_vector(llvm::seq<int64_t>(0, loopRank), [&](int64_t dim) {
+ return rewriter.getAffineDimExpr(dim);
+ });
+ SmallVector<AffineExpr> aMap(dimensions.begin(),
+ dimensions.begin() + batchRank);
+ aMap.push_back(dimensions[batchRank]);
+ aMap.push_back(dimensions[batchRank + 2]);
+ SmallVector<AffineExpr> bMap(dimensions.begin(),
+ dimensions.begin() + batchRank);
+ bMap.push_back(dimensions[batchRank + 1]);
+ bMap.push_back(dimensions[batchRank + 2]);
+ SmallVector<AffineExpr> outputMap(dimensions.begin(),
+ dimensions.begin() + batchRank);
+ outputMap.push_back(dimensions[batchRank]);
+ outputMap.push_back(dimensions[batchRank + 1]);
+ SmallVector<AffineMap> indexingMaps = {
+ AffineMap::get(loopRank, 0, aMap, rewriter.getContext()),
+ AffineMap::get(loopRank, 0, bMap, rewriter.getContext()),
+ AffineMap::get(loopRank, 0, outputMap, rewriter.getContext())};
+ SmallVector<utils::IteratorType> iteratorTypes(
+ loopRank, utils::IteratorType::parallel);
+ iteratorTypes.back() = utils::IteratorType::reduction;
+
+ auto generic = linalg::GenericOp::create(
+ rewriter, loc, outputType, ValueRange{a, 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;
@@ -1146,6 +1412,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 780901eab90ee..80f548877b206 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -62,6 +62,7 @@ struct TosaToLinalgNamed
target.addIllegalOp<tosa::AvgPool2dOp>();
target.addIllegalOp<tosa::AvgPool2dAdaptiveOp>();
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 2034db67e9e3d..9ccbf65aa61f5 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -95,6 +95,328 @@ 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_broadcast_both_rank4
+func.func @matmul_t_broadcast_both_rank4(%arg0: tensor<2x1x5x3xf32>, %arg1: tensor<1x3x6x3xf32>) -> tensor<2x3x5x6xf32> {
+ // CHECK: %[[B_SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0, 0] [1, 3, 6, 3] [1, 1, 1, 1] : tensor<1x3x6x3xf32> to tensor<3x6x3xf32>
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%[[B_SLICE]] : tensor<3x6x3xf32>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_SLICE:.+]] = tensor.extract_slice %arg0[0, 0, 0, 0] [2, 1, 5, 3] [1, 1, 1, 1] : tensor<2x1x5x3xf32> to tensor<2x5x3xf32>
+ // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%[[A_SLICE]] : tensor<2x5x3xf32>) {{.*}} dimensions = [1]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x5x3xf32> into tensor<6x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x6x3xf32> into tensor<6x6x3xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<6x5x3xf32>, tensor<6x6x3xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [2, 3, 5, 6] : tensor<6x5x6xf32> into tensor<2x3x5x6xf32>
+ %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<2x1x5x3xf32>, tensor<1x3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x5x6xf32>
+ return %0 : tensor<2x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_a_missing_batch_dim
+func.func @matmul_t_a_missing_batch_dim(%arg0: tensor<3x5x3xf32>, %arg1: tensor<4x3x6x3xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%arg0 : tensor<3x5x3xf32>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [4, 3, 5, 6] : tensor<12x5x6xf32> into tensor<4x3x5x6xf32>
+ %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<3x5x3xf32>, tensor<4x3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_b_missing_batch_dim
+func.func @matmul_t_b_missing_batch_dim(%arg0: tensor<4x3x5x3xf32>, %arg1: tensor<3x6x3xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%arg1 : tensor<3x6x3xf32>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [4, 3, 5, 6] : tensor<12x5x6xf32> into tensor<4x3x5x6xf32>
+ %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<4x3x5x3xf32>, tensor<3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_rank2_a
+func.func @matmul_t_rank2_a(%arg0: tensor<5x3xf32>, %arg1: tensor<4x3x6x3xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%arg0 : tensor<5x3xf32>) {{.*}} dimensions = [0, 1]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+ %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<5x3xf32>, tensor<4x3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_rank2_b
+func.func @matmul_t_rank2_b(%arg0: tensor<4x3x5x3xf32>, %arg1: tensor<6x3xf32>) -> tensor<4x3x5x6xf32> {
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%arg1 : tensor<6x3xf32>) {{.*}} dimensions = [0, 1]
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+ %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<4x3x5x3xf32>, tensor<6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+ return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_rank2
+func.func @matmul_t_rank2(%arg0: tensor<5x3xf32>, %arg1: tensor<6x3xf32>) -> tensor<5x6xf32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<5x6xf32>) -> tensor<5x6xf32>
+ // CHECK: linalg.matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<5x3xf32>, tensor<6x3xf32>) outs(%[[FILLED]] : tensor<5x6xf32>) -> tensor<5x6xf32>
+ %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<5x3xf32>, tensor<6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5x6xf32>
+ return %0 : tensor<5x6xf32>
+}
+
+// -----
+
+// CHECK-DAG: #[[$BCAST_MAP:.+]] = affine_map<(d0, d1, d2) -> (0, d1, d2)>
+// CHECK-DAG: #[[$IDENTITY_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: @matmul_t_dynamic_batch
+func.func @matmul_t_dynamic_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<?x6x3xf32>) -> tensor<?x5x6xf32> {
+ // CHECK: %[[A_BATCH:.+]] = tensor.dim %arg0, %{{.+}}
+ // CHECK: %[[B_BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[ONE:.+]] = arith.constant 1 : index
+ // CHECK: %[[A_IS_ONE:.+]] = arith.cmpi eq, %[[A_BATCH]], %[[ONE]] : index
+ // CHECK: %[[TARGET_BATCH:.+]] = arith.select %[[A_IS_ONE]], %[[B_BATCH]], %[[A_BATCH]] : index
+ // CHECK: %[[A:.+]] = scf.if %{{.+}} -> (tensor<?x5x3xf32>) {
+ // CHECK: %[[A_BCAST_INIT:.+]] = tensor.empty(%[[TARGET_BATCH]]) : tensor<?x5x3xf32>
+ // CHECK: %[[A_BCAST:.+]] = linalg.generic {indexing_maps = [#[[$BCAST_MAP]], #[[$IDENTITY_MAP]]], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0 : tensor<?x5x3xf32>) outs(%[[A_BCAST_INIT]] : tensor<?x5x3xf32>)
+ // CHECK: scf.yield %[[A_BCAST]] : tensor<?x5x3xf32>
+ // CHECK: } else
+ // CHECK: %[[B:.+]] = scf.if %{{.+}} -> (tensor<?x6x3xf32>) {
+ // CHECK: %[[B_BCAST_INIT:.+]] = tensor.empty(%[[TARGET_BATCH]]) : tensor<?x6x3xf32>
+ // CHECK: %[[B_BCAST:.+]] = linalg.generic {indexing_maps = [#[[$BCAST_MAP]], #[[$IDENTITY_MAP]]], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<?x6x3xf32>) outs(%[[B_BCAST_INIT]] : tensor<?x6x3xf32>)
+ // CHECK: scf.yield %[[B_BCAST]] : tensor<?x6x3xf32>
+ // CHECK: } else
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A]], %[[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>
+}
+
+// -----
+
+// A dynamic batch dimension may need broadcasting to a known static size.
+// CHECK-LABEL: @matmul_t_dynamic_static_batch
+func.func @matmul_t_dynamic_static_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<4x6x3xf32>) -> tensor<4x5x6xf32> {
+ // CHECK: %[[A:.+]] = scf.if %{{.+}} -> (tensor<?x5x3xf32>) {
+ // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x5x3xf32>
+ // CHECK: %[[BCAST:.+]] = linalg.generic {{.*}} ins(%arg0 : tensor<?x5x3xf32>) outs(%[[EMPTY]] : tensor<4x5x3xf32>)
+ // CHECK: %[[CAST:.+]] = tensor.cast %[[BCAST]] : tensor<4x5x3xf32> to tensor<?x5x3xf32>
+ // CHECK: scf.yield %[[CAST]] : tensor<?x5x3xf32>
+ // CHECK: } else {
+ // CHECK: scf.yield %arg0 : tensor<?x5x3xf32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A]], %arg1 : tensor<?x5x3xf32>, tensor<4x6x3xf32>)
+ %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<4x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x5x6xf32>
+ return %0 : tensor<4x5x6xf32>
+}
+
+// -----
+
+// The dynamic operand already determines the target size when its partner is 1.
+// CHECK-LABEL: @matmul_t_dynamic_singleton_batch
+func.func @matmul_t_dynamic_singleton_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<1x6x3xf32>) -> tensor<?x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %{{.+}}
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3] [1, 1, 1] : tensor<1x6x3xf32> to tensor<6x3xf32>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x6x3xf32>
+ // CHECK: %[[B:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xf32>) outs(%[[EMPTY]] : tensor<?x6x3xf32>) dimensions = [0]
+ // CHECK-NOT: scf.if
+ // 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<1x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x5x6xf32>
+ return %0 : tensor<?x5x6xf32>
+}
+
+// -----
+
+// Missing leading dimensions broadcast directly to the partner's runtime size.
+// CHECK-LABEL: @matmul_t_missing_dynamic_batch
+func.func @matmul_t_missing_dynamic_batch(%arg0: tensor<3x5x3xf32>, %arg1: tensor<?x3x6x3xf32>) -> tensor<?x3x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x3x5x3xf32>
+ // CHECK: %[[A:.+]] = linalg.broadcast ins(%arg0 : tensor<3x5x3xf32>) outs(%[[EMPTY]] : tensor<?x3x5x3xf32>) dimensions = [0]
+ // CHECK-NOT: scf.if
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A]] {{\[}}[0, 1], [2], [3]] : tensor<?x3x5x3xf32> into tensor<?x5x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<?x3x6x3xf32> into tensor<?x6x3xf32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : 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<3x5x3xf32>, tensor<?x3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x3x5x6xf32>
+ return %0 : tensor<?x3x5x6xf32>
+}
+
+// -----
+
+// Resolve dynamic batch dimensions one at a time while preserving matrix sizes.
+// CHECK-LABEL: @matmul_t_multiple_dynamic_batches
+func.func @matmul_t_multiple_dynamic_batches(%arg0: tensor<?x?x?x3xf32>, %arg1: tensor<?x?x6x3xf32>) -> tensor<?x?x?x6xf32> {
+ // CHECK: arith.select
+ // CHECK: %[[A0:.+]] = scf.if %{{.+}} -> (tensor<?x?x?x3xf32>)
+ // CHECK: %[[B0:.+]] = scf.if %{{.+}} -> (tensor<?x?x6x3xf32>)
+ // CHECK: arith.select
+ // CHECK: %[[A1:.+]] = scf.if %{{.+}} -> (tensor<?x?x?x3xf32>)
+ // CHECK: linalg.generic {{.*}} ins(%[[A0]] : tensor<?x?x?x3xf32>)
+ // CHECK: %[[B1:.+]] = scf.if %{{.+}} -> (tensor<?x?x6x3xf32>)
+ // CHECK: linalg.generic {{.*}} ins(%[[B0]] : tensor<?x?x6x3xf32>)
+ // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A1]] {{\[}}[0, 1], [2], [3]] : tensor<?x?x?x3xf32> into tensor<?x?x3xf32>
+ // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B1]] {{\[}}[0, 1], [2], [3]] : tensor<?x?x6x3xf32> into tensor<?x6x3xf32>
+ // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<?x?x3xf32>, tensor<?x6x3xf32>)
+ // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [%{{.+}}, %{{.+}}, %{{.+}}, 6] : tensor<?x?x6xf32> into tensor<?x?x?x6xf32>
+ %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<?x?x?x3xf32>, tensor<?x?x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x?x?x6xf32>
+ return %0 : tensor<?x?x?x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_i8_zero_point
+func.func @matmul_t_i8_zero_point(%arg0: tensor<2x5x3xi8>, %arg1: tensor<2x6x3xi8>) -> tensor<2x5x6xi32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : i32) outs(%[[INIT]] : tensor<2x5x6xi32>) -> tensor<2x5x6xi32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xi8>, tensor<2x6x3xi8>) outs(%[[FILLED]] : tensor<2x5x6xi32>) -> tensor<2x5x6xi32>
+ %a_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xi8>, tensor<2x6x3xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2x5x6xi32>
+ return %0 : tensor<2x5x6xi32>
+}
+
+// -----
+
+// 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-DAG: #[[$MATMUL_T_A_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d4)>
+// CHECK-DAG: #[[$MATMUL_T_B_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3, d4)>
+// CHECK-DAG: #[[$MATMUL_T_OUT_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)>
+// CHECK-LABEL: @matmul_t_zero_point_rank4
+func.func @matmul_t_zero_point_rank4(%arg0: tensor<2x1x5x3xi8>, %arg1: tensor<1x3x6x3xi8>) -> tensor<2x3x5x6xi32> {
+ // CHECK: %[[B_SLICE:.+]] = tensor.extract_slice %arg1
+ // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%[[B_SLICE]] : tensor<3x6x3xi8>) {{.*}} dimensions = [0]
+ // CHECK: %[[A_SLICE:.+]] = tensor.extract_slice %arg0
+ // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%[[A_SLICE]] : tensor<2x5x3xi8>) {{.*}} dimensions = [1]
+ // CHECK: linalg.generic {indexing_maps = [#[[$MATMUL_T_A_MAP]], #[[$MATMUL_T_B_MAP]], #[[$MATMUL_T_OUT_MAP]]], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction"]} ins(%[[A_BCAST]], %[[B_BCAST]] : tensor<2x3x5x3xi8>, tensor<2x3x6x3xi8>)
+ %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<2x1x5x3xi8>, tensor<1x3x6x3xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<2x3x5x6xi32>
+ return %0 : tensor<2x3x5x6xi32>
+}
+
+// -----
+
// 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