[Mlir-commits] [mlir] Feat/tosa variadic matmul sharding (PR #226180)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 24 07:20:56 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tosa
Author: Deeptanshu Sekhri (deeptanshusekhri)
<details>
<summary>Changes</summary>
---
Patch is 67.30 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/226180.diff
17 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td (+22-10)
- (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td (+8)
- (modified) mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp (+24-5)
- (modified) mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp (+91-14)
- (modified) mlir/lib/Dialect/Tosa/IR/TosaOps.cpp (+141-112)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp (+11)
- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp (+39-2)
- (modified) mlir/test/Dialect/Shard/sharding-propagation-failed.mlir (+11)
- (modified) mlir/test/Dialect/Shard/sharding-propagation.mlir (+117)
- (modified) mlir/test/Dialect/Tosa/error_if_check.mlir (+20)
- (modified) mlir/test/Dialect/Tosa/level_check.mlir (+36)
- (modified) mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir (+44)
- (modified) mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir (+88)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p0-invalid.mlir (+30)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p0-valid.mlir (+10)
- (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir (+40)
- (modified) mlir/test/Dialect/Tosa/verifier.mlir (+77-4)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 7cbe8e0b6c70e7..e55727ffdd2dd0 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -416,21 +416,28 @@ def Tosa_MatMulOp : Tosa_InferShapedTypeOp<"matmul", [NoMemoryEffect]> {
let summary = "Matrix multiplication operator.";
let description = [{
- Performs two dimensional matrix multiplications.
+ Performs batched matrix multiplication of `A` with shape `[..., H, C]` and
+ `B` with shape `[..., C, W]`. Their leading batch dimensions are aligned
+ from the right. Missing leading dimensions are treated as size 1, and each
+ aligned pair must either match or include a dimension of size 1. The output
+ shape is the resulting broadcast batch shape followed by `[H, W]`.
+
+ Both inputs must have rank 2 or greater. Rank 1 inputs are not accepted
+ directly and must first be reshaped to rank 2.
This operation is not pure. Undefined behaviour may occur if the accumulated
result overflows.
}];
let arguments = (ins
- Tosa_Tensor3D:$a,
- Tosa_Tensor3D:$b,
+ Tosa_TensorAtLeast2D:$a,
+ Tosa_TensorAtLeast2D:$b,
Tosa_ScalarIntOrFloatTensor:$a_zp,
Tosa_ScalarIntOrFloatTensor:$b_zp
);
let results = (outs
- Tosa_Tensor3D:$output
+ Tosa_TensorAtLeast2D:$output
);
let extraClassDeclaration = [{
@@ -454,16 +461,21 @@ def Tosa_MatMulTOp : Tosa_InferShapedTypeOp<"matmul_t", [NoMemoryEffect]> {
let summary = "Matrix multiplication with transposed B operator.";
let description = [{
- Performs two dimensional matrix multiplications. `A` matrix is of shape
- `N x H x C`. `B` matrix is of shape `D x W x C`. This is effectively a
- matrix multiply of `A` by the transposed `B` matrix. If the batched
- dimension of input `B` is of size 1, the `B` matrix is broadcast.
+ Performs batched matrix multiplication of `A` with shape `[..., H, C]` and
+ `B` with shape `[..., W, C]`, implicitly transposing only the two innermost
+ dimensions of `B`. Their leading batch dimensions are aligned from the
+ right. Missing leading dimensions are treated as size 1, and each aligned
+ pair must either match or include a dimension of size 1. The output shape
+ is the resulting broadcast batch shape followed by `[H, W]`.
+
+ Both inputs must have rank 2 or greater. Rank 1 inputs are not accepted
+ directly and must first be reshaped to rank 2.
}];
- let arguments = (ins Tosa_Tensor3D:$a, Tosa_Tensor3D:$b,
+ let arguments = (ins Tosa_TensorAtLeast2D:$a, Tosa_TensorAtLeast2D:$b,
Tosa_ScalarIntOrFloatTensor:$a_zp, Tosa_ScalarIntOrFloatTensor:$b_zp);
- let results = (outs Tosa_Tensor3D:$output);
+ let results = (outs Tosa_TensorAtLeast2D:$output);
let extraClassDeclaration = [{
FailureOr<int64_t> getAZeroPoint();
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
index b5701853aaf683..4d6727ed76c564 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td
@@ -171,6 +171,10 @@ def AtLeastRankOne : And<[
IsRankedTensorTypePred,
CPred<"::llvm::cast<::mlir::RankedTensorType>($_self).getRank() >= 1">]>;
+def AtLeastRankTwo : And<[
+ IsRankedTensorTypePred,
+ CPred<"::llvm::cast<::mlir::RankedTensorType>($_self).getRank() >= 2">]>;
+
def IsValidBlockScaledTensorType
: CPred<"::mlir::succeeded(::mlir::tosa::"
"verifyBlockScaledTensorType($_self))">;
@@ -257,6 +261,10 @@ def Tosa_TensorAtLeast1D
: TosaTensorOf<[Tosa_AnyNumber], [AtLeastRankOne],
"tosa-conformant tensor of at least rank 1">;
+def Tosa_TensorAtLeast2D
+ : TosaTensorOf<[Tosa_AnyNumber], [AtLeastRankTwo],
+ "tosa-conformant tensor of at least rank 2">;
+
def Tosa_MXFPDataTensor3D
: TosaTensorRankOf<[Tosa_MXFPValue], [3]>;
def Tosa_MXFPScaleTensor3D
diff --git a/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp b/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp
index cbdc274940d1e3..50d2bebe1853ac 100644
--- a/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp
+++ b/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp
@@ -73,6 +73,9 @@ checkOperandAffineExprRecursively(AffineExpr expr,
seenIds[position] = true;
return success();
}
+ case AffineExprKind::Constant:
+ // A constant indexing expression does not reference a loop dimension.
+ return success();
default:
return failure();
}
@@ -304,7 +307,14 @@ shard::detail::defaultGetShardingOption(Operation *op,
if (!shardAttr)
continue;
- anyShardingInResultsOrOperands = !shardAttr.getSplitAxes().empty();
+ // Record the grid before processing mapped dimensions. Constant-indexed
+ // dimensions represent replication and do not call fillShardingOption.
+ FlatSymbolRefAttr grid = shardAttr.getGridAttr();
+ if (shardingOption.grid && grid && shardingOption.grid != grid)
+ return failure();
+ if (grid)
+ shardingOption.grid = grid;
+ anyShardingInResultsOrOperands |= !shardAttr.getSplitAxes().empty();
AffineMap map = maps[shardingIt.index()];
unsigned numDims = map.getNumDims();
@@ -312,17 +322,26 @@ shard::detail::defaultGetShardingOption(Operation *op,
//
// TODO: Change to process the operands with single loop index first and
// then the operands with multiple loop indices.
- for (auto it : llvm::zip(map.getResults(), shardAttr.getSplitAxes())) {
- AffineExpr expr = std::get<0>(it);
- ArrayRef<GridAxis> axes = std::get<1>(it).asArrayRef();
+ for (auto [dim, exprAndAxes] : llvm::enumerate(
+ llvm::zip(map.getResults(), shardAttr.getSplitAxes()))) {
+ AffineExpr expr = std::get<0>(exprAndAxes);
+ ArrayRef<GridAxis> axes = std::get<1>(exprAndAxes).asArrayRef();
FailureOr<llvm::SmallSet<unsigned, 2>> loopIndices =
checkOperandAffineExpr(expr, numDims);
if (failed(loopIndices))
return op->emitOpError()
<< "operand's affine expression is restricted to const_i * "
"dim_i + const_j + dim_j + ...";
- if (loopIndices->empty())
+ // A constant expression selects the same element for every loop
+ // iteration. It can represent replication, but cannot carry split axes.
+ if (loopIndices->empty()) {
+ if (!axes.empty())
+ return op->emitOpError()
+ << "cannot shard operand " << shardingIt.index()
+ << " dimension " << dim + 1
+ << ": its indexing expression is constant";
continue;
+ }
if (loopIndices->size() == 1) {
unsigned loopIdx = *loopIndices->begin();
visitedLoopIndices.insert(loopIdx);
diff --git a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
index 45994a7ec679f1..b0ca3e355dd9b8 100644
--- a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
@@ -23,13 +23,14 @@ using namespace mlir::shard;
namespace {
-// loop types: [parallel, parallel, parallel, reduction_sum]
-// indexing maps:
-// (d0, d1, d2, d3) -> (d0, d1, d3)
-// (d0, d1, d2, d3) -> (d0, d3, d2)
-// (d0, d1, d2, d3) -> (d0, d1, d2)
-struct MatMulOpSharding
- : public ShardingInterface::ExternalModel<MatMulOpSharding, MatMulOp> {
+// For an output of rank R, use R parallel loops followed by one reduction
+// loop. Right align the input batch dimensions with the output batch loops and
+// map A[..., H, C] to [..., d(R-2), d(R)]. MATMUL maps B[..., C, W] to
+// [..., d(R), d(R-1)], while MATMUL_T maps B[..., W, C] to
+// [..., d(R-1), d(R)].
+template <typename OpType, bool TransposeB>
+struct MatMulSharding : public ShardingInterface::ExternalModel<
+ MatMulSharding<OpType, TransposeB>, OpType> {
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
auto tensorType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
if (!tensorType)
@@ -47,20 +48,95 @@ struct MatMulOpSharding
}
SmallVector<AffineMap> getIndexingMaps(Operation *op) const {
- auto tensorType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
- if (!tensorType)
+ auto aType = dyn_cast<RankedTensorType>(op->getOperand(0).getType());
+ auto bType = dyn_cast<RankedTensorType>(op->getOperand(1).getType());
+ auto outputType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
+ if (!aType || !bType || !outputType || aType.getRank() < 2 ||
+ bType.getRank() < 2 || outputType.getRank() < aType.getRank() ||
+ outputType.getRank() < bType.getRank())
return {};
+
MLIRContext *ctx = op->getContext();
+ const int64_t outputRank = outputType.getRank();
+ const int64_t loopRank = outputRank + 1;
+ const int64_t reductionLoop = outputRank;
+
+ auto getAlignedBatchDim = [&](RankedTensorType operandType,
+ int64_t outputDim) {
+ const int64_t operandDim =
+ outputDim - (outputRank - operandType.getRank());
+ return operandDim < 0 ? 1 : operandType.getDimSize(operandDim);
+ };
+ for (int64_t i = 0; i < outputRank - 2; ++i) {
+ const int64_t aDim = getAlignedBatchDim(aType, i);
+ const int64_t bDim = getAlignedBatchDim(bType, i);
+ // A dynamic dimension may be either one (broadcast) or the corresponding
+ // output extent. These cases require different indexing maps unless the
+ // other operand is known to be a singleton.
+ if ((ShapedType::isDynamic(aDim) && bDim != 1) ||
+ (ShapedType::isDynamic(bDim) && aDim != 1))
+ return {};
+ }
+
+ auto getOperandMap = [&](RankedTensorType operandType, bool isB) {
+ const int64_t operandRank = operandType.getRank();
+ const int64_t operandBatchRank = operandRank - 2;
+ const int64_t batchRankDiff = outputRank - operandRank;
+ SmallVector<AffineExpr> results;
+ results.reserve(operandRank);
+ for (int64_t i = 0; i < operandBatchRank; ++i) {
+ // A statically sized batch dimension of one is broadcast and must
+ // remain replicated rather than inherit the corresponding sharding.
+ if (operandType.getDimSize(i) == 1)
+ results.push_back(getAffineConstantExpr(0, ctx));
+ else
+ results.push_back(getAffineDimExpr(batchRankDiff + i, ctx));
+ }
+ if (isB) {
+ if constexpr (TransposeB) {
+ results.push_back(getAffineDimExpr(outputRank - 1, ctx));
+ results.push_back(getAffineDimExpr(reductionLoop, ctx));
+ } else {
+ results.push_back(getAffineDimExpr(reductionLoop, ctx));
+ results.push_back(getAffineDimExpr(outputRank - 1, ctx));
+ }
+ } else {
+ results.push_back(getAffineDimExpr(outputRank - 2, ctx));
+ results.push_back(getAffineDimExpr(reductionLoop, ctx));
+ }
+ return AffineMap::get(loopRank, 0, results, ctx);
+ };
+
+ SmallVector<unsigned> outputTargets;
+ outputTargets.reserve(outputRank);
+ for (int64_t i = 0; i < outputRank; ++i)
+ outputTargets.push_back(i);
+
SmallVector<AffineMap> maps;
- maps.push_back(AffineMap::getMultiDimMapWithTargets(4, {0, 1, 3}, ctx));
- maps.push_back(AffineMap::getMultiDimMapWithTargets(4, {0, 3, 2}, ctx));
- maps.push_back(AffineMap::get(0, 0, {}, ctx));
- maps.push_back(AffineMap::get(0, 0, {}, ctx));
- maps.push_back(AffineMap::getMultiDimMapWithTargets(4, {0, 1, 2}, ctx));
+ maps.push_back(getOperandMap(aType, /*isB=*/false));
+ maps.push_back(getOperandMap(bType, /*isB=*/true));
+ maps.push_back(AffineMap::get(loopRank, 0, {}, ctx));
+ maps.push_back(AffineMap::get(loopRank, 0, {}, ctx));
+ maps.push_back(
+ AffineMap::getMultiDimMapWithTargets(loopRank, outputTargets, ctx));
return maps;
}
+
+ FailureOr<ShardingOption>
+ getShardingOption(Operation *op, ArrayRef<Sharding> operandShardings,
+ ArrayRef<Sharding> resultShardings) const {
+ // Decline propagation when a dynamic batch dimension makes the operand
+ // indexing ambiguous between replication and the corresponding loop.
+ if (getIndexingMaps(op).empty())
+ return ShardingOption::makeEmpty();
+ return shard::detail::defaultGetShardingOption(op, operandShardings,
+ resultShardings);
+ }
};
+using MatMulOpSharding = MatMulSharding<MatMulOp, /*TransposeB=*/false>;
+using MatMulTOpSharding = MatMulSharding<MatMulTOp, /*TransposeB=*/true>;
+
struct NegateOpSharding
: public ShardingInterface::ExternalModel<NegateOpSharding, NegateOp> {
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
@@ -126,6 +202,7 @@ void mlir::tosa::registerShardingInterfaceExternalModels(
GreaterOp, GreaterEqualOp>(ctx);
MatMulOp::attachInterface<MatMulOpSharding>(*ctx);
+ MatMulTOp::attachInterface<MatMulTOpSharding>(*ctx);
NegateOp::attachInterface<NegateOpSharding>(*ctx);
});
}
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
index feac95c5a021a5..4cdd0b097e7b6a 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaOps.cpp
@@ -1899,32 +1899,86 @@ bool tosa::EqualOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
return succeeded(verifyCompatibleShape(l[0], r[0]));
}
-LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
- MLIRContext *context, ::std::optional<Location> location,
- MatMulOp::Adaptor adaptor,
- SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
- ShapeAdaptor lhsShape(adaptor.getA().getType());
- ShapeAdaptor rhsShape(adaptor.getB().getType());
-
- // All shapes are dynamic.
- SmallVector<int64_t> outShape;
- outShape.resize(3, ShapedType::kDynamic);
+// MATMUL batch shapes are right aligned and may have different ranks. Missing
+// leading dimensions are treated as one, while an unranked input remains
+// unknown because it may contain any number of batch dimensions.
+static int64_t getMatMulBatchDim(const ShapeAdaptor &shape, int64_t outputRank,
+ int64_t axis) {
+ if (!shape.hasRank())
+ return ShapedType::kDynamic;
+ const int64_t inputAxis = axis - (outputRank - shape.getRank());
+ return inputAxis < 0 ? 1 : shape.getDimSize(inputAxis);
+}
+
+static FailureOr<SmallVector<int64_t>>
+resolveMatMulOutputShape(const ShapeAdaptor &aShape, const ShapeAdaptor &bShape,
+ int64_t outputRank, bool transposeB) {
+ if (outputRank < 2 ||
+ (aShape.hasRank() &&
+ (aShape.getRank() < 2 || aShape.getRank() > outputRank)) ||
+ (bShape.hasRank() &&
+ (bShape.getRank() < 2 || bShape.getRank() > outputRank)))
+ return failure();
- if (lhsShape.hasRank()) {
- outShape[0] = lhsShape.getDimSize(0);
- outShape[1] = lhsShape.getDimSize(1);
+ SmallVector<int64_t> outputShape(outputRank, ShapedType::kDynamic);
+ for (int64_t axis = 0; axis < outputRank - 2; ++axis) {
+ const int64_t aDim = getMatMulBatchDim(aShape, outputRank, axis);
+ const int64_t bDim = getMatMulBatchDim(bShape, outputRank, axis);
+ FailureOr<int64_t> resolvedDim = resolveBroadcastDim(aDim, bDim);
+ if (failed(resolvedDim))
+ return failure();
+ outputShape[axis] = *resolvedDim;
}
- if (rhsShape.hasRank()) {
- outShape[0] = outShape[0] == ShapedType::kDynamic ? rhsShape.getDimSize(0)
- : outShape[0];
- outShape[2] = rhsShape.getDimSize(2);
+ if (aShape.hasRank())
+ outputShape[outputRank - 2] = aShape.getDimSize(aShape.getRank() - 2);
+ if (bShape.hasRank())
+ outputShape[outputRank - 1] =
+ bShape.getDimSize(bShape.getRank() - (transposeB ? 2 : 1));
+ return outputShape;
+}
+
+static LogicalResult inferMatMulReturnTypeComponents(
+ const ShapeAdaptor &aShape, const ShapeAdaptor &bShape, bool transposeB,
+ SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+ if ((aShape.hasRank() && aShape.getRank() < 2) ||
+ (bShape.hasRank() && bShape.getRank() < 2))
+ return failure();
+
+ // An unranked input can have an arbitrary batch prefix, so the output rank
+ // cannot yet be inferred.
+ if (!aShape.hasRank() || !bShape.hasRank()) {
+ inferredReturnShapes.emplace_back();
+ return success();
}
- inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
+ const int64_t aChannels = aShape.getDimSize(aShape.getRank() - 1);
+ const int64_t bChannels =
+ bShape.getDimSize(bShape.getRank() - (transposeB ? 1 : 2));
+ if (ShapedType::isStatic(aChannels) && ShapedType::isStatic(bChannels) &&
+ aChannels != bChannels)
+ return failure();
+
+ const int64_t outputRank = std::max(aShape.getRank(), bShape.getRank());
+ FailureOr<SmallVector<int64_t>> outputShape =
+ resolveMatMulOutputShape(aShape, bShape, outputRank, transposeB);
+ if (failed(outputShape))
+ return failure();
+
+ inferredReturnShapes.emplace_back(*outputShape);
return success();
}
+LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
+ MLIRContext *context, ::std::optional<Location> location,
+ MatMulOp::Adaptor adaptor,
+ SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+ return inferMatMulReturnTypeComponents(ShapeAdaptor(adaptor.getA().getType()),
+ ShapeAdaptor(adaptor.getB().getType()),
+ /*transposeB=*/false,
+ inferredReturnShapes);
+}
+
template <typename T>
static LogicalResult verifyMatMulQuantizedOperandsType(T op, Type aElementType,
Type bElementType) {
@@ -1978,11 +2032,68 @@ static LogicalResult verifyMatMulZeroPointType(T op, Value input, Value zp,
return diag;
}
+static SmallVector<int64_t> getMatMulBatchShape(const ShapeAdaptor &shape) {
+ SmallVector<int64_t> batchShape;
+ if (!shape.hasRank() || shape.getRank() < 2)
+ return batchShape;
+ batchShape.reserve(shape.getRank() - 2);
+ for (int64_t i = 0, e = shape.getRank() - 2; i < e; ++i)
+ batchShape.push_back(shape.getDimSize(i));
+ return batchShape;
+}
+
+template <typename T>
+static LogicalResult verifyMatMulShapes(T op, bool transposeB) {
+ const ShapeAdaptor aShape(op.getA().getType());
+ const ShapeAdaptor bShape(op.getB().getType());
+ const auto outputType = cast<ShapedType>(op.getResult().getType());
+
+ int64_t channels = aShape.hasRank() ? aShape.getDimSize(aShape.getRank() - 1)
+ : ShapedType::kDynamic;
+ if (bShape.hasRank() &&
+ failed(tryUpdateDimOrFailure(
+ op, channels,
+ bShape.getDimSize(bShape.getRank() - (transposeB ? 1 : 2)), "b",
+ "channels")))
+ return failure();
+
+ const int64_t minimumOutputRank =
+ std::max(aShape.hasRank() ? aShape.getRank() : 2,
+ bShape.hasRank() ? bShape.getRank() : 2);
+ if (outputType.hasRank() && outputType.getRank() < minimumOutputRank)
+ return op.emitOpError("expected output rank of at least ")
+ << minimumOutputRank << ", got " << outputType.getRank();
+
+ const bool bothInputsRanked = aShape.hasRank() && bShape.hasRank();
+ if (!bothInputsRanked && !outputType.hasRank())
+ return success();
+
+ // When an input is unranked, use the declared output rank to validate every
+ // result dimension constrained by the ranked input.
+ const int64_t expectedOutputRank =
+ bothInputsRanked ? minimumOutputRank : outputType.getRank();
+ FailureOr<SmallVector<int64_t>> expectedOutputShape =
+ resolveMatMulOutputShape(aShape, bShape, expectedOutputRank, transposeB);
+ if (failed(expectedOutputShape)) {
+ InFlightDiagnostic diag = op.emitOpError(
+ "expected batch dimensions of a and b to be broadcast compatible, "
+ "got a=[");
+ printShapeToDiagnostic(diag, getMatMulBatchShape(aShape));
+ diag << "] and b=[";
+ printShapeToDiagnostic(diag, getMatMulBatchShape(bShape));
+ diag << "]";
+ return diag;
+ ...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/226180
More information about the Mlir-commits
mailing list