This document is relevant for: Trn2, Trn3

nki.isa.NkiScalar#

class nki.isa.NkiScalar(state, dtype=None, is_comparison=False)[source]#

NkiScalar is the core scalar type in NKI. It holds one integer, computed on device, and can be used as an operand of arithmetic, as the scalar_offset of an access pattern, and as a bound of nki.language.fori_loop().

NkiScalars are instantiated from values in a NkiTensor where dtype is either nl.int32 or nl.int64, via nki.isa.scalar_load().

A NkiScalar can be used almost anywhere a Python int literal is used, and the other operand of an arithmetic operator may be a Python int or another NkiScalar so long as the dtypes match. Every operator returns a new NkiScalar.

Width is fixed at creation. A 32-bit scalar is never widened, so arithmetic that outgrows 32 bits wraps rather than promoting: 100000 * 100000 yields 1410065408, and adding 1 to the largest nl.int32 yields the smallest. Load from an nl.int64 tile when a value needs more range.

Not every operation exists at both widths. At 64 bits only addition, subtraction and multiplication are available:

Operation

nl.int32

nl.int64

+, - (binary and unary), *

Supported

Supported

<, <=, >, >=, ==, !=

Traced, but the result cannot be used

Rejected at compile time

Bound of nki.language.fori_loop()

Supported, including an arithmetic expression

Not supported; the trip-count comparison has no 64-bit form

scalar_offset on an HBM access pattern

Supported

Supported, including offsets beyond 4GB

scalar_offset on an SBUF or PSUM access pattern

Supported

Not supported; an on-chip offset must be 32 bits

A scalar has no Python value. bool(s), int(s), float(s), using a scalar as a Python index (range(s), some_list[s]), and using one in an if, while, or assert condition are all rejected, under the CPU simulator as well as the tracer, because the value arrives only when the kernel runs. Indexing a tensor with a scalar works through nki.language.ds(): t[nl.ds(s, 8)] and select() both take one, as does scalar_offset. A slice bound cannot be a scalar, because a slice needs both bounds while the kernel is built. nki.language.fori_loop() takes a scalar bound directly.

Choosing a width for addressing. scalar_offset counts elements of the access pattern’s indirect dimension, not bytes, and neuronx-cc scales that count by the dimension’s stride when it forms the address, 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 – when element-indexing a tensor of more than 2**31 elements, for instance.

Create scalars via nki.isa.scalar_load(). See NKI Scalar Values for a walkthrough with runnable examples.

Methods

Attributes

dtype

The NKI dtype name, e.g.
property dtype#

The NKI dtype name, e.g. "int32".

This document is relevant for: Trn2, Trn3