tenzr::execution::DeviceMesh¶

Logical grid of devices with optional named dimensions, mapping between device ids and mesh coordinates.

#include <collective.h>

Public Functions¶

Public Functions Documentation¶

function DeviceMesh¶

explicit DeviceMesh(
    std::vector< int64_t > shape,
    std::vector< std::string > dim_names ={}
)

Parameters:

• shape Mesh dimension sizes.
• dim_names Optional names for each mesh axis; missing names are generated.

Create a device mesh with the given shape and optional axis names.

function NumDevices¶

inline int64_t NumDevices() const

Total number of devices in the mesh.

function Ndim¶

inline int64_t Ndim() const

Number of mesh dimensions.

function shape¶

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

Size of each mesh dimension.

function DimNames¶

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

Name of each mesh dimension.

function DimSize¶

int64_t DimSize(
    std::string_view name
) const

Size of the named dimension.

Parameters:

• name Axis name.

Return: Size along that axis.

Size of the named mesh axis.

function DimIndex¶

int64_t DimIndex(
    std::string_view name
) const

Index of the named dimension, or -1 if absent.

Parameters:

• name Axis name.

Return: Zero-based axis index.

Index of the named mesh axis, throwing when absent.

function DeviceCoords¶

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

Coordinates of a device within the mesh.

Parameters:

• device_id Linear device index.

Return: Coordinates along each mesh axis.

Row-major coordinates of a linear device id.

function DeviceId¶

int64_t DeviceId(
    const std::vector< int64_t > & coords
) const

Device id located at the given mesh coordinates.

Parameters:

• coords Per-axis coordinates.

Return: Row-major linear device index.

Linear device id for the given mesh coordinates.