nki.isa.tensor_copy — AWS Neuron Documentation (original) (raw)

This document is relevant for: Inf2, Trn1, Trn2

nki.isa.tensor_copy#

nki.isa.tensor_copy(src, *, mask=None, dtype=None, engine=engine.unknown, **kwargs)[source]#

Create a copy of src tile within NeuronCore on-chip SRAMs using Vector, Scalar or GpSimd Engine.

The output tile has the same partition axis size and also the same number of elements per partition as the input tile src.

All three compute engines, Vector, Scalar and GpSimd Engine can perform tensor copy. However, their copy behavior is slightly different across engines:

• Scalar Engine on NeuronCore-v2 performs copy by first casting the input tile to FP32 internally and then casting from FP32 to the output dtype (dtype, or src.dtype if dtype is not specified). Therefore, users should be cautious with assigning this instruction to Scalar Engine when the input data type cannot be precisely cast to FP32 (e.g., INT32).
• Both GpSimd and Vector Engine can operate in two modes: (1) bit-accurate copy when input and output data types are the same or (2) intermediate FP32 cast when input and output data types differ, similar to Scalar Engine.

In addition, since GpSimd Engine cannot access PSUM in NeuronCore, Scalar or Vector Engine must be chosen when the input or output tile is in PSUM (see NeuronCore-v2 Compute Engines for details). By default, this API returns a tile in SBUF, unless the returned value is assigned to a pre-declared PSUM tile.

Estimated instruction cost:

max(MIN_II, N) engine cycles, where N is the number of elements per partition in the input tile, and MIN_II is the minimum instruction initiation interval for small input tiles.MIN_II is roughly 64 engine cycles.

Parameters:

• src – the source of copy, must be a tile in SBUF or PSUM.
• mask – (optional) a compile-time constant predicate that controls whether/how this instruction is executed (see NKI API Masking for details)
• dtype – (optional) data type to cast the output type to (see Supported Data Types for more information); if not specified, it will default to be the same as the data type of the input tile.
• engine – (optional) the engine to use for the operation: nki.isa.vector_engine, nki.isa.scalar_engine,nki.isa.gpsimd_engine or nki.isa.unknown_engine (default, compiler selects best engine based on engine workload).

Returns:

a tile with the same content and partition axis size as the src tile.

Example:

import neuronxcc.nki.isa as nisa import neuronxcc.nki.language as nl ...

############################################################################

Example 1: Copy over the tensor to another tensor using the Vector engine.

############################################################################ x = nl.load(in_tensor) x_copy = nisa.tensor_copy(x, engine=nisa.vector_engine) nl.store(out_tensor, value=x_copy)

This document is relevant for: Inf2, Trn1, Trn2