This document is relevant for: Trn2, Trn3

nki.isa.tensor_scalar_reduce#

nki.isa.tensor_scalar_reduce(dst: NkiTensor, data: NkiTensor, op0, operand0: float | nki.language.NkiTensor, reduce_op, reduce_res: Optional[NkiTensor] = None, reverse0=False, reduce_cmd: reduce_cmd = reduce_cmd.reset_reduce, reduce_init: Optional[Union[float, NkiTensor]] = None, reduce_res_negate=False, name=None)[source]#

Perform the same computation as nisa.tensor_scalar with one math operator and also a reduction along the free dimension of the nisa.tensor_scalar result using Vector Engine.

Refer to nisa.tensor_scalar for semantics of data/op0/operand0. Unlike regular nisa.tensor_scalar where two operators are supported, only one operator is supported in this API. Also, op0 can only be arithmetic operation in Supported Math Operators for NKI ISA. Bitvec operators are not supported in this API.

In addition to nisa.tensor_scalar computation, this API also performs a reduction along the free dimension(s) of the nisa.tensor_scalar result, at a small additional performance cost. The reduction result is written into reduce_res, which must be a SBUF/PSUM tile with the same partition axis size as the input tile data and one element per partition. The reduce_op can be any of nl.add, nl.multiply, nl.max or nl.min.

Reduction axis is not configurable in this API. If the input tile has multiple free axis, the API will reduce across all of them.

\[\begin{split}result = data <op0> operand0 \\ reduce\_res = reduce\_op(dst, axis=<FreeAxis>)\end{split}\]

Accumulator behavior:

The Vector Engine maintains internal accumulator registers that can be controlled via the reduce_cmd parameter:

  • reduce_cmd.reset_reduce: (default) Reset accumulators to the identity value for reduce_op, then accumulate the current results. The identity is 0 for nl.add, 1 for nl.multiply, -inf for nl.max and +inf for nl.min.

  • reduce_cmd.reduce: Continue accumulating without resetting (useful for multi-step reductions across tiles).

  • reduce_cmd.load_reduce: Load the values from reduce_init into the accumulator, then accumulate the current result on top of it.

Note

reduce_init should only be set when reduce_cmd is load_reduce.

Note

The accumulator registers are shared across Vector Engine accumulation instructions including nki.isa.exponential, nki.isa.range_select, nki.isa.select_reduce, and nki.isa.tensor_scalar_cumulative.

Tensor indirection.

On NeuronCore-v4 and later, dst and data support tensor indirection (gather/scatter) by passing a view created with .indirect(index). operand0, reduce_res, and reduce_init do 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 data uses .indirect(), data starts on the same partition as its index;

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

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

Parameters:

  • dst – an output tile of (data <op0> operand0) computation

  • data – the input tile

  • op0 – the math operator used with operand0 (any arithmetic operator in Supported Math Operators for NKI ISA is allowed).

  • operand0 – a scalar constant or a tile of shape (data.shape[0], 1), where data.shape[0] is the partition axis size of the input data tile. Must be None or 0 when op0 is a unary operator (e.g., nl.abs).

  • reverse0 – (not supported yet) reverse ordering of inputs to op0; if false, operand0 is the rhs of op0; if true, operand0 is the lhs of op0. <– currently not supported yet.

  • reduce_op – the reduce operation to perform on the free dimension of data <op0> operand0

  • reduce_res –

    a tile of shape (data.shape[0], 1), where data.shape[0] is the partition axis size of the input data tile. The result of reduce_op(data <op0> operand0) is written into the tile.

    Pass None to keep the reduction result in the Vector Engine’s internal accumulator without writing it out. This is useful when chaining multiple calls that reduce into the same accumulator — only the final call needs to pass a tile to retrieve the accumulated result.

  • reduce_cmd – Control the state of reduction registers for accumulating reduction results. Supported: reset_reduce (default), reduce, load_reduce.

  • reduce_init – Initial value for reduction when using reduce_cmd.load_reduce. Must be provided when reduce_cmd is load_reduce. Supported dtypes: float32.

  • reduce_res_negate – Negate the reduction result before writing it to reduce_res, does not modify the internal reduction registers.

This document is relevant for: Trn2, Trn3