This document is relevant for: Inf2, Trn1, Trn2

nki.isa.nc_matmul#

nki.isa.nc_matmul(stationary, moving, *, is_stationary_onezero=False, is_moving_onezero=False, is_transpose=False, mask=None, **kwargs)[source]#

Compute stationary.T @ moving matrix multiplication using Tensor Engine.

The nc_matmul instruction must read inputs from SBUF and write outputs to PSUM. Therefore, the stationary and moving must be SBUF tiles, and the result tile is a PSUM tile.

The nc_matmul instruction currently supports float8_e4m3/float8_e5m2/bfloat16/float16/tfloat32/float32 input data types as listed in Supported Data Types. The matmul accumulation and results are always in float32.

The Tensor Engine imposes special layout constraints on the input tiles. First, the partition axis sizes of the stationary and moving tiles must be identical and <=128, which corresponds to the contraction dimension of the matrix multiplication. Second, the free axis sizes of stationary and moving tiles must be <= 128 and <=512, respectively, For example, stationary.shape = (128, 126); moving.shape = (128, 512) and nc_matmul(stationary,moving) returns a tile of shape = (126, 512). For more information about the matmul layout, see Tensor Engine.

../../../../_images/matmul.png — Fig. 9 MxKxN Matrix Multiplication Visualization.#

If the contraction dimension of the matrix multiplication exceeds 128, you may accumulate multiple nc_matmul instruction output tiles into the same PSUM tile. See example code snippet below.

Estimated instruction cost:

The Tensor Engine has complex performance characteristics given its data flow and pipeline design. The below formula is the average nc_matmul cost assuming many nc_matmul instructions of the same shapes running back-to-back on the engine:

Cost (Tensor Engine Cycles)	Condition
`max(min(64, N_stationary), N_moving)`	input data type is one of `float8_e4m3/float8_e5m2/bfloat16/float16/tfloat32`
`4 * max(min(64, N_stationary), N_moving)`	input data type is `float32`

where,

N_stationary is the number of elements per partition in stationary tile.
N_moving is the number of elements per partition in moving tile.

Parameters:

stationary – the stationary operand on SBUF; layout: (partition axis <= 128, free axis <= 128)
moving – the moving operand on SBUF; layout: (partition axis <= 128, free axis <= 512)
mask – (optional) a compile-time constant predicate that controls whether/how this instruction is executed (see NKI API Masking for details)
is_stationary_onezero – hints to the compiler whether the stationary operand is a tile with ones/zeros only; setting this field explicitly could lead to 2x better performance if stationary tile is in float32; the field has no impact for non-float32 stationary.
is_moving_onezero – hints to the compiler if the moving operand is a tile with ones/zeros only; setting this field explicitly could lead to 2x better performance if moving tile is in float32; the field has no impact for non-float32 moving.
is_transpose – hints to the compiler that this is a transpose operation with moving as an identity matrix.

Returns:

a tile on PSUM that has the result of matrix multiplication of stationary and moving tiles; layout: partition axis comes from free axis of stationary, while free axis comes from free axis of moving.

Example:

import neuronxcc.nki.isa as nisa
import neuronxcc.nki.language as nl

##################################################################
# Example 1:
# multiply matrix a of shape (128, 128) and matrix b of shape (128, 512)
# to get matrix c in PSUM of shape (128, 512)
##################################################################
i_p_a = nl.arange(128)[:, None]
i_f_a = nl.arange(128)[None, :]
i_p_b = nl.arange(128)[:, None]
i_f_b = nl.arange(512)[None, :]
a = nl.load(a_tensor[i_p_a, i_f_a])
b = nl.load(b_tensor[i_p_b, i_f_b])

c_psum = nisa.nc_matmul(a[i_p_a, i_f_a], b[i_p_b, i_f_b])

nl.store(c_tensor[i_p_a, i_f_b], c_psum)

##################################################################
# Example 2:
# multiply matrix d of shape (256, 128) and matrix e of shape (256, 512)
# to get matrix f in PSUM of shape (128, 512) using psum accumulation
##################################################################
f_psum = nl.zeros((128, 512), nl.float32, buffer=nl.psum)

i_p_d = nl.arange(128)[:, None]
i_f_d = nl.arange(128)[None, :]
i_p_e = nl.arange(128)[:, None]
i_f_e = nl.arange(512)[None, :]

for i_contract in nl.affine_range(2):
  d = nl.load(d_tensor[i_contract * 128 + i_p_d, i_f_d])
  e = nl.load(e_tensor[i_contract * 128 + i_p_e, i_f_e])
  f_psum += nisa.nc_matmul(d[i_p_d, i_f_d],
                            e[i_p_e, i_f_e])
  
nl.store(f_tensor[i_p_d, i_f_e], f_psum)
return c_tensor, f_tensor

This document is relevant for: Inf2, Trn1, Trn2

nki.isa.nc_matmul

Contents

nki.isa.nc_matmul#