MLIR 23.0.0git
mlir::shard::MoveSplitAxisPattern Class Reference

Move a split axis between tensor dimensions: e.g. More...

Inheritance diagram for mlir::shard::MoveSplitAxisPattern:

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 a split axis between tensor dimensions: e.g.

[[0], []] -> [[], [0]].

Definition at line 265 of file Partition.cpp.

Member Function Documentation

◆ tryApply()

std::optional< std::tuple< TypedValue< ShapedType >, Sharding > > mlir::shard::MoveSplitAxisPattern::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.

Implements mlir::shard::ReshardingPattern.

Definition at line 350 of file Partition.cpp.

References mlir::shard::ReshardingPattern::hasStaticOffsetsOrHalos().


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