tenzr::Mesh¶

Named grid of devices defining axes used by sharding specifications.

#include <mesh.h>

Public Functions¶

Public Functions Documentation¶

function Mesh¶

inline Mesh()

function Mesh¶

inline Mesh(
    std::vector< int64_t > devices,
    std::vector< int64_t > shape,
    std::vector< std::string > axis_names
)

function Mesh¶

inline Mesh(
    std::vector< int64_t > shape,
    std::vector< std::string > axis_names
)

function Ndim¶

inline size_t Ndim() const

Number of grid dimensions.

function NumDevices¶

inline int64_t NumDevices() const

Total number of devices in the mesh.

function Size¶

inline int64_t Size() const

Alias for num_devices().

function Empty¶

inline bool Empty() const

Whether the mesh holds no devices.

function shape¶

inline const std::vector< int64_t > & shape() const

Extent of each grid axis.

function AxisNames¶

inline const std::vector< std::string > & AxisNames() const

Name of each grid axis.

function Devices¶

inline const std::vector< int64_t > & Devices() const

Device ids in row-major order.

function AxisSize¶

inline int64_t AxisSize(
    std::string_view axis_name
) const

Parameters:

• axis_name Axis name.

Return: Size of the axis.

Extent of a named axis.

function AxisIndex¶

inline int64_t AxisIndex(
    std::string_view axis_name
) const

Parameters:

• axis_name Axis name.

Return: Zero-based axis index, or -1 if absent.

Index of a named axis.

function HasAxis¶

inline bool HasAxis(
    std::string_view axis_name
) const

Whether a named axis exists.

function DeviceCoords¶

inline std::vector< int64_t > DeviceCoords(
    int64_t device_id
) const

Parameters:

• device_id Device id.

Return: Row-major coordinates of the device.

Coordinates of a device within the grid.

function DeviceAt¶

inline int64_t DeviceAt(
    const std::vector< int64_t > & coords
) const

Parameters:

• coords Row-major grid coordinates.

Return: Device id at those coordinates.

Device id at the given grid coordinates.

function ToString¶

inline std::string ToString() const

Render the mesh as a string.

function ToShardyMeshDecl¶

inline std::string ToShardyMeshDecl(
    std::string_view mesh_sym ="mesh"
) const

Emit Shardy MLIR dialect mesh declaration: sdy.mesh @mesh = <["data"=4, "model"=2]>

function operator==¶

inline bool operator==(
    const Mesh & other
) const

function operator!=¶

inline bool operator!=(
    const Mesh & other
) const