[Mlir-commits] [mlir] [mlir][shard] Empowering resharding (PR #180962)
Frank Schlimbach
llvmlistbot at llvm.org
Wed Feb 11 09:11:41 PST 2026
================
@@ -44,484 +46,468 @@ static bool arePartialAxesCompatible(const SourceAxes &sourceAxes,
});
}
-static Sharding targetShardingInSplitLastAxis(MLIRContext *ctx,
- const Sharding &sourceSharding,
- int64_t splitTensorAxis,
- GridAxis splitGridAxis) {
- SmallVector<GridAxesAttr> targetShardingSplitAxes =
- llvm::to_vector(sourceSharding.getSplitAxes());
- while (static_cast<int64_t>(targetShardingSplitAxes.size()) <=
- splitTensorAxis) {
- targetShardingSplitAxes.push_back(GridAxesAttr::get(ctx, {}));
- }
- auto targetSplitAxes =
- llvm::to_vector(targetShardingSplitAxes[splitTensorAxis].asArrayRef());
- targetSplitAxes.push_back(splitGridAxis);
- targetShardingSplitAxes[splitTensorAxis] =
- GridAxesAttr::get(ctx, targetSplitAxes);
- return Sharding::get(sourceSharding.getGridAttr(), targetShardingSplitAxes);
-}
+/// Base class for resharding patterns.
+/// Subclasses implement `tryApply` to detect and apply a specific resharding.
+class ReshardingPattern {
+public:
+ virtual ~ReshardingPattern() = default;
+
+ /// Try to apply this resharding pattern. Returns the resharded value and
+ /// resulting sharding on success, or std::nullopt if the pattern doesn't
+ /// match.
+ virtual std::optional<std::tuple<TypedValue<ShapedType>, Sharding>>
+ tryApply(ImplicitLocOpBuilder &builder, GridOp grid, int64_t tensorDim,
+ const Sharding &srcSharding, const Sharding &tgtSharding,
+ ShapedType srcUnshardedType, TypedValue<ShapedType> srcShard) = 0;
+
+protected:
+ /// Returns true if either sharding has non-empty static sharded dims offsets.
+ static bool hasStaticOffsets(const Sharding &srcSharding,
+ const Sharding &tgtSharding) {
+ return !srcSharding.getStaticShardedDimsOffsets().empty() ||
+ !tgtSharding.getStaticShardedDimsOffsets().empty();
+ }
-// Split a replicated tensor along a grid axis.
-// E.g. [[0, 1]] -> [[0, 1, 2]].
-// Returns the partitioned target value with its sharding.
-static std::tuple<TypedValue<ShapedType>, Sharding>
-splitLastAxisInResharding(ImplicitLocOpBuilder &builder,
- Sharding sourceSharding,
- TypedValue<ShapedType> sourceShard, GridOp grid,
- int64_t splitTensorAxis, GridAxis splitGridAxis) {
- TypedValue<ShapedType> targetShard =
- AllSliceOp::create(builder, sourceShard, grid,
- ArrayRef<GridAxis>(splitGridAxis), splitTensorAxis)
- .getResult();
- Sharding targetSharding = targetShardingInSplitLastAxis(
- builder.getContext(), std::move(sourceSharding), splitTensorAxis,
- splitGridAxis);
- return {targetShard, targetSharding};
-}
+ /// Returns true if either sharding has non-empty static sharded dims offsets
+ /// or non-empty static halo sizes.
+ static bool hasStaticOffsetsOrHalos(const Sharding &srcSharding,
+ const Sharding &tgtSharding) {
+ return hasStaticOffsets(srcSharding, tgtSharding) ||
+ !srcSharding.getStaticHaloSizes().empty() ||
+ !tgtSharding.getStaticHaloSizes().empty();
+ }
+};
-// Detect if the resharding is of type e.g.
-// [[0, 1]] -> [[0, 1, 2]].
-// If detected, returns the corresponding tensor axis grid axis pair.
-// Does not detect insertions like
-// [[0, 1]] -> [[0, 2, 1]].
-static std::optional<std::tuple<int64_t, GridAxis>>
-detectSplitLastAxisInResharding(const Sharding &sourceSharding,
- const Sharding &targetSharding) {
- for (size_t tensorAxis = 0; tensorAxis < targetSharding.getSplitAxes().size();
- ++tensorAxis) {
- if (sourceSharding.getSplitAxes().size() > tensorAxis) {
- if (sourceSharding.getSplitAxes()[tensorAxis].size() + 1 !=
- targetSharding.getSplitAxes()[tensorAxis].size()) {
- continue;
- }
- if (!llvm::equal(
- sourceSharding.getSplitAxes()[tensorAxis].asArrayRef(),
- llvm::make_range(
- targetSharding.getSplitAxes()[tensorAxis]
- .asArrayRef()
- .begin(),
- targetSharding.getSplitAxes()[tensorAxis].asArrayRef().end() -
- 1))) {
- continue;
- }
- } else {
- if (targetSharding.getSplitAxes()[tensorAxis].size() != 1) {
- continue;
- }
+/// Split a replicated axis: e.g. [[0, 1]] -> [[0, 1, 2]].
+class SplitLastAxisPattern : public ReshardingPattern {
+ static Sharding tgtSharding(MLIRContext *ctx, const Sharding &srcSharding,
+ int64_t splitTensorDim, GridAxis splitGridAxis) {
+ SmallVector<GridAxesAttr> tgtShardingSplitAxes =
+ llvm::to_vector(srcSharding.getSplitAxes());
+ while (static_cast<int64_t>(tgtShardingSplitAxes.size()) <=
+ splitTensorDim) {
+ tgtShardingSplitAxes.push_back(GridAxesAttr::get(ctx, {}));
}
- return std::make_tuple(
- tensorAxis,
- targetSharding.getSplitAxes()[tensorAxis].asArrayRef().back());
+ auto tgtSplitAxes =
+ llvm::to_vector(tgtShardingSplitAxes[splitTensorDim].asArrayRef());
+ tgtSplitAxes.push_back(splitGridAxis);
+ tgtShardingSplitAxes[splitTensorDim] = GridAxesAttr::get(ctx, tgtSplitAxes);
+ return Sharding::get(srcSharding.getGridAttr(), tgtShardingSplitAxes);
}
- return std::nullopt;
-}
-static std::optional<std::tuple<TypedValue<ShapedType>, Sharding>>
-trySplitLastAxisInResharding(ImplicitLocOpBuilder &builder, GridOp grid,
- const Sharding &sourceSharding,
- Sharding targetSharding,
- TypedValue<ShapedType> sourceShard) {
- if (auto detectRes = detectSplitLastAxisInResharding(
- sourceSharding, std::move(targetSharding))) {
- auto [tensorAxis, gridAxis] = detectRes.value();
- return splitLastAxisInResharding(builder, sourceSharding, sourceShard, grid,
- tensorAxis, gridAxis);
+ // Split a replicated tensor along a grid axis.
+ // E.g. [[0, 1]] -> [[0, 1, 2]].
+ // Returns the partitioned target value with its sharding.
+ static std::tuple<TypedValue<ShapedType>, Sharding>
+ apply(ImplicitLocOpBuilder &builder, Sharding srcSharding,
+ TypedValue<ShapedType> srcShard, GridOp grid, int64_t splitTensorDim,
+ GridAxis splitGridAxis) {
+ TypedValue<ShapedType> tgtShard =
+ AllSliceOp::create(builder, srcShard, grid,
+ ArrayRef<GridAxis>(splitGridAxis), splitTensorDim)
+ .getResult();
+ Sharding resultSharding =
+ tgtSharding(builder.getContext(), std::move(srcSharding),
+ splitTensorDim, splitGridAxis);
+ return {tgtShard, resultSharding};
}
- return std::nullopt;
-}
+ // Detect if the resharding is of type e.g.
+ // [[0, 1]] -> [[0, 1, 2]].
+ // If detected, returns the corresponding grid axis.
+ // Does not detect insertions like
+ // [[0, 1]] -> [[0, 2, 1]].
+ static std::optional<GridAxis> detect(const Sharding &srcSharding,
+ const Sharding &tgtSharding,
+ int64_t tensorDim) {
+ if (static_cast<size_t>(tensorDim) >= tgtSharding.getSplitAxes().size())
+ return std::nullopt;
+ auto tgtAxes = tgtSharding.getSplitAxes()[tensorDim].asArrayRef();
+ if (srcSharding.getSplitAxes().size() > static_cast<size_t>(tensorDim)) {
+ auto srcAxes = srcSharding.getSplitAxes()[tensorDim].asArrayRef();
+ if (srcAxes.size() + 1 != tgtAxes.size())
+ return std::nullopt;
+ if (!llvm::equal(srcAxes,
+ llvm::make_range(tgtAxes.begin(), tgtAxes.end() - 1)))
+ return std::nullopt;
+ } else {
+ if (tgtAxes.size() != 1)
+ return std::nullopt;
+ }
+ return tgtAxes.back();
+ }
+
+public:
+ std::optional<std::tuple<TypedValue<ShapedType>, Sharding>>
+ tryApply(ImplicitLocOpBuilder &builder, GridOp grid, int64_t tensorDim,
+ const Sharding &srcSharding, const Sharding &tgtSharding,
+ ShapedType srcUnshardedType,
+ TypedValue<ShapedType> srcShard) override {
+ if (hasStaticOffsetsOrHalos(srcSharding, tgtSharding))
+ return std::nullopt;
+ if (auto gridAxis = detect(srcSharding, tgtSharding, tensorDim))
+ return apply(builder, srcSharding, srcShard, grid, tensorDim,
+ gridAxis.value());
+ return std::nullopt;
+ }
+};
-// Detect if the resharding removes trailing split Axes along a tensor
-// dimension, e.g.
-// [[0, 1, 2]] -> [[0, 1]], [[0, 1, 2]] -> [0] or [[0, 1, 2]] -> [].
-// If detected, returns the corresponding (tensor dim, grid axes) pair, where
-// the "grid axes" are the removed trailing split axes.
-static std::optional<std::tuple<int64_t, SmallVector<GridAxis>>>
-detectUnsplitLastAxesInResharding(const Sharding &srcSharding,
- const Sharding &tgtSharding) {
- size_t dimOff = 0;
- size_t srcSize = srcSharding.getSplitAxes().size();
- for (size_t tensorDim = 0; tensorDim < srcSize; ++tensorDim) {
+/// Unsplit trailing axes: e.g. [[0, 1, 2]] -> [[0, 1]] or [[0, 1, 2]] -> [].
+class UnsplitLastAxesPattern : public ReshardingPattern {
+ // Detect if the resharding removes trailing split axes along a tensor
+ // dimension, e.g.
+ // [[0, 1, 2]] -> [[0, 1]], [[0, 1, 2]] -> [0] or [[0, 1, 2]] -> [].
+ // If detected, returns the removed trailing split axes (grid axes).
+ static std::optional<SmallVector<GridAxis>>
+ detect(const Sharding &srcSharding, const Sharding &tgtSharding,
+ int64_t tensorDim) {
+ if (static_cast<size_t>(tensorDim) >= srcSharding.getSplitAxes().size())
+ return std::nullopt;
+ size_t dimOff = 0;
auto srcSplitAxes = srcSharding.getSplitAxes()[tensorDim].asArrayRef();
- if (tgtSharding.getSplitAxes().size() > tensorDim) {
+ if (tgtSharding.getSplitAxes().size() > static_cast<size_t>(tensorDim)) {
auto tgtSplitAxes = tgtSharding.getSplitAxes()[tensorDim].asArrayRef();
- // No match if the target sharding does not have less split axes than the
- // source sharding along the current tensor dimension.
+ // No match if the target sharding does not have less split axes than
+ // the source sharding along the current tensor dimension.
if (srcSplitAxes.size() <= tgtSplitAxes.size())
- continue;
+ return std::nullopt;
// No match if the split axes of the target sharding are different from
// the first split axes of the source sharding.
if (!std::equal(tgtSplitAxes.begin(), tgtSplitAxes.end(),
srcSplitAxes.begin()))
- continue;
+ return std::nullopt;
dimOff = tgtSplitAxes.size();
} else {
- // Here the target dimension is replicated; there is nothing to do if the
- // source dimension is also replicated.
+ // Here the target dimension is replicated; there is nothing to do if
+ // the source dimension is also replicated.
if (srcSplitAxes.size() == 0)
- continue;
+ return std::nullopt;
dimOff = 0;
}
- // This is a match. Return the current tensor dimension and the trailing
- // grid axis of the source sharding along this dimension.
+ // This is a match. Return the trailing grid axes of the source sharding
+ // along this dimension.
ArrayRef<GridAxis> trailingAxes = srcSplitAxes.drop_front(dimOff);
SmallVector<GridAxis> unsplitAxes(trailingAxes.begin(), trailingAxes.end());
- return std::make_tuple(tensorDim, unsplitAxes);
+ return unsplitAxes;
}
- return std::nullopt;
-}
-// Return the resulting Sharding if the unsplit last axes resharding is applied.
-static Sharding targetShardingInUnsplitLastAxes(MLIRContext *ctx,
- const Sharding &sourceSharding,
- int64_t splitTensorDim,
- size_t numUnsplitAxes) {
- SmallVector<GridAxesAttr> resSplitAxes =
- llvm::to_vector(sourceSharding.getSplitAxes());
- assert(static_cast<int64_t>(resSplitAxes.size()) > splitTensorDim);
- ArrayRef<GridAxis> srcSplitAxes = resSplitAxes[splitTensorDim].asArrayRef();
- assert(srcSplitAxes.size() >= numUnsplitAxes);
- size_t numSplitAxes = srcSplitAxes.size() - numUnsplitAxes;
- SmallVector<GridAxis> newSplitAxes(srcSplitAxes.begin(),
- srcSplitAxes.begin() + numSplitAxes);
- resSplitAxes[splitTensorDim] = GridAxesAttr::get(ctx, newSplitAxes);
- return Sharding::get(sourceSharding.getGridAttr(), resSplitAxes);
-}
+ // Return the resulting Sharding if the unsplit last axes resharding is
+ // applied.
+ static Sharding tgtSharding(MLIRContext *ctx, const Sharding &srcSharding,
+ int64_t splitTensorDim, size_t numUnsplitAxes) {
+ SmallVector<GridAxesAttr> resSplitAxes =
+ llvm::to_vector(srcSharding.getSplitAxes());
+ assert(static_cast<int64_t>(resSplitAxes.size()) > splitTensorDim);
+ ArrayRef<GridAxis> srcSplitAxes = resSplitAxes[splitTensorDim].asArrayRef();
+ assert(srcSplitAxes.size() >= numUnsplitAxes);
+ size_t numSplitAxes = srcSplitAxes.size() - numUnsplitAxes;
+ SmallVector<GridAxis> newSplitAxes(srcSplitAxes.begin(),
+ srcSplitAxes.begin() + numSplitAxes);
+ resSplitAxes[splitTensorDim] = GridAxesAttr::get(ctx, newSplitAxes);
+ return Sharding::get(srcSharding.getGridAttr(), resSplitAxes);
+ }
-// Return the resulting Tensor type after applying the unsplit last axes
-// resharding.
-static ShapedType allGatherResultTypeInUnsplitLastAxes(
- ShapedType sourceType, int64_t splitTensorDim, ArrayRef<int64_t> gridShape,
- ArrayRef<GridAxis> unsplitAxes) {
- SmallVector<int64_t> targetShape = llvm::to_vector(sourceType.getShape());
- for (GridAxis gridAxis : unsplitAxes)
- targetShape[splitTensorDim] =
- gatherDimension(targetShape[splitTensorDim], gridShape[gridAxis]);
- return sourceType.cloneWith(targetShape, sourceType.getElementType());
-}
+ // Return the resulting Tensor type after applying the unsplit last axes
+ // resharding.
+ static ShapedType allGatherResultType(ShapedType srcType,
+ int64_t splitTensorDim,
+ ArrayRef<int64_t> gridShape,
+ ArrayRef<GridAxis> unsplitAxes) {
+ SmallVector<int64_t> tgtShape = llvm::to_vector(srcType.getShape());
+ for (GridAxis gridAxis : unsplitAxes)
+ tgtShape[splitTensorDim] =
+ gatherDimension(tgtShape[splitTensorDim], gridShape[gridAxis]);
+ return srcType.cloneWith(tgtShape, srcType.getElementType());
+ }
-// Perform the resharding for the unsplit last axes case.
-// This basically performs an all-gather along the unsplit grid axes.
-static std::tuple<TypedValue<ShapedType>, Sharding> unsplitLastAxesInResharding(
- ImplicitLocOpBuilder &builder, Sharding sourceSharding,
- ShapedType sourceUnshardedShape, TypedValue<ShapedType> sourceShard,
- GridOp grid, int64_t splitTensorDim, ArrayRef<GridAxis> unsplitAxes) {
- MLIRContext *ctx = builder.getContext();
- builder.setInsertionPointAfterValue(sourceShard);
-
- Sharding targetSharding = targetShardingInUnsplitLastAxes(
- ctx, std::move(sourceSharding), splitTensorDim, unsplitAxes.size());
- ShapedType allGatherResultType = allGatherResultTypeInUnsplitLastAxes(
- sourceShard.getType(), splitTensorDim, grid.getShape(), unsplitAxes);
- Value allGatherResult = AllGatherOp::create(
- builder,
- RankedTensorType::get(allGatherResultType.getShape(),
- allGatherResultType.getElementType()),
- grid.getSymName(), unsplitAxes, sourceShard, APInt(64, splitTensorDim));
- ShapedType targetType =
- shardShapedType(sourceUnshardedShape, grid, targetSharding);
- TypedValue<ShapedType> targetShard =
- tensor::CastOp::create(builder, targetType, allGatherResult).getResult();
- return {targetShard, targetSharding};
-}
+ // Perform the resharding for the unsplit last axes case.
+ // This basically performs an all-gather along the unsplit grid axes.
+ static std::tuple<TypedValue<ShapedType>, Sharding>
+ apply(ImplicitLocOpBuilder &builder, Sharding srcSharding,
+ ShapedType srcUnshardedType, TypedValue<ShapedType> srcShard,
+ GridOp grid, int64_t splitTensorDim, ArrayRef<GridAxis> unsplitAxes) {
+ MLIRContext *ctx = builder.getContext();
+ builder.setInsertionPointAfterValue(srcShard);
+
+ Sharding resultSharding = tgtSharding(ctx, std::move(srcSharding),
+ splitTensorDim, unsplitAxes.size());
+ ShapedType agResultType = allGatherResultType(
+ srcShard.getType(), splitTensorDim, grid.getShape(), unsplitAxes);
+ Value allGatherResult = AllGatherOp::create(
+ builder,
+ RankedTensorType::get(agResultType.getShape(),
+ agResultType.getElementType()),
+ grid.getSymName(), unsplitAxes, srcShard, APInt(64, splitTensorDim));
+ ShapedType tgtType =
+ shardShapedType(srcUnshardedType, grid, resultSharding);
+ TypedValue<ShapedType> tgtShard =
+ tensor::CastOp::create(builder, tgtType, allGatherResult).getResult();
+ return {tgtShard, resultSharding};
+ }
-static std::optional<std::tuple<TypedValue<ShapedType>, Sharding>>
-tryUnsplitLastAxesInResharding(ImplicitLocOpBuilder &builder, GridOp grid,
- const Sharding &sourceSharding,
- Sharding targetSharding,
- ShapedType sourceUnshardedShape,
- TypedValue<ShapedType> sourceShard) {
- if (auto detectRes = detectUnsplitLastAxesInResharding(
- sourceSharding, std::move(targetSharding))) {
- auto [tensorDim, gridAxes] = detectRes.value();
- return unsplitLastAxesInResharding(builder, sourceSharding,
- sourceUnshardedShape, sourceShard, grid,
- tensorDim, gridAxes);
- }
-
- return std::nullopt;
-}
+public:
+ std::optional<std::tuple<TypedValue<ShapedType>, Sharding>>
+ tryApply(ImplicitLocOpBuilder &builder, GridOp grid, int64_t tensorDim,
+ const Sharding &srcSharding, const Sharding &tgtSharding,
+ ShapedType srcUnshardedType,
+ TypedValue<ShapedType> srcShard) override {
+ if (hasStaticOffsetsOrHalos(srcSharding, tgtSharding))
+ return std::nullopt;
+ if (auto gridAxes = detect(srcSharding, tgtSharding, tensorDim))
+ return apply(builder, srcSharding, srcUnshardedType, srcShard, grid,
+ tensorDim, gridAxes.value());
+ return std::nullopt;
+ }
+};
-// Detect if the resharding is of type e.g.
-// [[0, 1], [2]] -> [[0], [1, 2]].
-// Only moving the last axis counts.
-// If detected, returns the corresponding (source_tensor_axis,
-// target_tensor_axis, grid_axis) tuple.
-static std::optional<std::tuple<int64_t, int64_t, GridAxis>>
-detectMoveLastSplitAxisInResharding(const Sharding &sourceSharding,
- const Sharding &targetSharding) {
- for (size_t sourceTensorAxis = 0;
- sourceTensorAxis < sourceSharding.getSplitAxes().size();
- ++sourceTensorAxis) {
- for (size_t targetTensorAxis = 0;
- targetTensorAxis < targetSharding.getSplitAxes().size();
- ++targetTensorAxis) {
- if (sourceTensorAxis == targetTensorAxis)
+/// Move a split axis between tensor dimensions:
+/// e.g. [[0], []] -> [[], [0]].
+class MoveLastSplitAxisPattern : public ReshardingPattern {
+ // Detect if the resharding moves a single split axis from one tensor
+ // dimension to another tensor dimension. If detected, returns the
+ // corresponding (tgt_tensor_dim, grid_axis) pair.
+ static std::optional<std::tuple<int64_t, GridAxis>>
+ detect(const Sharding &srcSharding, const Sharding &tgtSharding,
+ int64_t srcTensorDim) {
+ if (static_cast<size_t>(srcTensorDim) >= srcSharding.getSplitAxes().size())
+ return std::nullopt;
+ auto srcAxes = srcSharding.getSplitAxes()[srcTensorDim].asArrayRef();
+ if (srcAxes.size() != 1)
+ return std::nullopt;
----------------
fschlimb wrote:
The orig detection only covered a single axis. Renamed to `MoveAxis`.
https://github.com/llvm/llvm-project/pull/180962
More information about the Mlir-commits
mailing list