Skip to content

tenzr::execution

Classes

Name
class tenzr::execution::Communicator
Process-group communication handle exposing its rank, its size, and a barrier.
class tenzr::execution::DeviceMesh
Logical grid of devices with optional named dimensions, mapping between device ids and mesh coordinates.
class tenzr::execution::ExecutionStream
Serialized work queue backed by a worker thread, used to submit and synchronize device work.
class tenzr::execution::StreamEvent
Handle to a recorded point in a stream's work, used to query completion or to make later work wait.
struct tenzr::execution::StreamEventState
Shared readiness state for a stream event, holding the ready flag, the PJRT event, and completion callbacks.

Types

Name
enum class ReductionOp

Functions

Name
Tensor AllReduce(const Tensor & input, ReductionOp op =ReductionOp::kSum, const Communicator * comm =nullptr)
std::vector< Tensor > AllReduce(const std::vector< Tensor > & replica_tensors, ReductionOp op =ReductionOp::kSum)
Tensor AllGather(const Tensor & input, int64_t axis =0, const Communicator * comm =nullptr)
Tensor AllGather(const std::vector< Tensor > & replica_tensors, int64_t axis =0)
std::vector< Tensor > ReduceScatter(const std::vector< Tensor > & replica_tensors, ReductionOp op =ReductionOp::kSum, int64_t axis =0)
std::vector< Tensor > Broadcast(const Tensor & root_tensor, int num_replicas, int root_rank =0)

Types Documentation

enum ReductionOp

Enumerator Value Description
kSum
kProduct
kMin
kMax

Functions Documentation

function AllReduce

Tensor AllReduce(
    const Tensor & input,
    ReductionOp op =ReductionOp::kSum,
    const Communicator * comm =nullptr
)

Parameters:

  • input Tensor to reduce.
  • op Reduction operation.
  • comm Communicator to use, or the world communicator when null.
  • input Tensor to reduce.
  • comm Communicator to reduce over; may be null.

Return:

  • Reduced tensor.
  • The reduced tensor.

Reduce a tensor across processes. Reduce input across the communicator (currently a clone for single-rank or absent communicators).

function AllReduce

std::vector< Tensor > AllReduce(
    const std::vector< Tensor > & replica_tensors,
    ReductionOp op =ReductionOp::kSum
)

Parameters:

  • replica_tensors Per-replica input tensors.
  • op Reduction operation.
  • replica_tensors Per-replica inputs, all with matching shape and dtype.
  • op Reduction operation.

Return:

  • Reduced tensors.
  • One reduced tensor per replica.

Reduce one tensor per replica into a combined result. Reduce a list of per-replica tensors elementwise and broadcast the result to each replica.

function AllGather

Tensor AllGather(
    const Tensor & input,
    int64_t axis =0,
    const Communicator * comm =nullptr
)

Parameters:

  • input Local tensor to contribute.
  • axis Axis to concatenate along.
  • comm Communicator to use, or the world communicator when null.
  • input Tensor to gather.
  • comm Communicator to gather over; may be null.

Return:

  • Gathered tensor.
  • The gathered tensor.

Gather tensors from all processes along an axis. Gather input across the communicator (currently a clone for single-rank or absent communicators).

function AllGather

Tensor AllGather(
    const std::vector< Tensor > & replica_tensors,
    int64_t axis =0
)

Parameters:

  • replica_tensors Per-replica input tensors.
  • axis Axis to concatenate along.
  • replica_tensors Per-replica inputs, matching on all non-gathered axes.
  • axis Concatenation axis, supporting negative indices.

Return:

  • Gathered tensor.
  • The gathered tensor.

Gather per-replica tensors along an axis. Concatenate per-replica tensors along axis.

function ReduceScatter

std::vector< Tensor > ReduceScatter(
    const std::vector< Tensor > & replica_tensors,
    ReductionOp op =ReductionOp::kSum,
    int64_t axis =0
)

Parameters:

  • replica_tensors Per-replica input tensors.
  • op Reduction operation.
  • axis Axis to scatter along.
  • replica_tensors Per-replica inputs.
  • op Reduction operation.
  • axis Axis divided between replicas.

Return:

  • One reduced shard per replica.
  • One chunk per replica, in replica order.

Reduce per-replica tensors, then split the result into per-replica shards. Reduce per-replica tensors then slice the result evenly across replicas.

function Broadcast

std::vector< Tensor > Broadcast(
    const Tensor & root_tensor,
    int num_replicas,
    int root_rank =0
)

Parameters:

  • root_tensor Tensor to broadcast.
  • num_replicas Number of replicas to receive it.
  • root_rank Rank holding the source tensor.
  • root_tensor Tensor to broadcast.
  • num_replicas Number of replicas, at least 1.

Return:

  • One tensor per replica.
  • One independent copy per replica.

Broadcast a tensor from the root rank to every replica. Copy a root tensor to every replica.