This document is relevant for: Inf2, Trn1, Trn2

nki.isa.dma_copy#

nki.isa.dma_copy(*, dst, src)[source]#

Copy data from src to dst using DMA engine. Both src and dst tiles can be in device memory (HBM) or SBUF. However, if both src and dst tiles are in SBUF, consider using nisa.tensor_copy instead for better performance.

Parameters:

  • src – the source of copy.

  • dst – the dst of copy

A cast will happen if the src and dst have different dtype.

Example:

import neuronxcc.nki.isa as nisa
...


############################################################################
# Example 1: Copy over the tensor to another tensor
############################################################################
nisa.dma_copy(dst=b, src=a)

This document is relevant for: Inf2, Trn1, Trn2