mlir::mesh Namespace Reference

Namespaces
	detail

Classes
struct	ShardingOption
struct	IndependentParallelIteratorDomainShardingInterface
struct	ElementwiseShardingInterface
class	MeshSharding
struct	OpRewritePatternWithSymbolTableCollection

Typedefs
using	ShardingArray = SmallVector< SmallVector< MeshAxis > >
using	ShardingArrayRef = ArrayRef< SmallVector< MeshAxis > >
using	MeshAxis = int16_t
using	MeshAxesAttr = DenseI16ArrayAttr
using	ShardShapeAttr = DenseI64ArrayAttr
using	HaloSizePairAttr = DenseI64ArrayAttr
using	UnshardedToShardedValueMap = DenseMap< Value, Value >

Enumerations
enum class	TraversalOrder { Forward , Backward , ForwardBackward , BackwardForward }
	This enum controls the traversal order for the sharding propagation. More...

Functions
FailureOr< std::pair< bool, MeshSharding > >	getMeshSharding (OpResult result)
FailureOr< std::pair< bool, MeshSharding > >	getMeshSharding (OpOperand &opOperand)
void	spmdizeFullyReplicatedOperation (Operation &op, ArrayRef< Value > spmdizedOperands, ArrayRef< MeshSharding > operandShardings, ArrayRef< MeshSharding > resultShardings, IRMapping &spmdizationMap, SymbolTableCollection &symbolTable, OpBuilder &builder)
ShardingArray	getMeshAxisAssignmentForLoopIterators (ArrayRef< MeshSharding > operandShardings, ArrayRef< MeshSharding > resultShardings, ArrayRef< utils::IteratorType > loopIteratorTypes, ArrayRef< AffineMap > indexingMaps)
bool	isAtLeastOneReductionIteratorSharded (ArrayRef< utils::IteratorType > loopIteratorTypes, ArrayRef< SmallVector< MeshAxis >> meshAxisAssignmentForLoopIterators)
SmallVector< MeshAxis >	getReductionMeshAxes (ArrayRef< utils::IteratorType > loopIteratorTypes, ArrayRef< SmallVector< MeshAxis >> meshAxisAssignmentForLoopIterators)
void	spmdizeTriviallyShardableOperation (Operation &op, ArrayRef< Value > spmdizedOperands, ArrayRef< MeshSharding > operandShardings, ArrayRef< MeshSharding > resultShardings, IRMapping &spmdizationMap, SymbolTableCollection &symbolTable, OpBuilder &builder)
bool	isReductionLoop (utils::IteratorType iType)
template<typename T >
void	removeTrailingEmptySubArray (SmallVector< SmallVector< T >> &array)
bool	isFullReplication (MeshSharding sharding)
mesh::MeshOp	getMeshOrNull (Operation *op, FlatSymbolRefAttr meshSymbol, SymbolTableCollection &symbolTableCollection)
mesh::MeshOp	getMesh (Operation *op, FlatSymbolRefAttr meshSymbol, SymbolTableCollection &symbolTableCollection)
template<typename Op >
mesh::MeshOp	getMesh (Op op, SymbolTableCollection &symbolTableCollection)
template<>
mesh::MeshOp	getMesh< ShardOp > (ShardOp op, SymbolTableCollection &symbolTableCollection)
template<typename MeshAxesRange , typename MeshShapeRange >
int64_t	collectiveProcessGroupSize (MeshAxesRange &&meshAxes, MeshShapeRange &&meshShape)
template<typename MeshAxesRange >
int64_t	collectiveProcessGroupSize (MeshAxesRange &&meshAxes, MeshOp mesh)
int64_t	shardDimension (int64_t dimSize, int64_t shardCount)
int64_t	gatherDimension (int64_t dimSize, int64_t shardCount)
ShapedType	shardShapedType (ShapedType shape, MeshOp mesh, MeshSharding sharding)
Type	shardType (Type type, MeshOp mesh, MeshSharding sharding)
void	maybeInsertTargetShardingAnnotation (MeshSharding sharding, OpResult result, OpBuilder &builder)
void	maybeInsertSourceShardingAnnotation (MeshSharding sharding, OpOperand &operand, OpBuilder &builder)
SmallVector< Value >	getMixedAsValues (OpBuilder b, const Location &loc, llvm::ArrayRef< int64_t > statics, ValueRange dynamics, Type type=Type())
	Converts a vector of OpFoldResults (ints) into vector of Values of the provided type. More...
