[Mlir-commits] [mlir] 4400e2d - [mlir][tosa] Lower MATMUL_T to Linalg (#222661)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 14 03:41:45 PDT 2026
Author: Jack Frankland
Date: 2026-10-14T11:41:39+01:00
New Revision: 4400e2d4ae4bfb672558060a8bddecafc769ef4a
URL: https://github.com/llvm/llvm-project/commit/4400e2d4ae4bfb672558060a8bddecafc769ef4a
DIFF: https://github.com/llvm/llvm-project/commit/4400e2d4ae4bfb672558060a8bddecafc769ef4a.diff
LOG: [mlir][tosa] Lower MATMUL_T to Linalg (#222661)
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
Added:
mlir/lib/Conversion/TosaToLinalg/TosaToLinalgUtils.h
Modified:
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-pipeline.mlir
mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 215b1da1a2782a..25b4b360723b22 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
diff er 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 2308f22dc3118e..a7f7414bf56ed3 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 6239c7dc62fd65..f2771a084f344f 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 00000000000000..2bc56c043a57b2
--- /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 5b9b33d6476dd9..4a052337335a7f 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 cd347153596660..93ca7ade0ce491 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 7031aa11720746..67a51e1e2e80f9 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