148 lines
5.7 KiB
C++
148 lines
5.7 KiB
C++
//===- ShardingInterfaceImpl.h ----------------------------------*- C++ -*-===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_DIALECT_SHARD_INTERFACES_SHARDINGINTERFACEIMPL_H_
|
|
#define MLIR_DIALECT_SHARD_INTERFACES_SHARDINGINTERFACEIMPL_H_
|
|
|
|
#include "mlir/Dialect/Shard/IR/ShardOps.h"
|
|
#include "mlir/Dialect/Shard/Interfaces/ShardingInterface.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Value.h"
|
|
|
|
namespace mlir {
|
|
|
|
class Operation;
|
|
class IRMapping;
|
|
class SymbolTableCollection;
|
|
|
|
namespace shard {
|
|
|
|
// Retrieve the grid axes corresponding to each operation loop iterator based
|
|
// on the provided shardings for the op's operands and results.
|
|
// Assumes that the indexingMaps are projected permutations.
|
|
ShardingArray getGridAxisAssignmentForLoopIterators(
|
|
ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings,
|
|
ArrayRef<utils::IteratorType> loopIteratorTypes,
|
|
ArrayRef<AffineMap> indexingMaps);
|
|
|
|
bool isAtLeastOneReductionIteratorSharded(
|
|
ArrayRef<utils::IteratorType> loopIteratorTypes,
|
|
ArrayRef<SmallVector<GridAxis>> gridAxisAssignmentForLoopIterators);
|
|
|
|
// Get the set of grid axes that correspond to reduction loop iterators.
|
|
SmallVector<GridAxis> getReductionGridAxes(
|
|
ArrayRef<utils::IteratorType> loopIteratorTypes,
|
|
ArrayRef<SmallVector<GridAxis>> gridAxisAssignmentForLoopIterators);
|
|
|
|
// Inserts a clone of the operation that has all ranked tensor
|
|
// arguments/results sharded.
|
|
void partitionTriviallyShardableOperation(Operation &op,
|
|
ArrayRef<Value> partitionedOperands,
|
|
ArrayRef<Sharding> operandShardings,
|
|
ArrayRef<Sharding> resultShardings,
|
|
IRMapping &partitionMap,
|
|
SymbolTableCollection &symbolTable,
|
|
OpBuilder &builder);
|
|
|
|
// All ranked tensor argument and result dimensions have
|
|
// independent parallel loop iterators.
|
|
template <typename Op>
|
|
struct IndependentParallelIteratorDomainShardingInterface
|
|
: public ShardingInterface::ExternalModel<
|
|
IndependentParallelIteratorDomainShardingInterface<Op>, Op> {
|
|
SmallVector<utils::IteratorType>
|
|
getLoopIteratorTypes(Operation *operation) const {
|
|
SmallVector<utils::IteratorType> iterTypes;
|
|
for (Type t : operation->getOperandTypes()) {
|
|
populateIteratorTypes(t, iterTypes);
|
|
}
|
|
for (Type t : operation->getResultTypes()) {
|
|
populateIteratorTypes(t, iterTypes);
|
|
}
|
|
return iterTypes;
|
|
}
|
|
|
|
SmallVector<AffineMap> getIndexingMaps(Operation *op) const {
|
|
// TODO: implement.
|
|
return SmallVector<AffineMap>();
|
|
}
|
|
|
|
LogicalResult partition(Operation *op, ArrayRef<Value> partitionedOperands,
|
|
ArrayRef<Sharding> operandShardings,
|
|
ArrayRef<Sharding> resultShardings,
|
|
IRMapping &partitionMap,
|
|
SymbolTableCollection &symbolTable,
|
|
OpBuilder &builder) const {
|
|
partitionTriviallyShardableOperation(*op, partitionedOperands,
|
|
operandShardings, resultShardings,
|
|
partitionMap, symbolTable, builder);
|
|
return success();
|
|
}
|
|
|
|
private:
|
|
void
|
|
populateIteratorTypes(Type t,
|
|
SmallVector<utils::IteratorType> &iterTypes) const {
|
|
RankedTensorType rankedTensorType = dyn_cast<RankedTensorType>(t);
|
|
if (!rankedTensorType) {
|
|
return;
|
|
}
|
|
|
|
iterTypes.reserve(iterTypes.size() + rankedTensorType.getRank());
|
|
for (int64_t i = 0; i < rankedTensorType.getRank(); ++i) {
|
|
iterTypes.push_back(utils::IteratorType::parallel);
|
|
}
|
|
}
|
|
};
|
|
|
|
// Sharding of elementwise operations like tensor addition and multiplication.
|
|
template <typename ElemwiseOp>
|
|
struct ElementwiseShardingInterface
|
|
: public ShardingInterface::ExternalModel<
|
|
ElementwiseShardingInterface<ElemwiseOp>, ElemwiseOp> {
|
|
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
|
|
Value val = op->getOperand(0);
|
|
auto type = dyn_cast<RankedTensorType>(val.getType());
|
|
if (!type)
|
|
return {};
|
|
SmallVector<utils::IteratorType> types(type.getRank(),
|
|
utils::IteratorType::parallel);
|
|
return types;
|
|
}
|
|
|
|
SmallVector<AffineMap> getIndexingMaps(Operation *op) const {
|
|
MLIRContext *ctx = op->getContext();
|
|
Value val = op->getOperand(0);
|
|
auto type = dyn_cast<RankedTensorType>(val.getType());
|
|
if (!type)
|
|
return {};
|
|
int64_t rank = type.getRank();
|
|
int64_t num = op->getNumOperands() + op->getNumResults();
|
|
SmallVector<AffineMap> maps(num,
|
|
AffineMap::getMultiDimIdentityMap(rank, ctx));
|
|
return maps;
|
|
}
|
|
|
|
LogicalResult partition(Operation *op, ArrayRef<Value> partitionedOperands,
|
|
ArrayRef<Sharding> operandShardings,
|
|
ArrayRef<Sharding> resultShardings,
|
|
IRMapping &partitionMap,
|
|
SymbolTableCollection &symbolTable,
|
|
OpBuilder &builder) const {
|
|
partitionTriviallyShardableOperation(*op, partitionedOperands,
|
|
operandShardings, resultShardings,
|
|
partitionMap, symbolTable, builder);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace shard
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_DIALECT_SHARD_INTERFACES_SHARDINGINTERFACEIMPL_H_
|