template<typename AlgebraicOp >
void	populateAllReduceEndomorphismSimplificationPatterns (RewritePatternSet &patterns, ReductionKind reduction)
void	populateSimplificationPatterns (RewritePatternSet &patterns, SymbolTableCollection &symbolTableCollection)
void	populateFoldingPatterns (RewritePatternSet &patterns, SymbolTableCollection &symbolTableCollection)
TypedValue< ShapedType >	reshard (OpBuilder &builder, MeshOp mesh, ShardOp source, ShardOp target, TypedValue< ShapedType > sourceShardValue)
TypedValue< ShapedType >	reshard (OpBuilder &builder, ShardOp source, ShardOp target, TypedValue< ShapedType > sourceShardValue, SymbolTableCollection &symbolTableCollection)
void	reshardingRegisterDependentDialects (DialectRegistry &registry)
void	populateProcessMultiIndexOpLoweringPatterns (RewritePatternSet &patterns, SymbolTableCollection &symbolTableCollection)
void	registerProcessMultiIndexOpLoweringDialects (DialectRegistry &registry)
void	populateAllSliceOpLoweringPatterns (RewritePatternSet &patterns, SymbolTableCollection &symbolTableCollection)
void	registerAllSliceOpLoweringDialects (DialectRegistry &registry)
void	populateAllOpLoweringPatterns (RewritePatternSet &patterns, SymbolTableCollection &symbolTableCollection)
void	registerAllOpLoweringDialects (DialectRegistry &registry)
TypedValue< IndexType >	createCollectiveProcessGroupSize (MeshOp mesh, ArrayRef< MeshAxis > axes, ImplicitLocOpBuilder &builder)
TypedValue< IndexType >	createProcessLinearIndex (StringRef mesh, ArrayRef< MeshAxis > meshAxes, ImplicitLocOpBuilder &builder)
TypedValue< IndexType >	createProcessLinearIndex (StringRef mesh, ValueRange processInGroupMultiIndex, ArrayRef< MeshAxis > meshAxes, ImplicitLocOpBuilder &builder)
template<typename SourceAxes , typename TargetAxes >
static bool	arePartialAxesCompatible (const SourceAxes &sourceAxes, const TargetAxes &targetAxes)
static std::tuple< TypedValue< ShapedType >, MeshSharding >	handlePartialAxesDuringResharding (OpBuilder &builder, MeshSharding sourceSharding, MeshSharding targetSharding, TypedValue< ShapedType > sourceShard)
static MeshSharding	targetShardingInSplitLastAxis (MLIRContext *ctx, MeshSharding sourceSharding, int64_t splitTensorAxis, MeshAxis splitMeshAxis)
static std::tuple< TypedValue< ShapedType >, MeshSharding >	splitLastAxisInResharding (ImplicitLocOpBuilder &builder, MeshSharding sourceSharding, TypedValue< ShapedType > sourceShard, MeshOp mesh, int64_t splitTensorAxis, MeshAxis splitMeshAxis)
static std::optional< std::tuple< int64_t, MeshAxis > >	detectSplitLastAxisInResharding (MeshSharding sourceSharding, MeshSharding targetSharding)
static std::optional< std::tuple< TypedValue< ShapedType >, MeshSharding > >	trySplitLastAxisInResharding (ImplicitLocOpBuilder &builder, MeshOp mesh, MeshSharding sourceSharding, MeshSharding targetSharding, TypedValue< ShapedType > sourceShard)
static std::optional< std::tuple< int64_t, MeshAxis > >	detectUnsplitLastAxisInResharding (MeshSharding sourceSharding, MeshSharding targetSharding)
static MeshSharding	targetShardingInUnsplitLastAxis (MLIRContext *ctx, MeshSharding sourceSharding, int64_t splitTensorAxis)
static ShapedType	allGatherResultShapeInUnsplitLastAxis (ShapedType sourceShape, int64_t splitCount, int64_t splitTensorAxis)
static std::tuple< TypedValue< ShapedType >, MeshSharding >	unsplitLastAxisInResharding (ImplicitLocOpBuilder &builder, MeshSharding sourceSharding, ShapedType sourceUnshardedShape, TypedValue< ShapedType > sourceShard, MeshOp mesh, int64_t splitTensorAxis, MeshAxis splitMeshAxis)
static std::optional< std::tuple< TypedValue< ShapedType >, MeshSharding > >	tryUnsplitLastAxisInResharding (ImplicitLocOpBuilder &builder, MeshOp mesh, MeshSharding sourceSharding, MeshSharding targetSharding, ShapedType sourceUnshardedShape, TypedValue< ShapedType > sourceShard)
static std::optional< std::tuple< int64_t, int64_t, MeshAxis > >	detectMoveLastSplitAxisInResharding (MeshSharding sourceSharding, MeshSharding targetSharding)
static MeshSharding	targetShardingInMoveLastAxis (MLIRContext *ctx, MeshSharding sourceSharding, int64_t sourceTensorAxis, int64_t targetTensorAxis)
static ShapedType	allToAllResultShapeInMoveLastAxis (ShapedType sourceShape, int64_t splitCount, int64_t sourceTensorAxis, int64_t targetTensorAxis)
static std::tuple< TypedValue< ShapedType >, MeshSharding >	moveLastSplitAxisInResharding (ImplicitLocOpBuilder &builder, MeshOp mesh, MeshSharding sourceSharding, ShapedType sourceUnshardedShape, TypedValue< ShapedType > sourceShard, int64_t sourceTensorAxis, int64_t targetTensorAxis, MeshAxis meshAxis)
static std::optional< std::tuple< TypedValue< ShapedType >, MeshSharding > >	tryMoveLastSplitAxisInResharding (ImplicitLocOpBuilder &builder, MeshOp mesh, MeshSharding sourceSharding, MeshSharding targetSharding, ShapedType sourceUnshardedShape, TypedValue< ShapedType > sourceShard)
static std::optional< std::tuple< TypedValue< ShapedType >, MeshSharding > >	tryUpdateHaloInResharding (ImplicitLocOpBuilder &builder, MeshOp mesh, MeshSharding sourceSharding, MeshSharding targetSharding, ShapedType sourceUnshardedShape, TypedValue< ShapedType > sourceShard)
static TypedValue< ShapedType >	reshardOn1DMesh (ImplicitLocOpBuilder &builder, MeshOp mesh, MeshSharding sourceSharding, MeshSharding targetSharding, TypedValue< ShapedType > sourceUnshardedValue, TypedValue< ShapedType > sourceShard)
TypedValue< ShapedType >	reshard (ImplicitLocOpBuilder &builder, MeshOp mesh, MeshSharding sourceSharding, MeshSharding targetSharding, TypedValue< ShapedType > sourceUnshardedValue, TypedValue< ShapedType > sourceShard)
SmallVector< Type >	shardedBlockArgumentTypes (Block &block, SymbolTableCollection &symbolTableCollection)
static LogicalResult	spmdizeOperation (Operation &op, ArrayRef< Value > spmdizedOperands, ArrayRef< MeshSharding > operandShardings, ArrayRef< MeshSharding > resultShardings, IRMapping &spmdizationMap, SymbolTableCollection &symbolTableCollection, OpBuilder &builder)
static std::vector< MeshSharding >	getOperandShardings (Operation &op)
static std::vector< MeshSharding >	getResultShardings (Operation &op)
static LogicalResult	spmdizeOperation (ShardOp shardOp, IRMapping &spmdizationMap, SymbolTableCollection &symbolTableCollection, OpBuilder &builder)
static LogicalResult	spmdizeOperation (Operation &op, IRMapping &spmdizationMap, SymbolTableCollection &symbolTableCollection, OpBuilder &builder)
static LogicalResult	spmdizeBlock (Block &block, IRMapping &spmdizationMap, SymbolTableCollection &symbolTableCollection, OpBuilder &builder)
static LogicalResult	spmdizeFuncOp (FunctionOpInterface op, IRMapping &spmdizationMap, SymbolTableCollection &symbolTableCollection)

