mlir::mesh::MeshSharding Class Reference

#include "mlir/Dialect/Mesh/IR/MeshOps.h"

Public Member Functions
	MeshSharding ()=default
	MeshSharding (Value rhs)
::mlir::FlatSymbolRefAttr	getMeshAttr () const
::llvm::StringRef	getMesh () const
ArrayRef< MeshAxesAttr >	getSplitAxes () const
ArrayRef< MeshAxis >	getPartialAxes () const
ReductionKind	getPartialType () const
ArrayRef< int64_t >	getStaticHaloSizes () const
ArrayRef< int64_t >	getStaticShardedDimsOffsets () const
ArrayRef< Value >	getDynamicHaloSizes () const
ArrayRef< Value >	getDynamicShardedDimsOffsets () const
	operator bool () const
bool	operator== (Value rhs) const
bool	operator!= (Value rhs) const
bool	operator== (const MeshSharding &rhs) const
bool	operator!= (const MeshSharding &rhs) const
bool	equalSplitAndPartialAxes (const MeshSharding &rhs) const
bool	equalHaloAndShardSizes (const MeshSharding &rhs) const
bool	equalHaloSizes (const MeshSharding &rhs) const
bool	equalShardSizes (const MeshSharding &rhs) const

Static Public Member Functions
static MeshSharding	get (::mlir::FlatSymbolRefAttr mesh_, ArrayRef< MeshAxesAttr > split_axes_, ArrayRef< MeshAxis > partial_axes_={}, ReductionKind partial_type_=ReductionKind::Sum, ArrayRef< int64_t > static_halo_sizes_={}, ArrayRef< int64_t > static_sharded_dims_offsets_={}, ArrayRef< Value > dynamic_halo_sizes_={}, ArrayRef< Value > dynamic_sharded_dims_offsets_={})

Detailed Description

Definition at line 42 of file MeshOps.h.

Constructor & Destructor Documentation

MeshSharding() [1/2]

mlir::mesh::MeshSharding::MeshSharding ( )

default

MeshSharding() [2/2]

MeshSharding::MeshSharding

(

Value

rhs

)

Definition at line 710 of file MeshOps.cpp.

References get(), and mlir::Value::getDefiningOp().

Member Function Documentation

equalHaloAndShardSizes()

bool MeshSharding::equalHaloAndShardSizes

(

const MeshSharding &

rhs

)

const

Definition at line 653 of file MeshOps.cpp.

References equalHaloSizes(), and equalShardSizes().

Referenced by operator==().

equalHaloSizes()

bool MeshSharding::equalHaloSizes

(

const MeshSharding &

rhs

)

const

Definition at line 678 of file MeshOps.cpp.

References getDynamicHaloSizes(), and getStaticHaloSizes().

Referenced by equalHaloAndShardSizes(), and mlir::mesh::tryUpdateHaloInResharding().

equalShardSizes()

bool MeshSharding::equalShardSizes

(

const MeshSharding &

rhs

)

const

Definition at line 657 of file MeshOps.cpp.

References getDynamicShardedDimsOffsets(), and getStaticShardedDimsOffsets().

Referenced by equalHaloAndShardSizes().

equalSplitAndPartialAxes()

bool MeshSharding::equalSplitAndPartialAxes

(

const MeshSharding &

rhs

)

const

Definition at line 623 of file MeshOps.cpp.

References getMesh(), getPartialAxes(), getPartialType(), getSplitAxes(), and min().

Referenced by operator==(), and mlir::mesh::tryUpdateHaloInResharding().

get()

MeshSharding MeshSharding::get ( ::mlir::FlatSymbolRefAttr  mesh_, ArrayRef< MeshAxesAttr >  split_axes_, ArrayRef< MeshAxis >  partial_axes_ = {}, ReductionKind  partial_type_ = ReductionKind::Sum, ArrayRef< int64_t >  static_halo_sizes_ = {}, ArrayRef< int64_t >  static_sharded_dims_offsets_ = {}, ArrayRef< Value >  dynamic_halo_sizes_ = {}, ArrayRef< Value >  dynamic_sharded_dims_offsets_ = {}  )

static

Definition at line 722 of file MeshOps.cpp.

References mlir::clone(), copy(), mlir::detail::enumerate(), and mlir::detail::DenseArrayAttrImpl< T >::get().

Referenced by getSharding(), mlir::mesh::handlePartialAxesDuringResharding(), MeshSharding(), mlir::mesh::targetShardingInMoveLastAxis(), mlir::mesh::targetShardingInSplitLastAxis(), and mlir::mesh::targetShardingInUnsplitLastAxis().

getDynamicHaloSizes()

ArrayRef<Value> mlir::mesh::MeshSharding::getDynamicHaloSizes ( ) const

inline

Definition at line 73 of file MeshOps.h.

Referenced by equalHaloSizes(), and mlir::mesh::tryUpdateHaloInResharding().

getDynamicShardedDimsOffsets()

