This document is relevant for: Trn2, Trn3

nki.isa.tensor_copy_predicated#

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

Conditionally copy elements from the src tile to the destination 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.

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.

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.

This document is relevant for: Trn2, Trn3