Typedef Documentation

HaloSizePairAttr

using mlir::mesh::HaloSizePairAttr = typedef DenseI64ArrayAttr

Definition at line 29 of file MeshOps.h.

MeshAxesAttr

using mlir::mesh::MeshAxesAttr = typedef DenseI16ArrayAttr

Definition at line 27 of file MeshOps.h.

MeshAxis

using mlir::mesh::MeshAxis = typedef int16_t

Definition at line 26 of file MeshOps.h.

ShardingArray

using mlir::mesh::ShardingArray = typedef SmallVector<SmallVector<MeshAxis> >

Definition at line 25 of file ShardingInterface.h.

ShardingArrayRef

using mlir::mesh::ShardingArrayRef = typedef ArrayRef<SmallVector<MeshAxis> >

Definition at line 26 of file ShardingInterface.h.

ShardShapeAttr

using mlir::mesh::ShardShapeAttr = typedef DenseI64ArrayAttr

Definition at line 28 of file MeshOps.h.

UnshardedToShardedValueMap

using mlir::mesh::UnshardedToShardedValueMap = typedef DenseMap<Value, Value>

Definition at line 610 of file Spmdization.cpp.

Enumeration Type Documentation

TraversalOrder

enum mlir::mesh::TraversalOrder

strong

This enum controls the traversal order for the sharding propagation.

Enumerator
Forward	Forward traversal.
Backward	Backward traversal.
ForwardBackward	Forward then backward traversal.
BackwardForward	Backward then forward traversal.

Definition at line 23 of file Passes.h.

Function Documentation

allGatherResultShapeInUnsplitLastAxis()

static ShapedType mlir::mesh::allGatherResultShapeInUnsplitLastAxis ( ShapedType  sourceShape, int64_t  splitCount, int64_t  splitTensorAxis  )

static

Definition at line 247 of file Spmdization.cpp.

References gatherDimension().

Referenced by unsplitLastAxisInResharding().

allToAllResultShapeInMoveLastAxis()

static ShapedType mlir::mesh::allToAllResultShapeInMoveLastAxis ( ShapedType  sourceShape, int64_t  splitCount, int64_t  sourceTensorAxis, int64_t  targetTensorAxis  )

static

Definition at line 374 of file Spmdization.cpp.

References gatherDimension(), and shardDimension().

Referenced by moveLastSplitAxisInResharding().

arePartialAxesCompatible()

template<typename SourceAxes , typename TargetAxes >

static bool mlir::mesh::arePartialAxesCompatible ( const SourceAxes &  sourceAxes, const TargetAxes &  targetAxes  )

static

Definition at line 42 of file Spmdization.cpp.

Referenced by handlePartialAxesDuringResharding().

collectiveProcessGroupSize() [1/2]

template<typename MeshAxesRange >

int64_t mlir::mesh::collectiveProcessGroupSize	(	MeshAxesRange &&	meshAxes,
		MeshOp	mesh
	)

Definition at line 169 of file MeshOps.h.

References collectiveProcessGroupSize().

collectiveProcessGroupSize() [2/2]

template<typename MeshAxesRange , typename MeshShapeRange >

int64_t mlir::mesh::collectiveProcessGroupSize	(	MeshAxesRange &&	meshAxes,
		MeshShapeRange &&	meshShape
	)

Definition at line 153 of file MeshOps.h.

Referenced by collectiveProcessGroupSize(), sliceResultType(), verifyAllToAllOperandAndResultShape(), verifyGatherOperandAndResultShape(), and verifyScatterOrSliceOperandAndResultShape().

createCollectiveProcessGroupSize()

TypedValue< IndexType > mlir::mesh::createCollectiveProcessGroupSize	(	MeshOp	mesh,
		ArrayRef< MeshAxis >	axes,
		ImplicitLocOpBuilder &	builder
	)

Definition at line 201 of file Transforms.cpp.

References mlir::ImplicitLocOpBuilder::create(), mlir::arith::createProduct(), mlir::Builder::getIndexType(), and mlir::ImplicitLocOpBuilder::getLoc().