ArrayRef<Value> mlir::mesh::MeshSharding::getDynamicShardedDimsOffsets ( ) const

inline

Definition at line 74 of file MeshOps.h.

Referenced by equalShardSizes().

getMesh()

::llvm::StringRef mlir::mesh::MeshSharding::getMesh ( ) const

inline

Definition at line 65 of file MeshOps.h.

References mlir::FlatSymbolRefAttr::getValue().

Referenced by mlir::linalg::createAllReduceForResultWithoutPartialSharding(), and equalSplitAndPartialAxes().

getMeshAttr()

::mlir::FlatSymbolRefAttr mlir::mesh::MeshSharding::getMeshAttr ( ) const

inline

Definition at line 64 of file MeshOps.h.

Referenced by mlir::mesh::detail::defaultGetShardingOption(), mlir::mesh::handlePartialAxesDuringResharding(), mlir::mesh::targetShardingInMoveLastAxis(), mlir::mesh::targetShardingInSplitLastAxis(), and mlir::mesh::targetShardingInUnsplitLastAxis().

getPartialAxes()

ArrayRef<MeshAxis> mlir::mesh::MeshSharding::getPartialAxes ( ) const

inline

Definition at line 67 of file MeshOps.h.

Referenced by mlir::linalg::createAllReduceForResultWithoutPartialSharding(), mlir::mesh::detail::defaultGetShardingOption(), equalSplitAndPartialAxes(), mlir::mesh::handlePartialAxesDuringResharding(), mlir::mesh::isFullReplication(), mlir::mesh::targetShardingInMoveLastAxis(), mlir::mesh::targetShardingInSplitLastAxis(), mlir::mesh::targetShardingInUnsplitLastAxis(), and mlir::mesh::tryUpdateHaloInResharding().

getPartialType()

ReductionKind mlir::mesh::MeshSharding::getPartialType ( ) const

inline

Definition at line 68 of file MeshOps.h.

Referenced by equalSplitAndPartialAxes(), mlir::mesh::handlePartialAxesDuringResharding(), mlir::mesh::targetShardingInMoveLastAxis(), mlir::mesh::targetShardingInSplitLastAxis(), and mlir::mesh::targetShardingInUnsplitLastAxis().

getSplitAxes()

ArrayRef<MeshAxesAttr> mlir::mesh::MeshSharding::getSplitAxes ( ) const

inline

Definition at line 66 of file MeshOps.h.

Referenced by mlir::mesh::detail::defaultGetShardingOption(), mlir::mesh::detectMoveLastSplitAxisInResharding(), mlir::mesh::detectSplitLastAxisInResharding(), mlir::mesh::detectUnsplitLastAxisInResharding(), equalSplitAndPartialAxes(), mlir::mesh::handlePartialAxesDuringResharding(), mlir::mesh::isFullReplication(), mlir::mesh::shardShapedType(), mlir::mesh::targetShardingInMoveLastAxis(), mlir::mesh::targetShardingInSplitLastAxis(), mlir::mesh::targetShardingInUnsplitLastAxis(), and mlir::mesh::tryUpdateHaloInResharding().

getStaticHaloSizes()

ArrayRef<int64_t> mlir::mesh::MeshSharding::getStaticHaloSizes ( ) const

inline

Definition at line 69 of file MeshOps.h.

Referenced by equalHaloSizes(), mlir::mesh::reshardOn1DMesh(), mlir::mesh::shardShapedType(), and mlir::mesh::tryUpdateHaloInResharding().

getStaticShardedDimsOffsets()

ArrayRef<int64_t> mlir::mesh::MeshSharding::getStaticShardedDimsOffsets ( ) const

inline

Definition at line 70 of file MeshOps.h.

Referenced by equalShardSizes(), mlir::mesh::reshardOn1DMesh(), mlir::mesh::shardShapedType(), and mlir::mesh::tryUpdateHaloInResharding().

operator bool()

mlir::mesh::MeshSharding::operator bool ( ) const

inline

Definition at line 77 of file MeshOps.h.

operator!=() [1/2]

bool MeshSharding::operator!=

(

const MeshSharding &

rhs

)

const

Definition at line 706 of file MeshOps.cpp.

operator!=() [2/2]

bool MeshSharding::operator!=

(

Value

rhs

)

const

Definition at line 700 of file MeshOps.cpp.

operator==() [1/2]

bool MeshSharding::operator==

(

const MeshSharding &

rhs

)

const

Definition at line 702 of file MeshOps.cpp.

References equalHaloAndShardSizes(), and equalSplitAndPartialAxes().

operator==() [2/2]

bool MeshSharding::operator==

(

Value

rhs

)

const

Definition at line 696 of file MeshOps.cpp.

References equalHaloAndShardSizes(), and equalSplitAndPartialAxes().

---

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

• include/mlir/Dialect/Mesh/IR/MeshOps.h
• lib/Dialect/Mesh/IR/MeshOps.cpp