[Mlir-commits] [mlir] Feat/tosa variadic matmul sharding (PR #226180)
Deeptanshu Sekhri
llvmlistbot at llvm.org
Tue Oct 6 07:47:53 PDT 2026
https://github.com/deeptanshusekhri updated https://github.com/llvm/llvm-project/pull/226180
>From 84040f0293c3cb28b8c587c5e8f39ba9c8f2e9d3 Mon Sep 17 00:00:00 2001
From: Deeptanshu Sekhri <deeptanshu.sekhri at arm.com>
Date: Wed, 23 Sep 2026 13:31:03 +0100
Subject: [PATCH] [mlir][tosa] Support variadic matmul sharding
Build MATMUL and MATMUL_T indexing maps from operand and result ranks,
right align batch dimensions, and keep broadcast singleton dimensions
replicated. Decline propagation when dynamic batch dimensions make the
indexing ambiguous.
Signed-off-by: Deeptanshu Sekhri <deeptanshu.sekhri at arm.com>
---
.../Shard/Interfaces/ShardingInterface.cpp | 29 ++++-
.../Dialect/Tosa/IR/ShardingInterfaceImpl.cpp | 106 +++++++++++++---
.../Shard/sharding-propagation-failed.mlir | 11 ++
.../Dialect/Shard/sharding-propagation.mlir | 117 ++++++++++++++++++
4 files changed, 243 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp b/mlir/lib/Dialect/Shard/Interfaces/ShardingInterface.cpp
index cbdc274940d1e..50d2bebe1853a 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 45994a7ec679f..2af4fac3f6eba 100644
--- a/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/ShardingInterfaceImpl.cpp
@@ -22,14 +22,14 @@ using namespace mlir::tosa;
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 +47,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 +201,7 @@ void mlir::tosa::registerShardingInterfaceExternalModels(
GreaterOp, GreaterEqualOp>(ctx);
MatMulOp::attachInterface<MatMulOpSharding>(*ctx);
+ MatMulTOp::attachInterface<MatMulTOpSharding>(*ctx);
NegateOp::attachInterface<NegateOpSharding>(*ctx);
});
}
diff --git a/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir b/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir
index c80bf754a7bb3..4e8af9417a75a 100644
--- a/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir
+++ b/mlir/test/Dialect/Shard/sharding-propagation-failed.mlir
@@ -1,6 +1,7 @@
// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation))" %s -verify-diagnostics
shard.grid "private" @grid(shape = 1)
+shard.grid "private" @grid_2(shape = 2)
// expected-error @+1 {{'func.func' op only one block is supported!}}
func.func @multi_block_function(%arg0 : tensor<6x6xi32>) -> tensor<6x6xi32> {
%sharding = shard.sharding @grid split_axes = [[0]] : !shard.sharding
@@ -9,3 +10,13 @@ func.func @multi_block_function(%arg0 : tensor<6x6xi32>) -> tensor<6x6xi32> {
^bb1:
return %sharded : tensor<6x6xi32>
}
+
+// -----
+
+func.func @matmul_shard_broadcast_dimension(%a: tensor<1x4x8xf32>, %b: tensor<2x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x16xf32> {
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %a_sharded = shard.shard %a to %sharding annotate_for_users : tensor<1x4x8xf32>
+ // expected-error @+1 {{'tosa.matmul' op cannot shard operand 0 dimension 1: its indexing expression is constant}}
+ %0 = tosa.matmul %a_sharded, %b, %zp, %zp : (tensor<1x4x8xf32>, tensor<2x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x16xf32>
+ return %0 : tensor<2x4x16xf32>
+}
diff --git a/mlir/test/Dialect/Shard/sharding-propagation.mlir b/mlir/test/Dialect/Shard/sharding-propagation.mlir
index 34aaf0598b3f0..e1953526162e2 100644
--- a/mlir/test/Dialect/Shard/sharding-propagation.mlir
+++ b/mlir/test/Dialect/Shard/sharding-propagation.mlir
@@ -5,6 +5,123 @@ shard.grid @grid_1d(shape = ?)
shard.grid @grid_2d(shape = 2x4)
shard.grid @grid_3d(shape = ?x?x?)
+// CHECK-LABEL: func.func @matmul_broadcast_batches
+// CHECK-SAME: %[[A:.*]]: tensor<2x1x4x8xf32>, %[[B:.*]]: tensor<1x4x8x16xf32>
+func.func @matmul_broadcast_batches(%a: tensor<2x1x4x8xf32>, %b: tensor<1x4x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}0]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}], [1]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<2x1x4x8xf32>, tensor<1x4x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2d split_axes = [[0], [1], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_missing_batch
+// CHECK-SAME: %[[A:.*]]: tensor<4x4x8xf32>, %[[B:.*]]: tensor<2x4x8x16xf32>
+func.func @matmul_missing_batch(%a: tensor<4x4x8xf32>, %b: tensor<2x4x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}0]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<4x4x8xf32>, tensor<2x4x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0], [], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_t_rank2
+// CHECK-SAME: %[[A:.*]]: tensor<4x8xf32>, %[[B:.*]]: tensor<16x8xf32>
+func.func @matmul_t_rank2(%a: tensor<4x8xf32>, %b: tensor<16x8xf32>, %zp: tensor<1xf32>) -> tensor<4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}0]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}1]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul_t %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul_t %a, %b, %zp, %zp : (tensor<4x8xf32>, tensor<16x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x16xf32>
+ %sharding = shard.sharding @grid_2d split_axes = [[0], [1]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<4x16xf32>
+ return %sharded : tensor<4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_t_broadcast_batches
+// CHECK-SAME: %[[A:.*]]: tensor<2x1x4x8xf32>, %[[B:.*]]: tensor<1x4x16x8xf32>
+func.func @matmul_t_broadcast_batches(%a: tensor<2x1x4x8xf32>, %b: tensor<1x4x16x8xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}0]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2d split_axes = {{\[\[}}], [1]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul_t %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul_t %a, %b, %zp, %zp : (tensor<2x1x4x8xf32>, tensor<1x4x16x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2d split_axes = [[0], [1], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_t_missing_batch
+// CHECK-SAME: %[[A:.*]]: tensor<4x4x8xf32>, %[[B:.*]]: tensor<2x4x16x8xf32>
+func.func @matmul_t_missing_batch(%a: tensor<4x4x8xf32>, %b: tensor<2x4x16x8xf32>, %zp: tensor<1xf32>) -> tensor<2x4x4x16xf32> {
+ // CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+ // CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+ // CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}0]]
+ // CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+ // CHECK: tosa.matmul_t %[[A_SHARDED]], %[[B_SHARDED]]
+ %0 = tosa.matmul_t %a, %b, %zp, %zp : (tensor<4x4x8xf32>, tensor<2x4x16x8xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0], [], [], []] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x4x16xf32>
+ return %sharded : tensor<2x4x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_dynamic_lhs_batch
+// CHECK-NOT: annotate_for_users
+// CHECK: tosa.matmul
+func.func @matmul_dynamic_lhs_batch(%a: tensor<?x4x8xf32>, %b: tensor<2x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x16xf32> {
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<?x4x8xf32>, tensor<2x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x16xf32>
+ return %sharded : tensor<2x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_dynamic_rhs_batch
+// CHECK-NOT: annotate_for_users
+// CHECK: tosa.matmul
+func.func @matmul_dynamic_rhs_batch(%a: tensor<2x4x8xf32>, %b: tensor<?x8x16xf32>, %zp: tensor<1xf32>) -> tensor<2x4x16xf32> {
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<2x4x8xf32>, tensor<?x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<2x4x16xf32>
+ return %sharded : tensor<2x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_dynamic_and_singleton_batch
+// CHECK-SAME: %[[A:.*]]: tensor<?x4x8xf32>, %[[B:.*]]: tensor<1x8x16xf32>
+// CHECK: %[[A_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}0]]
+// CHECK: %[[A_SHARDED:.*]] = shard.shard %[[A]] to %[[A_SHARDING]] annotate_for_users
+// CHECK: %[[B_SHARDING:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+// CHECK: %[[B_SHARDED:.*]] = shard.shard %[[B]] to %[[B_SHARDING]] annotate_for_users
+// CHECK: tosa.matmul %[[A_SHARDED]], %[[B_SHARDED]]
+func.func @matmul_dynamic_and_singleton_batch(%a: tensor<?x4x8xf32>, %b: tensor<1x8x16xf32>, %zp: tensor<1xf32>) -> tensor<?x4x16xf32> {
+ %0 = tosa.matmul %a, %b, %zp, %zp : (tensor<?x4x8xf32>, tensor<1x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<?x4x16xf32>
+ %sharding = shard.sharding @grid_2 split_axes = [[0]] : !shard.sharding
+ %sharded = shard.shard %0 to %sharding : tensor<?x4x16xf32>
+ return %sharded : tensor<?x4x16xf32>
+}
+
+// CHECK-LABEL: func.func @matmul_operand_replication
+// CHECK-SAME: %[[A:.*]]: tensor<1x4x8xf32>, %[[B:.*]]: tensor<1x8x16xf32>
+// CHECK: %[[REPLICATION:.*]] = shard.sharding @grid_2 split_axes = {{\[\[}}]]
+// CHECK: %[[A_REPLICATED:.*]] = shard.shard %[[A]] to %[[REPLICATION]] annotate_for_users
+// CHECK: %[[B_REPLICATED:.*]] = shard.shard %[[B]] to %[[REPLICATION]] annotate_for_users
+// CHECK: tosa.matmul %[[A_REPLICATED]], %[[B_REPLICATED]]
+func.func @matmul_operand_replication(%a: tensor<1x4x8xf32>, %b: tensor<1x8x16xf32>, %zp: tensor<1xf32>) -> tensor<1x4x16xf32> {
+ %replication = shard.sharding @grid_2 split_axes = [[]] : !shard.sharding
+ %a_replicated = shard.shard %a to %replication annotate_for_users : tensor<1x4x8xf32>
+ %0 = tosa.matmul %a_replicated, %b, %zp, %zp : (tensor<1x4x8xf32>, tensor<1x8x16xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x4x16xf32>
+ return %0 : tensor<1x4x16xf32>
+}
+
// CHECK-LABEL: func.func @element_wise_empty_sharding_info
func.func @element_wise_empty_sharding_info(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
// CHECK-NEXT: tosa.sigmoid
More information about the Mlir-commits
mailing list