This document is relevant for: Trn2, Trn3

nki.isa.tensor_copy_predicated#

nki.isa.tensor_copy_predicated(dst: NkiTensor, src: float | NkiTensor, predicate: NkiTensor, reverse_pred=False, name=None)[source]#

Conditionally copy elements from the src tile to the dst tile on SBUF / PSUM based on a predicate using Vector Engine.

This instruction provides low-level control over conditional data movement on NeuronCores, optimized for scenarios where only selective copying of elements is needed. Either src or predicate may be in PSUM, but not both simultaneously. Both src and predicate are permitted to be in SBUF. reverse_pred can only be set to true if src is a scalar.

Shape and data type constraints:

  1. src (if it is a tensor), dst, and predicate must occupy the same number of partitions and same number of elements per partition.

  2. predicate must be of type uint8, uint16, or uint32.

  3. src and dst must share the same data type.

  4. dst, predicate, and a tensor src can have at most three free dimensions (maximum rank 4, including the partition dimension).

  5. if src is a scalar, i.e. a Python int or float, only the kind must match dst: a Python float requires a float dst dtype and a Python int requires an integer dst dtype. The bit widths need not match – the scalar is materialized in dst’s dtype, so e.g. a Python int may be copied into an int8.

Behavior:

  • Where predicate is True: The corresponding elements from src are copied to dst tile. If src is a scalar, the scalar is copied to the dst tile.

  • Where predicate is False: The corresponding values in dst tile are unmodified

Tensor indirection.

On NeuronCore-v4 and later, dst and predicate support tensor indirection (gather/scatter) by passing a view created with .indirect(index). src does not support tensor indirection. Runs on the Vector engine.

When operands are manually allocated, their base partitions must satisfy:

  • the index of every .indirect() view starts on a quadrant boundary (a multiple of 32);

  • if predicate uses .indirect(), predicate starts on the same partition as its index;

  • if dst uses .indirect() and predicate is in SBUF, dst’s index starts on the same partition as predicate;

  • if dst uses .indirect() and predicate is in PSUM and uses .indirect(), dst’s index starts on the same partition as predicate’s index.

Parameters:

  • src – The source tile or number to copy elements from when predicate is True

  • dst – The destination tile to copy elements to

  • predicate – A tile that determines which elements to copy

  • reverse_pred – A boolean that reverses the effect of predicate. Only available when src is a scalar.

This document is relevant for: Trn2, Trn3