mlir::shard::ReshardingPattern Class Reference

Base class for resharding patterns. More...

Inheritance diagram for mlir::shard::ReshardingPattern:

Public Member Functions
virtual	~ReshardingPattern ()=default
virtual 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)=0
	Try to apply this resharding pattern.

Static Protected Member Functions
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

Base class for resharding patterns.

Subclasses implement tryApply to detect and apply a specific resharding.

Definition at line 51 of file Partition.cpp.

Constructor & Destructor Documentation

~ReshardingPattern()

virtual mlir::shard::ReshardingPattern::~ReshardingPattern ( )

virtualdefault

Member Function Documentation

hasStaticOffsets()

bool mlir::shard::ReshardingPattern::hasStaticOffsets ( const Sharding & srcSharding, const Sharding & tgtSharding )

inlinestaticprotected

Returns true if either sharding has non-empty static sharded dims offsets.

Definition at line 65 of file Partition.cpp.

References mlir::shard::Sharding::getStaticShardedDimsOffsets().

Referenced by mlir::shard::UpdateHaloPattern::tryApply().

hasStaticOffsetsOrHalos()

bool mlir::shard::ReshardingPattern::hasStaticOffsetsOrHalos ( const Sharding & srcSharding, const Sharding & tgtSharding )

inlinestaticprotected

Returns true if either sharding has non-empty static sharded dims offsets or non-empty static halo sizes.

Definition at line 73 of file Partition.cpp.

Referenced by mlir::shard::MoveSplitAxisPattern::tryApply(), mlir::shard::SplitLastAxisPattern::tryApply(), and mlir::shard::UnsplitLastAxesPattern::tryApply().

tryApply()

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

pure virtual

Try to apply this resharding pattern.

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

Implemented in mlir::shard::MoveSplitAxisPattern, mlir::shard::SplitLastAxisPattern, mlir::shard::UnsplitLastAxesPattern, and mlir::shard::UpdateHaloPattern.

---

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

• lib/Dialect/Shard/Transforms/Partition.cpp