createProcessLinearIndex() [1/2]

TypedValue< IndexType > mlir::mesh::createProcessLinearIndex	(	StringRef	mesh,
		ArrayRef< MeshAxis >	meshAxes,
		ImplicitLocOpBuilder &	builder
	)

Definition at line 226 of file Transforms.cpp.

References mlir::ImplicitLocOpBuilder::create().

Referenced by mlir::linalg::createDestinationPassingStyleInitOperand().

createProcessLinearIndex() [2/2]

TypedValue< IndexType > mlir::mesh::createProcessLinearIndex	(	StringRef	mesh,
		ValueRange	processInGroupMultiIndex,
		ArrayRef< MeshAxis >	meshAxes,
		ImplicitLocOpBuilder &	builder
	)

Definition at line 211 of file Transforms.cpp.

References mlir::ImplicitLocOpBuilder::create(), and mlir::affine::linearizeIndex().

detectMoveLastSplitAxisInResharding()

static std::optional<std::tuple<int64_t, int64_t, MeshAxis> > mlir::mesh::detectMoveLastSplitAxisInResharding ( MeshSharding  sourceSharding, MeshSharding  targetSharding  )

static

Definition at line 303 of file Spmdization.cpp.

References mlir::mesh::MeshSharding::getSplitAxes().

Referenced by tryMoveLastSplitAxisInResharding().

detectSplitLastAxisInResharding()

static std::optional<std::tuple<int64_t, MeshAxis> > mlir::mesh::detectSplitLastAxisInResharding ( MeshSharding  sourceSharding, MeshSharding  targetSharding  )

static

Definition at line 151 of file Spmdization.cpp.

References mlir::mesh::MeshSharding::getSplitAxes().

Referenced by trySplitLastAxisInResharding().

detectUnsplitLastAxisInResharding()

static std::optional<std::tuple<int64_t, MeshAxis> > mlir::mesh::detectUnsplitLastAxisInResharding ( MeshSharding  sourceSharding, MeshSharding  targetSharding  )

static

Definition at line 201 of file Spmdization.cpp.

References mlir::mesh::MeshSharding::getSplitAxes().

Referenced by tryUnsplitLastAxisInResharding().

gatherDimension()

int64_t mlir::mesh::gatherDimension ( int64_t  dimSize, int64_t  shardCount  )

inline

Definition at line 184 of file MeshOps.h.

Referenced by allGatherResultShapeInUnsplitLastAxis(), and allToAllResultShapeInMoveLastAxis().

getMesh() [1/2]

template<typename Op >

mesh::MeshOp mlir::mesh::getMesh	(	Op	op,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 137 of file MeshOps.h.

References getMesh(), and mlir::Op< ConcreteType, Traits >::getOperation().

getMesh() [2/2]

mesh::MeshOp mlir::mesh::getMesh ( Operation *  op, FlatSymbolRefAttr  meshSymbol, SymbolTableCollection &  symbolTableCollection  )

inline

Definition at line 128 of file MeshOps.h.

References getMeshOrNull().

Referenced by getMesh(), mlir::linalg::getMesh(), getMesh< ShardOp >(), reshard(), and shardedBlockArgumentTypes().

getMesh< ShardOp >()

template<>

mesh::MeshOp mlir::mesh::getMesh< ShardOp > ( ShardOp  op, SymbolTableCollection &  symbolTableCollection  )

inline

Definition at line 142 of file MeshOps.h.

References getMesh().

getMeshAxisAssignmentForLoopIterators()

ShardingArray mlir::mesh::getMeshAxisAssignmentForLoopIterators	(	ArrayRef< MeshSharding >	operandShardings,
		ArrayRef< MeshSharding >	resultShardings,
		ArrayRef< utils::IteratorType >	loopIteratorTypes,
		ArrayRef< AffineMap >	indexingMaps
	)

Definition at line 642 of file ShardingInterface.cpp.

References updateMeshAxisAssignmentForLoopIterators().

getMeshOrNull()

mesh::MeshOp mlir::mesh::getMeshOrNull ( Operation *  op, FlatSymbolRefAttr  meshSymbol, SymbolTableCollection &  symbolTableCollection  )

inline

Definition at line 120 of file MeshOps.h.

References mlir::SymbolTableCollection::lookupNearestSymbolFrom().

Referenced by getMesh(), getMeshAndVerify(), and spmdizeTriviallyShardableOperation().

getMeshSharding() [1/2]

FailureOr< std::pair< bool, MeshSharding > > mlir::mesh::getMeshSharding

(

OpOperand &

opOperand

)

Definition at line 154 of file ShardingInterface.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), and mlir::Value::getDefiningOp().

getMeshSharding() [2/2]

FailureOr< std::pair< bool, MeshSharding > > mlir::mesh::getMeshSharding

(

OpResult

result

)

Definition at line 110 of file ShardingInterface.cpp.

References getSharding(), mlir::Value::getUsers(), and mlir::Value::hasOneUse().

Referenced by visitOp().

getMixedAsValues()

SmallVector< Value > mlir::mesh::getMixedAsValues	(	OpBuilder	b,
		const Location &	loc,
		llvm::ArrayRef< int64_t >	statics,
		ValueRange	dynamics,
		Type	type = Type()
	)

Converts a vector of OpFoldResults (ints) into vector of Values of the provided type.

Definition at line 77 of file MeshOps.cpp.

References mlir::OpBuilder::create(), mlir::Builder::getI64IntegerAttr(), mlir::Builder::getI64Type(), mlir::Builder::getIndexAttr(), and mlir::Builder::getIndexType().

