[Mlir-commits] [mlir] [mlir][shard] Empowering resharding (PR #180962)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Feb 11 07:56:02 PST 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Frank Schlimbach (fschlimb)
<details>
<summary>Changes</summary>
Previously, resharding was performed by patterns, which do the full resharding over all tensor dimensions or nothing.
This was a unnecessarily limiting approach, in particular because the patterns operated on a single dimension only.
The new approach implemented in this PR operates dimension by dimension, try-applying each pattern on a single dimension.
This enables many more resharding cases, in particular simple things like resharding a multi-dim sharding to replication (which is something that, for example, facilitates correctness checks, like it enables the 2d grid case in https://github.com/llvm/lighthouse/pull/50).
To make this possible in a reasonable way, the resharding patterns (`UpdateHalo`, `MoveLastSplitAxis`, `SplitLastAxis`, `UnsplitLastAxes`) were refactored into pattern classes which implement a `tryApply` method operating on a single dimension. This makes it a rather hard to review PR, I know (probably somewhat easier without whitespace diffs).
This PR also cleans up/unifies variable naming and doc-strings in the refactored code.
For debugging and diagnostic messages the `operator<<` got introduced for class `Sharding`.
---
Patch is 55.07 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/180962.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/Shard/IR/ShardOps.h (+12-2)
- (modified) mlir/lib/Dialect/Shard/IR/ShardOps.cpp (+23)
- (modified) mlir/lib/Dialect/Shard/Transforms/Partition.cpp (+420-429)
- (modified) mlir/test/Dialect/Shard/partition.mlir (+81)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Shard/IR/ShardOps.h b/mlir/include/mlir/Dialect/Shard/IR/ShardOps.h
index 457fe6f6b8d0a..da3dca8de7b09 100644
--- a/mlir/include/mlir/Dialect/Shard/IR/ShardOps.h
+++ b/mlir/include/mlir/Dialect/Shard/IR/ShardOps.h
@@ -12,6 +12,7 @@
#include "mlir/Bytecode/BytecodeOpInterface.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
@@ -79,6 +80,15 @@ class Sharding {
bool equalShardSizes(const Sharding &rhs) const;
};
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const Sharding &sharding);
+
+inline Diagnostic &operator<<(Diagnostic &diag, const Sharding &sharding) {
+ std::string str;
+ llvm::raw_string_ostream os(str);
+ os << sharding;
+ return diag << str;
+}
+
} // namespace shard
} // namespace mlir
@@ -181,8 +191,8 @@ inline int64_t gatherDimension(int64_t dimSize, int64_t shardCount) {
return dimSize * shardCount;
}
-// Return the sharded shape `shape` according ot sharding `sharding`.
-// The shape for the tensor on each device in the grid.
+// Return the per-device sharded type for `shape` based on `sharding`.
+// This is the tensor shape on each grid partition.
// Example:
// On a 2x4x? grid with split axes = [[0], [1], [2]] the shape ?x5x1 would
// result in a shape for each shard of ?x2x?.
diff --git a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
index 5941f7da6ed42..98234bada09e4 100644
--- a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
+++ b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
@@ -748,6 +748,29 @@ bool Sharding::operator==(const Sharding &rhs) const {
bool Sharding::operator!=(const Sharding &rhs) const { return !(*this == rhs); }
+llvm::raw_ostream &mlir::shard::operator<<(llvm::raw_ostream &os,
+ const Sharding &sharding) {
+ os << "Sharding<grid=" << sharding.getGrid() << ", split_axes=[";
+ llvm::interleaveComma(sharding.getSplitAxes(), os, [&](GridAxesAttr axes) {
+ os << "[";
+ llvm::interleaveComma(axes.asArrayRef(), os);
+ os << "]";
+ });
+ os << "]";
+ if (!sharding.getStaticHaloSizes().empty()) {
+ os << ", halo_sizes=[";
+ llvm::interleaveComma(sharding.getStaticHaloSizes(), os);
+ os << "]";
+ }
+ if (!sharding.getStaticShardedDimsOffsets().empty()) {
+ os << ", sharded_dims_offsets=[";
+ llvm::interleaveComma(sharding.getStaticShardedDimsOffsets(), os);
+ os << "]";
+ }
+ os << ">";
+ return os;
+}
+
Sharding::Sharding(::mlir::FlatSymbolRefAttr grid) : grid(grid) {}
Sharding::Sharding(Value rhs) {
diff --git a/mlir/lib/Dialect/Shard/Transforms/Partition.cpp b/mlir/lib/Dialect/Shard/Transforms/Partition.cpp
index a2b3c86cac28d..eb77de69619d0 100644
--- a/mlir/lib/Dialect/Shard/Transforms/Partition.cpp
+++ b/mlir/lib/Dialect/Shard/Transforms/Partition.cpp
@@ -29,7 +29,9 @@
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
+#include <array>
#include <iterator>
+#include <memory>
#include <optional>
#include <tuple>
#include <utility>
@@ -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(
+ ...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/180962
More information about the Mlir-commits
mailing list