This document is relevant for: Inf2, Trn1, Trn1n

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:

"""
Copyright (C) 2024, Amazon.com. All Rights Reserved

"""
import unittest

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

nki_jit = nki.trace

@nki_jit
def print_kernel(a):
  i0 = nl.arange(4)[:, None]
  i1 = nl.arange(4)[None, :]

  # 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[i0, i1], value=y)

class TestNkiIsaExamplesSimulateKernel(unittest.TestCase):
  def test_simulate_kernel(self):
    np.random.seed(0)
    a = np.random.random_sample([4, 4]).astype(np.float32)

    nki.simulate_kernel(print_kernel, a)

    self.assertTrue(np.allclose(a, np.zeros([4, 4])))

This document is relevant for: Inf2, Trn1, Trn1n