getOperandShardings()

static std::vector<MeshSharding> mlir::mesh::getOperandShardings ( Operation &  op )

static

Definition at line 667 of file Spmdization.cpp.

References mlir::Value::getDefiningOp(), mlir::Operation::getNumOperands(), and mlir::Operation::getOperands().

Referenced by spmdizeOperation().

getReductionMeshAxes()

SmallVector< MeshAxis > mlir::mesh::getReductionMeshAxes	(	ArrayRef< utils::IteratorType >	loopIteratorTypes,
		ArrayRef< SmallVector< MeshAxis >>	meshAxisAssignmentForLoopIterators
	)

Definition at line 696 of file ShardingInterface.cpp.

Referenced by mlir::linalg::spmdizeLinalgOpWithShardedReduction().

getResultShardings()

static std::vector<MeshSharding> mlir::mesh::getResultShardings ( Operation &  op )

static

Definition at line 687 of file Spmdization.cpp.

References mlir::Operation::getNumResults(), mlir::Operation::getOperands(), mlir::Operation::getResults(), and mlir::Value::getUsers().

Referenced by spmdizeOperation().

handlePartialAxesDuringResharding()

static std::tuple<TypedValue<ShapedType>, MeshSharding> mlir::mesh::handlePartialAxesDuringResharding ( OpBuilder &  builder, MeshSharding  sourceSharding, MeshSharding  targetSharding, TypedValue< ShapedType >  sourceShard  )

static

Definition at line 56 of file Spmdization.cpp.

References arePartialAxesCompatible(), mlir::OpBuilder::create(), mlir::mesh::MeshSharding::get(), mlir::mesh::MeshSharding::getMeshAttr(), mlir::mesh::MeshSharding::getPartialAxes(), mlir::mesh::MeshSharding::getPartialType(), mlir::mesh::MeshSharding::getSplitAxes(), and mlir::OpBuilder::setInsertionPointAfterValue().

Referenced by reshardOn1DMesh().

isAtLeastOneReductionIteratorSharded()

bool mlir::mesh::isAtLeastOneReductionIteratorSharded	(	ArrayRef< utils::IteratorType >	loopIteratorTypes,
		ArrayRef< SmallVector< MeshAxis >>	meshAxisAssignmentForLoopIterators
	)

Definition at line 683 of file ShardingInterface.cpp.

isFullReplication()

bool mlir::mesh::isFullReplication ( MeshSharding  sharding )

inline

Definition at line 112 of file MeshOps.h.

References mlir::mesh::MeshSharding::getPartialAxes(), and mlir::mesh::MeshSharding::getSplitAxes().

Referenced by isValueCompatibleWithFullReplicationSharding(), maybeInsertSourceShardingAnnotation(), and reshard().

isReductionLoop()

bool mlir::mesh::isReductionLoop ( utils::IteratorType  iType )

inline

Definition at line 100 of file MeshOps.h.

Referenced by mlir::mesh::detail::defaultGetShardingOption(), and getSharding().

maybeInsertSourceShardingAnnotation()

void mlir::mesh::maybeInsertSourceShardingAnnotation	(	MeshSharding	sharding,
		OpOperand &	operand,
		OpBuilder &	builder
	)

Definition at line 353 of file MeshOps.cpp.

