366 lines
23 KiB
C++
366 lines
23 KiB
C++
/*===- TableGen'erated file -------------------------------------*- C++ -*-===*\
|
|
|* *|
|
|
|* Interface Declarations *|
|
|
|* *|
|
|
|* Automatically generated file, do not edit! *|
|
|
|* *|
|
|
\*===----------------------------------------------------------------------===*/
|
|
|
|
namespace mlir::shard {
|
|
|
|
/// Interface for allowing operations to expose information needed to
|
|
/// shard them.
|
|
class ShardingInterface;
|
|
|
|
} // namespace mlir::shard
|
|
namespace mlir::shard {
|
|
|
|
/// Interface for allowing operations to expose information needed to
|
|
/// shard them.
|
|
namespace detail {
|
|
struct ShardingInterfaceInterfaceTraits {
|
|
struct Concept {
|
|
/// The methods defined by the interface.
|
|
SmallVector<mlir::utils::IteratorType> (*getLoopIteratorTypes)(const Concept *impl, ::mlir::Operation *);
|
|
SmallVector<ReductionKind> (*getReductionLoopIteratorKinds)(const Concept *impl, ::mlir::Operation *);
|
|
SmallVector<AffineMap> (*getIndexingMaps)(const Concept *impl, ::mlir::Operation *);
|
|
FailureOr<ShardingOption> (*getShardingOption)(const Concept *impl, ::mlir::Operation *, ArrayRef<Sharding>, ArrayRef<Sharding>);
|
|
FailureOr<std::vector<Sharding>> (*getShardingAnnotations)(const Concept *impl, ::mlir::Operation *, const ShardingOption &);
|
|
LogicalResult (*addShardingAnnotations)(const Concept *impl, ::mlir::Operation *, OpBuilder &, const ShardingOption &);
|
|
LogicalResult (*partition)(const Concept *impl, ::mlir::Operation *, ArrayRef<Value>, ArrayRef<Sharding>, ArrayRef<Sharding>, IRMapping&, SymbolTableCollection &, OpBuilder &);
|
|
};
|
|
template<typename ConcreteOp>
|
|
class Model : public Concept {
|
|
public:
|
|
using Interface = ::mlir::shard::ShardingInterface;
|
|
Model() : Concept{getLoopIteratorTypes, getReductionLoopIteratorKinds, getIndexingMaps, getShardingOption, getShardingAnnotations, addShardingAnnotations, partition} {}
|
|
|
|
static inline SmallVector<mlir::utils::IteratorType> getLoopIteratorTypes(const Concept *impl, ::mlir::Operation *tablegen_opaque_val);
|
|
static inline SmallVector<ReductionKind> getReductionLoopIteratorKinds(const Concept *impl, ::mlir::Operation *tablegen_opaque_val);
|
|
static inline SmallVector<AffineMap> getIndexingMaps(const Concept *impl, ::mlir::Operation *tablegen_opaque_val);
|
|
static inline FailureOr<ShardingOption> getShardingOption(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings);
|
|
static inline FailureOr<std::vector<Sharding>> getShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, const ShardingOption & shardingOption);
|
|
static inline LogicalResult addShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, OpBuilder & b, const ShardingOption & shardingOption);
|
|
static inline LogicalResult partition(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping& partitionMap, SymbolTableCollection & symbolTableCollection, OpBuilder & builder);
|
|
};
|
|
template<typename ConcreteOp>
|
|
class FallbackModel : public Concept {
|
|
public:
|
|
using Interface = ::mlir::shard::ShardingInterface;
|
|
FallbackModel() : Concept{getLoopIteratorTypes, getReductionLoopIteratorKinds, getIndexingMaps, getShardingOption, getShardingAnnotations, addShardingAnnotations, partition} {}
|
|
|
|
static inline SmallVector<mlir::utils::IteratorType> getLoopIteratorTypes(const Concept *impl, ::mlir::Operation *tablegen_opaque_val);
|
|
static inline SmallVector<ReductionKind> getReductionLoopIteratorKinds(const Concept *impl, ::mlir::Operation *tablegen_opaque_val);
|
|
static inline SmallVector<AffineMap> getIndexingMaps(const Concept *impl, ::mlir::Operation *tablegen_opaque_val);
|
|
static inline FailureOr<ShardingOption> getShardingOption(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings);
|
|
static inline FailureOr<std::vector<Sharding>> getShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, const ShardingOption & shardingOption);
|
|
static inline LogicalResult addShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, OpBuilder & b, const ShardingOption & shardingOption);
|
|
static inline LogicalResult partition(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping& partitionMap, SymbolTableCollection & symbolTableCollection, OpBuilder & builder);
|
|
};
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
class ExternalModel : public FallbackModel<ConcreteModel> {
|
|
public:
|
|
using ConcreteEntity = ConcreteOp;
|
|
SmallVector<mlir::utils::IteratorType> getLoopIteratorTypes(::mlir::Operation *tablegen_opaque_val) const;
|
|
SmallVector<ReductionKind> getReductionLoopIteratorKinds(::mlir::Operation *tablegen_opaque_val) const;
|
|
SmallVector<AffineMap> getIndexingMaps(::mlir::Operation *tablegen_opaque_val) const;
|
|
FailureOr<ShardingOption> getShardingOption(::mlir::Operation *tablegen_opaque_val, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings) const;
|
|
FailureOr<std::vector<Sharding>> getShardingAnnotations(::mlir::Operation *tablegen_opaque_val, const ShardingOption &shardingOption) const;
|
|
LogicalResult addShardingAnnotations(::mlir::Operation *tablegen_opaque_val, OpBuilder &b, const ShardingOption &shardingOption) const;
|
|
LogicalResult partition(::mlir::Operation *tablegen_opaque_val, ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping&partitionMap, SymbolTableCollection &symbolTableCollection, OpBuilder &builder) const;
|
|
};
|
|
};
|
|
template <typename ConcreteOp>
|
|
struct ShardingInterfaceTrait;
|
|
|
|
} // namespace detail
|
|
class ShardingInterface : public ::mlir::OpInterface<ShardingInterface, detail::ShardingInterfaceInterfaceTraits> {
|
|
public:
|
|
using ::mlir::OpInterface<ShardingInterface, detail::ShardingInterfaceInterfaceTraits>::OpInterface;
|
|
template <typename ConcreteOp>
|
|
struct Trait : public detail::ShardingInterfaceTrait<ConcreteOp> {};
|
|
|
|
/// Returns a list of iterator types that describe the number of loops.
|
|
/// The iterator types determine how the operation traverses its input and
|
|
/// output tensors.
|
|
///
|
|
/// Example 1: A gemm op has 3 loops, M, N and K. Their loop iterator
|
|
/// types are parallel, parallel, reduction. This indicates that M and
|
|
/// N are traversed in parallel, while the K dimension is used for
|
|
/// reduction.
|
|
SmallVector<mlir::utils::IteratorType> getLoopIteratorTypes();
|
|
|
|
/// Return the kind of all reduction loop iterators.
|
|
/// The order is the same as the same as the result from
|
|
/// `getLoopIteratorTypes`.
|
|
///
|
|
/// Example 1:
|
|
/// iterator types = (parallel, reduction, parallel, reduction)
|
|
/// || ||
|
|
/// reduction kinds = ( sum, max)
|
|
///
|
|
/// Example 2:
|
|
/// A softmax op's loop iterator types are parallel and
|
|
/// reduction.
|
|
/// The reduction iterator will be of kind `generic`, since it is non of
|
|
/// the available presets.
|
|
SmallVector<ReductionKind> getReductionLoopIteratorKinds();
|
|
|
|
/// Return the indexing maps attribute within the current operation.
|
|
/// Indexing maps determine how indices in the iteration space map to
|
|
/// tensor indices. They are specified using `affine_map` in MLIR, which
|
|
/// provides an affine transformation of indices.
|
|
SmallVector<AffineMap> getIndexingMaps();
|
|
|
|
/// Given that certain operands or results of the operation may have
|
|
/// sharding annotations, this method leverages this information to
|
|
/// deduce how the operation should be sharded.
|
|
/// The passed sharding may be incomplete, this gives freedom for the
|
|
/// op to select the most appropriate shardings for all the operands
|
|
/// and results and the op itself.
|
|
FailureOr<ShardingOption> getShardingOption(ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings);
|
|
|
|
/// Based on a given ShardingOption, get the operand and result
|
|
/// operations for the operands and results sharding annotations.
|
|
/// This is what shardings the operands and results need to have in order
|
|
/// to shard the op according to shardingOption.
|
|
FailureOr<std::vector<Sharding>> getShardingAnnotations(const ShardingOption & shardingOption);
|
|
|
|
/// Based on a given ShardingOption, this method adds `shard.shard`
|
|
/// operations for the operands and results that previously lacked
|
|
/// sharding annotations.
|
|
LogicalResult addShardingAnnotations(OpBuilder & b, const ShardingOption & shardingOption);
|
|
|
|
/// Convert self to SPMD form.
|
|
/// This method is used during the partition pass of a program fully
|
|
/// annotated with shardings.
|
|
///
|
|
/// The partition algorithm would read the surrounding sharding
|
|
/// annotations from the IR for each argument/result and prepare
|
|
/// `operandShardings` and `resultShardings`.
|
|
/// Values that are not ranked tensors do not have sharding annotations.
|
|
/// In this case their corresponding Sharding is null.
|
|
///
|
|
/// For convenience it will also prepare `partitionedOperands`, although
|
|
/// they can be retrieved from the `partitionMap`.
|
|
///
|
|
/// The `partitionMap` contains a mapping from unsharded to
|
|
/// sharded/partitioned values that are constructed during the partition
|
|
/// pass. The interface implementation must populate `partitionMap`
|
|
/// with the mapping for this op's results.
|
|
///
|
|
/// `builder` is set to insert new operations in the appropriate point.
|
|
/// The implementation should not return the builder to the original
|
|
/// insertion point.
|
|
/// It should leave it as is after all insertions are done.
|
|
///
|
|
/// The default implementation does full replication.
|
|
/// This assumes that all sharding annotations are for full replication.
|
|
LogicalResult partition(ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping& partitionMap, SymbolTableCollection & symbolTableCollection, OpBuilder & builder);
|
|
|
|
LogicalResult verifyShardingInterfaceImpl();
|
|
|
|
void printLoopTypesAndIndexingMaps(raw_ostream &os);
|
|
};
|
|
|
|
} // namespace mlir::shard
|
|
namespace mlir::shard::detail {
|
|
|
|
template <typename ConcreteOp>
|
|
struct ShardingInterfaceTrait : public ::mlir::OpInterface<ShardingInterface, detail::ShardingInterfaceInterfaceTraits>::Trait<ConcreteOp> {
|
|
|
|
/// Returns a list of iterator types that describe the number of loops.
|
|
/// The iterator types determine how the operation traverses its input and
|
|
/// output tensors.
|
|
///
|
|
/// Example 1: A gemm op has 3 loops, M, N and K. Their loop iterator
|
|
/// types are parallel, parallel, reduction. This indicates that M and
|
|
/// N are traversed in parallel, while the K dimension is used for
|
|
/// reduction.
|
|
SmallVector<mlir::utils::IteratorType> getLoopIteratorTypes() {
|
|
return {};
|
|
}
|
|
|
|
/// Return the kind of all reduction loop iterators.
|
|
/// The order is the same as the same as the result from
|
|
/// `getLoopIteratorTypes`.
|
|
///
|
|
/// Example 1:
|
|
/// iterator types = (parallel, reduction, parallel, reduction)
|
|
/// || ||
|
|
/// reduction kinds = ( sum, max)
|
|
///
|
|
/// Example 2:
|
|
/// A softmax op's loop iterator types are parallel and
|
|
/// reduction.
|
|
/// The reduction iterator will be of kind `generic`, since it is non of
|
|
/// the available presets.
|
|
SmallVector<ReductionKind> getReductionLoopIteratorKinds() {
|
|
return {};
|
|
}
|
|
|
|
/// Return the indexing maps attribute within the current operation.
|
|
/// Indexing maps determine how indices in the iteration space map to
|
|
/// tensor indices. They are specified using `affine_map` in MLIR, which
|
|
/// provides an affine transformation of indices.
|
|
SmallVector<AffineMap> getIndexingMaps() {
|
|
return {};
|
|
}
|
|
|
|
/// Given that certain operands or results of the operation may have
|
|
/// sharding annotations, this method leverages this information to
|
|
/// deduce how the operation should be sharded.
|
|
/// The passed sharding may be incomplete, this gives freedom for the
|
|
/// op to select the most appropriate shardings for all the operands
|
|
/// and results and the op itself.
|
|
FailureOr<ShardingOption> getShardingOption(ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings) {
|
|
return detail::defaultGetShardingOption(
|
|
(*static_cast<ConcreteOp *>(this)).getOperation(), operandShardings, resultShardings);
|
|
}
|
|
|
|
/// Based on a given ShardingOption, get the operand and result
|
|
/// operations for the operands and results sharding annotations.
|
|
/// This is what shardings the operands and results need to have in order
|
|
/// to shard the op according to shardingOption.
|
|
FailureOr<std::vector<Sharding>> getShardingAnnotations(const ShardingOption & shardingOption) {
|
|
return detail::defaultGetShardingAnnotations(
|
|
(*static_cast<ConcreteOp *>(this)).getOperation(), shardingOption);
|
|
}
|
|
|
|
/// Based on a given ShardingOption, this method adds `shard.shard`
|
|
/// operations for the operands and results that previously lacked
|
|
/// sharding annotations.
|
|
LogicalResult addShardingAnnotations(OpBuilder & b, const ShardingOption & shardingOption) {
|
|
return detail::defaultAddShardingAnnotations(
|
|
(*static_cast<ConcreteOp *>(this)).getOperation(), b, shardingOption);
|
|
}
|
|
|
|
/// Convert self to SPMD form.
|
|
/// This method is used during the partition pass of a program fully
|
|
/// annotated with shardings.
|
|
///
|
|
/// The partition algorithm would read the surrounding sharding
|
|
/// annotations from the IR for each argument/result and prepare
|
|
/// `operandShardings` and `resultShardings`.
|
|
/// Values that are not ranked tensors do not have sharding annotations.
|
|
/// In this case their corresponding Sharding is null.
|
|
///
|
|
/// For convenience it will also prepare `partitionedOperands`, although
|
|
/// they can be retrieved from the `partitionMap`.
|
|
///
|
|
/// The `partitionMap` contains a mapping from unsharded to
|
|
/// sharded/partitioned values that are constructed during the partition
|
|
/// pass. The interface implementation must populate `partitionMap`
|
|
/// with the mapping for this op's results.
|
|
///
|
|
/// `builder` is set to insert new operations in the appropriate point.
|
|
/// The implementation should not return the builder to the original
|
|
/// insertion point.
|
|
/// It should leave it as is after all insertions are done.
|
|
///
|
|
/// The default implementation does full replication.
|
|
/// This assumes that all sharding annotations are for full replication.
|
|
LogicalResult partition(ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping& partitionMap, SymbolTableCollection & symbolTableCollection, OpBuilder & builder) {
|
|
partitionFullyReplicatedOperation(
|
|
*(*static_cast<ConcreteOp *>(this)).getOperation(), partitionedOperands, operandShardings,
|
|
resultShardings, partitionMap, symbolTableCollection, builder);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
} // namespace mlir::shard::detail
|
|
namespace mlir::shard {
|
|
|
|
template<typename ConcreteOp>
|
|
SmallVector<mlir::utils::IteratorType> detail::ShardingInterfaceInterfaceTraits::Model<ConcreteOp>::getLoopIteratorTypes(const Concept *impl, ::mlir::Operation *tablegen_opaque_val) {
|
|
return (llvm::cast<ConcreteOp>(tablegen_opaque_val)).getLoopIteratorTypes();
|
|
}
|
|
template<typename ConcreteOp>
|
|
SmallVector<ReductionKind> detail::ShardingInterfaceInterfaceTraits::Model<ConcreteOp>::getReductionLoopIteratorKinds(const Concept *impl, ::mlir::Operation *tablegen_opaque_val) {
|
|
return (llvm::cast<ConcreteOp>(tablegen_opaque_val)).getReductionLoopIteratorKinds();
|
|
}
|
|
template<typename ConcreteOp>
|
|
SmallVector<AffineMap> detail::ShardingInterfaceInterfaceTraits::Model<ConcreteOp>::getIndexingMaps(const Concept *impl, ::mlir::Operation *tablegen_opaque_val) {
|
|
return (llvm::cast<ConcreteOp>(tablegen_opaque_val)).getIndexingMaps();
|
|
}
|
|
template<typename ConcreteOp>
|
|
FailureOr<ShardingOption> detail::ShardingInterfaceInterfaceTraits::Model<ConcreteOp>::getShardingOption(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings) {
|
|
return (llvm::cast<ConcreteOp>(tablegen_opaque_val)).getShardingOption(operandShardings, resultShardings);
|
|
}
|
|
template<typename ConcreteOp>
|
|
FailureOr<std::vector<Sharding>> detail::ShardingInterfaceInterfaceTraits::Model<ConcreteOp>::getShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, const ShardingOption & shardingOption) {
|
|
return (llvm::cast<ConcreteOp>(tablegen_opaque_val)).getShardingAnnotations(shardingOption);
|
|
}
|
|
template<typename ConcreteOp>
|
|
LogicalResult detail::ShardingInterfaceInterfaceTraits::Model<ConcreteOp>::addShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, OpBuilder & b, const ShardingOption & shardingOption) {
|
|
return (llvm::cast<ConcreteOp>(tablegen_opaque_val)).addShardingAnnotations(b, shardingOption);
|
|
}
|
|
template<typename ConcreteOp>
|
|
LogicalResult detail::ShardingInterfaceInterfaceTraits::Model<ConcreteOp>::partition(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping& partitionMap, SymbolTableCollection & symbolTableCollection, OpBuilder & builder) {
|
|
return (llvm::cast<ConcreteOp>(tablegen_opaque_val)).partition(partitionedOperands, operandShardings, resultShardings, partitionMap, symbolTableCollection, builder);
|
|
}
|
|
template<typename ConcreteOp>
|
|
SmallVector<mlir::utils::IteratorType> detail::ShardingInterfaceInterfaceTraits::FallbackModel<ConcreteOp>::getLoopIteratorTypes(const Concept *impl, ::mlir::Operation *tablegen_opaque_val) {
|
|
return static_cast<const ConcreteOp *>(impl)->getLoopIteratorTypes(tablegen_opaque_val);
|
|
}
|
|
template<typename ConcreteOp>
|
|
SmallVector<ReductionKind> detail::ShardingInterfaceInterfaceTraits::FallbackModel<ConcreteOp>::getReductionLoopIteratorKinds(const Concept *impl, ::mlir::Operation *tablegen_opaque_val) {
|
|
return static_cast<const ConcreteOp *>(impl)->getReductionLoopIteratorKinds(tablegen_opaque_val);
|
|
}
|
|
template<typename ConcreteOp>
|
|
SmallVector<AffineMap> detail::ShardingInterfaceInterfaceTraits::FallbackModel<ConcreteOp>::getIndexingMaps(const Concept *impl, ::mlir::Operation *tablegen_opaque_val) {
|
|
return static_cast<const ConcreteOp *>(impl)->getIndexingMaps(tablegen_opaque_val);
|
|
}
|
|
template<typename ConcreteOp>
|
|
FailureOr<ShardingOption> detail::ShardingInterfaceInterfaceTraits::FallbackModel<ConcreteOp>::getShardingOption(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings) {
|
|
return static_cast<const ConcreteOp *>(impl)->getShardingOption(tablegen_opaque_val, operandShardings, resultShardings);
|
|
}
|
|
template<typename ConcreteOp>
|
|
FailureOr<std::vector<Sharding>> detail::ShardingInterfaceInterfaceTraits::FallbackModel<ConcreteOp>::getShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, const ShardingOption & shardingOption) {
|
|
return static_cast<const ConcreteOp *>(impl)->getShardingAnnotations(tablegen_opaque_val, shardingOption);
|
|
}
|
|
template<typename ConcreteOp>
|
|
LogicalResult detail::ShardingInterfaceInterfaceTraits::FallbackModel<ConcreteOp>::addShardingAnnotations(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, OpBuilder & b, const ShardingOption & shardingOption) {
|
|
return static_cast<const ConcreteOp *>(impl)->addShardingAnnotations(tablegen_opaque_val, b, shardingOption);
|
|
}
|
|
template<typename ConcreteOp>
|
|
LogicalResult detail::ShardingInterfaceInterfaceTraits::FallbackModel<ConcreteOp>::partition(const Concept *impl, ::mlir::Operation *tablegen_opaque_val, ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping& partitionMap, SymbolTableCollection & symbolTableCollection, OpBuilder & builder) {
|
|
return static_cast<const ConcreteOp *>(impl)->partition(tablegen_opaque_val, partitionedOperands, operandShardings, resultShardings, partitionMap, symbolTableCollection, builder);
|
|
}
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
SmallVector<mlir::utils::IteratorType> detail::ShardingInterfaceInterfaceTraits::ExternalModel<ConcreteModel, ConcreteOp>::getLoopIteratorTypes(::mlir::Operation *tablegen_opaque_val) const {
|
|
return {};
|
|
}
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
SmallVector<ReductionKind> detail::ShardingInterfaceInterfaceTraits::ExternalModel<ConcreteModel, ConcreteOp>::getReductionLoopIteratorKinds(::mlir::Operation *tablegen_opaque_val) const {
|
|
return {};
|
|
}
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
SmallVector<AffineMap> detail::ShardingInterfaceInterfaceTraits::ExternalModel<ConcreteModel, ConcreteOp>::getIndexingMaps(::mlir::Operation *tablegen_opaque_val) const {
|
|
return {};
|
|
}
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
FailureOr<ShardingOption> detail::ShardingInterfaceInterfaceTraits::ExternalModel<ConcreteModel, ConcreteOp>::getShardingOption(::mlir::Operation *tablegen_opaque_val, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings) const {
|
|
return detail::defaultGetShardingOption(
|
|
(llvm::cast<ConcreteOp>(tablegen_opaque_val)).getOperation(), operandShardings, resultShardings);
|
|
}
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
FailureOr<std::vector<Sharding>> detail::ShardingInterfaceInterfaceTraits::ExternalModel<ConcreteModel, ConcreteOp>::getShardingAnnotations(::mlir::Operation *tablegen_opaque_val, const ShardingOption &shardingOption) const {
|
|
return detail::defaultGetShardingAnnotations(
|
|
(llvm::cast<ConcreteOp>(tablegen_opaque_val)).getOperation(), shardingOption);
|
|
}
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
LogicalResult detail::ShardingInterfaceInterfaceTraits::ExternalModel<ConcreteModel, ConcreteOp>::addShardingAnnotations(::mlir::Operation *tablegen_opaque_val, OpBuilder &b, const ShardingOption &shardingOption) const {
|
|
return detail::defaultAddShardingAnnotations(
|
|
(llvm::cast<ConcreteOp>(tablegen_opaque_val)).getOperation(), b, shardingOption);
|
|
}
|
|
template<typename ConcreteModel, typename ConcreteOp>
|
|
LogicalResult detail::ShardingInterfaceInterfaceTraits::ExternalModel<ConcreteModel, ConcreteOp>::partition(::mlir::Operation *tablegen_opaque_val, ArrayRef<Value> partitionedOperands, ArrayRef<Sharding> operandShardings, ArrayRef<Sharding> resultShardings, IRMapping&partitionMap, SymbolTableCollection &symbolTableCollection, OpBuilder &builder) const {
|
|
partitionFullyReplicatedOperation(
|
|
*(llvm::cast<ConcreteOp>(tablegen_opaque_val)).getOperation(), partitionedOperands, operandShardings,
|
|
resultShardings, partitionMap, symbolTableCollection, builder);
|
|
return success();
|
|
}
|
|
|
|
} // namespace mlir::shard
|