This document is relevant for: Trn2, Trn3

nki.isa.sendrecv_compute#

nki.isa.sendrecv_compute(src: NkiTensor, dst: NkiTensor, send_to_rank, recv_from_rank, reduce_op, name=None)[source]#

Perform point-to-point communication between NeuronCores and reduce the local and received tiles into a single result, using the DMA engine’s reduction capability.

Note

Available only when running with LNC (Logical NeuronCore) degree greater than 1.

This instruction combines a bidirectional data exchange (like nisa.sendrecv) with a reduction. The current NeuronCore sends its src tile to the target NeuronCore specified by send_to_rank, simultaneously receives the peer tile from recv_from_rank, and writes reduce_op(src, peer) to dst.

The use case is fusing an inter-core exchange with an element-wise reduction, such as cross-core all-reduce-style accumulation, without a separate nisa.tensor_tensor step after nisa.sendrecv.

The communication is between the physical NeuronCores that make up the same Logical NeuronCore (LNC), so this instruction requires an LNC degree greater than 1, and send_to_rank / recv_from_rank must be a valid rank within the LNC (0 or 1 for LNC degree 2).

The reduce_op parameter selects the reduction applied to the local src and the peer tile. Supported operations are nl.add, nl.maximum, and nl.minimum.

Constraints.

  • src and dst tiles must be in the same memory space: both in SBUF or both in private HBM (mixing SBUF and HBM, or using shared HBM, is not supported).

  • src and dst must have the same data type.

  • src and dst must have the same shape and layout.

  • src and dst must have the same partition dimension size and the same number of elements per partition.

Parameters:

  • src – the source tile on the current NeuronCore to be sent and reduced

  • dst – the destination tile on the current NeuronCore where the reduced result is stored

  • send_to_rank – rank ID of the target NeuronCore to send data to

  • recv_from_rank – rank ID of the source NeuronCore to receive data from

  • reduce_op – the reduction operation to apply: nl.add, nl.maximum, or nl.minimum

Example:

# All-reduce-style accumulation across two cores in an LNC
current_rank = nl.program_id()
peer_rank = 1 - current_rank

send_data = nl.ndarray((batch_size, hidden_dim), dtype=nl.float32, buffer=nl.sbuf)
result = nl.ndarray((batch_size, hidden_dim), dtype=nl.float32, buffer=nl.sbuf)

# result = send_data + peer's send_data
sendrecv_compute(
    src=send_data,
    dst=result,
    send_to_rank=peer_rank,
    recv_from_rank=peer_rank,
    reduce_op=nl.add,
)

This document is relevant for: Trn2, Trn3