References mlir::OpBuilder::create(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Value::getDefiningOp(), mlir::Value::getLoc(), mlir::detail::IROperandBase::getOwner(), mlir::Value::getType(), mlir::Operation::hasTrait(), isFullReplication(), mlir::Operation::replaceUsesWithIf(), and mlir::OpBuilder::setInsertionPoint().

Referenced by addShardOp().

maybeInsertTargetShardingAnnotation()

void mlir::mesh::maybeInsertTargetShardingAnnotation	(	MeshSharding	sharding,
		OpResult	result,
		OpBuilder &	builder
	)

Definition at line 339 of file MeshOps.cpp.

References mlir::Value::getUses(), and maybeInsertTargetShardingAnnotationImpl().

Referenced by addShardOp().

moveLastSplitAxisInResharding()

static std::tuple<TypedValue<ShapedType>, MeshSharding> mlir::mesh::moveLastSplitAxisInResharding ( ImplicitLocOpBuilder &  builder, MeshOp  mesh, MeshSharding  sourceSharding, ShapedType  sourceUnshardedShape, TypedValue< ShapedType >  sourceShard, int64_t  sourceTensorAxis, int64_t  targetTensorAxis, MeshAxis  meshAxis  )

static

Definition at line 387 of file Spmdization.cpp.

References allToAllResultShapeInMoveLastAxis(), mlir::ImplicitLocOpBuilder::create(), mlir::get(), mlir::Builder::getContext(), mlir::OpBuilder::setInsertionPointAfterValue(), shardShapedType(), and targetShardingInMoveLastAxis().

Referenced by tryMoveLastSplitAxisInResharding().

populateAllOpLoweringPatterns()

void mlir::mesh::populateAllOpLoweringPatterns	(	RewritePatternSet &	patterns,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 189 of file Transforms.cpp.

References mlir::patterns, populateAllSliceOpLoweringPatterns(), and populateProcessMultiIndexOpLoweringPatterns().

populateAllReduceEndomorphismSimplificationPatterns()

template<typename AlgebraicOp >

void mlir::mesh::populateAllReduceEndomorphismSimplificationPatterns	(	RewritePatternSet &	patterns,
		ReductionKind	reduction
	)

Definition at line 40 of file Simplifications.h.

References mlir::patterns.

populateAllSliceOpLoweringPatterns()

void mlir::mesh::populateAllSliceOpLoweringPatterns	(	RewritePatternSet &	patterns,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 177 of file Transforms.cpp.

References mlir::patterns.

Referenced by populateAllOpLoweringPatterns().

populateFoldingPatterns()

void mlir::mesh::populateFoldingPatterns	(	RewritePatternSet &	patterns,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 115 of file Simplifications.cpp.

References mlir::patterns.

Referenced by populateSimplificationPatterns().

populateProcessMultiIndexOpLoweringPatterns()

void mlir::mesh::populateProcessMultiIndexOpLoweringPatterns	(	RewritePatternSet &	patterns,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 167 of file Transforms.cpp.

References mlir::patterns.

Referenced by populateAllOpLoweringPatterns().

populateSimplificationPatterns()

void mlir::mesh::populateSimplificationPatterns	(	RewritePatternSet &	patterns,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 24 of file Simplifications.cpp.

References mlir::patterns, and populateFoldingPatterns().

registerAllOpLoweringDialects()

void mlir::mesh::registerAllOpLoweringDialects

(

DialectRegistry &

registry

)

Definition at line 195 of file Transforms.cpp.

References registerAllSliceOpLoweringDialects(), and registerProcessMultiIndexOpLoweringDialects().

registerAllSliceOpLoweringDialects()

void mlir::mesh::registerAllSliceOpLoweringDialects

(

DialectRegistry &

registry

)

Definition at line 183 of file Transforms.cpp.

References mlir::DialectRegistry::insert().

Referenced by registerAllOpLoweringDialects().

registerProcessMultiIndexOpLoweringDialects()

void mlir::mesh::registerProcessMultiIndexOpLoweringDialects

(

DialectRegistry &

registry

)

Definition at line 173 of file Transforms.cpp.

References mlir::DialectRegistry::insert().

Referenced by registerAllOpLoweringDialects().

removeTrailingEmptySubArray()

template<typename T >

void mlir::mesh::removeTrailingEmptySubArray

(

SmallVector< SmallVector< T >> &

array

)

Definition at line 106 of file MeshOps.h.

Referenced by mlir::mesh::detail::defaultGetShardingOption(), and getSharding().

reshard() [1/3]

TypedValue<ShapedType> mlir::mesh::reshard	(	ImplicitLocOpBuilder &	builder,
		MeshOp	mesh,
		MeshSharding	sourceSharding,
		MeshSharding	targetSharding,
		TypedValue< ShapedType >	sourceUnshardedValue,
		TypedValue< ShapedType >	sourceShard
	)

Definition at line 556 of file Spmdization.cpp.

References isFullReplication(), reshardOn1DMesh(), and tryUpdateHaloInResharding().

Referenced by spmdizeOperation().

reshard() [2/3]

TypedValue< ShapedType > mlir::mesh::reshard	(	OpBuilder &	builder,
		MeshOp	mesh,
		ShardOp	source,
		ShardOp	target,
		TypedValue< ShapedType >	sourceShardValue
	)

Definition at line 582 of file Spmdization.cpp.

Referenced by reshard().

reshard() [3/3]

TypedValue< ShapedType > mlir::mesh::reshard	(	OpBuilder &	builder,
		ShardOp	source,
		ShardOp	target,
		TypedValue< ShapedType >	sourceShardValue,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 594 of file Spmdization.cpp.

References getMesh(), and reshard().

reshardingRegisterDependentDialects()

void mlir::mesh::reshardingRegisterDependentDialects

(

DialectRegistry &

registry

)

Definition at line 603 of file Spmdization.cpp.

References mlir::DialectRegistry::insert().

reshardOn1DMesh()

static TypedValue<ShapedType> mlir::mesh::reshardOn1DMesh ( ImplicitLocOpBuilder &  builder, MeshOp  mesh, MeshSharding  sourceSharding, MeshSharding  targetSharding, TypedValue< ShapedType >  sourceUnshardedValue, TypedValue< ShapedType >  sourceShard  )

static

Definition at line 511 of file Spmdization.cpp.

References mlir::mesh::MeshSharding::getStaticHaloSizes(), mlir::mesh::MeshSharding::getStaticShardedDimsOffsets(), handlePartialAxesDuringResharding(), shardShapedType(), tryMoveLastSplitAxisInResharding(), trySplitLastAxisInResharding(), and tryUnsplitLastAxisInResharding().

Referenced by reshard().

shardDimension()

int64_t mlir::mesh::shardDimension ( int64_t  dimSize, int64_t  shardCount  )

inline

Definition at line 175 of file MeshOps.h.

Referenced by allToAllResultShapeInMoveLastAxis().

shardedBlockArgumentTypes()

SmallVector<Type> mlir::mesh::shardedBlockArgumentTypes	(	Block &	block,
		SymbolTableCollection &	symbolTableCollection
	)

Definition at line 616 of file Spmdization.cpp.

References mlir::Block::getArguments(), getMesh(), mlir::Operation::getUsers(), and shardShapedType().

Referenced by spmdizeBlock().

shardShapedType()

ShapedType mlir::mesh::shardShapedType	(	ShapedType	shape,
		MeshOp	mesh,
		MeshSharding	sharding
	)

Definition at line 282 of file MeshOps.cpp.

References mlir::mesh::MeshSharding::getSplitAxes(), mlir::mesh::MeshSharding::getStaticHaloSizes(), mlir::mesh::MeshSharding::getStaticShardedDimsOffsets(), and shardShape().

Referenced by moveLastSplitAxisInResharding(), reshardOn1DMesh(), shardedBlockArgumentTypes(), shardType(), and unsplitLastAxisInResharding().

shardType()

Type mlir::mesh::shardType	(	Type	type,
		MeshOp	mesh,
		MeshSharding	sharding
	)

Definition at line 292 of file MeshOps.cpp.

References shardShapedType().

Referenced by spmdizeTriviallyShardableOperation().

splitLastAxisInResharding()

static std::tuple<TypedValue<ShapedType>, MeshSharding> mlir::mesh::splitLastAxisInResharding ( ImplicitLocOpBuilder &  builder, MeshSharding  sourceSharding, TypedValue< ShapedType >  sourceShard, MeshOp  mesh, int64_t  splitTensorAxis, MeshAxis  splitMeshAxis  )

static

Definition at line 130 of file Spmdization.cpp.

References mlir::ImplicitLocOpBuilder::create(), mlir::Builder::getContext(), and targetShardingInSplitLastAxis().

Referenced by trySplitLastAxisInResharding().

spmdizeBlock()

static LogicalResult mlir::mesh::spmdizeBlock ( Block &  block, IRMapping &  spmdizationMap, SymbolTableCollection &  symbolTableCollection, OpBuilder &  builder  )

static

Definition at line 779 of file Spmdization.cpp.

References mlir::OpBuilder::createBlock(), mlir::Block::getArguments(), mlir::Block::getOperations(), mlir::Block::getParent(), mlir::IRMapping::map(), mlir::OpBuilder::setInsertionPointToEnd(), shardedBlockArgumentTypes(), and spmdizeOperation().

Referenced by spmdizeFuncOp().

spmdizeFullyReplicatedOperation()

void mlir::mesh::spmdizeFullyReplicatedOperation	(	Operation &	op,
		ArrayRef< Value >	spmdizedOperands,
		ArrayRef< MeshSharding >	operandShardings,
		ArrayRef< MeshSharding >	resultShardings,
		IRMapping &	spmdizationMap,
		SymbolTableCollection &	symbolTable,
		OpBuilder &	builder
	)

Definition at line 613 of file ShardingInterface.cpp.

References areValuesCompatibleWithFullReplicationShardings(), mlir::OpBuilder::clone(), mlir::Operation::getOperands(), and mlir::Operation::getResults().

Referenced by spmdizeOperation().

spmdizeFuncOp()

static LogicalResult mlir::mesh::spmdizeFuncOp ( FunctionOpInterface  op, IRMapping &  spmdizationMap, SymbolTableCollection &  symbolTableCollection  )

static

Definition at line 807 of file Spmdization.cpp.

References mlir::get(), mlir::Operation::getOperandTypes(), and spmdizeBlock().

spmdizeOperation() [1/3]

static LogicalResult mlir::mesh::spmdizeOperation ( Operation &  op, ArrayRef< Value >  spmdizedOperands, ArrayRef< MeshSharding >  operandShardings, ArrayRef< MeshSharding >  resultShardings, IRMapping &  spmdizationMap, SymbolTableCollection &  symbolTableCollection, OpBuilder &  builder  )

static

Definition at line 638 of file Spmdization.cpp.

References mlir::Operation::getResults(), and spmdizeFullyReplicatedOperation().

spmdizeOperation() [2/3]

static LogicalResult mlir::mesh::spmdizeOperation ( Operation &  op, IRMapping &  spmdizationMap, SymbolTableCollection &  symbolTableCollection, OpBuilder &  builder  )

static

Definition at line 746 of file Spmdization.cpp.

References mlir::OpBuilder::clone(), mlir::Operation::emitError(), mlir::Operation::getOperands(), getOperandShardings(), mlir::Operation::getResult(), getResultShardings(), and mlir::IRMapping::map().

Referenced by spmdizeBlock().

spmdizeOperation() [3/3]

static LogicalResult mlir::mesh::spmdizeOperation ( ShardOp  shardOp, IRMapping &  spmdizationMap, SymbolTableCollection &  symbolTableCollection, OpBuilder &  builder  )

static

Definition at line 721 of file Spmdization.cpp.

References mlir::IRMapping::contains(), mlir::IRMapping::lookup(), mlir::IRMapping::map(), and reshard().

spmdizeTriviallyShardableOperation()

void mlir::mesh::spmdizeTriviallyShardableOperation	(	Operation &	op,
		ArrayRef< Value >	spmdizedOperands,
		ArrayRef< MeshSharding >	operandShardings,
		ArrayRef< MeshSharding >	resultShardings,
		IRMapping &	spmdizationMap,
		SymbolTableCollection &	symbolTable,
		OpBuilder &	builder
	)

Definition at line 709 of file ShardingInterface.cpp.

References mlir::OpBuilder::clone(), getMeshOrNull(), mlir::Operation::getResults(), and shardType().

Referenced by mlir::mesh::IndependentParallelIteratorDomainShardingInterface< Op >::spmdize(), mlir::mesh::ElementwiseShardingInterface< ElemwiseOp >::spmdize(), and mlir::linalg::spmdizeLinalgOpWithShardedReduction().

targetShardingInMoveLastAxis()

static MeshSharding mlir::mesh::targetShardingInMoveLastAxis ( MLIRContext *  ctx, MeshSharding  sourceSharding, int64_t  sourceTensorAxis, int64_t  targetTensorAxis  )

static

Definition at line 344 of file Spmdization.cpp.

References mlir::mesh::MeshSharding::get(), mlir::detail::DenseArrayAttrImpl< T >::get(), mlir::mesh::MeshSharding::getMeshAttr(), mlir::mesh::MeshSharding::getPartialAxes(), mlir::mesh::MeshSharding::getPartialType(), and mlir::mesh::MeshSharding::getSplitAxes().

Referenced by moveLastSplitAxisInResharding().

targetShardingInSplitLastAxis()

static MeshSharding mlir::mesh::targetShardingInSplitLastAxis ( MLIRContext *  ctx, MeshSharding  sourceSharding, int64_t  splitTensorAxis, MeshAxis  splitMeshAxis  )

static

Definition at line 106 of file Spmdization.cpp.

References mlir::mesh::MeshSharding::get(), mlir::detail::DenseArrayAttrImpl< T >::get(), mlir::mesh::MeshSharding::getMeshAttr(), mlir::mesh::MeshSharding::getPartialAxes(), mlir::mesh::MeshSharding::getPartialType(), and mlir::mesh::MeshSharding::getSplitAxes().

Referenced by splitLastAxisInResharding().

targetShardingInUnsplitLastAxis()

static MeshSharding mlir::mesh::targetShardingInUnsplitLastAxis ( MLIRContext *  ctx, MeshSharding  sourceSharding, int64_t  splitTensorAxis  )

static

Definition at line 229 of file Spmdization.cpp.

References mlir::mesh::MeshSharding::get(), mlir::detail::DenseArrayAttrImpl< T >::get(), mlir::mesh::MeshSharding::getMeshAttr(), mlir::mesh::MeshSharding::getPartialAxes(), mlir::mesh::MeshSharding::getPartialType(), and mlir::mesh::MeshSharding::getSplitAxes().

Referenced by unsplitLastAxisInResharding().

tryMoveLastSplitAxisInResharding()

static std::optional<std::tuple<TypedValue<ShapedType>, MeshSharding> > mlir::mesh::tryMoveLastSplitAxisInResharding ( ImplicitLocOpBuilder &  builder, MeshOp  mesh, MeshSharding  sourceSharding, MeshSharding  targetSharding, ShapedType  sourceUnshardedShape, TypedValue< ShapedType >  sourceShard  )

static

Definition at line 414 of file Spmdization.cpp.

References detectMoveLastSplitAxisInResharding(), and moveLastSplitAxisInResharding().

Referenced by reshardOn1DMesh().

trySplitLastAxisInResharding()

static std::optional<std::tuple<TypedValue<ShapedType>, MeshSharding> > mlir::mesh::trySplitLastAxisInResharding ( ImplicitLocOpBuilder &  builder, MeshOp  mesh, MeshSharding  sourceSharding, MeshSharding  targetSharding, TypedValue< ShapedType >  sourceShard  )

static

Definition at line 183 of file Spmdization.cpp.

References detectSplitLastAxisInResharding(), and splitLastAxisInResharding().

Referenced by reshardOn1DMesh().

tryUnsplitLastAxisInResharding()

static std::optional<std::tuple<TypedValue<ShapedType>, MeshSharding> > mlir::mesh::tryUnsplitLastAxisInResharding ( ImplicitLocOpBuilder &  builder, MeshOp  mesh, MeshSharding  sourceSharding, MeshSharding  targetSharding, ShapedType  sourceUnshardedShape, TypedValue< ShapedType >  sourceShard  )

static

Definition at line 281 of file Spmdization.cpp.

References detectUnsplitLastAxisInResharding(), and unsplitLastAxisInResharding().

Referenced by reshardOn1DMesh().

tryUpdateHaloInResharding()

static std::optional<std::tuple<TypedValue<ShapedType>, MeshSharding> > mlir::mesh::tryUpdateHaloInResharding ( ImplicitLocOpBuilder &  builder, MeshOp  mesh, MeshSharding  sourceSharding, MeshSharding  targetSharding, ShapedType  sourceUnshardedShape, TypedValue< ShapedType >  sourceShard  )

static

Definition at line 435 of file Spmdization.cpp.

References mlir::ImplicitLocOpBuilder::create(), mlir::mesh::MeshSharding::equalHaloSizes(), mlir::mesh::MeshSharding::equalSplitAndPartialAxes(), mlir::get(), mlir::Builder::getContext(), mlir::mesh::MeshSharding::getDynamicHaloSizes(), mlir::mesh::MeshSharding::getPartialAxes(), mlir::mesh::MeshSharding::getSplitAxes(), mlir::mesh::MeshSharding::getStaticHaloSizes(), and mlir::mesh::MeshSharding::getStaticShardedDimsOffsets().

Referenced by reshard().

unsplitLastAxisInResharding()

static std::tuple<TypedValue<ShapedType>, MeshSharding> mlir::mesh::unsplitLastAxisInResharding ( ImplicitLocOpBuilder &  builder, MeshSharding  sourceSharding, ShapedType  sourceUnshardedShape, TypedValue< ShapedType >  sourceShard, MeshOp  mesh, int64_t  splitTensorAxis, MeshAxis  splitMeshAxis  )

static

Definition at line 256 of file Spmdization.cpp.

References allGatherResultShapeInUnsplitLastAxis(), mlir::ImplicitLocOpBuilder::create(), mlir::get(), mlir::Builder::getContext(), mlir::OpBuilder::setInsertionPointAfterValue(), shardShapedType(), and targetShardingInUnsplitLastAxis().

Referenced by tryUnsplitLastAxisInResharding().