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

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.