This document is relevant for: Inf2, Trn1, Trn2

nki.isa.scalar_tensor_tensor#

nki.isa.scalar_tensor_tensor(*, data, op0, operand0, op1, operand1, reverse0=False, reverse1=False, dtype=None, mask=None, **kwargs)[source]#

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

data input can be an SBUF or PSUM tile of 2D shape. operand0 can be SBUF or PSUM tile of shape (data.shape[0], 1), i.e., vector, or a compile-time constant scalar. operand1 can be SBUF or PSUM tile of shape (data.shape[0], data.shape[1]) (i.e., has to match data shape), note that operand1 and data can’t both be on PSUM.

Estimated instruction cost:

Cost (Vector Engine Cycles)

Condition

N

data and operand1 are both bfloat16, op0=nl.subtract and op1=nl.multiply, and N is even

2*N

otherwise

where,

  • N is the number of elements per partition in data.

Parameters:

  • 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 of shape with the same partition and free dimension as data input.

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

  • dtype – (optional) data type to cast the output type to (see Supported Data Types for more information); if not specified, it will default to be the same as the data type of the input tile.

  • mask – (optional) a compile-time constant predicate that controls whether/how this instruction is executed (see NKI API Masking for details)

Returns:

an output tile of (data <op0> operand0) <op1> operand1 computation

This document is relevant for: Inf2, Trn1, Trn2