This document is relevant for: Trn2, Trn3
NKI Scalar Values#
Most of NKI moves tiles: rectangles of data in SBUF, PSUM and HBM. Alongside
that, a kernel needs a few plain integers — an index to gather from, a trip
count read at run time, an offset computed on device. An
nki.isa.NkiScalar holds one such value.
This page covers where scalars come from, what arithmetic they support, how their dtype decides both range and available operations, and the two places a kernel puts them to work: dynamic addressing and for-loop bounds.
Note
Scalars require the tracer frontend; the parser frontend does not
support them. Every example below is run with nki.simulate(), and each
one produces the same value on Trainium.
Loading and Storing#
nki.isa.scalar_load() reads a 1x1 tile into a scalar, and
nki.isa.scalar_store() writes one back to a 1x1 tile.
Arithmetic on a scalar is traced, not evaluated. i * 4 + 1 computes
nothing at compile time — it records the operations and yields a new scalar.
import nki
import nki.isa as nisa
import nki.language as nl
import numpy as np
@nki.jit
def scalar_arithmetic(index: nl.NkiTensor):
index_tile = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=index_tile, src=index)
i = nisa.scalar_load(index_tile) # an NkiScalar
i = i * 4 + 1 # traced, not evaluated in Python
result = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.scalar_store(result, i)
out = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.shared_hbm)
nisa.dma_copy(dst=out, src=result)
return out
nki.simulate(scalar_arithmetic)(np.array([[10]], dtype=np.int32)) # [[41]]
The source tile must be nl.int32 or nl.int64. Unsigned dtypes are not
accepted; use the signed dtype of the same width. scalar_store takes a
scalar, not a constant.
Width and Range#
A scalar’s dtype is fixed when it is loaded and never changes. Nothing widens
automatically, so arithmetic that outgrows nl.int32 wraps:
@nki.jit
def scalar_wraps(value: nl.NkiTensor):
tile = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=tile, src=value)
v = nisa.scalar_load(tile)
squared = v * v # 10**10 does not fit 32 bits
result = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.scalar_store(result, squared)
out = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.shared_hbm)
nisa.dma_copy(dst=out, src=result)
return out
nki.simulate(scalar_wraps)(np.array([[100000]], dtype=np.int32)) # [[1410065408]]
Warning
Wrapping is silent, and identical on device and under nki.simulate(). A
scalar keeps the low bits of the true result and nothing else: overflow past
the largest value of its dtype continues from the smallest, and underflow the
other way. 100000 * 100000 is 10**10, whose low 32 bits are
1410065408. A value that needs more than 32 bits has to be loaded as
nl.int64 from the start.
For a wider value, load from an nl.int64 tile. Arithmetic then keeps all 64
bits, and a literal too large for 32 bits is added as written:
@nki.jit
def wide_scalar(value: nl.NkiTensor):
tile = nl.ndarray((1, 1), dtype=nl.int64, buffer=nl.sbuf)
nisa.dma_copy(dst=tile, src=value)
v = nisa.scalar_load(tile) # nl.int64 in, so a 64-bit scalar
v = (v + 12345) * 3
result = nl.ndarray((1, 1), dtype=nl.int64, buffer=nl.sbuf)
nisa.scalar_store(result, v)
out = nl.ndarray((1, 1), dtype=nl.int64, buffer=nl.shared_hbm)
nisa.dma_copy(dst=out, src=result)
return out
nki.simulate(wide_scalar)(np.array([[1 << 40]], dtype=np.int64)) # [[3298534920363]]
The two widths cannot be mixed. Combining an nl.int32 and an nl.int64
scalar in one operation is an error at compile time, and so is storing a scalar
into a tile of the other width. A Python int has no width of its own: it
takes the width of the scalar it is used with, and a literal that does not fit
that width will be rejected rather than truncated.
Supported Operations by Scalar Width#
nl.int64 supports fewer operations than nl.int32:
Operation |
|
|
|---|---|---|
|
Supported |
Supported |
|
Traced, but the result cannot be used |
Rejected at compile time |
Bound of |
Supported, including an arithmetic expression |
Not supported |
|
Supported |
Supported, including offsets beyond 4GB |
|
Supported |
Not supported |
The operation limits come from the hardware, not from a NKI check that could be relaxed: on Trainium 1, 2 and 3 the register ALU implements only addition, subtraction and multiplication at 64 bits, so a comparison has no 64-bit form, and neither does the compare-and-branch that a loop bound compiles into. The restriction on an SBUF or PSUM offset has a different cause, given under Choosing a Width for an Offset.
Every unsupported case above is reported at compile time. A comparison traces, but its result cannot be stored, used in further arithmetic, or used as an offset, and each of those is refused where it is written.
Dynamic Addressing#
An access pattern’s scalar_offset takes a scalar, which is how a kernel
reads a location it cannot know until run time. The offset counts elements of
the dimension named by indirect_dim, so with indirect_dim=0 a scalar
selects a row:
ROWS, COLS = 4, 8
@nki.jit
def gather_at_scalar_offset(data: nl.NkiTensor, index: nl.NkiTensor):
index_tile = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=index_tile, src=index)
row = nisa.scalar_load(index_tile)
tile = nl.ndarray((1, COLS), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(
dst=tile,
src=data.ap(
pattern=[[COLS, 1], [1, COLS]], offset=0,
scalar_offset=row, indirect_dim=0,
),
)
out = nl.ndarray((1, COLS), dtype=nl.float32, buffer=nl.shared_hbm)
nisa.dma_copy(dst=out, src=tile)
return out
data = np.arange(32, dtype=np.float32).reshape(4, 8)
nki.simulate(gather_at_scalar_offset)(data, np.array([[2]], dtype=np.int32))
# [[16. 17. 18. 19. 20. 21. 22. 23.]]
Warning
A scalar_offset is not bounds-checked. An index computed on device that
falls outside the tensor still forms an address, and the DMA reads that
address: the kernel returns another tensor’s data, or uninitialized HBM,
without an error from the compiler or from nki.simulate(). Check the
arithmetic that produces an index, not just the kernel that consumes it.
Choosing a Width for an Offset#
The offset counts elements of the indirect dimension rather than bytes, and is
scaled by that dimension’s stride when the address is formed, so a 32-bit index
reaches well past 4GB. Rows of 512 nl.float32 values are 2KB apart, so
2**31 of them span roughly 4TB. An nl.int64 offset is needed when the
index value itself passes 2**31 — element-indexing a tensor of more than
2**31 elements, for instance.
An offset into SBUF or PSUM must be nl.int32. A 64-bit one is rejected:
SBUF and PSUM today are too small for an offset that needs 64 bits to encode.
Loop Bounds and the Loop Variable#
nki.language.fori_loop() accepts a scalar as a bound, which is how a
kernel loops a number of times it reads at run time. The loop variable handed
to the body is itself a scalar, usable in arithmetic and as an offset:
@nki.jit
def sum_rows_to_bound(data: nl.NkiTensor, count: nl.NkiTensor):
count_tile = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=count_tile, src=count)
trip_count = nisa.scalar_load(count_tile)
total = nl.zeros((1, COLS), dtype=nl.float32, buffer=nl.sbuf)
def body(i): # i is an NkiScalar
row = nl.ndarray((1, COLS), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(
dst=row,
src=data.ap(
pattern=[[COLS, 1], [1, COLS]], offset=0,
scalar_offset=i, indirect_dim=0,
),
)
nisa.tensor_tensor(dst=total, data1=total, data2=row, op=nl.add)
nl.fori_loop(0, trip_count, body)
out = nl.ndarray((1, COLS), dtype=nl.float32, buffer=nl.shared_hbm)
nisa.dma_copy(dst=out, src=total)
return out
nki.simulate(sum_rows_to_bound)(data, np.array([[3]], dtype=np.int32))
# [[24. 27. 30. 33. 36. 39. 42. 45.]]
The loop variable is always 32-bit. A 64-bit scalar can be loaded inside the body and used as an offset there, but it cannot be combined with the loop variable: mixing widths is refused, so a wide offset stays the same on every iteration unless the loop carries it.
nki.language.while_loop() does not take a scalar as its condition.
This document is relevant for: Trn2, Trn3