# AUTO GENERATED FILE - DO NOT MANUALLY MODIFY.
from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
from typing import *
if TYPE_CHECKING:
from nki.isa import NkiScalar
from nki.isa import VirtualRegister
DType = str
bool_: DType = ...
r"""Boolean (True or False) stored as a byte"""
int8: DType = ...
r"""8-bit signed integer number"""
int16: DType = ...
r"""16-bit signed integer number"""
int32: DType = ...
r"""32-bit signed integer number"""
int64: DType = ...
r"""64-bit signed integer number"""
uint8: DType = ...
r"""8-bit unsigned integer number"""
uint16: DType = ...
r"""16-bit unsigned integer number"""
uint32: DType = ...
r"""32-bit unsigned integer number"""
uint64: DType = ...
r"""64-bit unsigned integer number"""
float16: DType = ...
r"""16-bit floating-point number"""
float32: DType = ...
r"""32-bit floating-point number"""
bfloat16: DType = ...
r"""16-bit floating-point number (1S,8E,7M)"""
tfloat32: DType = ...
r"""32-bit floating-point number (1S,8E,10M)"""
float8_e4m3: DType = ...
r"""8-bit floating-point number (1S,4E,3M)"""
float8_e4m3fn: DType = ...
r"""8-bit floating-point number (1S,4E,3M), Extended range: no inf, NaN represented by 0bS111'1111"""
float8_e5m2: DType = ...
r"""8-bit floating-point number (1S,5E,2M)"""
float8_e5m2_x4: DType = ...
r"""4x packed float8_e5m2 elements, custom data type for nki.isa.nc_matmul_mx on NeuronCore-v4"""
float8_e4m3fn_x4: DType = ...
r"""4x packed float8_e4m3fn elements, custom data type for nki.isa.nc_matmul_mx on NeuronCore-v4"""
float8_e8m0fnu: DType = ...
r"""8-bit floating-point exponent type (0S,8E,0M) - unsigned, NaN represented by 0xFF. Used as power-of-two shared scale factor in MX quantization"""
float4_e2m1fn_x4: DType = ...
r"""4x packed float4_e2m1fn elements, custom data type for nki.isa.nc_matmul_mx on NeuronCore-v4"""
sbuf: MemoryRegion = ...
r"""State Buffer - Only visible to each individual kernel instance in the SPMD grid"""
psum: MemoryRegion = ...
r"""PSUM - Only visible to each individual kernel instance in the SPMD grid"""
hbm: MemoryRegion = ...
r"""HBM - Alias of private_hbm"""
private_hbm: MemoryRegion = ...
r"""HBM - Only visible to each individual kernel instance in the SPMD grid"""
shared_hbm: MemoryRegion = ...
r"""Shared HBM - Visible to all kernel instances in the SPMD grid"""
tile_size: _TileSize = ...
r"""Hardware tile size constants (pmax, psum_bank_fmax, gemm_stationary_fmax, etc.)"""
prelu: Op = ...
r"""Parametric ReLU activation function. Used as the ``op`` parameter in activation ISA instructions such as :doc:`nki.isa.activation </nki/api/generated/nki.isa.activation>`. The slope for negative inputs is controlled by the ``relu_param`` argument (see :ref:`nki-act-func`)."""
bypass: Op = ...
r"""No-op operator that passes data through unchanged. Used as the ``op0`` or ``op1`` parameter in tensor-scalar ISA instructions (e.g., :doc:`nki.isa.activation </nki/api/generated/nki.isa.activation>`) to skip a computation stage."""
class MemoryRegion(Enum):
r"""Memory region constants for NKI tensors."""
sbuf = 'sbuf'
psum = 'psum'
private_hbm = 'private_hbm'
shared_hbm = 'shared_hbm'
def __eq__(self, other):
...
...
class NKIObject():
r"""Base class for NKI kernel dataclasses and configuration objects."""
def __init__(self, **kwargs: Any) -> None:
...
...
class Op(NKIObject):
r"""A lightweight wrapper for NKI operator constants.
Use for type annotations for NKI operators."""
def __init__(self, name):
...
def __call__(self, x, y=None, dtype=None):
...
...
class NkiGrid(NKIObject):
r"""Base class for tensor views organized into tiles or blocks.
``NkiGrid`` is the common interface for :class:`NkiTileGrid` and
:class:`NkiBlockGrid`. A grid describes an existing tensor region through
logical tile or block coordinates while retaining its dtype, buffer, and
underlying storage.
Users do not construct ``NkiGrid`` directly. Grid views are created via
:func:`nki.language.tiles` and :func:`nki.language.blocks`:
.. code-block:: python
tile_grid = nl.tiles(src, tile_shape=(128, 256))
block_grid = nl.blocks(
src,
tile_shape=(128, 256),
block_shape=(2, 2),
)
:func:`nki.language.tiles` returns an :class:`NkiTileGrid`, and
:func:`nki.language.blocks` returns an :class:`NkiBlockGrid`. Both concrete
objects are instances of ``NkiGrid``. Refer to those factory functions for
the tiling, block-grouping, and boundary behavior of the returned view.
See Also:
:func:`nki.language.tiles`: Create a tile-coordinate view.
:func:`nki.language.blocks`: Group tiles into block coordinates."""
@property
def shape(self) -> tuple[int, ...]:
r"""Return the extents of the grid coordinates available for indexing.
Tile-grid dimensions count tiles and block-grid dimensions count
blocks. Untiled leading HBM dimensions appear first. Counts include
partial boundary cells. Integer indexing consumes a coordinate level;
slicing preserves it with the selected count.
:return: One positive count per currently indexable grid dimension."""
...
@property
def tile_shape(self) -> tuple[int, ...]:
r"""Return the nominal source-element shape of one tile.
:return: The configured per-tile shape in source-element units."""
...
@property
def element_shape(self) -> tuple[int, ...]:
r"""Return the number of selected elements along each source dimension.
:return: The selected element counts in source-dimension order."""
...
@property
def dtype(self) -> DType:
r"""Return the element dtype of the backing tensor.
:return: The grid's NKI element dtype."""
...
@property
def buffer(self) -> MemoryRegion:
r"""Return the memory region containing the backing tensor.
:return: The :class:`MemoryRegion` backing the view."""
...
@property
def is_remainder(self) -> bool:
r"""Return whether this selection contains a partial boundary tile.
:return: ``True`` when at least one selected tile extent is partial."""
...
@property
def is_data_representable(self) -> bool:
r"""Return whether :attr:`data` can represent this selection exactly.
:return: ``True`` when one native tensor operand can describe the
complete selection."""
...
@property
def index_step_elements(self) -> tuple[int, ...]:
r"""Return each grid-coordinate step in source-element units.
:return: The source-element stride of each current grid coordinate."""
...
@property
def data(self) -> NkiTensor:
r"""Expose the complete grid selection as one native tensor operand.
Grid navigation retains tile or block coordinates. ``data`` is the
explicit transition to the :class:`NkiTensor` accepted by NKI
instructions. The returned tensor preserves the selection's offsets,
strides, and physical gaps.
:return: An exact native tensor operand over the current grid region.
:raises ValueError: if one native access pattern cannot represent the
selected region, for example across an unsupported nonuniform
partition remainder or an overly complex access pattern. Split the
grid into representable regions in that case. Use
:attr:`is_data_representable` to test this before accessing
``data``.
.. note::
``data.shape`` describes the native access-pattern traversal. It may
differ from :attr:`element_shape` and from the private backing
allocation shape. For example, a complete two-dimensional tile grid can
expose ``(P, tile_rows, tile_columns, F)``, while a contiguous
free-dimension tail can compact to ``(P, valid_free_elements)``.
Capacity reserved behind an SBUF or PSUM remainder is never exposed.
Examples:
.. code-block:: python
import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def copy_grid_as_one_operand(src):
# src is tile-aligned 2D, for example (256, 512).
src_tiles = nl.tiles(src, tile_shape=(128, 128))
staging_tiles = nl.tiles(
nl.ndarray(
src_tiles.element_shape,
dtype=src.dtype,
buffer=nl.sbuf,
),
tile_shape=src_tiles.tile_shape,
)
dst = nl.ndarray(
src.shape, dtype=src.dtype, buffer=nl.shared_hbm
)
dst_tiles = nl.tiles(dst, tile_shape=(128, 128))
nisa.dma_copy(dst=staging_tiles.data, src=src_tiles.data)
nisa.dma_copy(dst=dst_tiles.data, src=staging_tiles.data)
return dst"""
...
def __getitem__(self, key: int | slice | tuple[int | slice, ...]) -> Self:
r"""Select tile or block coordinates while preserving the grid type.
Supported key forms:
* An ``int`` selects and consumes one grid coordinate. Negative
coordinates are resolved relative to the current grid extent.
* A ``slice`` selects a static positive-step coordinate range and
preserves that navigation level with its selected count.
* A ``tuple[GridIndex, ...]`` applies one key to each successive grid
dimension. Omitted trailing dimensions remain available for later
indexing.
:param key: One supported per-dimension key or a tuple of such keys.
:return: The same concrete grid type as ``self``, sharing the source
storage.
:raises TypeError: if a key or slice boundary has an unsupported type.
:raises IndexError: if an integer is outside the selected grid extent,
too many keys are supplied, or the grid has no remaining
coordinates at its current level.
:raises ValueError: if a slice is empty, has a nonpositive step, or
uses a non-static boundary."""
...
...
class NkiBlockGrid(NkiGrid):
r"""A tensor view navigated in block coordinates.
A block groups a fixed number of tiles along every tiled dimension.
:attr:`NkiGrid.shape` counts blocks, :attr:`block_shape` gives tiles per
block, and :attr:`NkiGrid.tile_shape` gives source elements per tile. This
separates a coarser transfer or reuse grain from the tile grain consumed by
compute instructions.
Partial final blocks and tiles are retained and clamped to the exact source
extent. Creating and indexing a block view allocates no storage and moves
no data.
Examples:
.. code-block:: python
import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def copy_block_grid(src):
tile_shape = (128, 256)
block_shape = (2, 2)
src_blocks = nl.blocks(
src, tile_shape=tile_shape, block_shape=block_shape
)
dst = nl.ndarray(
src.shape, dtype=src.dtype, buffer=nl.shared_hbm
)
dst_blocks = nl.blocks(
dst, tile_shape=tile_shape, block_shape=block_shape
)
for i in range(src_blocks.shape[0]):
for j in range(src_blocks.shape[1]):
src_block = src_blocks[i, j]
block = nl.blocks(
nl.ndarray(
src_block.element_shape,
dtype=src.dtype,
buffer=nl.sbuf,
),
tile_shape=src_block.tile_shape,
block_shape=src_block.block_shape,
)
nisa.dma_copy(dst=block.data, src=src_block.data)
nisa.dma_copy(
dst=dst_blocks[i, j].data,
src=block.data,
)
return dst
See Also:
:func:`nki.language.blocks`: Create or regroup a block grid."""
@property
def block_shape(self) -> tuple[int, ...]:
r"""Return the nominal number of tiles grouped along each block axis.
A boundary block may contain fewer tiles. Use
:func:`nki.language.tiles` on a selected block to inspect its exact
interior tile counts.
:return: The configured tile count per block dimension."""
...
...
[docs]class NkiTensor(NKIObject):
r"""NKI tensor with shape-based view operations.
``NkiTensor`` is the core tensor type in NKI. It represents a multi-dimensional
array allocated on a specific memory buffer (SBUF, PSUM, or HBM) with a dtype
and shape. All view operations (``slice``, ``permute``, ``reshape``, etc.) return
new ``NkiTensor`` objects that share the same underlying storage — no data is
copied. Views are consumed by NKI ISA instructions
(e.g., ``nisa.tensor_copy``, ``nisa.dma_copy``).
Tensors are created via :func:`nl.ndarray`:
.. code-block:: python
sb = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
hbm = nl.ndarray((4, 128, 64), dtype=nl.float32, buffer=nl.shared_hbm)
**Partition dimension.**
On-chip tensors (SBUF, PSUM) have a partition dimension at dim 0 that maps to
the hardware's parallel partitions. Most view operations cannot modify this
dimension — see individual method docs for constraints."""
shape: tuple[int, ...] = ...
r"""Tuple giving the count of each dimension of the view. Each element is
an integer, or a :class:`~nki.isa.VirtualRegister` for dynamic-access
views constructed via :meth:`ap`."""
strides: tuple[int, ...] = ...
r"""Tuple with stride per dimension. Each stride is the element step between
consecutive indices along that dimension in the underlying storage.
Each entry is an integer, or a :class:`~nki.isa.VirtualRegister` for
dynamic-access views constructed via :meth:`ap`.
A stride of ``0`` means the dimension is broadcast."""
offset: int = ...
r"""Element offset (not byte offset) from the start of the underlying
storage at which this view begins."""
name: str = ...
r"""Optional debug label, used in compiler diagnostics and the profiler."""
dtype: DType = ...
r"""Element type, e.g. ``nl.float32`` or ``nl.bfloat16``."""
buffer: MemoryRegion | str = ...
r"""The memory region the storage lives in. One of ``nl.sbuf``,
``nl.psum``, ``nl.hbm``, ``nl.private_hbm``, or ``nl.shared_hbm``."""
ndim: int = ...
r"""Number of dimensions (``len(shape)``)."""
size: int = ...
r"""Total number of elements (``prod(shape)``). ``None`` when any element
of ``shape`` is a :class:`~nki.isa.VirtualRegister`."""
def __init__(self, shape: tuple[int, ...], dtype: DType, storage: Any, buffer: MemoryRegion | str=sbuf, name: str='', address=None):
r"""Create an ``NkiTensor`` bound to an existing storage handle.
Most code should use :func:`nl.ndarray` instead — it allocates fresh
storage and wraps it in a tensor. This constructor is useful when you
need to re-bind a storage handle you already have (e.g. from a
framework integration or test fixture).
:param shape: tuple of positive integers. Register-bound (dynamic) dims
are not accepted here — construct dynamic-access views via :meth:`ap`.
:param dtype: NKI dtype string (e.g. ``nl.float32``).
:param storage: backend-specific storage handle. May be ``None`` when
the tensor is used as a signature placeholder and never materialized.
:param buffer: memory region (``nl.sbuf``, ``nl.psum``, ``nl.hbm``,
``nl.shared_hbm``). Defaults to ``nl.sbuf``.
:param name: optional debug name, used in compiler diagnostics and
:ref:`scheduling <how-to-scheduling-apis>`.
:param address: fixed address the storage is pinned to, or ``None`` to
let the compiler choose one."""
...
[docs] def is_contiguous(self) -> bool:
r"""Return True if the view covers storage contiguously (row-major order).
Computed from the current strides: each non-size-1 dimension's stride
must equal the product of the shape sizes of inner dimensions.
For on-chip tensors (SBUF, PSUM), the partition dim (dim 0) is skipped:
partitions are physically independent memory banks, so the partition
stride does not represent physical contiguity. Contiguity is evaluated
per-partition over the free dims only."""
...
[docs] def get_pattern(self) -> list[list[int | VirtualRegister]]:
r"""Return the view's access pattern as ``[[stride, count], ...]``.
Useful as a starting point when constructing a new ``.ap()`` that
keeps most of the current layout intact. The returned pattern
pairs each of the view's dimensions with its current stride, in
the same order as :attr:`shape` / :attr:`strides`."""
...
[docs] def is_indirect(self) -> bool:
r"""Return True if this view already uses indirect addressing.
Indirect addressing is produced by dynamic :meth:`select`,
:meth:`vector_select`, or :meth:`ap` with ``scalar_offset`` /
``vector_offset``. Indirect views cannot be re-indirected, and the
dimension that participates in the indirection cannot be further
sliced or selected — use this query to guard against those chains."""
...
def __getitem__(self, key) -> NkiTensor:
r"""Index the tensor using NumPy-style slicing and selection.
Returns a new view sharing the same storage. Supports integers (select),
slices, ``Ellipsis``, negative indices, and :func:`nl.ds` (DynamicSlice).
.. code-block:: python
sb = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
a = sb[:, 0:32] # slice dim 1
b = sb[:, 0:32:2] # slice with step
c = sb[..., 0] # select last dim
d = sb[nl.ds(i, 8), :] # dynamic slice on partition dim
**Runtime indices.** An index may also be a runtime value (see
:meth:`select`), either on its own — selecting one element and removing
the dim — or as the ``start`` of :func:`nl.ds`, which keeps the dim and
slides a fixed-size window:
.. code-block:: python
idx = nl.ndarray((1, 1), dtype=nl.uint32, buffer=nl.sbuf)
t = nl.ndarray((128, 8, 64), dtype=nl.float32, buffer=nl.sbuf)
t[:, idx, :] # shape (128, 64) — dim 1 removed
t[:, nl.ds(idx, 2), :] # shape (128, 2, 64) — window of 2
See :meth:`select` for the constraints on runtime indices.
**On-chip tensors** (SBUF, PSUM): integer indexing on dim 0 produces a
size-1 slice (partition dim is preserved, not removed). Result is always
at least 2D.
**HBM tensors**: integer indexing removes the dimension. Result is at
least 1D.
:param key: index expression (int, slice, Ellipsis, tuple, DynamicSlice,
or a runtime ``NkiTensor`` / ``VirtualRegister`` / ``NkiScalar``)
:return: new ``NkiTensor`` view"""
...
[docs] def permute(self, dims: tuple[int, ...]) -> NkiTensor:
r"""Reorder tensor dimensions.
Returns a new view with dimensions rearranged according to ``dims``.
No data is copied.
.. code-block:: python
t = nl.ndarray((128, 4, 8), dtype=nl.float32, buffer=nl.sbuf)
t.permute((0, 2, 1)) # shape becomes (128, 8, 4)
**Constraints.**
- ``dims`` must be a permutation of ``range(t.ndim)``
- On-chip tensors: ``dims[0]`` must be 0 (partition dim stays outermost)
- After ``vector_select``: ``dims[0]`` must be 0 (indirect partition dim
stays outermost; see :meth:`is_indirect`)
:param dims: tuple of dimension indices in the desired order
:return: new ``NkiTensor`` view with reordered dimensions"""
...
[docs] def broadcast(self, dim: int, size: int) -> NkiTensor:
r"""Expand a size-1 dimension to ``size`` by repeating elements.
The dimension must have size 1 before broadcasting. No data is copied.
.. code-block:: python
t = nl.ndarray((128, 1, 64), dtype=nl.float32, buffer=nl.sbuf)
t.broadcast(1, 8) # shape becomes (128, 8, 64)
**Constraints.**
- ``t.shape[dim]`` must be 1
- On-chip tensors: ``dim`` must not be 0 (partition dim)
:param dim: dimension to broadcast
:param size: new size for the dimension
:return: new ``NkiTensor`` view with the broadcasted dimension"""
...
[docs] def expand_dim(self, dim: int) -> NkiTensor:
r"""Insert a new dimension of size 1 at position ``dim``.
.. code-block:: python
t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t.expand_dim(1) # shape becomes (128, 1, 64)
**Constraints.**
- On-chip tensors: ``dim`` must not be 0
- After ``vector_select``: ``dim`` must not be 0 (cannot insert before
the indirect partition dim; see :meth:`is_indirect`)
:param dim: position at which to insert the new dimension
:return: new ``NkiTensor`` view with an additional size-1 dimension"""
...
[docs] def squeeze_dim(self, dim: int) -> NkiTensor:
r"""Remove a dimension of size 1.
.. code-block:: python
t = nl.ndarray((128, 1, 64), dtype=nl.float32, buffer=nl.sbuf)
t.squeeze_dim(1) # shape becomes (128, 64)
**Constraints.**
- ``t.shape[dim]`` must be 1
- On-chip tensors: ``dim`` must not be 0
- After ``vector_select``: ``dim`` must not be 0
:param dim: dimension to remove (must have size 1)
:return: new ``NkiTensor`` view with the dimension removed"""
...
[docs] def reshape_dim(self, dim: int, shape: tuple[int, ...]) -> NkiTensor:
r"""Split a single dimension into multiple dimensions.
The product of ``shape`` must equal ``t.shape[dim]``. One element of
``shape`` may be -1, in which case its value is inferred.
.. code-block:: python
t = nl.ndarray((128, 24), dtype=nl.float32, buffer=nl.sbuf)
t.reshape_dim(1, (4, 6)) # shape becomes (128, 4, 6)
t.reshape_dim(1, (4, -1)) # same result, 6 is inferred
**Constraints.**
- ``prod(shape) == t.shape[dim]``
- On-chip tensors: ``dim`` must not be 0 (unless ``shape`` is trivial, e.g., ``(128,)``)
- After ``vector_select``: ``dim`` must not be 0
:param dim: dimension to split
:param shape: tuple of sizes for the new dimensions (may contain one -1)
:return: new ``NkiTensor`` view with the dimension split"""
...
[docs] def flatten_dims(self, start_dim: int, end_dim: int) -> NkiTensor:
r"""Merge a contiguous range of dimensions into one.
Dimensions ``start_dim`` through ``end_dim`` (inclusive) are merged into
a single dimension. The dimensions must already be contiguous in memory
(no ``permute`` or non-contiguous slicing across them) so the merged view
is itself a valid view of storage.
.. code-block:: python
t = nl.ndarray((128, 2, 3, 4), dtype=nl.float32, buffer=nl.sbuf)
t.flatten_dims(1, 2) # shape becomes (128, 6, 4)
**Constraints.**
- Dimensions ``start_dim..end_dim`` must be contiguous in memory
- On-chip tensors: ``start_dim`` must be > 0
- After ``vector_select``: ``start_dim`` must be > 0
:param start_dim: first dimension to merge (inclusive)
:param end_dim: last dimension to merge (inclusive)
:return: new ``NkiTensor`` view with the merged dimension"""
...
[docs] def reshape(self, shape: tuple[int, ...]) -> NkiTensor:
r"""Reshape the tensor to a new shape without copying data.
The total number of elements must remain the same. Fails if the current
memory layout is incompatible with the requested shape (e.g. after a
non-contiguous slice or permute).
.. code-block:: python
t = nl.ndarray((128, 4, 6), dtype=nl.float32, buffer=nl.sbuf)
t.reshape((128, 24)) # merge last two dims
t.reshape((128, 2, 12)) # split differently
**Constraints.**
- ``prod(shape) == prod(t.shape)``
- On-chip tensors: ``shape[0]`` must equal ``t.shape[0]`` (partition dim preserved)
- After ``vector_select``: ``shape[0]`` must equal ``t.shape[0]`` (indirect
partition dim preserved)
- Fails if the current layout is incompatible with the requested shape
:param shape: tuple of new dimension sizes
:return: new ``NkiTensor`` view with the requested shape"""
...
[docs] def rearrange(self, src_pattern: tuple, dst_pattern: tuple, fixed_sizes: dict[str, int] | None=None) -> NkiTensor:
r"""Rearrange tensor dimensions using einops-style patterns.
Combines splitting, reordering, and merging dimensions into a single
named operation. Patterns are tuples of strings (dimension names) or
tuples of strings (grouped dimensions that are split or merged).
.. code-block:: python
t = nl.ndarray((128, 24), dtype=nl.float32, buffer=nl.sbuf)
# Split dim 1 into (h, w), then reorder to (b, w, h):
t.rearrange(('b', ('h', 'w')), ('b', 'w', 'h'), {'h': 4})
# Result shape: (128, 6, 4)
:param src_pattern: source dimension pattern (tuple of str or tuple-of-str)
:param dst_pattern: destination dimension pattern (same dimension names)
:param fixed_sizes: dict mapping dimension names to known sizes (for -1 inference)
:return: new ``NkiTensor`` view with rearranged dimensions"""
...
def __setitem__(self, key, value):
...
[docs] def slice(self, dim: int, start: int, end: int, step: int=1) -> NkiTensor:
r"""Slice along a single dimension.
Returns a view selecting elements from ``start`` to ``end`` (exclusive)
with the given ``step``. Equivalent to ``t[:, start:end:step, :]`` when
``dim=1``.
.. code-block:: python
t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t.slice(1, 8, 24, 1) # shape becomes (128, 16)
t.slice(1, 0, 64, 2) # shape becomes (128, 32)
**Constraints.**
- ``0 <= start < end <= t.shape[dim]``
- ``step >= 1``
- On an indirect view (see :meth:`is_indirect`), cannot slice a
dimension that participates in the indirection.
:param dim: dimension to slice
:param start: start index (inclusive)
:param end: end index (exclusive)
:param step: step size (default 1)
:return: new ``NkiTensor`` view with the sliced dimension"""
...
[docs] def select(self, dim: int, index: Union[int, NkiTensor, VirtualRegister, NkiScalar]) -> NkiTensor:
r"""Select a single element along a dimension, removing it.
The index is either a compile-time ``int`` (equivalent to
``t[:, index, :]`` when ``dim=1``) or a runtime value resolved on
device: an ``NkiTensor`` 1-element tile, a ``VirtualRegister``, or an
``NkiScalar``.
.. code-block:: python
t = nl.ndarray((128, 8, 64), dtype=nl.float32, buffer=nl.sbuf)
t.select(1, 3) # static: shape becomes (128, 64)
# Dynamic select (HBM tensor, index resolved at runtime):
idx = nl.ndarray((1, 1), dtype=nl.uint32, buffer=nl.sbuf)
hbm_t = nl.ndarray((4, 128, 8), dtype=nl.float32, buffer=nl.shared_hbm)
hbm_t.select(0, idx) # shape becomes (128, 8)
# Same thing driven by a register:
reg = nisa.register_alloc()
nisa.register_load(reg, idx)
hbm_t.select(0, reg)
**Constraints.**
- Static: ``0 <= index < t.shape[dim]``; runtime indices are not
bounds-checked, and out-of-range is undefined behaviour
- One runtime index per view (no chaining) — see :meth:`is_indirect`
- Runtime on-chip: ``dim`` must not be 0 (partition dim)
- Static selection cannot target a dimension already indirected
:param dim: dimension to select from
:param index: compile-time ``int``, or a runtime ``NkiTensor`` /
``VirtualRegister`` / ``NkiScalar``
:return: new ``NkiTensor`` view with the dimension removed"""
...
[docs] def vector_select(self, dim: int, vector_offset: NkiTensor) -> NkiTensor:
r"""Per-partition indirect addressing using a vector of offsets.
Each partition uses its own index from ``vector_offset`` as the base
address for the selected dimension. The output partition count is
``vector_offset.shape[0]``.
This is used for gather-style operations where different partitions
read from different locations in the source tensor.
.. code-block:: python
hbm_t = nl.ndarray((64, 128, 8), dtype=nl.float32, buffer=nl.shared_hbm)
offsets = nl.ndarray((128, 1), dtype=nl.int32, buffer=nl.sbuf)
# Each of 128 partitions reads from a different row of hbm_t
hbm_t.vector_select(0, offsets) # shape becomes (128, 128, 8)
**Constraints.**
- ``dim`` must be 0
- Only supported on HBM tensors
- Only one dynamic select per tensor (cannot combine with a prior
dynamic :meth:`select` or :meth:`vector_select`); check
:meth:`is_indirect` to guard
- The result is an indirect view: the selected dimension cannot be
further sliced or selected.
:param dim: dimension to apply indirect addressing (must be 0)
:param vector_offset: SBUF tensor with per-partition indices, shape ``(num_partitions, 1)``
:return: new ``NkiTensor`` view with dim 0 size set to ``vector_offset.shape[0]``"""
...
[docs] def ap(self, pattern: list[list[int | VirtualRegister]] | list[list[int]], offset: int | None=None, scalar_offset: NkiTensor | None=None, vector_offset: NkiTensor | None=None, indirect_dim: int=0, dtype: DType | None=None) -> NkiTensor:
r"""Low-level access pattern override (escape hatch).
Replaces shape and strides with an explicit
``[[stride, count], ...]`` pattern addressing the underlying
storage directly. Analogous to ``torch.as_strided``.
Each ``stride`` and ``count`` slot may be a Python ``int`` or dynamic
value only known at runtime — :class:`nki.isa.VirtualRegister`. Note
that when any slot is dynamic, ``scalar_offset`` /
``vector_offset`` / non-zero ``indirect_dim`` are not supported.
.. code-block:: python
sb = nl.ndarray((128, 32), dtype=nl.float32, buffer=nl.sbuf)
# Explicit 2D pattern: partition stride=32, free stride=1
sb.ap(pattern=[[32, 128], [1, 32]])
# With indirect access and dtype reinterpret:
sb.ap(pattern=[[64, 128], [1, 16]], dtype=nl.bfloat16,
scalar_offset=idx, indirect_dim=1)
# Register-bound stride (dynamic access pattern):
stride0 = nisa.register_alloc(...)
count1 = nisa.register_alloc(...)
sb.ap(pattern=[[stride0, 128], [1, count1]])
:param pattern: list of ``[stride, count]`` pairs defining the access
pattern. Each entry may be ``int`` or ``VirtualRegister``.
:param offset: element offset added to the view's base storage offset.
When ``None`` (default), inherits the current view's storage offset
unchanged. Pass an explicit integer to compose with the base offset
(e.g. offset=0 keeps the base, offset=N shifts by N additional elements).
:param scalar_offset: dynamic scalar index for indirect access, as a
1x1 index tensor or an :class:`~nki.isa.NkiScalar`. Must be
32-bit when ``self`` is in SBUF or PSUM; a 64-bit offset is
only valid on an HBM operand, where it reaches past 4GB
:param vector_offset: per-partition index tensor for indirect access
:param indirect_dim: dimension in ``self.shape`` whose stride scales
the indirect scalar/vector offset (default 0)
:param dtype: reinterpret storage as this dtype (default: tensor's dtype)
:return: new ``NkiTensor`` with the explicit access pattern"""
...
[docs] def view(self, dtype: DType) -> NkiTensor:
r"""Reinterpret the tensor's data as a different dtype.
No data is copied. Equivalent to ``numpy.ndarray.view(dtype)`` or
C++ ``reinterpret_cast``. The last dimension's size is scaled by the
ratio of dtype sizes: reinterpreting a float32 tensor as uint8
multiplies the last-dim count by 4; reinterpreting uint8 as float32
divides it by 4.
.. code-block:: python
t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t.view(nl.uint8) # shape becomes (128, 256), 4x expansion
t.view(nl.int32) # shape stays (128, 64), same-size reinterpret
u = nl.ndarray((128, 256), dtype=nl.uint8, buffer=nl.sbuf)
u.view(nl.float32) # shape becomes (128, 64), 4x contraction
**Constraints.**
- The last dimension must be contiguous in memory
- For contraction (larger dtype): last-dim size must be divisible by the ratio
- Not supported after dynamic / vector select
:param dtype: target NKI dtype to reinterpret as
:return: new ``NkiTensor`` view with the adjusted dtype and shape"""
...
[docs] def reinterpret_cast(self, dtype: DType) -> NkiTensor:
r"""Alias for :meth:`view`.
:param dtype: target NKI dtype to reinterpret as
:return: new ``NkiTensor`` view with the adjusted dtype and shape"""
...
[docs] def indirect(self, index: NkiTensor, num_elem: int | None=None) -> NkiTensor:
r"""Create an indirect tensor view for Tensor Indirection (TI).
Available on NeuronCore-v4 (Trn3) and later.
TI allows reading or writing a column of data at given free-dimension
offsets across contiguous partition dimensions. Can be used as input
(gather) or output (scatter) in nisa operations.
Offsets are stored in a snake pattern across partition groups: offset i
comes from ``index[i % G, i // G]`` where G is the group size (16 for
vector/scalar/gpsimd engines, 32 for tensor engine).
:param index: SBUF tensor containing free-dimension offsets, shape ``(P, K)``
where ``P == self.shape[0]``.
:param num_elem: number of offsets to use. Defaults to ``index.size``.
:return: new ``NkiTensor`` view with TI attached. Output shape is
``(P, num_elem)``."""
...
...
class NkiTileGrid(NkiGrid):
r"""A tensor view navigated in tile coordinates.
Each tiled entry in :attr:`NkiGrid.shape` counts tiles, while
:attr:`NkiGrid.tile_shape` gives the nominal source-element shape of one
tile. Untiled HBM dimensions form a leading prefix and are indexed one slab
at a time.
Integer indexing consumes coordinates but retains ``NkiTileGrid``.
:attr:`NkiGrid.data` exposes the exact :class:`NkiTensor` operand for the
selected view.
Instances come from :func:`nki.language.tiles` or a block-to-tile
transition. Creating and indexing the view allocates no storage and moves
no data.
Examples:
.. code-block:: python
import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def copy_tile_grid(src):
tile_shape = (128, 256)
src_tiles = nl.tiles(src, tile_shape=tile_shape)
dst = nl.ndarray(
src.shape, dtype=src.dtype, buffer=nl.shared_hbm
)
dst_tiles = nl.tiles(dst, tile_shape=tile_shape)
for i in range(src_tiles.shape[0]):
for j in range(src_tiles.shape[1]):
src_tile = src_tiles[i, j]
tile = nl.ndarray(
src_tile.element_shape,
dtype=src.dtype,
buffer=nl.sbuf,
)
nisa.dma_copy(dst=tile.data, src=src_tile.data)
nisa.dma_copy(
dst=dst_tiles[i, j].data,
src=tile.data,
)
return dst
See Also:
:class:`NkiBlockGrid`: A grid whose cells group multiple tiles.
:func:`nki.language.tiles`: Create or retile a tile grid."""
...
class _TileSize(NKIObject):
r"""Hardware tile size constants (pmax, psum_bank_fmax, gemm_stationary_fmax, etc.)."""
pmax: int = ...
r"""Maximum partition dimension of a tile."""
psum_fmax: int = ...
r"""Maximum free dimension per PSUM bank per partition, in FP32 elements.
.. deprecated:: 0.4.0
Despite the name, this does not represent the free-dim size per
partition. Use :py:attr:`psum_bank_fmax` instead."""
psum_bank_fmax: int = ...
r"""Maximum free dimension per PSUM bank per partition, in FP32 elements."""
psum_bank_fmax_bytes: int = ...
r"""Maximum free dimension per PSUM bank per partition, in bytes."""
gemm_stationary_fmax: int = ...
r"""Maximum free dimension of the stationary operand of General Matrix Multiplication on Tensor Engine"""
gemm_moving_fmax: int = ...
r"""Maximum free dimension of the moving operand of General Matrix Multiplication on Tensor Engine"""
bn_stats_fmax: int = ...
r"""Maximum free dimension of BN_STATS"""
psum_min_align: int = ...
r"""Minimum byte alignment requirement for PSUM free dimension address"""
sbuf_min_align: int = ...
r"""Minimum byte alignment requirement for SBUF free dimension address"""
@property
def psum_num_banks(self) -> int:
r"""Number of usable PSUM banks per partition"""
...
@property
def sbuf_size_bytes(self) -> int:
r"""Total SBUF capacity in bytes (all partitions combined)."""
...
@property
def sbuf_fmax(self) -> int:
r"""Maximum free dimension of a tile on SBUF buffer, in FP32 elements."""
...
@property
def sbuf_fmax_bytes(self) -> int:
r"""Maximum free dimension of a tile on SBUF buffer, in bytes."""
...
@property
def total_available_sbuf_size(self) -> int:
r"""Usable SBUF free dimension per partition, in bytes.
.. deprecated:: 0.4.0
Despite the name, this returns the usable SBUF capacity *per
partition*, not the total SBUF capacity. Use :py:attr:`sbuf_size_bytes`
for the total SBUF capacity across all partitions, or
:py:attr:`sbuf_fmax_bytes` for the usable per-partition size."""
...
...
[docs]def is_hbm(buffer):
r"""Check if buffer is any HBM type."""
...
[docs]def is_on_chip(buffer):
r"""Check if buffer is on-chip (SBUF or PSUM)."""
...
[docs]def is_psum(buffer):
r"""Check if buffer is PSUM."""
...
[docs]def is_sbuf(buffer):
r"""Check if buffer is SBUF."""
...
[docs]def add(x, y, dtype=None):
r"""Add the inputs, element-wise.
((Similar to `numpy.add <https://numpy.org/doc/stable/reference/generated/numpy.add.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has ``x + y``, element-wise.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.add -- element-wise addition of two tiles
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 512), 2.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.add(a, b)
expected = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.add -- adding a scalar to every element of a tile
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.add(a, 2.0)
expected = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def subtract(x, y, dtype=None):
r"""Subtract the inputs, element-wise.
((Similar to `numpy.subtract <https://numpy.org/doc/stable/reference/generated/numpy.subtract.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has ``x - y``, element-wise.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.subtract -- element-wise subtraction of two tiles
a = nl.full((128, 512), 10.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.subtract(a, b)
expected = nl.full((128, 512), 7.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.subtract -- subtracting a scalar from every element
a = nl.full((128, 512), 10.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.subtract(a, 3.0)
expected = nl.full((128, 512), 7.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def multiply(x, y, dtype=None):
r"""Multiply the inputs, element-wise.
((Similar to `numpy.multiply <https://numpy.org/doc/stable/reference/generated/numpy.multiply.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has ``x * y``, element-wise.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.multiply -- element-wise multiplication of two tiles
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 512), 4.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.multiply(a, b)
expected = nl.full((128, 512), 12.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.multiply -- scaling every element by a scalar
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.multiply(a, 4.0)
expected = nl.full((128, 512), 12.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
def divide(x, y, dtype=None):
r"""Divide the inputs, element-wise.
((Similar to `numpy.divide <https://numpy.org/doc/stable/reference/generated/numpy.divide.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has ``x / y``, element-wise.
"""
...
[docs]def maximum(x, y, dtype=None):
r"""Maximum of the inputs, element-wise.
((Similar to `numpy.maximum <https://numpy.org/doc/stable/reference/generated/numpy.maximum.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has the maximum of each element from x and y.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.maximum -- max(3.0, 5.0) = 5.0
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.maximum(a, b)
expected = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.maximum -- with a scalar operand
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.maximum(a, 5.0)
expected = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def minimum(x, y, dtype=None):
r"""Minimum of the inputs, element-wise.
((Similar to `numpy.minimum <https://numpy.org/doc/stable/reference/generated/numpy.minimum.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has the minimum of each element from x and y.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.minimum -- min(3.0, 5.0) = 3.0
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.minimum(a, b)
expected = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.minimum -- with a scalar operand
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.minimum(a, 5.0)
expected = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def abs_max(x, y, dtype=None):
r"""Maximum of the inputs compared by magnitude, element-wise.
Compares ``abs(x)`` and ``abs(y)`` and returns the **original (signed) value**
of whichever input has the larger absolute value.
For example, ``abs_max(-5, 3)`` returns ``-5`` because ``|-5| > |3|``.
.. note::
Available only on NeuronCore-v4 (trn3) and newer.
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile where each element is ``x`` if ``|x| > |y|``, else ``y``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.abs_max -- returns the input with the larger absolute value
a = nl.full((128, 512), -5.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 1), 3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.abs_max(a, b) # |−5| > |3| → returns -5.0 (original signed value)
expected = nl.full((128, 512), -5.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
.. code-block:: python
# nki.language.abs_max -- tie-breaking: returns y when |x| == |y|
a = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 1), -1.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.abs_max(a, b) # |1| == |-1| → tie, returns y = -1.0
expected = nl.full((128, 512), -1.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def abs_min(x, y, dtype=None):
r"""Minimum of the inputs compared by magnitude, element-wise.
Compares ``abs(x)`` and ``abs(y)`` and returns the **original (signed) value**
of whichever input has the smaller absolute value.
For example, ``abs_min(-5, 3)`` returns ``3`` because ``|3| < |-5|``.
.. note::
Available only on NeuronCore-v4 (trn3) and newer.
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile where each element is ``x`` if ``|x| < |y|``, else ``y``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.abs_min -- element-wise absolute minimum of two tiles
a = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 1), -3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.abs_min(a, b) # |-3| < |5|, so returns -3.0 (original signed value)
expected = nl.full((128, 512), -3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
.. code-block:: python
# nki.language.abs_min -- tie-breaking: returns y when |x| == |y|
a = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 1), -1.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.abs_min(a, b) # |1| == |-1| → tie, returns y = -1.0
expected = nl.full((128, 512), -1.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def abs(x, dtype=None):
r"""Absolute value of the input, element-wise.
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has absolute values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.abs
a = nl.full((128, 512), -1.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.abs(a)
expected = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(b, expected)
# nki.language.abs with explicit dtype
a = nl.full((128, 512), -1.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.abs(a, dtype=nl.float16)
expected = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def square(x, dtype=None):
r"""Square of the input, element-wise.
((Similar to `numpy.square <https://numpy.org/doc/stable/reference/generated/numpy.square.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has square of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.square -- square(3.0) = 9.0
a = nl.full((128, 512), 3.0, dtype=nl.float32,
buffer=nl.sbuf)
b = nl.square(a)
expected = nl.full((128, 512), 9.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def power(x, y, dtype=None):
r"""Elements of x raised to powers of y, element-wise.
((Similar to `numpy.power <https://numpy.org/doc/stable/reference/generated/numpy.power.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has values ``x`` to the power of ``y``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.power -- element-wise exponentiation of two tiles
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 512), 2.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.power(a, b)
expected = nl.full((128, 512), 9.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def equal(x, y, dtype=None):
r"""Return (x == y) element-wise.
((Similar to `numpy.equal <https://numpy.org/doc/stable/reference/generated/numpy.equal.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); Defaults to the input tile dtype.
Use ``dtype=nl.uint8`` for a boolean-like result.
:return: a tile with 1 where equal, 0 otherwise."""
...
[docs]def not_equal(x, y, dtype=None):
r"""Return (x != y) element-wise.
((Similar to `numpy.not_equal <https://numpy.org/doc/stable/reference/generated/numpy.not_equal.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); Defaults to the input tile dtype.
Use ``dtype=nl.uint8`` for a boolean-like result.
:return: a tile with 1 where not equal, 0 otherwise."""
...
[docs]def less(x, y, dtype=None):
r"""Return (x < y) element-wise.
((Similar to `numpy.less <https://numpy.org/doc/stable/reference/generated/numpy.less.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); Defaults to the input tile dtype.
Use ``dtype=nl.uint8`` for a boolean-like result.
:return: a tile with 1 where x < y, 0 otherwise."""
...
[docs]def less_equal(x, y, dtype=None):
r"""Return (x <= y) element-wise.
((Similar to `numpy.less_equal <https://numpy.org/doc/stable/reference/generated/numpy.less_equal.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); Defaults to the input tile dtype.
Use ``dtype=nl.uint8`` for a boolean-like result.
:return: a tile with 1 where x <= y, 0 otherwise."""
...
[docs]def greater(x, y, dtype=None):
r"""Return (x > y) element-wise.
((Similar to `numpy.greater <https://numpy.org/doc/stable/reference/generated/numpy.greater.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); Defaults to the input tile dtype.
Use ``dtype=nl.uint8`` for a boolean-like result.
:return: a tile with 1 where x > y, 0 otherwise."""
...
[docs]def greater_equal(x, y, dtype=None):
r"""Return (x >= y) element-wise.
((Similar to `numpy.greater_equal <https://numpy.org/doc/stable/reference/generated/numpy.greater_equal.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); Defaults to the input tile dtype.
Use ``dtype=nl.uint8`` for a boolean-like result.
:return: a tile with 1 where x >= y, 0 otherwise."""
...
[docs]def bitwise_and(x, y, dtype=None):
r"""Compute the bitwise AND of two tiles element-wise.
((Similar to `numpy.bitwise_and <https://numpy.org/doc/stable/reference/generated/numpy.bitwise_and.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs must be integer typed.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the bitwise AND result."""
...
[docs]def bitwise_or(x, y, dtype=None):
r"""Compute the bitwise OR of two tiles element-wise.
((Similar to `numpy.bitwise_or <https://numpy.org/doc/stable/reference/generated/numpy.bitwise_or.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs must be integer typed.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the bitwise OR result."""
...
[docs]def bitwise_xor(x, y, dtype=None):
r"""Compute the bitwise XOR of two tiles element-wise.
((Similar to `numpy.bitwise_xor <https://numpy.org/doc/stable/reference/generated/numpy.bitwise_xor.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs must be integer typed.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the bitwise XOR result."""
...
[docs]def invert(x, dtype=None):
r"""Compute the bitwise NOT element-wise.
((Similar to `numpy.invert <https://numpy.org/doc/stable/reference/generated/numpy.invert.html>`_))
.. warning::
This API is experimental and may change in future releases.
Input must be integer typed. Implemented as XOR with all-ones.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile with the bitwise NOT result."""
...
[docs]def left_shift(x, y, dtype=None):
r"""Left shift the bits of x by y positions element-wise.
((Similar to `numpy.left_shift <https://numpy.org/doc/stable/reference/generated/numpy.left_shift.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs must be integer typed.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the left-shifted result."""
...
[docs]def right_shift(x, y, dtype=None):
r"""Right shift the bits of x by y positions element-wise.
((Similar to `numpy.right_shift <https://numpy.org/doc/stable/reference/generated/numpy.right_shift.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs must be integer typed.
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the right-shifted result."""
...
[docs]def logical_and(x, y, dtype=None):
r"""Compute the logical AND of two tiles element-wise.
((Similar to `numpy.logical_and <https://numpy.org/doc/stable/reference/generated/numpy.logical_and.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs should be boolean-like (0 or 1 values).
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the logical AND result."""
...
[docs]def logical_or(x, y, dtype=None):
r"""Compute the logical OR of two tiles element-wise.
((Similar to `numpy.logical_or <https://numpy.org/doc/stable/reference/generated/numpy.logical_or.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs should be boolean-like (0 or 1 values).
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the logical OR result."""
...
[docs]def logical_xor(x, y, dtype=None):
r"""Compute the logical XOR of two tiles element-wise.
((Similar to `numpy.logical_xor <https://numpy.org/doc/stable/reference/generated/numpy.logical_xor.html>`_))
.. warning::
This API is experimental and may change in future releases.
Inputs should be boolean-like (0 or 1 values).
:param x: a tile or a scalar value.
:param y: a tile or a scalar value. At least one of x, y must be a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile with the logical XOR result."""
...
[docs]def logical_not(x, dtype=None):
r"""Compute the logical NOT element-wise.
((Similar to `numpy.logical_not <https://numpy.org/doc/stable/reference/generated/numpy.logical_not.html>`_))
.. warning::
This API is experimental and may change in future releases.
Implemented as XOR with 1, so inputs should be boolean-like (0 or 1 values).
For non-boolean inputs, use ``nl.equal(x, 0)`` instead.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile with the logical NOT result."""
...
[docs]def exp(x, dtype=None):
r"""Exponential of the input, element-wise.
((Similar to `numpy.exp <https://numpy.org/doc/stable/reference/generated/numpy.exp.html>`_))
.. warning::
This API is experimental and may change in future releases.
The ``exp(x)`` is ``e^x`` where ``e`` is the Euler's number = 2.718281...
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has exponential values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.exp -- exp(0.0) = 1.0
a = nl.full((128, 512), 0.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.exp(a)
expected = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def log(x, dtype=None):
r"""Natural logarithm of the input, element-wise.
((Similar to `numpy.log <https://numpy.org/doc/stable/reference/generated/numpy.log.html>`_))
.. warning::
This API is experimental and may change in future releases.
It is the inverse of the exponential function, such that: ``log(exp(x)) = x`` .
The natural logarithm base is ``e``.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has natural logarithm values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.log -- log(1.0) = 0.0
a = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.log(a)
expected = nl.full((128, 512), 0.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def sqrt(x, dtype=None):
r"""Non-negative square-root of the input, element-wise.
((Similar to `numpy.sqrt <https://numpy.org/doc/stable/reference/generated/numpy.sqrt.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has square-root values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.sqrt -- sqrt(4.0) = 2.0
a = nl.full((128, 512), 4.0, dtype=nl.float32,
buffer=nl.sbuf)
b = nl.sqrt(a)
expected = nl.full((128, 512), 2.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def rsqrt(x, dtype=None):
r"""Reciprocal of the square-root of the input, element-wise.
((Similar to `torch.rsqrt <https://pytorch.org/docs/stable/generated/torch.rsqrt.html>`_))
.. warning::
This API is experimental and may change in future releases.
``rsqrt(x) = 1 / sqrt(x)``
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has reciprocal square-root values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.rsqrt -- rsqrt(4.0) = 0.5
a = nl.full((128, 512), 4.0, dtype=nl.float32,
buffer=nl.sbuf)
b = nl.rsqrt(a)
expected = nl.full((128, 512), 0.5, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def tanh(x, dtype=None):
r"""Hyperbolic tangent, element-wise."""
...
[docs]def sigmoid(x, dtype=None):
r"""Sigmoid activation, element-wise."""
...
[docs]def relu(x, dtype=None):
r"""ReLU activation, element-wise."""
...
[docs]def gelu(x, dtype=None):
r"""GELU activation, element-wise."""
...
[docs]def gelu_apprx_sigmoid(x, dtype=None):
r"""GELU approximation using sigmoid, element-wise."""
...
[docs]def gelu_apprx_tanh(x, dtype=None):
r"""GELU approximation using tanh, element-wise."""
...
[docs]def gelu_dx(x, dtype=None):
r"""Derivative of GELU activation, element-wise."""
...
[docs]def gelu_apprx_sigmoid_dx(x, dtype=None):
r"""Derivative of sigmoid-approximated GELU, element-wise."""
...
[docs]def silu(x, dtype=None):
r"""SiLU (Swish) activation, element-wise."""
...
[docs]def silu_dx(x, dtype=None):
r"""Derivative of SiLU activation, element-wise."""
...
[docs]def softplus(x, dtype=None):
r"""Softplus activation, element-wise."""
...
def log1p(x, dtype=None):
r"""``ln(1 + x)`` activation, element-wise.
Requires NeuronCore-v4 or later."""
...
[docs]def mish(x, dtype=None):
r"""Mish activation, element-wise."""
...
[docs]def erf(x, dtype=None):
r"""Error function, element-wise."""
...
[docs]def erf_dx(x, dtype=None):
r"""Derivative of error function, element-wise."""
...
[docs]def reciprocal(x, dtype=None):
r"""Reciprocal of the input, element-wise.
((Similar to `numpy.reciprocal <https://numpy.org/doc/stable/reference/generated/numpy.reciprocal.html>`_))
.. warning::
This API is experimental and may change in future releases.
``reciprocal(x) = 1 / x``
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has reciprocal values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.reciprocal -- reciprocal(4.0) = 0.25
a = nl.full((128, 512), 4.0, dtype=nl.float32,
buffer=nl.sbuf)
b = nl.reciprocal(a)
expected = nl.full((128, 512), 0.25, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def copy(x, dtype=None):
r"""Create a copy of the input tile on SBUF.
.. note::
Using a ``dtype`` different from the input's data type may result in
precision loss.
:param x: the source of copy, must be a tile in SBUF or PSUM.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a new tile with the same layout as ``x``, allocated on SBUF.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.copy -- copy an nl.matmul result from PSUM back to SBUF.
# nl.matmul writes its result to PSUM; nl.copy brings it to SBUF.
x = nl.load(x_tensor)
y = nl.load(y_tensor)
result_psum = nl.matmul(x, y, transpose_x=True) # x.T @ y, in PSUM
result = nl.copy(result_psum) # PSUM -> SBUF"""
...
[docs]def sin(x, dtype=None):
r"""Sine of the input, element-wise.
((Similar to `numpy.sin <https://numpy.org/doc/stable/reference/generated/numpy.sin.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has sine values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.sin -- sin(0.0) = 0.0
a = nl.full((128, 512), 0.0, dtype=nl.float32,
buffer=nl.sbuf)
b = nl.sin(a)
expected = nl.full((128, 512), 0.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def cos(x, dtype=None):
r"""Cosine of the input, element-wise.
((Similar to `numpy.cos <https://numpy.org/doc/stable/reference/generated/numpy.cos.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has cosine values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.cos -- cos(0.0) = 1.0
a = nl.full((128, 512), 0.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.cos(a)
expected = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def tan(x, dtype=None):
r"""Tangent of the input, element-wise.
((Similar to `numpy.tan <https://numpy.org/doc/stable/reference/generated/numpy.tan.html>`_))
.. warning::
This API is experimental and may change in future releases.
There is no ``tan`` operator on the Scalar Engine. It is computed as
``sin(x) / cos(x)`` from the ``sin`` activation, so the result can diverge
from the true tangent near its poles (odd multiples of ``pi/2``).
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has tangent values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.tan -- tan(0.0) = 0.0
a = nl.full((128, 512), 0.0, dtype=nl.float32,
buffer=nl.sbuf)
b = nl.tan(a)
expected = nl.full((128, 512), 0.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def arctan(x, dtype=None):
r"""Inverse tangent of the input, element-wise.
((Similar to `numpy.arctan <https://numpy.org/doc/stable/reference/generated/numpy.arctan.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has inverse tangent values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.arctan -- arctan(0.0) = 0.0
a = nl.full((128, 512), 0.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.arctan(a)
expected = nl.full((128, 512), 0.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
[docs]def ceil(x, dtype=None):
r"""Ceiling of the input, element-wise.
((Similar to `numpy.ceil <https://numpy.org/doc/stable/reference/generated/numpy.ceil.html>`_))
.. warning::
This API is experimental and may change in future releases.
The ceil of the scalar x is the smallest integer i, such that i >= x.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has ceiling values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.ceil -- rounds 3.2 up to 4.0
a = nl.full((128, 512), 3.2, dtype=nl.float32, buffer=nl.sbuf)
c = nl.ceil(a)
expected = nl.full((128, 512), 4.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.ceil -- rounds -3.7 up to -3.0
a = nl.full((128, 512), -3.7, dtype=nl.float32, buffer=nl.sbuf)
c = nl.ceil(a)
expected = nl.full((128, 512), -3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def floor(x, dtype=None):
r"""Floor of the input, element-wise.
((Similar to `numpy.floor <https://numpy.org/doc/stable/reference/generated/numpy.floor.html>`_))
.. warning::
This API is experimental and may change in future releases.
The floor of the scalar x is the largest integer i, such that i <= x.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has floor values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.floor -- rounds 3.7 down to 3.0
a = nl.full((128, 512), 3.7, dtype=nl.float32, buffer=nl.sbuf)
c = nl.floor(a)
expected = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.floor -- rounds -3.2 down to -4.0
a = nl.full((128, 512), -3.2, dtype=nl.float32, buffer=nl.sbuf)
c = nl.floor(a)
expected = nl.full((128, 512), -4.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def trunc(x, dtype=None):
r"""Truncated value of the input, element-wise.
((Similar to `numpy.trunc <https://numpy.org/doc/stable/reference/generated/numpy.trunc.html>`_))
.. warning::
This API is experimental and may change in future releases.
The truncated value of the scalar x is the nearest integer i which is closer to zero than x is.
In short, the fractional part of the signed number x is discarded.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has truncated values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.trunc -- truncates 3.7 toward zero to 3.0
a = nl.full((128, 512), 3.7, dtype=nl.float32, buffer=nl.sbuf)
c = nl.trunc(a)
expected = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.trunc -- truncates -3.7 toward zero to -3.0
a = nl.full((128, 512), -3.7, dtype=nl.float32, buffer=nl.sbuf)
c = nl.trunc(a)
expected = nl.full((128, 512), -3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def negative(x, dtype=None):
r"""Numerical negative of the input, element-wise.
((Similar to `numpy.negative <https://numpy.org/doc/stable/reference/generated/numpy.negative.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has numerical negative values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.negative -- negates 5.0 to -5.0
a = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.negative(a)
expected = nl.full((128, 512), -5.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)
# nki.language.negative -- negates -3.0 to 3.0
a = nl.full((128, 512), -3.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.negative(a)
expected = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def sign(x, dtype=None):
r"""Sign of the numbers of the input, element-wise.
((Similar to `numpy.sign <https://numpy.org/doc/stable/reference/generated/numpy.sign.html>`_))
.. warning::
This API is experimental and may change in future releases.
The sign function returns ``-1`` if ``x < 0``, ``0`` if ``x==0``, ``1`` if ``x > 0``.
:param x: a tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has sign values of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.sign -- sign(-5.0) = -1.0
a = nl.full((128, 512), -5.0, dtype=nl.float32,
buffer=nl.sbuf)
b = nl.sign(a)
expected = nl.full((128, 512), -1.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(b, expected)"""
...
def fmod(x, y, dtype=None):
r"""Floating-point remainder of ``x / y``, element-wise.
The remainder has the same sign as the dividend x.
It is equivalent to the Matlab(TM) rem function and should not be confused with the Python modulus operator x % y.
((Similar to `numpy.fmod <https://numpy.org/doc/stable/reference/generated/numpy.fmod.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile. If x is a scalar value it will be broadcast to the shape of y.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has values ``x fmod y``.
"""
...
def mod(x, y, dtype=None):
r"""Remainder of ``x / y``, element-wise.
Computes the remainder complementary to the floor_divide function.
It is equivalent to the Python modulus x % y and has the same sign as the divisor y.
((Similar to `numpy.mod <https://numpy.org/doc/stable/reference/generated/numpy.mod.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile. If x is a scalar value it will be broadcast to the shape of y.
:param y: a tile or a scalar value.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information);
:return: a tile that has values ``x mod y``.
"""
...
[docs]def sum(x, axis, dtype=None, keepdims=False):
r"""Sum of elements along the specified axis (or axes) of the input.
((Similar to `numpy.sum <https://numpy.org/doc/stable/reference/generated/numpy.sum.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. The axis (or axes) along which to operate;
must be free dimensions, not partition dimension (0); can only be the
last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4].
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:param keepdims: if True, the reduced axes are kept as size-one dimensions.
:return: a tile with the sum along the provided axis."""
...
[docs]def max(x, axis, dtype=None, keepdims=False):
r"""Maximum of elements along the specified axis (or axes) of the input.
((Similar to `numpy.max <https://numpy.org/doc/stable/reference/generated/numpy.max.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. The axis (or axes) along which to operate;
must be free dimensions, not partition dimension (0); can only be the
last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4].
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:param keepdims: if True, the reduced axes are kept as size-one dimensions.
:return: a tile with the maximum along the provided axis."""
...
[docs]def min(x, axis, dtype=None, keepdims=False):
r"""Minimum of elements along the specified axis (or axes) of the input.
((Similar to `numpy.min <https://numpy.org/doc/stable/reference/generated/numpy.min.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. The axis (or axes) along which to operate;
must be free dimensions, not partition dimension (0); can only be the
last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4].
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:param keepdims: if True, the reduced axes are kept as size-one dimensions.
:return: a tile with the minimum along the provided axis."""
...
[docs]def mean(x, axis, dtype=None, keepdims=False):
r"""Arithmetic mean along the specified axis (or axes) of the input.
((Similar to `numpy.mean <https://numpy.org/doc/stable/reference/generated/numpy.mean.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. Free dimensions can only be the last
contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. The partition
dimension (0) is also supported on its own, reducing across partitions; it
cannot be mixed with free dimensions in one call.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:param keepdims: if True, the reduced axes are kept as size-one dimensions. It
has no effect for ``axis=0``: dimension 0 is structurally the partition
dimension of a tile, so it is always retained as size one rather than dropped.
:return: a tile with the average of elements along the provided axis. Float32
intermediate values are used for the computation."""
...
[docs]def prod(x, axis, dtype=None, keepdims=False):
r"""Product of elements along the specified axis (or axes) of the input.
((Similar to `numpy.prod <https://numpy.org/doc/stable/reference/generated/numpy.prod.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. The axis (or axes) along which to operate;
must be free dimensions, not partition dimension (0); can only be the
last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4].
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:param keepdims: if True, the reduced axes are kept as size-one dimensions.
:return: a tile with the product along the provided axis."""
...
[docs]def var(x, axis, dtype=None, keepdims=False):
r"""Variance along the specified axis (or axes) of the input.
((Similar to `numpy.var <https://numpy.org/doc/stable/reference/generated/numpy.var.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. The axis (or axes) along which to operate;
must be free dimensions, not partition dimension (0); can only be the
last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4].
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:param keepdims: currently ignored; result always has keepdims=True shape.
:return: a tile with the variance of the elements along the provided axis."""
...
[docs]def all(x, axis, dtype=None):
r"""Whether all elements along the specified axis (or axes) evaluate to True.
((Similar to `numpy.all <https://numpy.org/doc/stable/reference/generated/numpy.all.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. The axis (or axes) along which to operate;
must be free dimensions, not partition dimension (0); can only be the
last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4].
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile with the logical AND reduction along the provided axis."""
...
[docs]def softmax(x, axis=-1, dtype=None):
r"""Softmax activation function on the input, element-wise.
((Similar to `torch.nn.functional.softmax <https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param axis: int or tuple/list of ints. The axis (or axes) along which to operate; must be free dimensions, not partition dimension (0); can only be the last contiguous dim(s) of the tile: ``[1], [1,2], [1,2,3], [1,2,3,4]``
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has softmax of ``x``.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.softmax -- uniform input produces uniform output
a = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
result = nl.softmax(a, axis=1)"""
...
[docs]def rms_norm(x, w, axis, n, epsilon=1e-06, dtype=None, compute_dtype=None):
r"""Apply Root Mean Square Layer Normalization.
.. warning::
This API is experimental and may change in future releases.
:param x: input tile.
:param w: weight tile.
:param axis: axis along which to compute the root mean square (rms) value.
:param n: total number of values to calculate rms.
:param epsilon: epsilon value used by rms calculation to avoid divide-by-zero.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:param compute_dtype: (optional) dtype for the internal computation.
:return: ``x / RMS(x) * w``
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.rms_norm -- normalize with unit weights
x = nl.full((128, 512), 2.0, dtype=nl.float32, buffer=nl.sbuf)
w = nl.full((128, 512), 1.0, dtype=nl.float32, buffer=nl.sbuf)
result = nl.rms_norm(x, w, axis=1, n=512)"""
...
[docs]def where(condition, x, y, dtype=None):
r"""Return elements chosen from x or y depending on condition.
((Similar to `numpy.where <https://numpy.org/doc/stable/reference/generated/numpy.where.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param condition: condition tile with float values (1.0 for True, 0.0 for False).
:param x: tensor from which to take elements where condition is True.
:param y: tensor from which to take elements where condition is False.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: tensor with elements from x or y based on condition.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.where -- select 10.0 where condition is 1, else 0.0
cond = nl.full((128, 512), 1.0, dtype=nl.float32,
buffer=nl.sbuf)
x = nl.full((128, 512), 10.0, dtype=nl.float32,
buffer=nl.sbuf)
y = nl.full((128, 512), 0.0, dtype=nl.float32,
buffer=nl.sbuf)
result = nl.where(cond, x, y)
expected = nl.full((128, 512), 10.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(result, expected)
# nki.language.where -- select 5.0 where condition is 0
cond = nl.full((128, 512), 0.0, dtype=nl.float32,
buffer=nl.sbuf)
x = nl.full((128, 512), 10.0, dtype=nl.float32,
buffer=nl.sbuf)
y = nl.full((128, 512), 5.0, dtype=nl.float32,
buffer=nl.sbuf)
result = nl.where(cond, x, y)
expected = nl.full((128, 512), 5.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(result, expected)"""
...
[docs]def dropout(x, rate, dtype=None):
r"""Randomly zeroes some of the elements of the input tile given a probability rate.
.. warning::
This API is experimental and may change in future releases.
:param x: a tile.
:param rate: the probability of zeroing each element. Can be a scalar constant
or a tile of shape ``(x.shape[0], 1)`` for per-partition drop probabilities.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile with randomly zeroed elements of ``x``."""
...
def alloc_blocks(*, tile_shape: tuple[int, ...], block_shape: tuple[int, ...], element_shape: tuple[int, ...], buffer: MemoryRegion, dtype: DType, address: tuple | None=None, name: str='') -> NkiBlockGrid:
r"""Allocate new storage organized as a grid of tile blocks.
``alloc_blocks`` creates the same tile geometry as
:func:`nki.language.alloc_tiles`, then adds a block-coordinate level above
it. ``tile_shape`` defines the compute grain, ``block_shape`` defines the
number of tiles in one transfer or reuse grain, and ``element_shape``
remains the exact logical tensor extent.
Block grouping does not change physical tile order or add per-block
allocations. HBM uses the compact logical extent as backing storage; SBUF
and PSUM use the same dense rank-two capacity layout as
:func:`nki.language.alloc_tiles`.
The returned :class:`NkiBlockGrid` owns one independent uninitialized allocation.
:param tile_shape: The positive source-element shape of one compute tile. It
must have at least two dimensions and applies to the trailing
dimensions of ``element_shape``.
:param block_shape: The positive tile counts grouped along each block
dimension. Its rank must equal ``tile_shape``.
:param element_shape: The exact positive logical element extent. Any
dimensions preceding ``tile_shape`` remain untiled.
:param buffer: Destination :class:`MemoryRegion`, such as
``nl.shared_hbm``, ``nl.sbuf``, or ``nl.psum``.
:param dtype: Destination NKI element dtype.
:param address: Optional manual address
``(partition_offset, free_offset)`` with the same contract as
:func:`nki.language.ndarray`.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:return: A new :class:`NkiBlockGrid` over independent, uninitialized
storage. Partial boundary blocks and tiles retain exact logical extents;
aggregate ``data`` access requires one exact native access pattern.
:raises TypeError: if a shape is not a tuple of exact Python integers or
``buffer`` is not an NKI :class:`MemoryRegion`.
:raises ValueError: if a shape entry is nonpositive, ``tile_shape`` has fewer
than two dimensions or greater rank than ``element_shape``,
``block_shape`` and ``tile_shape`` ranks differ, or ``buffer`` is
unsupported.
:raises AssertionError: if the native allocation violates a
:func:`nki.language.ndarray` constraint, including PSUM address
alignment.
Examples:
.. code-block:: python
import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def initialize_block_grid():
staging = nl.alloc_blocks(
tile_shape=(128, 128),
block_shape=(2, 2),
element_shape=(512, 512),
buffer=nl.sbuf,
dtype=nl.bfloat16,
)
dst = nl.ndarray(
(512, 512), dtype=nl.bfloat16, buffer=nl.shared_hbm
)
dst_blocks = nl.blocks(
dst,
tile_shape=(128, 128),
block_shape=(2, 2),
)
nisa.memset(dst=staging.data, value=0.0)
nisa.dma_copy(dst=dst_blocks.data, src=staging.data)
return dst
See Also:
:func:`nki.language.blocks`: Group tiles in existing storage.
:func:`nki.language.alloc_tiles`: Allocate a tile grid without blocks."""
...
def alloc_tiles(*, tile_shape: tuple[int, ...], element_shape: tuple[int, ...], buffer: MemoryRegion, dtype: DType, address: tuple | None=None, name: str='') -> NkiTileGrid:
r"""Allocate independent, uninitialized storage organized as a grid of tiles.
The returned :class:`NkiTileGrid` owns one backing allocation and provides
both tile navigation and an aggregate :attr:`NkiGrid.data` operand.
:param tile_shape: The positive source-element shape of one tile. It must be
a tuple containing at least two exact Python integers and applies to
the trailing dimensions of ``element_shape``.
:param element_shape: The exact positive logical element extent.
:param buffer: Destination :class:`MemoryRegion`, such as
``nl.shared_hbm``, ``nl.sbuf``, or ``nl.psum``.
:param dtype: Destination NKI element dtype.
:param address: Optional manual address
``(partition_offset, free_offset)`` with the same contract as
:func:`nki.language.ndarray`.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:return: A new :class:`NkiTileGrid` over independent, uninitialized
storage. Its ``shape`` counts tiles and its boundary selections expose
their exact clamped extents.
:raises TypeError: if a shape is not a tuple of exact Python integers or
``buffer`` is not an NKI :class:`MemoryRegion`.
:raises ValueError: if a shape entry is nonpositive, ``tile_shape`` has fewer
than two dimensions or greater rank than ``element_shape``, or
``buffer`` is unsupported.
:raises AssertionError: if the native allocation violates a
:func:`nki.language.ndarray` constraint, including PSUM address
alignment.
Examples:
.. code-block:: python
import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def initialize_tile_grid():
staging = nl.alloc_tiles(
tile_shape=(128, 256),
element_shape=(256, 512),
buffer=nl.sbuf,
dtype=nl.bfloat16,
)
dst = nl.ndarray(
(256, 512), dtype=nl.bfloat16, buffer=nl.shared_hbm
)
dst_tiles = nl.tiles(dst, tile_shape=(128, 256))
nisa.memset(dst=staging.data, value=0.0)
nisa.dma_copy(dst=dst_tiles.data, src=staging.data)
return dst
See Also:
:func:`nki.language.tiles`: View existing storage as a tile grid.
:func:`nki.language.alloc_blocks`: Allocate tiles grouped into blocks."""
...
def blocks(source: NkiTensor | NkiTileGrid | NkiBlockGrid, *, tile_shape: tuple[int, ...] | None=None, block_shape: tuple[int, ...]) -> NkiBlockGrid:
r"""Create a block view over a tensor where each block contains multiple tiles.
``blocks()`` adds a second level of grouping above the tile grid. A block is
a rectangular group of tiles intended as a coarser transfer or reuse grain,
while ``tile_shape`` remains the compute grain presented to ``nisa.*``
operations. When one native access pattern can represent the complete
block, ``block.data`` allows it to participate in one aggregate operation.
``block_shape`` is expressed in tile counts. For a ``4 x 4`` tile grid,
``block_shape=(2, 2)`` creates a ``2 x 2`` block grid:
.. code-block:: text
block grid: 2 x 2
+-------------------+-------------------+
| tile (0,0) (0,1) | tile (0,2) (0,3) |
| tile (1,0) (1,1) | tile (1,2) (1,3) |
+-------------------+-------------------+
| tile (2,0) (2,1) | tile (2,2) (2,3) |
| tile (3,0) (3,1) | tile (3,2) (3,3) |
+-------------------+-------------------+
Blocks are logical views over the source. Creating, regrouping, slicing,
or indexing a block grid allocates no storage, copies no data, and emits no
instruction.
.. note::
Indexing preserves :class:`NkiBlockGrid`. Call
:func:`nki.language.tiles` on a selected block to expose its interior
tile coordinates, then use :attr:`NkiGrid.data` when an instruction
requires one native tensor operand.
A block containing nonuniform boundary tiles may need to be split into
uniform tile regions first.
:param source: A direct :class:`NkiTensor`, :class:`NkiTileGrid`, or
:class:`NkiBlockGrid` whose storage and selected addresses are shared by
the returned view.
:param tile_shape: The nominal source-element shape of one compute tile.
Required for an :class:`NkiTensor`; optional for an existing grid.
:param block_shape: The positive tile counts grouped along each block
dimension. It must be a tuple with the same rank as the effective
``tile_shape``.
:return: An :class:`NkiBlockGrid` over the same storage. Its ``shape``
counts blocks, while ``tile_shape`` and ``block_shape`` preserve both
levels of the hierarchy. Partial boundary blocks and tiles are retained
at their exact logical extents; aggregate ``data`` access still requires
one exact native access pattern.
:raises TypeError: if ``source`` is unsupported, ``tile_shape`` is omitted
for an :class:`NkiTensor`, or a shape is not a tuple of exact Python
integers.
:raises ValueError: if a shape entry is nonpositive, ``block_shape`` and
``tile_shape`` have different ranks, the source cannot be tiled, or
retiling or grouping cannot preserve the selected addresses.
Examples:
.. code-block:: python
import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def scale_by_block(src):
dst = nl.ndarray(
src.shape, dtype=src.dtype, buffer=nl.shared_hbm
)
src_blocks = nl.blocks(
src,
tile_shape=(128, 256),
block_shape=(2, 2),
)
dst_blocks = nl.blocks(
dst,
tile_shape=(128, 256),
block_shape=(2, 2),
)
for i in range(src_blocks.shape[0]):
for j in range(src_blocks.shape[1]):
src_block = src_blocks[i, j]
dst_block = dst_blocks[i, j]
block = nl.blocks(
nl.ndarray(
src_block.element_shape,
dtype=src.dtype,
buffer=nl.sbuf,
),
tile_shape=src_block.tile_shape,
block_shape=src_block.block_shape,
)
# copy entire block in single instruction
nisa.dma_copy(dst=block.data, src=src_block.data)
# demote to tile-grid to access individual tiles
block_tiles = nl.tiles(block)
for k in range(block_tiles.shape[0]):
for p in range(block_tiles.shape[1]):
tile = block_tiles[k, p].data
nisa.tensor_scalar(
dst=tile,
data=tile,
op0=nl.multiply,
operand0=2.0,
)
nisa.dma_copy(
dst=dst_block.data,
src=block.data,
)
return dst
**Source kinds.**
The source can be ``NkiTensor`` or ``NkiTileGrid`` or ``NkiBlockGrid``.
.. code-block:: text
source tile_shape behavior
------------------ ----------------- ------------------------------
NkiTensor required tile, then group
NkiTileGrid omitted group existing tiles
NkiTileGrid supplied retile, then group
NkiBlockGrid omitted remove old groups, regroup
NkiBlockGrid supplied remove groups, retile, regroup
**Remainder blocks.**
Tile counts are ceiling-divided from the source extent, and block counts
are ceiling-divided from the tile counts. A final block can therefore
contain fewer tiles, and its final tile can contain fewer source elements.
Both levels are clamped to the exact source extent.
See Also:
:func:`nki.language.tiles`: Create a tile grid or expose block tiles.
:class:`NkiBlockGrid`: Block-grid metadata and indexing behavior."""
...
[docs]def ones(shape, dtype, buffer=sbuf, name=''):
r"""Create a new tensor of given shape and dtype on the specified buffer, filled with ones.
((Similar to `numpy.ones <https://numpy.org/doc/stable/reference/generated/numpy.ones.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param shape: the shape of the tensor.
:param dtype: the data type of the tensor.
:param buffer: the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:return: a new :class:`NkiTensor` allocated on the buffer."""
...
[docs]def full(shape, fill_value, dtype, buffer=sbuf, name=''):
r"""Create a new tensor of given shape and dtype on the specified buffer, filled with initial value.
.. warning::
This API is experimental and may change in future releases.
:param shape: the shape of the tensor.
:param fill_value: the value to fill the tensor with.
:param dtype: the data type of the tensor.
:param buffer: the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:return: a new :class:`NkiTensor` allocated on the buffer."""
...
def tiles(source: NkiTensor | NkiTileGrid | NkiBlockGrid, *, tile_shape: tuple[int, ...] | None=None) -> NkiTileGrid:
r"""Decompose a tensor into a logical grid of fixed-size tiles.
The returned grid is a view over ``source``, indexed by
tile coordinate rather than by manually calculated element offsets and
strides. A kernel can select ``src_tiles[i, j]`` to address one tile,
``src_tiles[i, :]`` to address a row of tiles, or ``src_tiles[:, j]`` to
address a column.
For example, ``tile_shape=(128, 256)`` over a ``(512, 1024)`` tensor
produces a ``4 x 4`` grid:
.. code-block:: text
F (free dimension)
256 256 256 256
+---------+---------+---------+---------+
P 128 | (0,0) | (0,1) | (0,2) | (0,3) |
(partition) +---------+---------+---------+---------+
128 | (1,0) | (1,1) | (1,2) | (1,3) |
+---------+---------+---------+---------+
128 | (2,0) | (2,1) | (2,2) | (2,3) |
+---------+---------+---------+---------+
128 | (3,0) | (3,1) | (3,2) | (3,3) |
+---------+---------+---------+---------+
.. note::
Grid indexing preserves the grid abstraction. ``src_tiles[i, j]``
returns a one-tile :class:`NkiTileGrid`, not an :class:`NkiTensor`.
Use ``src_tiles[i, j].data`` when an NKI instruction requires the
underlying tensor operand. Accessing ``data`` is also a view operation.
:param source: The tensor region to tile. It may be a direct
:class:`NkiTensor`, an existing :class:`NkiTileGrid`, or an
:class:`NkiBlockGrid`. The returned grid shares its storage, dtype,
buffer, offsets, and physical strides.
:param tile_shape: The nominal shape of one tile in source-element units,
with at least two dimensions. Every entry must be a positive, exact
Python integer. It is required when ``source`` is an
:class:`NkiTensor` and optional when ``source`` is already a grid.
For an HBM tensor, fewer dimensions than ``source`` leaves the leading
dimensions as batch coordinates. A tile extent may be larger than its
source extent, producing one partial tile.
:return: An :class:`NkiTileGrid` over ``source``. Its ``shape`` counts
addressable tiles, its ``tile_shape`` records the nominal
``tile_shape``, and its ``element_shape`` records the number of selected
elements in each source dimension. Partial boundary tiles expose only
valid elements, although independently allocated on-chip storage may
reserve nominal tile capacity behind them.
:raises TypeError: if ``source`` is unsupported, ``tile_shape`` is omitted
for an :class:`NkiTensor`, a shape is not a tuple of exact Python
integers, or the source buffer is not an NKI :class:`MemoryRegion`.
:raises ValueError: if ``tile_shape`` has fewer than two dimensions,
contains a nonpositive extent, has a rank incompatible with the source,
the source is an indirect :class:`NkiTensor`, the source buffer is
unsupported, or an existing grid cannot be retiled while preserving
its selected addresses.
Examples:
.. code-block:: python
import nki
import nki.isa as nisa
import nki.language as nl
@nki.jit
def scale_by_tile(src):
dst = nl.ndarray(
src.shape, dtype=src.dtype, buffer=nl.shared_hbm
)
src_tiles = nl.tiles(src, tile_shape=(128, 256))
dst_tiles = nl.tiles(dst, tile_shape=(128, 256))
for i in range(src_tiles.shape[0]):
for j in range(src_tiles.shape[1]):
src_tile = src_tiles[i, j]
dst_tile = dst_tiles[i, j]
tile = nl.ndarray(
src_tile.element_shape,
dtype=src.dtype,
buffer=nl.sbuf,
)
nisa.dma_copy(dst=tile.data, src=src_tile.data)
nisa.tensor_scalar(
dst=tile.data, data=tile.data, op0=nl.multiply, operand0=2.0
)
nisa.dma_copy(dst=dst_tile.data, src=tile.data)
return dst
**Source kinds.**
The source can be ``NkiTensor`` or ``NkiTileGrid`` or ``NkiBlockGrid``.
.. code-block:: text
source tile_shape result
------------------ ----------------- ------------------------------
HBM NkiTensor required tile grid over HBM
SBUF NkiTensor required tile grid over SBUF
PSUM NkiTensor required tile grid over PSUM
NkiTileGrid omitted same tile-grid view
NkiTileGrid supplied address-preserving retile
NkiBlockGrid omitted remove block grouping
NkiBlockGrid supplied remove blocks, then retile
Both :class:`NkiBlockGrid` forms remove the block-coordinate level.
Omitting ``tile_shape`` preserves the existing tile structure. Supplying it
rebuilds the tile grid at the requested compute grain. See **Retiling** section
below.
**Remainder tiles.**
Grid dimensions use ceiling division. When a source extent is not an exact
multiple of its tile extent, the final tile is clamped to the remaining
elements. For ``source.shape == (300, 600)`` and
``tile_shape == (128, 256)``:
.. code-block:: text
column 0 column 1 column 2
256 elements 256 elements 88 elements
+--------------+--------------+-----------+
row 0: 128 elements | 128 x 256 | 128 x 256 | 128 x 88 |
+--------------+--------------+-----------+
row 1: 128 elements | 128 x 256 | 128 x 256 | 128 x 88 |
+--------------+--------------+-----------+
row 2: 44 elements | 44 x 256 | 44 x 256 | 44 x 88 |
+--------------+--------------+-----------+
The nominal ``tile_shape`` remains ``(128, 256)``, while
``src_tiles[-1, -1].data.shape`` is ``(44, 88)``. The selected boundary
tile reports ``is_remainder == True``. Its logical extent and ``data``
operand contain no padding.
.. code-block:: python
# check if last row/col tiles are partials
src_tiles[-1].is_remainder
src_tiles[:, -1].is_remainder
src_tiles[-1, -1].is_remainder
See Also:
:func:`nki.language.blocks`: Group tiles into blocks.
:class:`NkiTileGrid`: Tile-grid metadata and indexing.
:attr:`NkiGrid.data`: Access a selected region as an NKI tensor."""
...
[docs]def zeros_like(x, dtype=None, buffer=None, name=''):
r"""Create a new tensor of zeros with the same shape and type as a given tensor.
((Similar to `numpy.zeros_like <https://numpy.org/doc/stable/reference/generated/numpy.zeros_like.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param x: the tensor.
:param dtype: the data type of the tensor.
:param buffer: the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:return: a new :class:`NkiTensor` of zeros with the same shape as ``x``."""
...
[docs]def empty_like(x: NkiTensor | NkiTileGrid | NkiBlockGrid, dtype: DType | None=None, buffer: MemoryRegion | None=None, name: str='', *, transpose_axes: tuple[int, ...] | None=None, address: tuple | None=None) -> NkiTensor | NkiTileGrid | NkiBlockGrid:
r"""Allocate independent, uninitialized storage with matching geometry.
The destination uses the source dtype and buffer unless either is
specified explicitly. Tensor inputs preserve the source shape. Grid inputs
preserve the concrete grid type, navigation shape, tile shape, block
grouping, and exact remainder extents.
A positive-step grid slice can select cells separated by physical gaps.
The destination packs those cells densely while preserving their logical
hierarchy and boundary extents.
``transpose_axes`` prepares an SBUF destination for a later
:func:`nki.isa.dma_transpose`; allocation does not issue the instruction.
Supported tiled-dimension orders are ``(1, 0)``, ``(2, 1, 0)``, and
``(3, 1, 2, 0)`` for ranks two through four. The source free extent that
becomes the destination partition extent must not exceed 128.
:param x: The :class:`NkiTensor`, :class:`NkiTileGrid`, or
:class:`NkiBlockGrid` whose geometry defines the allocation.
:param dtype: Destination element dtype. Defaults to ``x.dtype``.
:param buffer: Destination :class:`MemoryRegion`. Defaults to ``x.buffer``.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:param transpose_axes: Optional tiled-dimension order for a native
:func:`nki.isa.dma_transpose` destination.
:param address: Optional manual address
``(partition_offset, free_offset)`` with the same contract as
:func:`nki.language.ndarray`.
:return: New uninitialized storage with the matching tensor or grid
abstraction. Transpose allocation permutes the tiled dimensions.
:raises TypeError: if ``x`` is unsupported, ``buffer`` is not an NKI
:class:`MemoryRegion`, or ``transpose_axes`` is not a tuple.
:raises ValueError: if transpose rank, order, destination buffer, or
partition extent is unsupported.
:raises AssertionError: if manual ``address`` validation in
:func:`nki.language.ndarray` fails.
See Also:
:func:`nki.language.alloc_tiles`: Allocate from explicit tile geometry.
:func:`nki.language.alloc_blocks`: Allocate from explicit block geometry."""
...
[docs]def random_seed(seed):
r"""Set the random seed for random number generation.
Using the same seed will generate the same sequence of random numbers
when used with ``rand()``.
.. warning::
This API is experimental and may change in future releases.
:param seed: a [1,1] tensor on SBUF or PSUM with a 32-bit seed value.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.random_seed -- set seed for reproducible random values
seed = nl.full((1, 1), 42, dtype=nl.int32, buffer=nl.sbuf)
nl.random_seed(seed)
a = nl.rand((128, 512), dtype=nl.float32)
# nki.language.random_seed -- same seed produces same values
seed = nl.full((1, 1), 42, dtype=nl.int32, buffer=nl.sbuf)
nl.random_seed(seed)
a = nl.rand((128, 512), dtype=nl.float32)
nl.random_seed(seed)
b = nl.rand((128, 512), dtype=nl.float32)
assert nl.equal(a, b)"""
...
[docs]def expand_dims(x, axis):
r"""Expand the shape of a tile.
((Similar to `numpy.expand_dims <https://numpy.org/doc/stable/reference/generated/numpy.expand_dims.html>`_))
.. warning::
This API is experimental and may change in future releases.
Insert a new axis that will appear at the axis position in the expanded tile shape.
:param x: a tile.
:param axis: position in the expanded axes where the new axis is placed.
:return: a tile with view of input data with the number of dimensions increased."""
...
[docs]def broadcast_to(x, shape, dtype=None):
r"""Broadcast a tile to a new shape following numpy broadcasting rules.
((Similar to `numpy.broadcast_to <https://numpy.org/doc/stable/reference/generated/numpy.broadcast_to.html>`_))
.. warning::
This API is experimental and may change in future releases.
If ``x.shape`` is already the same as ``shape``, returns ``x`` unchanged
(or a dtype-cast copy if ``dtype`` differs).
:param x: the source tile in SBUF or PSUM.
:param shape: the target shape. Must have the same rank as ``x``.
Each dimension must either match or be broadcast from size 1.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile with the target shape containing broadcast values from ``x``."""
...
[docs]def transpose(x, dtype=None):
r"""Transposes a 2D tile between its partition and free dimension.
.. warning::
This API is experimental and may change in future releases.
:param x: 2D input tile.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a tile that has the values of the input tile with its partition and free
dimensions swapped.
Examples:
.. code-block:: python
import nki.isa as nisa
import nki.language as nl
# nki.language.transpose -- transpose of identity is identity
x = nl.shared_identity_matrix(n=128, dtype=nl.float32)
result_psum = nl.transpose(x)
result = nl.ndarray((128, 128), dtype=nl.float32, buffer=nl.sbuf)
nisa.tensor_copy(result, result_psum)
assert nl.equal(result, x)"""
...
[docs]def matmul(x, y, transpose_x=False):
r"""x @ y matrix multiplication of x and y.
.. warning::
This API is experimental and may change in future releases.
:param x: a tile on SBUF (partition dimension <= 128, free dimension <= 128),
x's free dimension must match y's partition dimension.
:param y: a tile on SBUF (partition dimension <= 128, free dimension <= 512).
:param transpose_x: defaults to False. If True, x is treated as already transposed.
If False, an additional transpose will be inserted to make x's partition
dimension the contract dimension of the matmul to align with the Tensor Engine.
:return: x @ y or x.T @ y if transpose_x=True.
Examples:
.. code-block:: python
import nki.isa as nisa
import nki.language as nl
# nki.language.matmul -- identity.T @ ones = ones
x = nl.shared_identity_matrix(n=128, dtype=nl.float32)
y = nl.full((128, 128), 1.0, dtype=nl.float32, buffer=nl.sbuf)
result_psum = nl.matmul(x, y, transpose_x=True)
result = nl.ndarray((128, 128), dtype=nl.float32, buffer=nl.sbuf)
nisa.tensor_copy(result, result_psum)
expected = nl.full((128, 128), 1.0, dtype=nl.float32,
buffer=nl.sbuf)
assert nl.equal(result, expected)"""
...
[docs]def no_reorder():
r"""Prevent the scheduler from reordering operations in this region.
Use as a context manager (``with nl.no_reorder():``) to guarantee that
operations inside the block execute in program order. Without this
directive, the compiler scheduler is free to reorder independent
operations for better hardware utilization.
Register-based dynamic loops (``nl.fori_loop``, ``nl.while_loop``) are
not supported inside a ``no_reorder`` block. Static loops
(``nl.affine_range``, ``nl.sequential_range``, ``nl.static_range``) are
allowed.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.no_reorder -- guarantee execution order
with nl.no_reorder():
a = nl.full((128, 512), 3.0, dtype=nl.float32, buffer=nl.sbuf)
b = nl.full((128, 512), 2.0, dtype=nl.float32, buffer=nl.sbuf)
c = nl.add(a, b)
expected = nl.full((128, 512), 5.0, dtype=nl.float32, buffer=nl.sbuf)
assert nl.equal(c, expected)"""
...
[docs]def load_transpose2d(src: NkiTensor, dtype=None):
r"""Load a tensor from device memory (HBM) and 2D-transpose the data before storing into on-chip memory (SBUF).
.. warning::
This API is experimental and may change in future releases.
:param src: HBM tensor to load the data from.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a new tile on SBUF with values from ``src`` 2D-transposed."""
...
[docs]def gather_flattened(data: NkiTensor, indices: NkiTensor, axis=0, dtype=None):
r"""Gather elements from data tensor using indices after flattening.
This instruction gathers elements from the data tensor using integer indices
provided in the indices tensor. For each element in the indices tensor, it
retrieves the corresponding value from the data tensor using the index value
to select from the free dimension of data.
.. warning::
This API is experimental and may change in future releases.
:param data: input tensor to gather from.
:param indices: indices to gather.
:param axis: axis along which to gather.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: gathered tensor.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.gather_flattened -- gather elements by index
data = nl.load(data_tensor[0:128, 0:512])
indices = nl.load(indices_tensor[0:128, 0:512])
result = nl.gather_flattened(data, indices)
nl.store(actual_tensor[0:128, 0:512], result)"""
...
[docs]def ds(start: Union[int, NkiTensor, VirtualRegister, NkiScalar], size: int):
r"""A ``size``-element window starting at ``start``.
Used inside an index expression: ``t[nl.ds(start, size), :]`` is equivalent
to ``t[start:start + size, :]`` for a compile-time ``start``. ``size`` must
be a compile-time integer (the access pattern needs a static element count);
``start`` may be an ``int`` or a runtime index (see
:meth:`NkiTensor.select`). A runtime ``start`` is only spelled this way,
because a slice needs both bounds while the kernel is built.
.. code-block:: python
t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t[:, nl.ds(8, 16)] # static: columns 8..23
idx = nl.ndarray((1, 1), dtype=nl.uint32, buffer=nl.sbuf)
t[:, nl.ds(idx, 16)] # runtime: 16 columns starting at idx
With a runtime ``start`` the window slides at execution time and the
resulting view is indirect (see :meth:`NkiTensor.is_indirect`).
:param start: window start — compile-time ``int`` or a runtime
``NkiTensor`` / ``VirtualRegister`` / ``NkiScalar``
:param size: window length, must be a compile-time integer
:return: a ``DynamicSlice`` to use inside an index expression"""
...
[docs]def affine_range(start, stop=None, step=1):
r"""Create a sequence for fully unrolled loop iteration.
Create a sequence of numbers for use as loop iterators in NKI, resulting in
a fully unrolled loop. This function is an alias for Python's ``range()``
function. Prefer using ``range()`` directly instead.
.. warning::
This API is deprecated and will be removed in future releases.
:param start: start value (or stop if ``stop`` is None).
:param stop: stop value (exclusive).
:param step: step size.
:return: an iterator yielding integer values from start to stop.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.affine_range
for i in nl.affine_range(input_tensor.shape[1] // 512):
offset = i * 512
tile = nl.load(input_tensor[0:128, offset:offset+512])
result = nl.multiply(tile, tile)
nl.store(out_tensor[0:128, offset:offset+512], result)"""
...
[docs]def sequential_range(start, stop=None, step=1):
r"""Create a sequence for fully unrolled loop iteration.
Create a sequence of numbers for use as loop iterators in NKI, resulting in
a fully unrolled loop. This function is an alias for Python's ``range()``
function. Prefer using ``range()`` directly instead.
.. warning::
This API is deprecated and will be removed in future releases.
:param start: start value (or stop if ``stop`` is None).
:param stop: stop value (exclusive).
:param step: step size.
:return: an iterator yielding integer values from start to stop.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.sequential_range
for i in nl.sequential_range(input_tensor.shape[1] // 512):
offset = i * 512
tile = nl.load(input_tensor[0:128, offset:offset+512])
result = nl.multiply(tile, tile)
nl.store(out_tensor[0:128, offset:offset+512], result)"""
...
[docs]def static_range(start, stop=None, step=1):
r"""Create a sequence for fully unrolled loop iteration.
Create a sequence of numbers for use as loop iterators in NKI, resulting in
a fully unrolled loop. This function is an alias for Python's ``range()``
function. Prefer using ``range()`` directly instead.
.. warning::
This API is deprecated and will be removed in future releases.
:param start: start value (or stop if ``stop`` is None).
:param stop: stop value (exclusive).
:param step: step size.
:return: an iterator yielding integer values from start to stop.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.static_range -- fully unrolled iteration over tiles
for i in nl.static_range(input_tensor.shape[1] // 512):
offset = i * 512
tile = nl.load(input_tensor[0:128, offset:offset+512])
result = nl.multiply(tile, tile)
nl.store(out_tensor[0:128, offset:offset+512], result)"""
...
[docs]def dynamic_range(start, stop=None, step=1):
r"""Create a sequence for **dynamic** loop iteration.
Create a sequence of numbers for use as **dynamic** loop iterators in NKI.
The loop runs on device with dynamic bounds.
:param start: start value (or stop if ``stop`` is None), can be VirtualRegister.
:param stop: stop value (exclusive), can be VirtualRegister.
:param step: step size, must be a compile-time positive integer (not VirtualRegister).
:return: an iterator yielding integer values from start to stop.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.dynamic_range -- dynamic iteration with runtime bounds
for _ in nl.dynamic_range(1):
tile = nl.load(input_tensor[0:128, 0:512])
result = nl.multiply(tile, tile)
nl.store(out_tensor[0:128, 0:512], result)"""
...
[docs]def fori_loop(lower, upper, body_fun, step=1):
r"""Structured for loop with dynamic bounds.
Executes ``body_fun(i)`` for each iteration value from ``lower`` to
``upper`` (exclusive) with the given ``step``. The body is a callable that
receives the current iteration value.
The loop is roughly equivalent to the following Python code:
.. code-block:: python
for i in range(lower, upper, step):
body_fun(i)
The body receives the current iteration value ``i`` as a VirtualRegister or NkiScalar,
passed by value. Read it inside the body (for example, as a ``scalar_offset``
into a tensor). The loop controls iteration via ``lower``/``upper``/``step``,
so writing to ``i`` has no effect on the loop; carry any other loop state
through SBUF/HBM.
:param lower: start value (int, VirtualRegister, or NkiScalar).
:param upper: end value (exclusive) (int, VirtualRegister, or NkiScalar).
:param body_fun: function ``(i) -> None`` called each iteration with the
current iteration value. ``i`` is a ``NkiScalar`` unless
either ``lower`` or ``upper`` is a ``VirtualRegister``,
then it will be a ``VirtualRegister``
:param step: step size, must be a compile-time positive integer.
:return: None. Side effects in ``body_fun`` persist after the loop.
A bound must be 32-bit: the loop compiles into a compare-and-branch that has
no 64-bit form, so an ``nl.int64`` scalar cannot bound it. A 64-bit scalar can
still be computed inside the body and used as an offset there.
Examples:
.. code-block:: python
import nki.isa as nisa
import nki.language as nl
# nki.language.fori_loop -- counted loop with a runtime (dynamic) upper bound
ub_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=ub_sb, src=ub_input)
ub_reg = nisa.register_alloc()
nisa.register_load(ub_reg, ub_sb)
zeros = nl.zeros((1, N), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(dst=output, src=zeros)
temp = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf)
def body(i):
# `i` is the loop induction register; read it as a dynamic offset
nisa.dma_copy(
dst=temp,
src=data.ap([[1, 1], [1, 1]], scalar_offset=i, indirect_dim=1),
)
nisa.dma_copy(
dst=output.ap([[1, 1], [1, 1]], scalar_offset=i, indirect_dim=1),
src=temp,
)
nl.fori_loop(0, ub_reg, body)"""
...
[docs]def while_loop(init, body_fun):
r"""Structured while loop with a register condition.
Loops while the condition register is nonzero, checking the condition before
each iteration (a true ``while``, not ``do-while``: if ``init`` is zero the
body never runs).
The loop is roughly equivalent to the following Python code:
.. code-block:: python
r = init
while r != 0:
r = body_fun(r)
The body receives the current condition value ``r`` as a VirtualRegister,
passed by value. Read it inside the body (for example, materialize it with
``register_store`` and use it as a ``scalar_offset``). The loop's next
condition is the register the body **returns**; carry any other loop state
through SBUF/HBM.
:param init: Initial condition register (VirtualRegister).
:param body_fun: function ``(r: VirtualRegister) -> VirtualRegister`` called
each iteration with the current condition value; returns the
next condition register.
:return: None. Side effects in ``body_fun`` persist after the loop.
Examples:
.. code-block:: python
import nki.isa as nisa
import nki.language as nl
# nki.language.while_loop -- flag-driven loop with a register condition
val = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf)
nisa.dma_copy(dst=val, src=data)
acc = nl.zeros((1, 1), dtype=nl.float32, buffer=nl.sbuf)
count_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.dma_copy(dst=count_sb, src=count_input)
one_sb = nl.ndarray((1, 1), dtype=nl.int32, buffer=nl.sbuf)
nisa.memset(dst=one_sb, value=1)
reg = nisa.register_alloc()
nisa.register_load(reg, count_sb)
def body(r):
nisa.tensor_tensor(dst=acc, data1=acc, data2=val, op=nl.add)
nisa.tensor_tensor(dst=count_sb, data1=count_sb, data2=one_sb, op=nl.subtract)
nisa.register_load(r, count_sb)
return r
nl.while_loop(reg, body)"""
...
[docs]def ndarray(shape, dtype, buffer=sbuf, name='', address=None):
r"""Create a new tensor of given shape and dtype on the specified buffer.
:param shape: the shape of the tensor.
:param dtype: the data type of the tensor.
:param buffer: the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:param address: optional memory address. For SBUF/PSUM this is a 2-tuple
``(partition_offset, free_offset)``; for HBM it is a 1-tuple
``(free_offset,)``.
:return: a new :class:`NkiTensor` allocated on the buffer."""
...
[docs]def zeros(shape, dtype, buffer=sbuf, name=''):
r"""Create a new tensor of given shape and dtype on the specified buffer, filled with zeros.
((Similar to `numpy.zeros <https://numpy.org/doc/stable/reference/generated/numpy.zeros.html>`_))
.. warning::
This API is experimental and may change in future releases.
:param shape: the shape of the tensor.
:param dtype: the data type of the tensor.
:param buffer: the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:return: a new :class:`NkiTensor` allocated on the buffer."""
...
[docs]def load(src: NkiTensor, dtype=None):
r"""Load a tensor from device memory (HBM) into on-chip memory (SBUF).
.. warning::
This API is experimental and may change in future releases.
:param src: HBM tensor to load the data from.
:param dtype: (optional) data type to cast the output type to (see :ref:`nki-dtype` for more information); if not specified, it will default to be the same as the data type of the input tile.
:return: a new tile on SBUF with values from ``src``."""
...
[docs]def store(dst: NkiTensor, value: NkiTensor):
r"""Store into a tensor on device memory (HBM) from on-chip memory (SBUF).
.. warning::
This API is experimental and may change in future releases.
:param dst: HBM tensor to store the data into.
:param value: an SBUF tile that contains the values to store."""
...
[docs]def program_id(axis=0):
r"""Index of the current SPMD program along the given axis in the launch grid.
:param axis: the axis of the launch grid.
:return: the program id along ``axis``."""
...
[docs]def num_programs(axes=0):
r"""Number of SPMD programs along the given axes in the launch grid.
:param axes: the axes of the launch grid. If not provided, returns the total
number of programs along the entire launch grid.
:return: the number of SPMD programs along ``axes`` in the launch grid."""
...
[docs]def program_ndim():
r"""Number of dimensions in the SPMD launch grid.
:return: the number of dimensions in the launch grid, i.e. the number of axes. 0 if no grid."""
...
[docs]def shared_constant(constant):
r"""Create a tensor in shared HBM initialized with constant data.
The constant is embedded in the compiled binary and loaded to HBM
at model load time. With LNC=2, both cores share the same constant;
the data must not diverge across cores.
Supported element types: float32, float16, bfloat16, int32, int16,
int8, uint32, uint16, uint8, float8_e4m3fn, float8_e5m2,
float8_e8m0fnu.
Packed types (float8_e4m3fn_x4, float8_e5m2_x4, float4_e2m1fn_x4)
and tfloat32 are supported at the MLIR level but not yet tested
end-to-end on hardware.
:param constant: the constant data. Can be a numpy array or a file path
to a ``.npy`` file.
:return: an NkiTensor in shared_hbm containing the constant data."""
...
[docs]def shared_identity_matrix(n, dtype=uint8, dst: NkiTensor | None=None):
r"""Create an identity matrix in SBUF with the specified data type.
This function has the same behavior to :doc:`nki.language.shared_constant <nki.language.shared_constant>` but
is preferred if the constant matrix is an identity matrix. The
compiler will reuse all the identity matrices of the same
dtype in the graph to save space.
:param n: the number of rows (and columns) of the returned identity matrix
:param dtype: the data type of the tensor, default to be ``nl.uint8`` (see :ref:`nki-dtype` for more information).
:return: a new :class:`NkiTensor` which contains the identity tensor
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.shared_identity_matrix -- 128x128 identity matrix
identity = nl.shared_identity_matrix(n=128, dtype=nl.float32)
expected = nl.load(expected_tensor[0:128, 0:128])
assert nl.equal(identity, expected)
nl.store(actual_tensor[0:128, 0:128], identity)"""
...
[docs]def device_print(print_prefix, tensor: NkiTensor):
r"""Print a message with a string prefix followed by the value of a tile.
During kernel execution on hardware, the Neuron Runtime (NRT) exports device-printed tensors
via the NRT debug stream API. By default, setting the environment variable
``NEURON_RT_DEBUG_OUTPUT_DIR`` to a directory path enables the default stream consumer,
which dumps tensor data to that directory. The output is organized as:
``<output_dir>/<print_prefix>/core_<logical_core_id>/<iteration>/``.
In CPU simulation, this prints immediately to stdout.
:param print_prefix: prefix of the print message. Evaluated at trace time; must be a constant string.
:param tensor: tensor to print out. Can be in SBUF or HBM."""
...
[docs]def rand(shape, dtype, buffer=sbuf, name=''):
r"""Create a new tensor of given shape and dtype on the specified buffer, filled with random values.
Values are sampled from a uniform distribution between 0 and 1.
.. warning::
This API is experimental and may change in future releases.
:param shape: the shape of the tensor.
:param dtype: the data type of the tensor (see :ref:`nki-dtype` for more information).
:param buffer: the specific buffer (ie, :doc:`sbuf<nki.language.sbuf>`, :doc:`psum<nki.language.psum>`, :doc:`hbm<nki.language.hbm>`), defaults to :doc:`sbuf<nki.language.sbuf>`.
:param name: the name of the tensor, used in :ref:`scheduling <how-to-scheduling-apis>`.
:return: a new :class:`NkiTensor` allocated on the buffer with random values.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.rand -- generate random values in [0, 1)
a = nl.rand((128, 512), dtype=nl.float32)"""
...
def profiler_region_open(label, color=''):
r"""Mark the start of a named annotation region.
Annotates all instructions between :py:func:`profiler_region_open` and
:py:func:`profiler_region_close` with the given label. These annotations
appear in Neuron Explorer as named regions on the execution timeline.
This API has no execution time cost — it only adds debug metadata during
compilation.
All instructions within a profiler region are associated with that region,
even if the compiler instruction scheduler reorders them to execute before
an instruction that appears in program order prior to
``profiler_region_open()``. The start of a region in the timeline is set
to the execution time of the first instruction associated with the region,
and therefore may cover instructions that were not associated with the
region.
Nested regions are supported: opening a region inside another region
associates instructions with both the inner and outer labels.
Opening a region without closing it is a compilation error.
:param label: name of the annotation region. Must not contain ``,``, ``@``,
``#``, or ``|`` characters.
:param color: optional CSS color string (e.g., ``"#4CAF50"``) for the region
in Neuron Explorer.
Examples:
.. code-block:: python
import nki.language as nl
# nki.language.profiler_region_open -- annotate regions for Neuron Explorer
nl.profiler_region_open("load_phase", color="#4CAF50")
tile = nl.ndarray((128, 128), dtype=nl.float16, buffer=nl.sbuf)
nisa.dma_copy(tile, a[0:128, 0:128])
nl.profiler_region_close()
nl.profiler_region_open("compute_phase", color="#2196F3")
dst = nl.ndarray((128, 128), dtype=nl.float16, buffer=nl.sbuf)
nisa.tensor_tensor(dst, tile, tile, op=nl.add)
nl.profiler_region_close()"""
...
def profiler_region_close():
r"""Mark the end of the most recently opened annotation region.
Must be paired with a preceding :py:func:`profiler_region_open` call.
Failing to close a region will result in a compilation error."""
...
