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 >

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, OpOperand &operand, OpBuilder &builder, ShardOp &newShardOp)
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 612 of file Spmdization.cpp.

Function Documentation

allGatherResultShapeInUnsplitLastAxis()

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

static

Definition at line 249 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 376 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 44 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 305 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 153 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 203 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 78 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 669 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 689 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 58 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 351 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() [1/2]

void mlir::mesh::maybeInsertTargetShardingAnnotation	(	MeshSharding	sharding,
		OpOperand &	operand,
		OpBuilder &	builder,
		ShardOp &	newShardOp
	)

Definition at line 301 of file MeshOps.cpp.

References mlir::OpBuilder::create(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Value::getLoc(), mlir::detail::IROperandBase::getOwner(), mlir::RewriterBase::replaceAllUsesExcept(), mlir::RewriterBase::replaceUsesWithIf(), and mlir::OpBuilder::setInsertionPointAfterValue().

Referenced by addShardOp(), and maybeInsertTargetShardingAnnotation().

maybeInsertTargetShardingAnnotation() [2/2]

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

Definition at line 342 of file MeshOps.cpp.

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

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 389 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 116 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 25 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 558 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 584 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 596 of file Spmdization.cpp.

References getMesh(), and reshard().

reshardingRegisterDependentDialects()

void mlir::mesh::reshardingRegisterDependentDialects

(

DialectRegistry &

registry

)

Definition at line 605 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 513 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 618 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 283 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 293 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 132 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 781 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 809 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 640 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 748 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 723 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 346 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 108 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 231 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 416 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 185 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 283 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 437 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 258 of file Spmdization.cpp.

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

Referenced by tryUnsplitLastAxisInResharding().