This document is relevant for: Trn2, Trn3

nki.isa.scalar_tensor_tensor#

nki.isa.scalar_tensor_tensor(dst: NkiTensor, data: NkiTensor, op0, operand0: float | NkiTensor, op1, operand1: NkiTensor, reverse0=False, reverse1=False, reduce_cmd: reduce_cmd = reduce_cmd.reset_reduce, reduce_op=Ellipsis, reduce_res: NkiTensor | None = None, name=None)[source]#

Apply two math operators in sequence using Vector Engine: (data <op0> operand0) <op1> operand1.

This instruction is equivalent to running two operations back-to-back: 1. temp_result = tensor_scalar(data, op0, operand0) - broadcast operand0 and apply op0 2. dst = tensor_tensor(temp_result, op1, operand1) - element-wise operation with operand1

The operand0 can be either a compile-time constant scalar for broadcast across all elements of data or a tile of shape (data.shape[0], 1) for broadcast along the free dimension. The operand1 tile must have the same shape as data for element-wise operation.

The scalar broadcasting in the first operation is performed at no additional performance cost, making this instruction have approximately the same latency as a regular tensor_tensor instruction.

Both op0 and op1 must be arithmetic operators (see Supported Math Operators for NKI ISA for supported operators). Bitvec operators are not supported. When the operators are non-commutative (e.g., subtract), operand ordering can be reversed using reverse0 and reverse1 flags.

Memory types.

The input data tile can be an SBUF or PSUM tile. The operand0 can be an SBUF or PSUM tile or a compile-time constant scalar. The operand1 must be an SBUF or PSUM tile. However, data and operand1 cannot both reside in PSUM. The output dst tile can be written to either SBUF or PSUM.

Data types.

All input tiles can be any supported NKI data type (see Supported Data Types for more information). The Vector Engine automatically casts input data types to float32 and performs all computations in float32 math. The float32 results are cast to the data type of output dst. See nki-type-casting for more information about type casting behavior.

Layout.

The parallel computation dimension of nisa.scalar_tensor_tensor is along the partition dimension.

Tile size.

The partition dimension size of input data, operand1, and output dst tiles must be the same and must not exceed 128. The total number of elements per partition of input data, operand1, and output dst tiles must be the same and must not exceed the physical size of each SBUF partition. If operand0 is not a scalar, the partition dimension size of operand0 must be the same as that of data and the number of elements per partition of operand0 must be 1. The data, operand1, and dst tiles can have at most two free dimensions (maximum rank 3, including the partition dimension).

In addition to the elementwise computation above, this API can optionally perform a reduction along the free dimension(s) of the elementwise result using Vector Engine, 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. reduce_op must be nl.add; bitvec operators do not support accumulation.

\[\begin{split}dst = (data <op0> operand0) <op1> operand1 \\ 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 0, then accumulate the current results.

  • reduce_cmd.reduce: Continue accumulating without resetting.

Note

The accumulator registers are shared across Vector Engine accumulation instructions including exponential(), range_select(), scalar_tensor_tensor(), select_reduce(), tensor_scalar_cumulative(), and tensor_scalar_reduce().

Parameters:

  • dst – the output tile

  • data – the input tile

  • op0 – the first math operator used with operand0 (see Supported Math Operators for NKI ISA for supported operators)

  • 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

  • reverse0 – reverse ordering of inputs to op0; if false, operand0 is the rhs of op0; if true, operand0 is the lhs of op0

  • op1 – the second math operator used with operand1 (see Supported Math Operators for NKI ISA for supported operators)

  • operand1 – a tile with the same size as data for element-wise operation

  • reverse1 – reverse ordering of inputs to op1; if false, operand1 is the rhs of op1; if true, operand1 is the lhs of op1

  • reduce_cmd – (optional) Controls the accumulator behavior, supports: reset_reduce (default) and reduce.

  • reduce_op – (optional) Reduction operator to apply (only nl.add is supported)

  • reduce_res – (optional) (data.shape[0], 1) SBUF/PSUM tile that receives the DVE accumulator value after reduce_cmd is applied, via a fused non-destructive read out (the accumulator is preserved, so a following DVE reduction continues to accumulate into it). Pass None to skip the read out when chaining multiple instructions.

This document is relevant for: Trn2, Trn3