Unsplit trailing axes: e.g. [[0, 1, 2]] -> [[0, 1]] or [[0, 1, 2]] -> []. More...

Inheritance diagram for mlir::shard::UnsplitLastAxesPattern:

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

Unsplit trailing axes: e.g. [[0, 1, 2]] -> [[0, 1]] or [[0, 1, 2]] -> [].

Definition at line 156 of file Partition.cpp.

Member Function Documentation

std::optional< std::tuple< TypedValue< ShapedType >, Sharding > > mlir::shard::UnsplitLastAxesPattern::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 250 of file Partition.cpp.

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