This document is relevant for: Trn2, Trn3

nki.simulate#

nki.simulate(kernel)[source]#

Create a CPU-simulated version of an NKI kernel.

Warning

This API is experimental and may change in future releases. It has not been tested or confirmed to work on all hardware platforms and operating systems.

Currently, Neuron confirms support for nki.simulate on these 2 operating systems:

  • Ubuntu 22.04

  • Amazon Linux 2023

See NKI CPU Simulator for full documentation including target platform selection, precise floating-point mode, debugging, and known limitations.

Example:

@nki.jit
def my_kernel(a, b): ...

# Explicit simulation
result = nki.simulate(my_kernel)(a_np, b_np)

# With LNC2
result = nki.simulate(my_kernel[2])(a_np, b_np)
Parameters:

kernel – NKI kernel function, typically decorated with @nki.jit. If a plain function is passed, it is automatically wrapped.

Returns:

A callable that, when invoked with NumPy arrays or torch Tensors, executes the kernel on CPU and returns results in the same format (NumPy arrays or torch Tensors respectively).

This document is relevant for: Trn2, Trn3