[Mlir-commits] [mlir] [mlir][tosa] Lower MATMUL_T to Linalg (PR #222661)
Jack Frankland
llvmlistbot at llvm.org
Thu Sep 10 06:53:17 PDT 2026
https://github.com/FranklandJack created https://github.com/llvm/llvm-project/pull/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.
>From 5a116abfe1f11f72406785c3d318d0cd30ab6bdc Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Thu, 10 Sep 2026 14:09:26 +0100
Subject: [PATCH] [mlir][tosa] Lower MATMUL_T to Linalg
Broadcast a singleton RHS batch explicitly before lowering
floating-point operations to linalg.batch_matmul_transpose_b. Lower
quantized operations with linalg.generic so zero-point adjustment is
represented in the reduction.
---
.../Conversion/TosaToLinalg/CMakeLists.txt | 1 +
.../TosaToLinalg/TosaToLinalgNamed.cpp | 171 ++++++++++++++++++
.../TosaToLinalg/TosaToLinalgNamedPass.cpp | 1 +
.../TosaToLinalg/tosa-to-linalg-named.mlir | 79 ++++++++
4 files changed, 252 insertions(+)
diff --git a/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt b/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt
index 5123c2a7cf916..98d8637686ba2 100644
--- a/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt
+++ b/mlir/lib/Conversion/TosaToLinalg/CMakeLists.txt
@@ -20,6 +20,7 @@ add_mlir_conversion_library(MLIRTosaToLinalg
MLIRLinalgUtils
MLIRMathDialect
MLIRPass
+ MLIRSCFDialect
MLIRSupport
MLIRTensorDialect
MLIRTosaDialect
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index ee57a5a781c5f..48b6f86aba278 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -14,6 +14,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
@@ -677,6 +678,175 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
}
};
+// Broadcast the batch dimension of B from one to the batch dimension of A.
+// linalg.broadcast adds dimensions rather than expanding dimensions of size
+// one, so first rank-reduce B from 1xWxC to WxC.
+static Value broadcastMatMulTRhs(PatternRewriter &rewriter, Location loc,
+ Value a, Value b) {
+ auto bType = cast<RankedTensorType>(b.getType());
+ auto rankReducedBType =
+ RankedTensorType::get(bType.getShape().drop_front(),
+ bType.getElementType(), bType.getEncoding());
+
+ SmallVector<OpFoldResult> offsets(3, rewriter.getIndexAttr(0));
+ SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1),
+ tensor::getMixedSize(rewriter, loc, b, 1),
+ tensor::getMixedSize(rewriter, loc, b, 2)};
+ SmallVector<OpFoldResult> strides(3, rewriter.getIndexAttr(1));
+ Value rankReducedB = tensor::ExtractSliceOp::create(
+ rewriter, loc, rankReducedBType, b, offsets, sizes, strides);
+
+ SmallVector<OpFoldResult> broadcastShape = {
+ tensor::getMixedSize(rewriter, loc, a, 0),
+ tensor::getMixedSize(rewriter, loc, b, 1),
+ tensor::getMixedSize(rewriter, loc, b, 2)};
+ Value broadcastInit =
+ tensor::EmptyOp::create(rewriter, loc, broadcastShape,
+ bType.getElementType(), bType.getEncoding());
+ return linalg::BroadcastOp::create(rewriter, loc, rankReducedB, broadcastInit,
+ ArrayRef<int64_t>{0})
+ .getResult()
+ .front();
+}
+
+// Materialize MATMUL_T's optional B batch broadcast. When the B batch is
+// dynamic, select the broadcasted value at runtime if its size is one.
+static Value normalizeMatMulTRhsBatch(PatternRewriter &rewriter, Location loc,
+ Value a, Value b) {
+ auto aType = cast<RankedTensorType>(a.getType());
+ auto bType = cast<RankedTensorType>(b.getType());
+ const int64_t aBatch = aType.getDimSize(0);
+ const int64_t bBatch = bType.getDimSize(0);
+
+ if (bBatch == 1 && aBatch != 1)
+ return broadcastMatMulTRhs(rewriter, loc, a, b);
+
+ if (ShapedType::isStatic(bBatch) || aBatch == 1)
+ return b;
+
+ Value bBatchSize = tensor::DimOp::create(rewriter, loc, b, 0);
+ Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ Value requiresBroadcast = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, bBatchSize, one);
+
+ auto thenBuilder = [&](OpBuilder &builder, Location nestedLoc) {
+ Value broadcastedB = broadcastMatMulTRhs(rewriter, nestedLoc, a, b);
+ Value castB =
+ builder.createOrFold<tensor::CastOp>(nestedLoc, bType, broadcastedB);
+ scf::YieldOp::create(builder, nestedLoc, castB);
+ };
+ auto elseBuilder = [&](OpBuilder &builder, Location nestedLoc) {
+ scf::YieldOp::create(builder, nestedLoc, b);
+ };
+ return scf::IfOp::create(rewriter, loc, requiresBroadcast, thenBuilder,
+ elseBuilder)
+ .getResult(0);
+}
+
+class MatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
+public:
+ using OpConversionPattern<tosa::MatMulTOp>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(tosa::MatMulTOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const final {
+ Location loc = op.getLoc();
+ auto aType = dyn_cast<RankedTensorType>(adaptor.getA().getType());
+ auto bType = dyn_cast<RankedTensorType>(adaptor.getB().getType());
+ auto outputType = dyn_cast<RankedTensorType>(op.getType());
+ if (!aType || !bType || !outputType)
+ return rewriter.notifyMatchFailure(op, "only supports ranked tensors");
+
+ FailureOr<int64_t> maybeAZp = op.getAZeroPoint();
+ FailureOr<int64_t> maybeBZp = op.getBZeroPoint();
+ if (failed(maybeAZp))
+ return rewriter.notifyMatchFailure(
+ op, "input a zero point cannot be statically determined");
+ if (failed(maybeBZp))
+ return rewriter.notifyMatchFailure(
+ op, "input b zero point cannot be statically determined");
+
+ const int64_t aZpVal = *maybeAZp;
+ const int64_t bZpVal = *maybeBZp;
+ if (op.verifyAZeroPoint(aZpVal).failed())
+ return rewriter.notifyMatchFailure(
+ op, "input a zero point must be zero for non-int8 integer types");
+ if (op.verifyBZeroPoint(bZpVal).failed())
+ return rewriter.notifyMatchFailure(
+ op, "input b zero point must be zero for non-int8 integer types");
+
+ Value b =
+ normalizeMatMulTRhsBatch(rewriter, loc, adaptor.getA(), adaptor.getB());
+
+ SmallVector<Value> dynamicDims(3);
+ if (outputType.isDynamicDim(0))
+ dynamicDims[0] = tensor::DimOp::create(rewriter, loc, adaptor.getA(), 0);
+ if (outputType.isDynamicDim(1))
+ dynamicDims[1] = tensor::DimOp::create(rewriter, loc, adaptor.getA(), 1);
+ if (outputType.isDynamicDim(2))
+ dynamicDims[2] = tensor::DimOp::create(rewriter, loc, adaptor.getB(), 1);
+
+ Value zero = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getZeroAttr(outputType.getElementType()));
+ Value emptyTensor = tensor::EmptyOp::create(
+ rewriter, loc, outputType.getShape(), outputType.getElementType(),
+ condenseValues(dynamicDims));
+ Value zeroTensor = linalg::FillOp::create(rewriter, loc, ValueRange{zero},
+ ValueRange{emptyTensor})
+ .result();
+
+ Type outputElementType = outputType.getElementType();
+ if (!isa<IntegerType>(outputElementType)) {
+ rewriter.replaceOpWithNewOp<linalg::BatchMatmulTransposeBOp>(
+ op, TypeRange{outputType}, ValueRange{adaptor.getA(), b},
+ ValueRange{zeroTensor});
+ return success();
+ }
+
+ AffineExpr batch, height, width, channels;
+ bindDims(rewriter.getContext(), batch, height, width, channels);
+ SmallVector<AffineMap> indexingMaps = {
+ AffineMap::get(4, 0, {batch, height, channels}, rewriter.getContext()),
+ AffineMap::get(4, 0, {batch, width, channels}, rewriter.getContext()),
+ AffineMap::get(4, 0, {batch, height, width}, rewriter.getContext())};
+ SmallVector<utils::IteratorType> iteratorTypes = {
+ utils::IteratorType::parallel, utils::IteratorType::parallel,
+ utils::IteratorType::parallel, utils::IteratorType::reduction};
+
+ auto generic = linalg::GenericOp::create(
+ rewriter, loc, outputType, ValueRange{adaptor.getA(), b}, zeroTensor,
+ indexingMaps, iteratorTypes,
+ [=](OpBuilder &builder, Location nestedLoc, ValueRange args) {
+ auto extendToOutputType = [&](Value value) {
+ if (value.getType() == outputElementType)
+ return value;
+ return arith::ExtSIOp::create(builder, nestedLoc, outputElementType,
+ value)
+ .getResult();
+ };
+
+ Value aValue = extendToOutputType(args[0]);
+ Value bValue = extendToOutputType(args[1]);
+ Value aZp = arith::ConstantOp::create(
+ builder, nestedLoc,
+ builder.getIntegerAttr(outputElementType, aZpVal));
+ Value bZp = arith::ConstantOp::create(
+ builder, nestedLoc,
+ builder.getIntegerAttr(outputElementType, bZpVal));
+ Value adjustedA =
+ arith::SubIOp::create(builder, nestedLoc, aValue, aZp);
+ Value adjustedB =
+ arith::SubIOp::create(builder, nestedLoc, bValue, bZp);
+ Value product =
+ arith::MulIOp::create(builder, nestedLoc, adjustedA, adjustedB);
+ Value sum =
+ arith::AddIOp::create(builder, nestedLoc, args[2], product);
+ linalg::YieldOp::create(builder, nestedLoc, sum);
+ });
+ rewriter.replaceOp(op, generic.getResult(0));
+ return success();
+ }
+};
+
class MaxPool2dConverter : public OpConversionPattern<tosa::MaxPool2dOp> {
public:
using OpConversionPattern::OpConversionPattern;
@@ -1132,6 +1302,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 3733d043f9ebc..051c9f8f7a624 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamedPass.cpp
@@ -60,6 +60,7 @@ struct TosaToLinalgNamed
target.addIllegalOp<tosa::MaxPool2dOp>();
target.addIllegalOp<tosa::AvgPool2dOp>();
target.addIllegalOp<tosa::MatMulOp>();
+ target.addIllegalOp<tosa::MatMulTOp>();
target.addIllegalOp<tosa::TransposeOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
index 801144c01829a..1e320cf06c7fe 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir
@@ -95,6 +95,85 @@ 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_broadcast
+func.func @matmul_t_broadcast(%arg0: tensor<4x5x3xf32>, %arg1: tensor<1x6x3xf32>) -> tensor<4x5x6xf32> {
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3] [1, 1, 1] : tensor<1x6x3xf32> to tensor<6x3xf32>
+ // CHECK: %[[BCAST_INIT:.+]] = tensor.empty() : tensor<4x6x3xf32>
+ // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xf32>) outs(%[[BCAST_INIT]] : tensor<4x6x3xf32>) dimensions = [0]
+ // CHECK: %[[C0:.+]] = arith.constant 0
+ // CHECK: %[[INIT:.+]] = tensor.empty()
+ // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %[[BCAST]] : tensor<4x5x3xf32>, tensor<4x6x3xf32>) outs(%[[FILLED]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<4x5x3xf32>, tensor<1x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x5x6xf32>
+ return %0 : tensor<4x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_dynamic_batch
+func.func @matmul_t_dynamic_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<?x6x3xf32>) -> tensor<?x5x6xf32> {
+ // CHECK: %[[BATCH:.+]] = tensor.dim %arg1, %{{.+}}
+ // CHECK: %[[ONE:.+]] = arith.constant 1 : index
+ // CHECK: %[[NEEDS_BCAST:.+]] = arith.cmpi eq, %[[BATCH]], %[[ONE]]
+ // CHECK: %[[B:.+]] = scf.if %[[NEEDS_BCAST]] -> (tensor<?x6x3xf32>) {
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3]
+ // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xf32>)
+ // CHECK: scf.yield %[[BCAST]] : tensor<?x6x3xf32>
+ // CHECK: } else {
+ // CHECK: scf.yield %arg1 : tensor<?x6x3xf32>
+ // CHECK: }
+ // CHECK: linalg.batch_matmul indexing_maps =
+ // CHECK-SAME: ins(%arg0, %[[B]] : tensor<?x5x3xf32>, tensor<?x6x3xf32>)
+ %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<?x5x3xf32>, tensor<?x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x5x6xf32>
+ return %0 : tensor<?x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_quantized_broadcast
+func.func @matmul_t_quantized_broadcast(%arg0: tensor<4x5x3xi8>, %arg1: tensor<1x6x3xi8>) -> tensor<4x5x6xi32> {
+ // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3] [1, 1, 1] : tensor<1x6x3xi8> to tensor<6x3xi8>
+ // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xi8>)
+ // CHECK: %[[RESULT:.+]] = linalg.generic
+ // CHECK-SAME: ins(%arg0, %[[BCAST]] : tensor<4x5x3xi8>, tensor<4x6x3xi8>)
+ // CHECK: ^bb0(%[[A:.+]]: i8, %[[BVAL:.+]]: i8, %[[ACC:.+]]: i32):
+ // CHECK: %[[EXT_A:.+]] = arith.extsi %[[A]] : i8 to i32
+ // CHECK: %[[EXT_B:.+]] = arith.extsi %[[BVAL]] : i8 to i32
+ // CHECK: %[[AZP:.+]] = arith.constant 1 : i32
+ // CHECK: %[[BZP:.+]] = arith.constant 2 : i32
+ // CHECK: %[[ADJUSTED_A:.+]] = arith.subi %[[EXT_A]], %[[AZP]] : i32
+ // CHECK: %[[ADJUSTED_B:.+]] = arith.subi %[[EXT_B]], %[[BZP]] : i32
+ // CHECK: %[[PRODUCT:.+]] = arith.muli %[[ADJUSTED_A]], %[[ADJUSTED_B]] : i32
+ // CHECK: %[[SUM:.+]] = arith.addi %[[ACC]], %[[PRODUCT]] : i32
+ // CHECK: linalg.yield %[[SUM]] : i32
+ %a_zp = "tosa.const"() <{values = dense<1> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %b_zp = "tosa.const"() <{values = dense<2> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<4x5x3xi8>, tensor<1x6x3xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<4x5x6xi32>
+ return %0 : tensor<4x5x6xi32>
+}
+
+// -----
+
// CHECK-LABEL: @max_pool
func.func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () {
// CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38
More information about the Mlir-commits
mailing list