Update halo sizes: handles cases where only the halo sizes differ between source and target sharding. More...
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
Update halo sizes: handles cases where only the halo sizes differ between source and target sharding.
Requires copying the "core" of the source tensor into the "core" of the destination tensor followed by an update halo op.
Definition at line 368 of file Partition.cpp.
Member Function Documentation
◆ tryApply()
|
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 371 of file Partition.cpp.
References mlir::shard::Sharding::equalHaloSizes(), mlir::shard::Sharding::equalSplitAxes(), mlir::Builder::getContext(), mlir::shard::Sharding::getDynamicHaloSizes(), mlir::ImplicitLocOpBuilder::getLoc(), mlir::shard::Sharding::getSplitAxes(), mlir::shard::Sharding::getStaticHaloSizes(), and mlir::shard::ReshardingPattern::hasStaticOffsets().
The documentation for this class was generated from the following file:
- lib/Dialect/Shard/Transforms/Partition.cpp
Generated on for MLIR by
1.14.0