[Mlir-commits] [mlir] [mlir][tosa] Lower MATMUL_T to Linalg (PR #222661)
Jack Frankland
llvmlistbot at llvm.org
Tue Oct 13 09:52:08 PDT 2026
================
@@ -679,6 +680,271 @@ class MatMulConverter : public OpConversionPattern<tosa::MatMulOp> {
}
};
+// Insert missing leading batch dimensions, taking their sizes from `other`.
+static Value broadcastMatMulTLeadingDimensions(PatternRewriter &rewriter,
+ Location loc, Value operand,
+ Value other) {
+ auto type = cast<RankedTensorType>(operand.getType());
+ int64_t missingDims =
+ cast<RankedTensorType>(other.getType()).getRank() - type.getRank();
+ if (missingDims <= 0)
+ return operand;
+
+ SmallVector<OpFoldResult> shape;
+ SmallVector<int64_t> dimensions;
+ for (int64_t dim = 0; dim < missingDims; ++dim) {
+ shape.push_back(tensor::getMixedSize(rewriter, loc, other, dim));
+ dimensions.push_back(dim);
+ }
+ llvm::append_range(shape, tensor::getMixedSizes(rewriter, loc, operand));
+ Value empty = tensor::EmptyOp::create(
+ rewriter, loc, shape, type.getElementType(), type.getEncoding());
+ return linalg::BroadcastOp::create(rewriter, loc, operand, empty, dimensions)
+ .getResult()
+ .front();
+}
+
+// Replace a known singleton batch dimension with `size`.
+static Value broadcastMatMulTSingletonDimension(PatternRewriter &rewriter,
+ Location loc, Value operand,
+ int64_t dim,
+ OpFoldResult size) {
+ auto type = cast<RankedTensorType>(operand.getType());
+ SmallVector<int64_t> sliceShape(type.getShape());
+ sliceShape.erase(sliceShape.begin() + dim);
+ auto sliceType = RankedTensorType::get(sliceShape, type.getElementType(),
+ type.getEncoding());
+ Value slice = tensor::createCanonicalRankReducingExtractSliceOp(
+ rewriter, loc, operand, sliceType);
+
+ SmallVector<OpFoldResult> shape =
+ tensor::getMixedSizes(rewriter, loc, operand);
+ shape[dim] = size;
+ Value empty = tensor::EmptyOp::create(
+ rewriter, loc, shape, type.getElementType(), type.getEncoding());
+ return linalg::BroadcastOp::create(rewriter, loc, slice, empty,
+ ArrayRef<int64_t>{dim})
+ .getResult()
+ .front();
+}
+
+// Align the ranks, then broadcast each shared batch dimension independently.
+static std::pair<Value, Value>
+normalizeMatMulTBatchDimensions(PatternRewriter &rewriter, Location loc,
+ Value a, Value b, int64_t batchRank) {
+ int64_t aRank = cast<RankedTensorType>(a.getType()).getRank();
+ int64_t bRank = cast<RankedTensorType>(b.getType()).getRank();
+ int64_t firstSharedDim = std::max(aRank, bRank) - std::min(aRank, bRank);
+ a = broadcastMatMulTLeadingDimensions(rewriter, loc, a, b);
+ b = broadcastMatMulTLeadingDimensions(rewriter, loc, b, a);
+
+ // The inserted leading dimensions already match their partner's sizes.
+ for (int64_t dim = firstSharedDim; dim < batchRank; ++dim) {
+ auto aType = cast<RankedTensorType>(a.getType());
+ auto bType = cast<RankedTensorType>(b.getType());
+ int64_t aSize = aType.getDimSize(dim);
+ int64_t bSize = bType.getDimSize(dim);
+ if (aSize == 1 && bSize == 1)
+ continue;
+ if (aSize == 1) {
+ a = broadcastMatMulTSingletonDimension(
+ rewriter, loc, a, dim, tensor::getMixedSize(rewriter, loc, b, dim));
+ continue;
+ }
+ if (bSize == 1) {
+ b = broadcastMatMulTSingletonDimension(
+ rewriter, loc, b, dim, tensor::getMixedSize(rewriter, loc, a, dim));
+ continue;
+ }
+ if (!aType.isDynamicDim(dim) && !bType.isDynamicDim(dim))
+ continue;
+
+ OpFoldResult size;
+ if (!aType.isDynamicDim(dim)) {
+ size = rewriter.getIndexAttr(aSize);
+ } else if (!bType.isDynamicDim(dim)) {
+ size = rewriter.getIndexAttr(bSize);
+ } else {
+ Value aDim = tensor::DimOp::create(rewriter, loc, a, dim);
+ Value bDim = tensor::DimOp::create(rewriter, loc, b, dim);
+ Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ Value aIsSingleton = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, aDim, one);
+ // If A is a singleton, B determines the size; otherwise A does.
+ size = arith::SelectOp::create(rewriter, loc, aIsSingleton, bDim, aDim)
+ .getResult();
+ }
+ a = broadcastDynamicDimension(rewriter, loc, a, dim, size);
+ b = broadcastDynamicDimension(rewriter, loc, b, dim, size);
+ }
+ return {a, b};
+}
+
+static SmallVector<ReassociationIndices>
+getMatMulTBatchCollapseReassociation(int64_t rank) {
+ SmallVector<ReassociationIndices> reassociation(3);
+ for (int64_t dim = 0; dim < rank - 2; ++dim)
+ reassociation[0].push_back(dim);
+ reassociation[1].push_back(rank - 2);
+ reassociation[2].push_back(rank - 1);
+ return reassociation;
+}
+
+class MatMulTConverter : public OpConversionPattern<tosa::MatMulTOp> {
+public:
+ using OpConversionPattern<tosa::MatMulTOp>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(tosa::MatMulTOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const final {
+ Location loc = op.getLoc();
+ auto aType = dyn_cast<RankedTensorType>(adaptor.getA().getType());
+ auto bType = dyn_cast<RankedTensorType>(adaptor.getB().getType());
+ auto outputType = dyn_cast<RankedTensorType>(op.getType());
+ if (!aType || !bType || !outputType)
+ return rewriter.notifyMatchFailure(op, "only supports ranked tensors");
+
+ FailureOr<int64_t> maybeAZp = op.getAZeroPoint();
+ FailureOr<int64_t> maybeBZp = op.getBZeroPoint();
+ if (failed(maybeAZp))
+ return rewriter.notifyMatchFailure(
+ op, "input A zero point cannot be statically determined");
+ if (failed(maybeBZp))
+ return rewriter.notifyMatchFailure(
+ op, "input B zero point cannot be statically determined");
+
+ const int64_t aZpVal = *maybeAZp;
+ const int64_t bZpVal = *maybeBZp;
+
+ const int64_t outputRank = outputType.getRank();
+ const int64_t batchRank = outputRank - 2;
+ auto [a, b] = normalizeMatMulTBatchDimensions(rewriter, loc, adaptor.getA(),
+ adaptor.getB(), batchRank);
+
+ SmallVector<OpFoldResult> outputShape;
+ outputShape.reserve(outputRank);
+ for (int64_t batchDim = 0; batchDim < batchRank; ++batchDim)
+ outputShape.push_back(tensor::getMixedSize(rewriter, loc, a, batchDim));
+ outputShape.push_back(
+ tensor::getMixedSize(rewriter, loc, a, outputRank - 2));
+ outputShape.push_back(
+ tensor::getMixedSize(rewriter, loc, b, outputRank - 2));
+ SmallVector<Value> dynamicDims;
+ for (auto [dim, size] : llvm::enumerate(outputShape))
+ if (outputType.isDynamicDim(dim))
+ dynamicDims.push_back(
+ getValueOrCreateConstantIndexOp(rewriter, loc, size));
+
+ Value zero = arith::ConstantOp::create(
+ rewriter, loc, rewriter.getZeroAttr(outputType.getElementType()));
+ Value emptyTensor = tensor::EmptyOp::create(
+ rewriter, loc, outputType.getShape(), outputType.getElementType(),
+ condenseValues(dynamicDims));
+ Value zeroTensor = linalg::FillOp::create(rewriter, loc, ValueRange{zero},
+ ValueRange{emptyTensor})
+ .result();
+
+ Type outputElementType = outputType.getElementType();
+ if (aZpVal == 0 && bZpVal == 0) {
+ if (outputRank == 2) {
+ rewriter.replaceOpWithNewOp<linalg::MatmulTransposeBOp>(
+ op, TypeRange{outputType}, ValueRange{a, b},
+ ValueRange{zeroTensor});
+ return success();
+ }
+
+ if (outputRank == 3) {
+ rewriter.replaceOpWithNewOp<linalg::BatchMatmulTransposeBOp>(
+ op, TypeRange{outputType}, ValueRange{a, b},
+ ValueRange{zeroTensor});
+ return success();
+ }
+
+ SmallVector<ReassociationIndices> reassociation =
+ getMatMulTBatchCollapseReassociation(outputRank);
+ auto collapse = [&](Value value) {
+ auto type = cast<RankedTensorType>(value.getType());
+ auto collapsedType =
+ tensor::CollapseShapeOp::inferCollapsedType(type, reassociation);
+ return tensor::CollapseShapeOp::create(rewriter, loc, collapsedType,
+ value, reassociation);
+ };
+ Value collapsedA = collapse(a);
+ Value collapsedB = collapse(b);
+ Value collapsedInit = collapse(zeroTensor);
+ auto collapsedOutputType =
+ cast<RankedTensorType>(collapsedInit.getType());
+ Value collapsedResult =
+ linalg::BatchMatmulTransposeBOp::create(
+ rewriter, loc, TypeRange{collapsedOutputType},
+ ValueRange{collapsedA, collapsedB}, ValueRange{collapsedInit})
+ .getResult(0);
+ Value result = tensor::ExpandShapeOp::create(
+ rewriter, loc, outputType, collapsedResult, reassociation,
+ tensor::getMixedSizes(rewriter, loc, zeroTensor));
+ rewriter.replaceOp(op, result);
+ return success();
+ }
+
+ int64_t loopRank = outputRank + 1;
+ SmallVector<AffineExpr> dimensions =
+ llvm::map_to_vector(llvm::seq<int64_t>(0, loopRank), [&](int64_t dim) {
+ return rewriter.getAffineDimExpr(dim);
+ });
+ SmallVector<AffineExpr> aMap(dimensions.begin(),
+ dimensions.begin() + batchRank);
+ aMap.push_back(dimensions[batchRank]);
+ aMap.push_back(dimensions[batchRank + 2]);
+ SmallVector<AffineExpr> bMap(dimensions.begin(),
+ dimensions.begin() + batchRank);
+ bMap.push_back(dimensions[batchRank + 1]);
+ bMap.push_back(dimensions[batchRank + 2]);
+ SmallVector<AffineExpr> outputMap(dimensions.begin(),
+ dimensions.begin() + batchRank);
+ outputMap.push_back(dimensions[batchRank]);
+ outputMap.push_back(dimensions[batchRank + 1]);
+ SmallVector<AffineMap> indexingMaps = {
+ AffineMap::get(loopRank, 0, aMap, rewriter.getContext()),
+ AffineMap::get(loopRank, 0, bMap, rewriter.getContext()),
+ AffineMap::get(loopRank, 0, outputMap, rewriter.getContext())};
+ SmallVector<utils::IteratorType> iteratorTypes(
+ loopRank, utils::IteratorType::parallel);
+ iteratorTypes.back() = utils::IteratorType::reduction;
+
+ auto generic = linalg::GenericOp::create(
----------------
FranklandJack wrote:
Yeah good point, I've moved it and some of the shared utilities.
https://github.com/llvm/llvm-project/pull/222661
More information about the Mlir-commits
mailing list