This document is relevant for: Inf2, Trn1, Trn2

nki.simulate_kernel#

nki.simulate_kernel(kernel, *args, **kwargs)[source]#

Simulate a nki kernel on CPU using a built-in simulator in Neuron Compiler. This simulation mode is especially useful for inspecting intermediate tensor values using nki.language.device_print (see code example below).

Note

All input and output tensors to the kernel must be numpy.ndarray when using this simulate_kernel API.

To run the kernel on a NeuronCore instead, please refer to Getting Started with NKI.

Parameters:

  • kernel – The kernel to be simulated

  • args – The args of the kernel

  • kwargs – The kwargs of the kernel

Returns:

Examples:

import neuronxcc.nki as nki
import neuronxcc.nki.language as nl
import numpy as np


@nki.jit
def print_kernel():
  a = nl.ndarray([4, 4], dtype=nl.float32, buffer=nl.shared_hbm)

  # Create (4, 4) tensor in sbuf
  y = nl.zeros([4, 4], dtype=np.float32)

  # Print tensor y
  nl.device_print("value of y:", y)

  # Directly store tensor y as a single tile
  nl.store(a, value=y)

  return a

np.random.seed(0)

a = nki.simulate_kernel(print_kernel)

assert np.allclose(a, np.zeros([4, 4]))

This document is relevant for: Inf2, Trn1, Trn2