This document is relevant for: Trn2, Trn3

nki.isa.scalar_load#

nki.isa.scalar_load(src: NkiTensor) → NkiScalar[source]#

Load a scalar value from a 1x1 tensor.

Returns an NkiScalar, whose dtype matches src and fixes the width of every value derived from it. See that class for which operations each width supports, and NKI Scalar Values for a walkthrough.

The source tensor must use nl.int32 or nl.int64 dtype. Unsigned dtypes are not accepted; use the signed dtype of the same width.

Only the tracer frontend supports this. The parser frontend rejects the call.

Unlike the compute instructions, scalar_load accepts nl.int64, and the value is not truncated. An int64 scalar can be written back with nki.isa.scalar_store(), and can be passed as scalar_offset to address an HBM tensor beyond the 4GB reach of a 32-bit byte offset. It cannot bound an nki.language.fori_loop(), because the loop’s trip-count comparison has no 64-bit form.

A scalar of either dtype can be used as a scalar_offset argument to access-pattern methods for dynamic addressing, with one restriction: on an SBUF or PSUM operand the offset must be nl.int32. On-chip buffers are megabytes, so a 64-bit offset there is rejected rather than silently narrowed.

Parameters:

src – A 1x1 tensor in SBUF or HBM with dtype nl.int32 or nl.int64.

Returns:

The loaded scalar, with dtype matching src.

Return type:

NkiScalar

Example:

buf = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=buf, src=src)
idx = nisa.scalar_load(buf)
idx = idx + 1  # arithmetic on the scalar

This document is relevant for: Trn2, Trn3