derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
1555 lines
56 KiB
C++
1555 lines
56 KiB
C++
//===- ShardOps.cpp - Shard Dialect Operations ----------------------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Shard/IR/ShardOps.h"
|
|
|
|
#include "mlir/Dialect/Arith/IR/Arith.h"
|
|
#include "mlir/Dialect/Shard/IR/ShardDialect.h"
|
|
#include "mlir/Dialect/Utils/StaticValueUtils.h"
|
|
#include "mlir/IR/Attributes.h"
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypeInterfaces.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
#include "mlir/IR/Diagnostics.h"
|
|
#include "mlir/IR/DialectImplementation.h"
|
|
#include "mlir/IR/IRMapping.h"
|
|
#include "mlir/IR/Location.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/IR/TypeUtilities.h"
|
|
#include "mlir/IR/Value.h"
|
|
#include "mlir/Interfaces/ViewLikeInterface.h"
|
|
#include "mlir/Support/LLVM.h"
|
|
#include "mlir/Transforms/InliningUtils.h"
|
|
#include "llvm/ADT/ArrayRef.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/ADT/SmallSet.h"
|
|
#include "llvm/ADT/SmallVector.h"
|
|
#include "llvm/ADT/TypeSwitch.h"
|
|
#include <algorithm>
|
|
#include <functional>
|
|
#include <iterator>
|
|
#include <numeric>
|
|
#include <optional>
|
|
#include <utility>
|
|
|
|
#define DEBUG_TYPE "shard-ops"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::shard;
|
|
|
|
#include "mlir/Dialect/Shard/IR/ShardDialect.cpp.inc"
|
|
|
|
namespace {
|
|
|
|
struct DimensionSize {
|
|
static DimensionSize dynamic() { return DimensionSize(ShapedType::kDynamic); }
|
|
DimensionSize(int64_t val) : val(val) {}
|
|
int64_t value() const { return val; }
|
|
operator int64_t() const { return val; }
|
|
bool isDynamic() const { return ShapedType::isDynamic(val); }
|
|
|
|
private:
|
|
int64_t val;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
static DimensionSize operator/(DimensionSize lhs, DimensionSize rhs) {
|
|
if (lhs.isDynamic() || rhs.isDynamic()) {
|
|
return DimensionSize::dynamic();
|
|
}
|
|
return lhs.value() / rhs.value();
|
|
}
|
|
|
|
static DimensionSize operator*(DimensionSize lhs, DimensionSize rhs) {
|
|
if (lhs.isDynamic() || rhs.isDynamic()) {
|
|
return DimensionSize::dynamic();
|
|
}
|
|
return lhs.value() * rhs.value();
|
|
}
|
|
|
|
SmallVector<Value>
|
|
mlir::shard::getMixedAsValues(OpBuilder b, const Location &loc,
|
|
llvm::ArrayRef<int64_t> statics,
|
|
ValueRange dynamics, Type type) {
|
|
SmallVector<Value> values;
|
|
auto dyn = dynamics.begin();
|
|
Type i64 = b.getI64Type();
|
|
if (!type)
|
|
type = i64;
|
|
assert((i64 == type || b.getIndexType() == type) &&
|
|
"expected an i64 or an intex type");
|
|
for (auto s : statics) {
|
|
if (s == ShapedType::kDynamic) {
|
|
values.emplace_back(*(dyn++));
|
|
} else {
|
|
TypedAttr val = type == i64 ? b.getI64IntegerAttr(s) : b.getIndexAttr(s);
|
|
values.emplace_back(arith::ConstantOp::create(b, loc, type, val));
|
|
}
|
|
}
|
|
return values;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Inliner
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
struct ShardInlinerinterface : public DialectInlinerInterface {
|
|
using DialectInlinerInterface::DialectInlinerInterface;
|
|
// Currently no restrictions are encoded for inlining.
|
|
bool isLegalToInline(Operation *, Operation *, bool) const final {
|
|
return true;
|
|
}
|
|
bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
|
|
return true;
|
|
}
|
|
bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
|
|
return true;
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Shard dialect
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void ShardDialect::initialize() {
|
|
addOperations<
|
|
#define GET_OP_LIST
|
|
#include "mlir/Dialect/Shard/IR/ShardOps.cpp.inc"
|
|
>();
|
|
addAttributes<
|
|
#define GET_ATTRDEF_LIST
|
|
#include "mlir/Dialect/Shard/IR/ShardAttributes.cpp.inc"
|
|
>();
|
|
addTypes<
|
|
#define GET_TYPEDEF_LIST
|
|
#include "mlir/Dialect/Shard/IR/ShardTypes.cpp.inc"
|
|
>();
|
|
addInterface<ShardInlinerinterface>();
|
|
}
|
|
|
|
Operation *ShardDialect::materializeConstant(OpBuilder &builder,
|
|
Attribute value, Type type,
|
|
Location loc) {
|
|
return arith::ConstantOp::materialize(builder, value, type, loc);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Shard utilities
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
static FailureOr<GridOp> getGridAndVerify(Operation *op,
|
|
FlatSymbolRefAttr gridSymbol,
|
|
SymbolTableCollection &symbolTable) {
|
|
shard::GridOp grid = getGridOrNull(op, gridSymbol, symbolTable);
|
|
if (!grid) {
|
|
return op->emitError() << "Undefined required grid symbol \""
|
|
<< gridSymbol.getValue() << "\".";
|
|
}
|
|
|
|
return grid;
|
|
}
|
|
|
|
template <typename It>
|
|
static bool isUnique(It begin, It end) {
|
|
if (begin == end) {
|
|
return true;
|
|
}
|
|
It next = std::next(begin);
|
|
if (next == end) {
|
|
return true;
|
|
}
|
|
for (; next != end; ++next, ++begin) {
|
|
if (*begin == *next) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static LogicalResult verifyGridAxes(Location loc, ArrayRef<GridAxis> axes,
|
|
GridOp grid) {
|
|
SmallVector<GridAxis> sorted = llvm::to_vector(axes);
|
|
llvm::sort(sorted);
|
|
if (!isUnique(sorted.begin(), sorted.end())) {
|
|
return emitError(loc) << "Grid axes contains duplicate elements.";
|
|
}
|
|
|
|
GridAxis rank = grid.getRank();
|
|
for (auto axis : axes) {
|
|
if (axis >= rank || axis < 0) {
|
|
return emitError(loc)
|
|
<< "0-based grid axis index " << axis
|
|
<< " is out of bounds. The referenced grid \"" << grid.getSymName()
|
|
<< "\" is of rank " << rank << ".";
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
template <typename Op>
|
|
static FailureOr<GridOp>
|
|
getGridAndVerifyAxes(Op op, SymbolTableCollection &symbolTable) {
|
|
auto grid =
|
|
::getGridAndVerify(op.getOperation(), op.getGridAttr(), symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyGridAxes(op.getLoc(), op.getGridAxes(), grid.value()))) {
|
|
return failure();
|
|
}
|
|
return grid;
|
|
}
|
|
|
|
template <typename InShape, typename GridShape, typename SplitAxes,
|
|
typename OutShape>
|
|
static void shardShape(const InShape &inShape, const GridShape &gridShape,
|
|
const SplitAxes &splitAxes, OutShape &outShape,
|
|
ArrayRef<int64_t> shardedDimsOffsets = {},
|
|
ArrayRef<int64_t> haloSizes = {}) {
|
|
// 0d tensors cannot be sharded and must get replicated
|
|
if (inShape.empty()) {
|
|
assert(outShape.empty());
|
|
return;
|
|
}
|
|
|
|
std::copy(llvm::adl_begin(inShape), llvm::adl_end(inShape),
|
|
llvm::adl_begin(outShape));
|
|
|
|
if (!shardedDimsOffsets.empty()) {
|
|
auto isDynShape = ShapedType::isDynamicShape(gridShape);
|
|
uint64_t pos = 1;
|
|
for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
|
|
if (!innerSplitAxes.empty()) {
|
|
auto sz = shardedDimsOffsets[pos];
|
|
bool same = !isDynShape;
|
|
if (same) {
|
|
// Find sharded dims in shardedDimsOffsets with same static size on
|
|
// all devices. Use kDynamic for dimensions with dynamic or
|
|
// non-uniform offs in shardedDimsOffsets.
|
|
uint64_t numShards = 0;
|
|
for (auto i : innerSplitAxes.asArrayRef()) {
|
|
numShards += gridShape[i];
|
|
}
|
|
for (size_t i = 1; i < numShards; ++i) {
|
|
if (shardedDimsOffsets[pos + i] - shardedDimsOffsets[pos + i - 1] !=
|
|
sz) {
|
|
same = false;
|
|
break;
|
|
}
|
|
}
|
|
pos += numShards + 1;
|
|
}
|
|
outShape[tensorAxis] = same ? sz : ShapedType::kDynamic;
|
|
}
|
|
}
|
|
} else {
|
|
for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
|
|
outShape[tensorAxis] = shardDimension(
|
|
inShape[tensorAxis],
|
|
collectiveProcessGroupSize(innerSplitAxes.asArrayRef(), gridShape));
|
|
}
|
|
|
|
if (!haloSizes.empty()) {
|
|
// add halo sizes if requested
|
|
int haloAxis = 0;
|
|
for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
|
|
if (ShapedType::isStatic(outShape[tensorAxis]) &&
|
|
!innerSplitAxes.empty()) {
|
|
if (haloSizes[haloAxis * 2] >= 0 &&
|
|
haloSizes[haloAxis * 2 + 1] >= 0) {
|
|
outShape[tensorAxis] +=
|
|
haloSizes[haloAxis * 2] + haloSizes[haloAxis * 2 + 1];
|
|
++haloAxis;
|
|
} else {
|
|
outShape[tensorAxis] = ShapedType::kDynamic;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
ShapedType shard::shardShapedType(ShapedType shape, GridOp grid,
|
|
Sharding sharding) {
|
|
using Dim = std::decay_t<decltype(shape.getDimSize(0))>;
|
|
SmallVector<Dim> resShapeArr(shape.getShape().size());
|
|
shardShape(shape.getShape(), grid.getShape(), sharding.getSplitAxes(),
|
|
resShapeArr, sharding.getStaticShardedDimsOffsets(),
|
|
sharding.getStaticHaloSizes());
|
|
return shape.clone(resShapeArr);
|
|
}
|
|
|
|
Type shard::shardType(Type type, GridOp grid, Sharding sharding) {
|
|
RankedTensorType rankedTensorType = dyn_cast<RankedTensorType>(type);
|
|
if (rankedTensorType && !rankedTensorType.getShape().empty()) {
|
|
return shardShapedType(rankedTensorType, grid, sharding);
|
|
}
|
|
return type;
|
|
}
|
|
|
|
static void maybeInsertTargetShardingAnnotationImpl(Sharding sharding,
|
|
Value &operandValue,
|
|
Operation *operandOp,
|
|
OpBuilder &builder,
|
|
ShardOp &newShardOp) {
|
|
OpBuilder::InsertionGuard insertionGuard(builder);
|
|
builder.setInsertionPointAfterValue(operandValue);
|
|
ShardOp shardOp = dyn_cast<ShardOp>(operandOp);
|
|
if (shardOp && sharding == shardOp.getSharding() &&
|
|
!shardOp.getAnnotateForUsers()) {
|
|
// No need for anything if the correct sharding is already set.
|
|
if (!newShardOp) {
|
|
newShardOp = shardOp;
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (!newShardOp) {
|
|
auto shardingOp =
|
|
ShardingOp::create(builder, operandValue.getLoc(), sharding);
|
|
newShardOp = ShardOp::create(builder, operandValue.getLoc(), operandValue,
|
|
shardingOp,
|
|
/*annotate_for_users*/ false);
|
|
}
|
|
operandValue.replaceUsesWithIf(
|
|
newShardOp, [operandOp, operandValue](OpOperand &use) {
|
|
return use.getOwner() == operandOp && use.get() == operandValue;
|
|
});
|
|
|
|
if (!shardOp || shardOp.getAnnotateForUsers()) {
|
|
return;
|
|
}
|
|
|
|
auto newShardOp2 = ShardOp::create(builder, operandValue.getLoc(), newShardOp,
|
|
newShardOp.getSharding(),
|
|
/*annotate_for_users*/ true);
|
|
newShardOp.getResult().replaceAllUsesExcept(newShardOp2, newShardOp2);
|
|
}
|
|
|
|
void mlir::shard::maybeInsertTargetShardingAnnotation(Sharding sharding,
|
|
OpResult result,
|
|
OpBuilder &builder) {
|
|
ShardOp newShardOp;
|
|
SmallVector<std::pair<Value, Operation *>> uses;
|
|
for (auto &use : result.getUses()) {
|
|
uses.emplace_back(use.get(), use.getOwner());
|
|
}
|
|
for (auto &[operandValue, operandOp] : uses) {
|
|
maybeInsertTargetShardingAnnotationImpl(sharding, operandValue, operandOp,
|
|
builder, newShardOp);
|
|
}
|
|
}
|
|
|
|
void mlir::shard::maybeInsertSourceShardingAnnotation(Sharding sharding,
|
|
OpOperand &operand,
|
|
OpBuilder &builder) {
|
|
OpBuilder::InsertionGuard insertionGuard(builder);
|
|
Value operandValue = operand.get();
|
|
Operation *operandSrcOp = operandValue.getDefiningOp();
|
|
bool isBlockArg = !operandSrcOp;
|
|
{
|
|
[[maybe_unused]] auto opType =
|
|
dyn_cast<mlir::RankedTensorType>(operandValue.getType());
|
|
assert(!opType || opType.getRank() > 0 || isFullReplication(sharding));
|
|
}
|
|
if (!isa<RankedTensorType>(operandValue.getType()) && operandSrcOp &&
|
|
operandSrcOp->hasTrait<OpTrait::ConstantLike>()) {
|
|
return;
|
|
}
|
|
|
|
Operation *operandOp = operand.getOwner();
|
|
ShardOp shardOp = dyn_cast_or_null<ShardOp>(operandSrcOp);
|
|
|
|
if (shardOp && sharding == shardOp.getSharding() &&
|
|
shardOp.getAnnotateForUsers()) {
|
|
// No need for anything the correct sharding is already set.
|
|
return;
|
|
}
|
|
|
|
builder.setInsertionPoint(operandOp);
|
|
auto shardingOp =
|
|
ShardingOp::create(builder, operand.get().getLoc(), sharding);
|
|
auto newShardOp =
|
|
ShardOp::create(builder, operandValue.getLoc(), operandValue, shardingOp,
|
|
/*annotate_for_users*/ true);
|
|
IRRewriter rewriter(builder);
|
|
rewriter.replaceUsesWithIf(
|
|
operandValue, newShardOp, [operandOp, operandValue](OpOperand &use) {
|
|
return use.getOwner() == operandOp && use.get() == operandValue;
|
|
});
|
|
|
|
if (isBlockArg || !shardOp || !shardOp.getAnnotateForUsers()) {
|
|
// No need for resharding.
|
|
return;
|
|
}
|
|
|
|
builder.setInsertionPoint(newShardOp);
|
|
auto newPreceedingShardOp =
|
|
ShardOp::create(builder, operandValue.getLoc(), operandValue, shardingOp,
|
|
/*annotate_for_users*/ false);
|
|
rewriter.replaceUsesWithIf(
|
|
newShardOp.getSrc(), newPreceedingShardOp, [&newShardOp](OpOperand &use) {
|
|
return use.getOwner() == newShardOp.getOperation();
|
|
});
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.grid op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult GridOp::verify() {
|
|
int64_t rank = getRank();
|
|
|
|
if (rank <= 0)
|
|
return emitOpError("rank of grid is expected to be a positive integer");
|
|
|
|
for (int64_t dimSize : getShape()) {
|
|
if (dimSize < 0 && ShapedType::isStatic(dimSize))
|
|
return emitOpError("dimension size of a grid is expected to be "
|
|
"non-negative or dynamic");
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.grid_shape op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
GridShapeOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = ::getGridAndVerify(getOperation(), getGridAttr(), symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyGridAxes(getLoc(), getAxes(), grid.value()))) {
|
|
return failure();
|
|
}
|
|
|
|
size_t expectedResultsCount =
|
|
getAxes().empty() ? grid->getRank() : getAxes().size();
|
|
if (getResult().size() != expectedResultsCount) {
|
|
return emitError() << "Unexpected number of results " << getResult().size()
|
|
<< ". Expected " << expectedResultsCount << ".";
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void GridShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
GridOp grid) {
|
|
build(odsBuilder, odsState, grid, SmallVector<GridAxis>());
|
|
}
|
|
|
|
void GridShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
GridOp grid, ArrayRef<GridAxis> axes) {
|
|
build(odsBuilder, odsState,
|
|
SmallVector<Type>(axes.empty() ? grid.getRank() : axes.size(),
|
|
odsBuilder.getIndexType()),
|
|
grid.getSymName(), GridAxesAttr::get(odsBuilder.getContext(), axes));
|
|
}
|
|
|
|
void GridShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
StringRef grid, ArrayRef<GridAxis> axes) {
|
|
assert(!axes.empty());
|
|
build(odsBuilder, odsState,
|
|
SmallVector<Type>(axes.size(), odsBuilder.getIndexType()), grid,
|
|
GridAxesAttr::get(odsBuilder.getContext(), axes));
|
|
}
|
|
|
|
void GridShapeOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResults()[0], "grid_shape");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.sharding
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void ShardingOp::build(::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
|
|
FlatSymbolRefAttr grid, ArrayRef<GridAxesAttr> splitAxes,
|
|
ArrayRef<int64_t> staticHalos,
|
|
ArrayRef<int64_t> staticOffsets) {
|
|
return build(
|
|
b, odsState, grid, GridAxesArrayAttr::get(b.getContext(), splitAxes),
|
|
::mlir::DenseI64ArrayAttr::get(b.getContext(), staticHalos), {},
|
|
::mlir::DenseI64ArrayAttr::get(b.getContext(), staticOffsets), {});
|
|
}
|
|
|
|
void ShardingOp::build(::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
|
|
llvm::StringRef grid, ArrayRef<GridAxesAttr> splitAxes,
|
|
ArrayRef<int64_t> staticHalos,
|
|
ArrayRef<int64_t> staticOffsets) {
|
|
return build(b, odsState, FlatSymbolRefAttr::get(b.getContext(), grid),
|
|
GridAxesArrayAttr::get(b.getContext(), splitAxes),
|
|
::mlir::DenseI64ArrayAttr::get(b.getContext(), staticHalos), {},
|
|
::mlir::DenseI64ArrayAttr::get(b.getContext(), staticOffsets),
|
|
{});
|
|
}
|
|
|
|
void ShardingOp::build(
|
|
::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
|
|
FlatSymbolRefAttr grid, ArrayRef<GridAxesAttr> splitAxes,
|
|
::mlir::ArrayRef<::mlir::OpFoldResult> haloSizes,
|
|
::mlir::ArrayRef<::mlir::OpFoldResult> shardedDimsOffsets) {
|
|
mlir::SmallVector<int64_t> staticHalos, staticDims;
|
|
mlir::SmallVector<mlir::Value> dynamicHalos, dynamicDims;
|
|
dispatchIndexOpFoldResults(haloSizes, dynamicHalos, staticHalos);
|
|
dispatchIndexOpFoldResults(shardedDimsOffsets, dynamicDims, staticDims);
|
|
return build(
|
|
b, odsState, grid, GridAxesArrayAttr::get(b.getContext(), splitAxes),
|
|
::mlir::DenseI64ArrayAttr::get(b.getContext(), staticHalos), dynamicHalos,
|
|
::mlir::DenseI64ArrayAttr::get(b.getContext(), staticDims), dynamicDims);
|
|
}
|
|
|
|
void ShardingOp::build(::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
|
|
mlir::shard::Sharding from) {
|
|
|
|
build(b, odsState, ShardingType::get(b.getContext()), from.getGridAttr(),
|
|
GridAxesArrayAttr::get(b.getContext(), from.getSplitAxes()),
|
|
from.getStaticShardedDimsOffsets().empty()
|
|
? DenseI64ArrayAttr()
|
|
: b.getDenseI64ArrayAttr(from.getStaticShardedDimsOffsets()),
|
|
from.getDynamicShardedDimsOffsets(),
|
|
from.getStaticHaloSizes().empty()
|
|
? DenseI64ArrayAttr()
|
|
: b.getDenseI64ArrayAttr(from.getStaticHaloSizes()),
|
|
from.getDynamicHaloSizes());
|
|
}
|
|
|
|
LogicalResult ShardingOp::verify() {
|
|
llvm::SmallSet<GridAxis, 4> visitedAxes;
|
|
|
|
auto checkGridAxis = [&](ArrayRef<GridAxis> axesArray) -> LogicalResult {
|
|
for (GridAxis axis : axesArray) {
|
|
if (axis < 0)
|
|
return emitError() << "grid axis is expected to be non-negative";
|
|
if (!visitedAxes.insert(axis).second)
|
|
return emitError() << "grid axis duplicated";
|
|
}
|
|
return success();
|
|
};
|
|
|
|
for (auto subAxes : getSplitAxes().getAxes()) {
|
|
ArrayRef<GridAxis> subAxesArray = subAxes.asArrayRef();
|
|
if (failed(checkGridAxis(subAxesArray)))
|
|
return failure();
|
|
}
|
|
|
|
if (!getStaticHaloSizes().empty() && !getStaticShardedDimsOffsets().empty()) {
|
|
return emitOpError("halo sizes and shard offsets are mutually exclusive");
|
|
}
|
|
|
|
if (!getStaticHaloSizes().empty()) {
|
|
auto numSplitAxes = getSplitAxes().getAxes().size();
|
|
for (auto splitAxis : getSplitAxes().getAxes()) {
|
|
if (splitAxis.empty()) {
|
|
--numSplitAxes;
|
|
}
|
|
}
|
|
if (getStaticHaloSizes().size() != numSplitAxes * 2) {
|
|
return emitError() << "halo sizes must be specified for all split axes.";
|
|
}
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void ShardingOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "sharding");
|
|
}
|
|
|
|
LogicalResult ShardingOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = ::getGridAndVerify(getOperation(), getGridAttr(), symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (mlir::ShapedType::isDynamicShape(grid->getShape()) &&
|
|
!getStaticShardedDimsOffsets().empty()) {
|
|
return emitError() << "sharded dims offsets are not allowed for "
|
|
"device grids with dynamic shape.";
|
|
}
|
|
|
|
auto shardedDimsOffsets = getStaticShardedDimsOffsets();
|
|
if (!shardedDimsOffsets.empty()) {
|
|
auto gridShape = grid.value().getShape();
|
|
assert(ShapedType::isStaticShape(gridShape));
|
|
uint64_t pos = 0;
|
|
for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(getSplitAxes())) {
|
|
if (!innerSplitAxes.empty()) {
|
|
int64_t numShards = 0, off = 0;
|
|
for (auto i : innerSplitAxes.asArrayRef()) {
|
|
numShards += gridShape[i];
|
|
}
|
|
for (int64_t i = 0; i <= numShards; ++i) {
|
|
if (shardedDimsOffsets.size() <= pos + i) {
|
|
return emitError() << "sharded dims offsets has wrong size.";
|
|
}
|
|
if (ShapedType::isStatic(shardedDimsOffsets[pos + i])) {
|
|
if (shardedDimsOffsets[pos + i] < off) {
|
|
return emitError()
|
|
<< "sharded dims offsets must be non-decreasing.";
|
|
}
|
|
off = shardedDimsOffsets[pos + i];
|
|
}
|
|
}
|
|
pos += numShards + 1;
|
|
}
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
namespace {
|
|
// Sharding annotations "halo sizes" and "sharded dims offsets"
|
|
// are a mix of attributes and dynamic values. This canonicalization moves
|
|
// constant values to the respective attribute lists, minimizing the number
|
|
// of values.
|
|
// It also removes sharded_dims_sizes and halos if they are effectively "empty".
|
|
class NormalizeSharding final : public OpRewritePattern<ShardingOp> {
|
|
public:
|
|
using OpRewritePattern<ShardingOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(ShardingOp op,
|
|
PatternRewriter &b) const override {
|
|
auto mixedHalos =
|
|
getMixedValues(op.getStaticHaloSizes(), op.getDynamicHaloSizes(), b);
|
|
auto mixedOffs = getMixedValues(op.getStaticShardedDimsOffsets(),
|
|
op.getDynamicShardedDimsOffsets(), b);
|
|
|
|
// No constant operands were folded, just return;
|
|
bool modified = succeeded(foldDynamicIndexList(mixedHalos, true)) ||
|
|
succeeded(foldDynamicIndexList(mixedOffs, true));
|
|
|
|
auto decomposedHalos = decomposeMixedValues(mixedHalos);
|
|
auto staticHalos = decomposedHalos.first;
|
|
auto dynamicHalos = decomposedHalos.second;
|
|
auto decomposedOffs = decomposeMixedValues(mixedOffs);
|
|
auto staticOffs = decomposedOffs.first;
|
|
auto dynamicOffs = decomposedOffs.second;
|
|
|
|
if (dynamicHalos.empty() && !staticHalos.empty()) {
|
|
if (staticHalos[0] == 0 && llvm::all_equal(staticHalos)) {
|
|
staticHalos.clear();
|
|
modified = true;
|
|
}
|
|
}
|
|
|
|
// Remove sharded dims offsets if they are effectively the default values,
|
|
// e.g. if they define equi-distance between all neighboring shards.
|
|
// Requires static-only offsets. Compares the first distance as the
|
|
// difference between the first two offsets. Only if all consecutive
|
|
// distances are the same, the offsets are removed.
|
|
if (dynamicOffs.empty() && !staticOffs.empty()) {
|
|
assert(staticOffs.size() >= 2);
|
|
auto diff = staticOffs[1] - staticOffs[0];
|
|
bool allSame = staticOffs.size() > 2;
|
|
for (auto i = 2u; i < staticOffs.size(); ++i) {
|
|
if (staticOffs[i] - staticOffs[i - 1] != diff) {
|
|
allSame = false;
|
|
break;
|
|
}
|
|
}
|
|
if (allSame) {
|
|
staticOffs.clear();
|
|
modified = true;
|
|
}
|
|
}
|
|
|
|
if (!modified) {
|
|
return failure();
|
|
}
|
|
|
|
b.modifyOpInPlace(op, [&]() {
|
|
op.setStaticHaloSizes(staticHalos);
|
|
op.getDynamicHaloSizesMutable().assign(dynamicHalos);
|
|
op.setStaticShardedDimsOffsets(staticOffs);
|
|
op.getDynamicShardedDimsOffsetsMutable().assign(dynamicOffs);
|
|
});
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void ShardingOp::getCanonicalizationPatterns(mlir::RewritePatternSet &results,
|
|
mlir::MLIRContext *context) {
|
|
results.add<NormalizeSharding>(context);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Sharding
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
bool Sharding::equalSplitAxes(const Sharding &rhs) const {
|
|
if (getGrid() != rhs.getGrid()) {
|
|
return false;
|
|
}
|
|
|
|
auto minSize = std::min(getSplitAxes().size(), rhs.getSplitAxes().size());
|
|
if (!llvm::equal(llvm::make_range(getSplitAxes().begin(),
|
|
getSplitAxes().begin() + minSize),
|
|
llvm::make_range(rhs.getSplitAxes().begin(),
|
|
rhs.getSplitAxes().begin() + minSize))) {
|
|
return false;
|
|
}
|
|
|
|
return llvm::all_of(llvm::drop_begin(getSplitAxes(), minSize),
|
|
std::mem_fn(&GridAxesAttr::empty)) &&
|
|
llvm::all_of(llvm::drop_begin(rhs.getSplitAxes(), minSize),
|
|
std::mem_fn(&GridAxesAttr::empty));
|
|
}
|
|
|
|
bool Sharding::equalHaloAndShardSizes(const Sharding &rhs) const {
|
|
return equalShardSizes(rhs) && equalHaloSizes(rhs);
|
|
}
|
|
|
|
bool Sharding::equalShardSizes(const Sharding &rhs) const {
|
|
if (rhs.getStaticShardedDimsOffsets().size() !=
|
|
getStaticShardedDimsOffsets().size() ||
|
|
!llvm::equal(getStaticShardedDimsOffsets(),
|
|
rhs.getStaticShardedDimsOffsets())) {
|
|
return false;
|
|
}
|
|
if (rhs.getDynamicShardedDimsOffsets().size() !=
|
|
getDynamicShardedDimsOffsets().size() ||
|
|
!llvm::equal(getDynamicShardedDimsOffsets(),
|
|
rhs.getDynamicShardedDimsOffsets())) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool Sharding::equalHaloSizes(const Sharding &rhs) const {
|
|
if (rhs.getStaticHaloSizes().size() != getStaticHaloSizes().size() ||
|
|
!llvm::equal(getStaticHaloSizes(), rhs.getStaticHaloSizes())) {
|
|
return false;
|
|
}
|
|
if (rhs.getDynamicHaloSizes().size() != getDynamicHaloSizes().size() ||
|
|
!llvm::equal(getDynamicHaloSizes(), rhs.getDynamicHaloSizes())) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool Sharding::operator==(Value rhs) const {
|
|
return equalSplitAxes(rhs) && equalHaloAndShardSizes(rhs);
|
|
}
|
|
|
|
bool Sharding::operator!=(Value rhs) const { return !(*this == rhs); }
|
|
|
|
bool Sharding::operator==(const Sharding &rhs) const {
|
|
return equalSplitAxes(rhs) && equalHaloAndShardSizes(rhs);
|
|
}
|
|
|
|
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) {
|
|
auto shardingOp = rhs.getDefiningOp<ShardingOp>();
|
|
assert(shardingOp && "expected sharding op");
|
|
auto splitAxes = shardingOp.getSplitAxes().getAxes();
|
|
// If splitAxes are empty, use "empty" constructor.
|
|
if (splitAxes.empty()) {
|
|
*this = Sharding(shardingOp.getGridAttr());
|
|
return;
|
|
}
|
|
*this =
|
|
get(shardingOp.getGridAttr(), splitAxes, shardingOp.getStaticHaloSizes(),
|
|
shardingOp.getStaticShardedDimsOffsets(),
|
|
SmallVector<Value>(shardingOp.getDynamicHaloSizes()),
|
|
SmallVector<Value>(shardingOp.getDynamicShardedDimsOffsets()));
|
|
}
|
|
|
|
Sharding Sharding::get(::mlir::FlatSymbolRefAttr grid,
|
|
ArrayRef<GridAxesAttr> splitAxes,
|
|
ArrayRef<int64_t> staticHaloSizes,
|
|
ArrayRef<int64_t> staticShardedDimsOffsets,
|
|
ArrayRef<Value> dynamicHaloSizes,
|
|
ArrayRef<Value> dynamicShardedDimsOffsets) {
|
|
Sharding res(grid);
|
|
if (splitAxes.empty()) {
|
|
return res;
|
|
}
|
|
|
|
res.split_axes.resize(splitAxes.size());
|
|
for (auto [i, axis] : llvm::enumerate(splitAxes)) {
|
|
res.split_axes[i] = GridAxesAttr::get(grid.getContext(), axis.asArrayRef());
|
|
}
|
|
|
|
auto clone = [](const auto src, auto &dst) {
|
|
dst.resize(src.size());
|
|
llvm::copy(src, dst.begin());
|
|
};
|
|
|
|
clone(staticHaloSizes, res.static_halo_sizes);
|
|
clone(staticShardedDimsOffsets, res.static_sharded_dims_offsets);
|
|
clone(dynamicHaloSizes, res.dynamic_halo_sizes);
|
|
clone(dynamicShardedDimsOffsets, res.dynamic_sharded_dims_offsets);
|
|
|
|
return res;
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.shard_shape
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void ShardShapeOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult()[0], "shard_shape");
|
|
}
|
|
|
|
void ShardShapeOp::build(::mlir::OpBuilder &odsBuilder,
|
|
::mlir::OperationState &odsState,
|
|
::llvm::ArrayRef<int64_t> dims,
|
|
ArrayRef<Value> dimsDyn, ::mlir::Value sharding,
|
|
::mlir::ValueRange device) {
|
|
SmallVector<mlir::Type> resType(dims.size(), odsBuilder.getIndexType());
|
|
build(odsBuilder, odsState, resType, dims, dimsDyn, sharding,
|
|
SmallVector<int64_t>(device.size(), ShapedType::kDynamic), device);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.shard op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
void ShardOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "sharding_annotated");
|
|
}
|
|
|
|
namespace {
|
|
// Determine if the given ShardOp is a duplicate of another ShardOp
|
|
// on the same value. This can happen if constant values are sharded.
|
|
class FoldDuplicateShardOp final : public OpRewritePattern<ShardOp> {
|
|
public:
|
|
using OpRewritePattern<ShardOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(ShardOp op, PatternRewriter &b) const override {
|
|
// Get the use-list of the value being sharded and check if it has more than
|
|
// one use.
|
|
Value value = op.getSrc();
|
|
if (value.hasOneUse() || value.getDefiningOp<ShardOp>()) {
|
|
return failure();
|
|
}
|
|
|
|
// Iterate through the uses of the value to find a duplicate ShardOp.
|
|
for (auto &use : value.getUses()) {
|
|
if (use.getOwner() != op.getOperation()) {
|
|
auto otherOp = dyn_cast<ShardOp>(use.getOwner());
|
|
if (!otherOp || !otherOp->isBeforeInBlock(op)) {
|
|
return failure();
|
|
}
|
|
// Create a Sharding object for the current and the other ShardOp
|
|
// If the two are equal replace current op with the other op.
|
|
Sharding currentSharding(op.getSharding());
|
|
Sharding otherSharding(otherOp.getSharding());
|
|
if (currentSharding == otherSharding) {
|
|
b.replaceAllUsesWith(op.getResult(), otherOp.getResult());
|
|
b.eraseOp(op.getOperation());
|
|
} else {
|
|
// use the other sharding as input for op
|
|
b.modifyOpInPlace(
|
|
op, [&]() { op.getSrcMutable().assign(otherOp.getResult()); });
|
|
}
|
|
return success();
|
|
}
|
|
}
|
|
|
|
return failure();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void ShardOp::getCanonicalizationPatterns(mlir::RewritePatternSet &results,
|
|
mlir::MLIRContext *context) {
|
|
results.add<FoldDuplicateShardOp>(context);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.process_multi_index op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
ProcessMultiIndexOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = ::getGridAndVerify(getOperation(), getGridAttr(), symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyGridAxes(getLoc(), getAxes(), grid.value()))) {
|
|
return failure();
|
|
}
|
|
|
|
size_t expectedResultsCount =
|
|
getAxes().empty() ? grid->getRank() : getAxes().size();
|
|
if (getResult().size() != expectedResultsCount) {
|
|
return emitError() << "Unexpected number of results " << getResult().size()
|
|
<< ". Expected " << expectedResultsCount << ".";
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void ProcessMultiIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
GridOp grid) {
|
|
build(odsBuilder, odsState,
|
|
SmallVector<Type>(grid.getRank(), odsBuilder.getIndexType()),
|
|
grid.getSymName(), ArrayRef<GridAxis>());
|
|
}
|
|
|
|
void ProcessMultiIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
StringRef grid, ArrayRef<GridAxis> axes) {
|
|
build(odsBuilder, odsState,
|
|
SmallVector<Type>(axes.size(), odsBuilder.getIndexType()), grid,
|
|
GridAxesAttr::get(odsBuilder.getContext(), axes));
|
|
}
|
|
|
|
void ProcessMultiIndexOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResults()[0], "proc_linear_idx");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.process_linear_index op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
ProcessLinearIndexOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = ::getGridAndVerify(getOperation(), getGridAttr(), symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
void ProcessLinearIndexOp::build(OpBuilder &odsBuilder,
|
|
OperationState &odsState, GridOp grid) {
|
|
build(odsBuilder, odsState, grid.getSymName());
|
|
}
|
|
|
|
void ProcessLinearIndexOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "proc_linear_idx");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.neighbors_linear_indices op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
NeighborsLinearIndicesOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = ::getGridAndVerify(getOperation(), getGridAttr(), symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
void NeighborsLinearIndicesOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getNeighborDown(), "down_linear_idx");
|
|
setNameFn(getNeighborUp(), "up_linear_idx");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// collective communication ops
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
template <typename Op>
|
|
struct EmptyGridAxesCanonicalizationPattern : OpRewritePattern<Op> {
|
|
using OpRewritePattern<Op>::OpRewritePattern;
|
|
LogicalResult matchAndRewrite(Op op,
|
|
PatternRewriter &rewriter) const override {
|
|
auto gridAxes = op.getGridAxes();
|
|
if (!gridAxes.empty()) {
|
|
return failure();
|
|
}
|
|
if (op.getInput().getType() != op.getResult().getType()) {
|
|
return failure();
|
|
}
|
|
|
|
rewriter.replaceAllUsesWith(op.getResult(), op.getInput());
|
|
rewriter.eraseOp(op.getOperation());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
static LogicalResult verifyInGroupDevice(Location loc, StringRef deviceName,
|
|
ArrayRef<int64_t> device,
|
|
Operation::operand_range deviceDynamic,
|
|
ArrayRef<GridAxis> gridAxes,
|
|
ArrayRef<int64_t> gridShape) {
|
|
if (device.size() != gridAxes.size()) {
|
|
return emitError(loc) << "In-group device \"" << deviceName
|
|
<< "\" has unexpected multi-index size "
|
|
<< device.size() << ". Expected " << gridAxes.size()
|
|
<< ".";
|
|
}
|
|
|
|
for (size_t i = 0; i < device.size(); ++i) {
|
|
if (ShapedType::isStatic(device[i]) &&
|
|
ShapedType::isStatic(gridShape[gridAxes[i]]) &&
|
|
gridShape[gridAxes[i]] <= device[i]) {
|
|
return emitError(loc)
|
|
<< "Out of bounds coordinate " << i << " for in-group device \""
|
|
<< deviceName << "\"."
|
|
<< " Got " << device[i] << ", but expected value in the range [0, "
|
|
<< (gridShape[gridAxes[i]] - 1) << "].";
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verifyDimensionCompatibility(Location loc,
|
|
int64_t expectedDimSize,
|
|
int64_t resultDimSize,
|
|
int64_t resultAxis) {
|
|
if (ShapedType::isStatic(resultDimSize) && expectedDimSize != resultDimSize) {
|
|
return emitError(loc) << "Dimension size mismatch for result axis "
|
|
<< resultAxis << ". Expected "
|
|
<< (ShapedType::isDynamic(expectedDimSize)
|
|
? Twine("dynamic")
|
|
: Twine(expectedDimSize))
|
|
<< ", but got " << resultDimSize << ".";
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verifyGatherOperandAndResultShape(
|
|
Value operand, Value result, int64_t gatherAxis,
|
|
ArrayRef<GridAxis> gridAxes, ArrayRef<int64_t> gridShape) {
|
|
auto resultRank = cast<ShapedType>(result.getType()).getRank();
|
|
if (gatherAxis < 0 || gatherAxis >= resultRank) {
|
|
return emitError(result.getLoc())
|
|
<< "Gather axis " << gatherAxis << " is out of bounds [0, "
|
|
<< resultRank << ").";
|
|
}
|
|
|
|
ShapedType operandType = cast<ShapedType>(operand.getType());
|
|
ShapedType resultType = cast<ShapedType>(result.getType());
|
|
auto deviceGroupSize =
|
|
DimensionSize(collectiveProcessGroupSize(gridAxes, gridShape));
|
|
for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
|
|
auto operandDimSize = DimensionSize(operandType.getDimSize(axis));
|
|
auto resultDimSize = DimensionSize(resultType.getDimSize(axis));
|
|
auto expectedResultDimSize =
|
|
axis == gatherAxis ? deviceGroupSize * operandDimSize : operandDimSize;
|
|
if (failed(verifyDimensionCompatibility(
|
|
result.getLoc(), expectedResultDimSize, resultDimSize, axis))) {
|
|
return failure();
|
|
}
|
|
}
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verifyAllToAllOperandAndResultShape(
|
|
Value operand, Value result, int64_t splitAxis, int64_t concatAxis,
|
|
ArrayRef<GridAxis> gridAxes, ArrayRef<int64_t> gridShape) {
|
|
ShapedType operandType = cast<ShapedType>(operand.getType());
|
|
ShapedType resultType = cast<ShapedType>(result.getType());
|
|
for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
|
|
if ((axis != splitAxis && axis != concatAxis) || splitAxis == concatAxis) {
|
|
if (failed(verifyDimensionCompatibility(
|
|
result.getLoc(), operandType.getDimSize(axis),
|
|
resultType.getDimSize(axis), axis))) {
|
|
return failure();
|
|
}
|
|
}
|
|
}
|
|
|
|
if (splitAxis == concatAxis) {
|
|
return success();
|
|
}
|
|
|
|
auto deviceGroupSize =
|
|
DimensionSize(collectiveProcessGroupSize(gridAxes, gridShape));
|
|
auto operandConcatDimSize = DimensionSize(operandType.getDimSize(concatAxis));
|
|
auto operandSplitDimSize = DimensionSize(operandType.getDimSize(splitAxis));
|
|
DimensionSize expectedResultConcatDimSize =
|
|
operandConcatDimSize * deviceGroupSize;
|
|
DimensionSize expectedResultSplitDimSize =
|
|
operandSplitDimSize / deviceGroupSize;
|
|
if (!expectedResultSplitDimSize.isDynamic() &&
|
|
int64_t(operandSplitDimSize) % int64_t(deviceGroupSize) != 0) {
|
|
expectedResultSplitDimSize = DimensionSize::dynamic();
|
|
}
|
|
if (failed(verifyDimensionCompatibility(
|
|
result.getLoc(), expectedResultConcatDimSize.value(),
|
|
resultType.getDimSize(concatAxis), concatAxis))) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyDimensionCompatibility(
|
|
result.getLoc(), expectedResultSplitDimSize.value(),
|
|
resultType.getDimSize(splitAxis), splitAxis))) {
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult verifyScatterOrSliceOperandAndResultShape(
|
|
Value operand, Value result, int64_t tensorAxis,
|
|
ArrayRef<GridAxis> gridAxes, ArrayRef<int64_t> gridShape) {
|
|
ShapedType operandType = cast<ShapedType>(operand.getType());
|
|
ShapedType resultType = cast<ShapedType>(result.getType());
|
|
for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
|
|
if (axis != tensorAxis) {
|
|
if (failed(verifyDimensionCompatibility(
|
|
result.getLoc(), operandType.getDimSize(axis),
|
|
resultType.getDimSize(axis), axis))) {
|
|
return failure();
|
|
}
|
|
}
|
|
}
|
|
|
|
auto deviceGroupSize =
|
|
DimensionSize(collectiveProcessGroupSize(gridAxes, gridShape));
|
|
auto operandScatterDimSize =
|
|
DimensionSize(operandType.getDimSize(tensorAxis));
|
|
if (!operandScatterDimSize.isDynamic() && !deviceGroupSize.isDynamic() &&
|
|
int64_t(operandScatterDimSize) % int64_t(deviceGroupSize) != 0) {
|
|
return emitError(result.getLoc())
|
|
<< "Operand dimension size " << int64_t(operandScatterDimSize)
|
|
<< " is not divisible by collective device group size "
|
|
<< int64_t(deviceGroupSize) << " for tensor axis " << tensorAxis
|
|
<< ".";
|
|
}
|
|
DimensionSize expectedResultTensorDimSize =
|
|
operandScatterDimSize / deviceGroupSize;
|
|
if (failed(verifyDimensionCompatibility(
|
|
result.getLoc(), expectedResultTensorDimSize.value(),
|
|
resultType.getDimSize(tensorAxis), tensorAxis))) {
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
static RankedTensorType sliceResultType(Type operandType, GridOp grid,
|
|
ArrayRef<GridAxis> gridAxes,
|
|
int64_t sliceAxis) {
|
|
RankedTensorType operandRankedTensorType =
|
|
cast<RankedTensorType>(operandType);
|
|
DimensionSize operandSliceAxisSize =
|
|
operandRankedTensorType.getShape()[sliceAxis];
|
|
SmallVector<int64_t> resultShape =
|
|
llvm::to_vector(operandRankedTensorType.getShape());
|
|
|
|
resultShape[sliceAxis] =
|
|
operandSliceAxisSize /
|
|
DimensionSize(collectiveProcessGroupSize(gridAxes, grid));
|
|
return operandRankedTensorType.clone(resultShape);
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.all_gather op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
AllGatherOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
auto gatherAxis = getGatherAxis().getSExtValue();
|
|
return verifyGatherOperandAndResultShape(getOperand(), getResult(),
|
|
gatherAxis, getGridAxes(),
|
|
grid.value().getShape());
|
|
}
|
|
|
|
void AllGatherOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<AllGatherOp>>(context);
|
|
}
|
|
|
|
void AllGatherOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "all_gather");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.all_reduce op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
AllReduceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
return getGridAndVerifyAxes(*this, symbolTable);
|
|
}
|
|
|
|
void AllReduceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<AllReduceOp>>(context);
|
|
}
|
|
|
|
void AllReduceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
Value input, StringRef grid,
|
|
ArrayRef<GridAxis> gridAxes, ReductionKind reduction) {
|
|
build(odsBuilder, odsState, input.getType(), grid, gridAxes, input,
|
|
reduction);
|
|
}
|
|
|
|
void AllReduceOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "all_reduce");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.all_slice op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult AllSliceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
return verifyScatterOrSliceOperandAndResultShape(
|
|
getOperand(), getResult(), getSliceAxis().getSExtValue(), getGridAxes(),
|
|
grid.value().getShape());
|
|
}
|
|
|
|
void AllSliceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<AllSliceOp>>(context);
|
|
}
|
|
|
|
void AllSliceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
Value input, GridOp grid, ArrayRef<GridAxis> gridAxes,
|
|
int64_t sliceAxis) {
|
|
Type resultType = sliceResultType(input.getType(), grid, gridAxes, sliceAxis);
|
|
build(odsBuilder, odsState, resultType, input, grid.getSymName(), gridAxes,
|
|
sliceAxis);
|
|
}
|
|
|
|
void AllSliceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
|
|
Type resultType, Value input, StringRef grid,
|
|
ArrayRef<GridAxis> gridAxes, int64_t sliceAxis) {
|
|
build(odsBuilder, odsState, resultType, grid, gridAxes, input,
|
|
APInt(sizeof(sliceAxis) * CHAR_BIT, sliceAxis));
|
|
}
|
|
|
|
void AllSliceOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "all_slice");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.all_to_all op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult AllToAllOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
|
|
return verifyAllToAllOperandAndResultShape(
|
|
getOperand(), getResult(), getSplitAxis().getSExtValue(),
|
|
getConcatAxis().getSExtValue(), getGridAxes(), grid.value().getShape());
|
|
}
|
|
|
|
void AllToAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<AllToAllOp>>(context);
|
|
}
|
|
|
|
void AllToAllOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "all_to_all");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.broadcast op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
BroadcastOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
|
|
getRootDynamic(), getGridAxes(),
|
|
grid.value().getShape()))) {
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<BroadcastOp>>(context);
|
|
}
|
|
|
|
void BroadcastOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "broadcast");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.gather op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult GatherOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
|
|
getRootDynamic(), getGridAxes(),
|
|
grid.value().getShape()))) {
|
|
return failure();
|
|
}
|
|
|
|
auto gatherAxis = getGatherAxis().getSExtValue();
|
|
return verifyGatherOperandAndResultShape(getInput(), getResult(), gatherAxis,
|
|
getGridAxes(),
|
|
grid.value().getShape());
|
|
}
|
|
|
|
void GatherOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<GatherOp>>(context);
|
|
}
|
|
|
|
void GatherOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "gather");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.recv op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult RecvOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (getSource() &&
|
|
failed(verifyInGroupDevice(getLoc(), getSourceAttrName(),
|
|
getSource().value(), getSourceDynamic(),
|
|
getGridAxes(), grid.value().getShape()))) {
|
|
return failure();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
void RecvOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<RecvOp>>(context);
|
|
}
|
|
|
|
void RecvOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "recv");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.reduce op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult ReduceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
|
|
getRootDynamic(), getGridAxes(),
|
|
grid.value().getShape()))) {
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void ReduceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<ReduceOp>>(context);
|
|
}
|
|
|
|
void ReduceOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "reduce");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.reduce_scatter op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
ReduceScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
|
|
return verifyScatterOrSliceOperandAndResultShape(
|
|
getOperand(), getResult(), getScatterDim().getSExtValue(), getGridAxes(),
|
|
grid.value().getShape());
|
|
}
|
|
|
|
void ReduceScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<ReduceScatterOp>>(context);
|
|
}
|
|
|
|
void ReduceScatterOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "reduce_scatter");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.scatter op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult ScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
|
|
getRootDynamic(), getGridAxes(),
|
|
grid.value().getShape()))) {
|
|
return failure();
|
|
}
|
|
|
|
auto scatterDim = getScatterDim().getSExtValue();
|
|
return verifyScatterOrSliceOperandAndResultShape(getInput(), getResult(),
|
|
scatterDim, getGridAxes(),
|
|
grid.value().getShape());
|
|
}
|
|
|
|
void ScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<ScatterOp>>(context);
|
|
}
|
|
|
|
void ScatterOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "scatter");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.send op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult SendOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
if (failed(verifyInGroupDevice(getLoc(), getDestinationAttrName(),
|
|
getDestination(), getDestinationDynamic(),
|
|
getGridAxes(), grid.value().getShape()))) {
|
|
return failure();
|
|
}
|
|
return success();
|
|
}
|
|
|
|
void SendOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
patterns.add<EmptyGridAxesCanonicalizationPattern<SendOp>>(context);
|
|
}
|
|
|
|
void SendOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "send");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.shift op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult ShiftOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerifyAxes(*this, symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
|
|
auto gridAxes = getGridAxes();
|
|
auto shiftAxis = getShiftAxis().getZExtValue();
|
|
if (!llvm::is_contained(gridAxes, shiftAxis)) {
|
|
return emitError() << "Invalid shift axis " << shiftAxis
|
|
<< ". It must be one of the grouping grid axes.";
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
void ShiftOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
|
|
MLIRContext *context) {
|
|
// TODO: remove op when offset is 0 or if it is a rotate with and
|
|
// offset % shift_axis_grid_dim_size == 0.
|
|
}
|
|
|
|
void ShiftOp::getAsmResultNames(
|
|
function_ref<void(Value, StringRef)> setNameFn) {
|
|
setNameFn(getResult(), "shift");
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// shard.update_halo op
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
LogicalResult
|
|
UpdateHaloOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
|
|
auto grid = getGridAndVerify(getOperation(), getGridAttr(), symbolTable);
|
|
if (failed(grid)) {
|
|
return failure();
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TableGen'd op method definitions
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#define GET_OP_CLASSES
|
|
#include "mlir/Dialect/Shard/IR/ShardOps.cpp.inc"
|
|
|
|
#define GET_ATTRDEF_CLASSES
|
|
#include "mlir/Dialect/Shard/IR/ShardAttributes.cpp.inc"
|
|
|
|
#define GET_TYPEDEF_CLASSES
|
|
#include "mlir/Dialect/Shard/IR/ShardTypes.cpp.inc"
|
|
|
|
#include "mlir/Dialect/Shard/IR/ShardEnums.cpp.inc"
|