[Mlir-commits] [mlir] [mlir][tosa] Lower MATMUL_T to Linalg (PR #222661)
Jack Frankland
llvmlistbot at llvm.org
Wed Oct 14 02:30:33 PDT 2026
https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/222661
>From bc60ba706bb732e3603f5d0837dacfe2ecd7959b 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
Normalize arbitrary leading batch dimensions in both ranked operands,
broadcasting missing and singleton dimensions as needed, including
dynamic batch dimensions.
Use named transpose-B matmul operations when both zero points are zero,
collapsing and restoring multiple batch dimensions around the batched
operation. Lower operations with at least one nonzero zero point to
linalg.generic so zero-point adjustment is represented in the reduction.
Assisted-by: Codex
---
.../Conversion/TosaToLinalg/TosaToLinalg.cpp | 226 +++++++++++++
.../TosaToLinalg/TosaToLinalgNamed.cpp | 84 +++++
.../TosaToLinalg/TosaToLinalgNamedPass.cpp | 8 +
.../TosaToLinalg/TosaToLinalgUtils.h | 34 ++
.../TosaToLinalg/tosa-to-linalg-named.mlir | 296 ++++++++++++++++++
.../TosaToLinalg/tosa-to-linalg-pipeline.mlir | 22 ++
.../TosaToLinalg/tosa-to-linalg.mlir | 87 +++++
7 files changed, 757 insertions(+)
create mode 100644 mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 215b1da1a2782..25b4b360723b2 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -11,6 +11,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
+#include "TosaToLinalgUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
@@ -1430,7 +1431,230 @@ static LogicalResult reduceMatchAndRewriteHelper(OpTy op, uint64_t axis,
return success();
}
+// 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();
+}
+
+// Expand an existing singleton batch dimension to `size`, preserving the rank.
+// Missing dimensions are handled by broadcastMatMulTLeadingDimensions.
+// linalg.broadcast inserts dimensions, so first remove the singleton dimension
+// with a rank-reducing slice, then reinsert it at the target 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);
+
+ IndexPool indexPool;
+ // 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 == bSize && ShapedType::isStatic(aSize))
+ continue;
+
+ // Compute the shared broadcast target only for batch dimensions.
+ // The matrix dimensions may differ between A and B.
+ auto [size, masterOperand] =
+ computeTargetSize(rewriter, loc, indexPool, ValueRange{a, b}, dim);
+ if (aSize == 1) {
+ a = broadcastMatMulTSingletonDimension(rewriter, loc, a, dim, size);
+ continue;
+ }
+ if (bSize == 1) {
+ b = broadcastMatMulTSingletonDimension(rewriter, loc, b, dim, size);
+ continue;
+ }
+ a = broadcastDynamicDimension(rewriter, loc, indexPool, a, dim, size,
+ masterOperand);
+ b = broadcastDynamicDimension(rewriter, loc, indexPool, b, dim, size,
+ masterOperand);
+ }
+ return {a, b};
+}
+
+MatMulTOperands
+mlir::tosa::prepareMatMulTOperands(PatternRewriter &rewriter, Location loc,
+ Value a, Value b,
+ RankedTensorType outputType) {
+ const int64_t outputRank = outputType.getRank();
+ const int64_t batchRank = outputRank - 2;
+ auto normalizedOperands =
+ normalizeMatMulTBatchDimensions(rewriter, loc, a, b, batchRank);
+ a = normalizedOperands.first;
+ b = normalizedOperands.second;
+
+ 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();
+
+ return {a, b, zeroTensor};
+}
+
namespace {
+// Linalg has quantized matmul and batch matmul operations, but neither has a
+// transpose-B variant. Lower nonzero zero points separately to linalg.generic
+// so the reduction can accumulate (A - aZp) * (B - bZp).
+class QuantizedMatMulTConverter : 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;
+
+ if (aZpVal == 0 && bZpVal == 0)
+ return rewriter.notifyMatchFailure(
+ op, "requires at least one nonzero zero point");
+
+ const int64_t outputRank = outputType.getRank();
+ const int64_t batchRank = outputRank - 2;
+ auto [a, b, zeroTensor] = prepareMatMulTOperands(
+ rewriter, loc, adaptor.getA(), adaptor.getB(), outputType);
+ Type outputElementType = outputType.getElementType();
+
+ 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();
+ }
+};
template <typename SrcOp>
class PointwiseConverter : public OpConversionPattern<SrcOp> {
@@ -3250,6 +3474,8 @@ void mlir::tosa::populateTosaToLinalgConversionPatterns(
PointwiseConverter<tosa::SigmoidOp>
>(converter, patterns->getContext());
+ patterns->add<QuantizedMatMulTConverter>(converter, patterns->getContext());
+
patterns->add<
IdentityNConverter<tosa::IdentityOp>,
GatherConverter,
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index 2308f22dc3118..a7f7414bf56ed 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -10,8 +10,10 @@
//
//===----------------------------------------------------------------------===//
+#include "TosaToLinalgUtils.h"
#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 +681,87 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
}
};
+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;
+
+ if (aZpVal != 0 || bZpVal != 0)
+ return rewriter.notifyMatchFailure(
+ op, "nonzero zero points require the generic MATMUL_T lowering");
+
+ const int64_t outputRank = outputType.getRank();
+ auto [a, b, zeroTensor] = prepareMatMulTOperands(
+ rewriter, loc, adaptor.getA(), adaptor.getB(), outputType);
+
+ 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();
+ }
+};
+
class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
public:
MaxPool2dConverter(const TypeConverter &typeConverter, MLIRContext *context,
@@ -1157,6 +1240,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 6239c7dc62fd6..f2771a084f344 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -62,6 +62,14 @@ struct TosaToLinalgNamed
target.addIllegalOp<tosa::AvgPool2dOp>();
target.addIllegalOp<tosa::AvgPool2dAdaptiveOp>();
target.addIllegalOp<tosa::MatMulOp>();
+ // Linalg has a quantized batch matmul, but no quantized transpose-B
+ // variant. Defer nonzero zero points to the generic pass's adjusted
+ // contraction.
+ target.addDynamicallyLegalOp<tosa::MatMulTOp>([](tosa::MatMulTOp op) {
+ FailureOr<int64_t> aZp = op.getAZeroPoint();
+ FailureOr<int64_t> bZp = op.getBZeroPoint();
+ return succeeded(aZp) && succeeded(bZp) && (*aZp != 0 || *bZp != 0);
+ });
target.addIllegalOp<tosa::TransposeOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h
new file mode 100644
index 0000000000000..2bc56c043a57b
--- /dev/null
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h
@@ -0,0 +1,34 @@
+//===- TosaToLinalgUtils.h - Shared TOSA to Linalg lowering utilities
+//-------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_LIB_CONVERSION_TOSATOLINALG_TOSATOLINALGUTILS_H
+#define MLIR_LIB_CONVERSION_TOSATOLINALG_TOSATOLINALGUTILS_H
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir {
+class PatternRewriter;
+
+namespace tosa {
+struct MatMulTOperands {
+ Value a;
+ Value b;
+ Value init;
+};
+
+/// Broadcast the batch dimensions of the ranked operands and create a
+/// zero-filled result tensor for either MATMUL_T lowering pattern.
+MatMulTOperands prepareMatMulTOperands(PatternRewriter &rewriter, Location loc,
+ Value a, Value b,
+ RankedTensorType outputType);
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_LIB_CONVERSION_TOSATOLINALG_TOSATOLINALGUTILS_H
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index 5b9b33d6476dd..4a052337335a7 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -96,6 +96,302 @@ 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: %[[TARGET_BATCH:.+]] = arith.maxui %[[A_BATCH]], %[[B_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>
+}
+
+// -----
+
+// The dynamic B operand determines the target size when A's batch size is 1.
+// CHECK-LABEL: @matmul_t_singleton_dynamic_batch
+func.func @matmul_t_singleton_dynamic_batch(%arg0: tensor<1x5x3xf32>, %arg1: tensor<?x6x3xf32>) -> tensor<?x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg0[0, 0, 0] [1, 5, 3] [1, 1, 1] : tensor<1x5x3xf32> to tensor<5x3xf32>
+ // CHECK: %[[EMPTY:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x5x3xf32>
+ // CHECK: %[[A:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<5x3xf32>) outs(%[[EMPTY]] : tensor<?x5x3xf32>) dimensions = [0]
+ // CHECK-NOT: scf.if
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%[[A]], %arg1 : 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<1x5x3xf32>, tensor<?x6x3xf32>, 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.maxui
+ // CHECK: %[[A0:.+]] = scf.if %{{.+}} -> (tensor<?x?x?x3xf32>)
+ // CHECK: %[[B0:.+]] = scf.if %{{.+}} -> (tensor<?x?x6x3xf32>)
+ // CHECK: arith.maxui
+ // 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: @max_pool
func.func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () {
// FINITE-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
index cd34715359666..93ca7ade0ce49 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
@@ -90,3 +90,25 @@ func.func @rescale_inexactround(%arg0: tensor<8x9x7x14xi32>) -> tensor<8x9x7x14x
%4 = tosa.rescale %arg0, %0, %1, %2, %3 scale32(true) rounding_mode<INEXACT_ROUND> per_channel(true) input_unsigned(false) output_unsigned(false) : (tensor<8x9x7x14xi32>, tensor<14xi32>, tensor<14xi8>, tensor<1xi32>, tensor<1xi8>) -> tensor<8x9x7x14xi8>
return %4 : tensor<8x9x7x14xi8>
}
+
+// -----
+
+// Check that the named pass defers nonzero MATMUL_T zero points to the generic
+// pass in the full pipeline.
+// CHECK-LABEL: func.func @matmul_t_nonzero_zero_points(
+func.func @matmul_t_nonzero_zero_points(%arg0: tensor<5x3xi8>, %arg1: tensor<6x3xi8>) -> tensor<5x6xi32> {
+ // CHECK: %[[RESULT:.+]] = linalg.generic
+ // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+ // CHECK: %[[EXT_A:.+]] = arith.extsi %{{.+}} : i8 to i32
+ // CHECK: %[[EXT_B:.+]] = arith.extsi %{{.+}} : 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: arith.muli %[[ADJUSTED_A]], %[[ADJUSTED_B]] : i32
+ // CHECK: return %[[RESULT]] : tensor<5x6xi32>
+ %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<5x3xi8>, tensor<6x3xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<5x6xi32>
+ return %0 : tensor<5x6xi32>
+}
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index 7031aa1172074..67a51e1e2e80f 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -2940,3 +2940,90 @@ func.func @argmax_f16(%arg0: tensor<6xf16>) -> () {
%0 = tosa.argmax %arg0 axis(0) : (tensor<6xf16>) -> tensor<i32>
return
}
+
+// -----
+
+// Nonzero zero points require a generic lowering because
+// linalg.quantized_batch_matmul has no transpose-B variant.
+
+// CHECK-DAG: #[[$MATMUL_T_RANK2_A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MATMUL_T_RANK2_B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$MATMUL_T_RANK2_OUT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+// CHECK-LABEL: @matmul_t_zero_point_rank2
+func.func @matmul_t_zero_point_rank2(%arg0: tensor<5x3xi8>, %arg1: tensor<6x3xi8>) -> tensor<5x6xi32> {
+ // CHECK: %[[C0:.+]] = arith.constant 0 : i32
+ // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<5x6xi32>
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : i32) outs(%[[INIT]] : tensor<5x6xi32>) -> tensor<5x6xi32>
+ // CHECK: %[[RESULT:.+]] = linalg.generic {indexing_maps = [#[[$MATMUL_T_RANK2_A_MAP]], #[[$MATMUL_T_RANK2_B_MAP]], #[[$MATMUL_T_RANK2_OUT_MAP]]], iterator_types = ["parallel", "parallel", "reduction"]}
+ // CHECK-SAME: ins(%arg0, %arg1 : tensor<5x3xi8>, tensor<6x3xi8>) outs(%[[FILLED]] : tensor<5x6xi32>)
+ // CHECK: ^bb0(%[[A:.+]]: i8, %[[B:.+]]: i8, %[[ACC:.+]]: i32):
+ // CHECK: %[[EXT_A:.+]] = arith.extsi %[[A]] : i8 to i32
+ // CHECK: %[[EXT_B:.+]] = arith.extsi %[[B]] : 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
+ // CHECK: return %[[RESULT]] : tensor<5x6xi32>
+ %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<5x3xi8>, tensor<6x3xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<5x6xi32>
+ return %0 : tensor<5x6xi32>
+}
+
+// -----
+
+// 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: @matmul_t_b_zero_point
+func.func @matmul_t_b_zero_point(%arg0: tensor<5x3xi8>, %arg1: tensor<6x3xi8>) -> tensor<5x6xi32> {
+ // CHECK: linalg.generic
+ // CHECK: %[[BZP:.+]] = arith.constant 2 : i32
+ // CHECK: arith.subi %{{.+}}, %[[BZP]] : i32
+ %a_zp = "tosa.const"() <{values = dense<0> : 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<5x3xi8>, tensor<6x3xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<5x6xi32>
+ return %0 : tensor<5x6xi32>
+}
More information about the Mlir-commits
mailing list