Move the last split axis of one tensor dimension to the front of another tensor dimension's split axes, e.g. More...

Inheritance diagram for mlir::shard::MoveLastSplitAxisPattern:

Public Member Functions
std::optional< std::tuple< TypedValue< ShapedType >, Sharding > >	tryApply (ImplicitLocOpBuilder &builder, GridOp grid, int64_t tensorDim, const Sharding &srcSharding, const Sharding &tgtSharding, ShapedType srcUnshardedType, TypedValue< ShapedType > srcShard) override
	Try to apply this resharding pattern.
Public Member Functions inherited from mlir::shard::ReshardingPattern
virtual	~ReshardingPattern ()=default

Additional Inherited Members
Static Protected Member Functions inherited from mlir::shard::ReshardingPattern
static bool	hasStaticOffsets (const Sharding &srcSharding, const Sharding &tgtSharding)
	Returns true if either sharding has non-empty static sharded dims offsets.
static bool	hasStaticOffsetsOrHalos (const Sharding &srcSharding, const Sharding &tgtSharding)
	Returns true if either sharding has non-empty static sharded dims offsets or non-empty static halo sizes.

Detailed Description

Move the last split axis of one tensor dimension to the front of another tensor dimension's split axes, e.g.

[[0], []] -> [[], [0]] or [[0, 1], [2]] -> [[0], [1, 2]].

Definition at line 277 of file Partition.cpp.

Member Function Documentation

std::optional< std::tuple< TypedValue< ShapedType >, Sharding > > mlir::shard::MoveLastSplitAxisPattern::tryApply	(	ImplicitLocOpBuilder &	builder,
		GridOp	grid,
		int64_t	tensorDim,
		const Sharding &	srcSharding,
		const Sharding &	tgtSharding,
		ShapedType	srcUnshardedType,
		TypedValue< ShapedType >	srcShard )

inlineoverridevirtual

Try to apply this resharding pattern.

Returns the resharded value and resulting sharding on success, or std::nullopt if the pattern doesn't match.

Definition at line 384 of file Partition.cpp.

The documentation for this class was generated from the following file: