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 >	getStaticShardedDimsSizes () const
ArrayRef< Value >	getDynamicHaloSizes () const
ArrayRef< Value >	getDynamicShardedDimsSizes () 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

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_sizes_={}, ArrayRef< Value > dynamic_halo_sizes_={}, ArrayRef< Value > dynamic_sharded_dims_sizes_={})

Detailed Description

Definition at line 41 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 625 of file MeshOps.cpp.

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

Member Function Documentation

equalHaloAndShardSizes()

bool MeshSharding::equalHaloAndShardSizes

(

const MeshSharding &

rhs

)

const

Definition at line 578 of file MeshOps.cpp.

References getDynamicHaloSizes(), getDynamicShardedDimsSizes(), getStaticHaloSizes(), and getStaticShardedDimsSizes().

Referenced by operator==().

equalSplitAndPartialAxes()

bool MeshSharding::equalSplitAndPartialAxes

(

const MeshSharding &

rhs

)

const

Definition at line 553 of file MeshOps.cpp.

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

Referenced by operator==().

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_sizes_ = {}, ArrayRef< Value >  dynamic_halo_sizes_ = {}, ArrayRef< Value >  dynamic_sharded_dims_sizes_ = {}  )

static

Definition at line 637 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 72 of file MeshOps.h.

Referenced by equalHaloAndShardSizes().

getDynamicShardedDimsSizes()

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

inline

Definition at line 73 of file MeshOps.h.

Referenced by equalHaloAndShardSizes().

getMesh()

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

inline

Definition at line 64 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 63 of file MeshOps.h.

Referenced by 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 66 of file MeshOps.h.

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

getPartialType()

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

inline

Definition at line 67 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 65 of file MeshOps.h.

Referenced by 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(), and mlir::mesh::targetShardingInUnsplitLastAxis().

getStaticHaloSizes()

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

inline

Definition at line 68 of file MeshOps.h.

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

getStaticShardedDimsSizes()

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

inline

Definition at line 69 of file MeshOps.h.

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

operator bool()

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

inline

Definition at line 76 of file MeshOps.h.

operator!=() [1/2]

bool MeshSharding::operator!=

(

const MeshSharding &

rhs

)

const

Definition at line 621 of file MeshOps.cpp.

operator!=() [2/2]

bool MeshSharding::operator!=

(

Value

rhs

)

const

Definition at line 615 of file MeshOps.cpp.

operator==() [1/2]

bool MeshSharding::operator==

(

const MeshSharding &

rhs

)

const

Definition at line 617 of file MeshOps.cpp.

References equalHaloAndShardSizes(), and equalSplitAndPartialAxes().

operator==() [2/2]

bool MeshSharding::operator==

(

Value

rhs

)

const

Definition at line 611 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