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_tensor):
  b = nl.empty_like(a_tensor, buffer=nl.hbm)

  # Load tensor into sbuf
  a = nl.load(a_tensor)

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

  # Directly store a into hbm
  nl.store(b, value=a)

  return b

np.random.seed(0)
a = np.random.random_sample([3, 4]).astype(np.float32) * 10

b = nki.simulate_kernel(print_kernel, a)

assert np.allclose(a, b)

This document is relevant for: Inf2, Trn1, Trn2