Axon integration

Axon is a neural processing unit (NPU) and software stack that accelerates neural networks and selected signal‑processing algorithms, providing a compiler, driver, simulator, and tools to run workloads either on a Zephyr‑based SoC or on a host‑based simulator.

Integration prerequisites

Integration overview

The Axon NPU is a peripheral processor that runs independently of the CPU. Use the Axon driver to control the Axon hardware and run inference workloads. The driver provides the following functionality:

  • Initializing the Axon hardware and driver.

  • Submitting jobs to the hardware and handling hardware events.

  • Running inference in synchronous or asynchronous mode.

  • Executing intrinsics, which are small, pre-compiled Axon code snippets that perform limited functions.

The driver also includes wrapper and test functions for managing compiled AI models, which are provided in source form. The Axon driver is implemented as a platform-independent library, while all platform-specific behavior is handled by the nRF Axon platform component. This separation allows the same driver library and related components, including samples, simulator, and tool chain files, to be built and used on Zephyr, in a software simulator running on the host machine, or in bare-metal environments.

The Axon NPU can accelerate neural networks as well as some feature extraction algorithms. Applications can target either the SoC running Zephyr or a host-based software simulator, depending on the development and testing workflow. The workflows for neural network models and algorithms differ, even though they share many of the same components.

The Axon compiler incorporates the user-supplied model_name into all model-specific symbols, macros, and file names.

Compiled model

The compiled model is placed in a header file of the name nrf_axon_model_<model_name>_h. This file declares the models parameters and compiled code, then encapsulates the model in a nrf_axon_nn_compiled_model_s instance of name nrf_axon_model_<model_name>. Structure nrf_axon_nn_compiled_model_s is declared in include/drivers/nrf_axon_nn_infer.h.

This structure provides all the model’s meta data:

  • Model name as text

  • Input and output dimensions, locations, quantization parameters

  • Pre-allocated output buffer

  • Memory

Memory Management

A global, common buffer nrf_axon_interlayer_buffer is used to store input, intermediate results, and output. This buffer is owned by whichever model is executing. This buffer must be sized to the requirement of largest need of all the models in the system. Models store their interlayer buffer requirement in nrf_axon_model_<model_name>.interlayer_buffer_needed and perform a run-time check during model initialization. Users can see this value declared in the model header file macro NRF_AXON_MODEL_<MODEL_NAME>MAX_IL_BUFFER_USED, then update the allocation in NRF_AXON_INTERLAYER_BUFFER_SIZE.

The synchronous and asynchronous inference APIs accept as parameters the input and output buffers, and will fill/drain the interlayer buffer with input/output in a thread safe manner.

Power Management

The Axon NPU is automatically put in a low power state when not in use. You do not need to complete any additional power management steps.

Other System Resources

The Axon NPU driver executes both in the caller’s thread and in a workqueue. Jobs are initiated in the caller’s thread, interrupts are processed in the workqueue, and in synchornous mode, job completion is signaled with a semaphore. In asynchronous mode, your callback is invoked on job completion; you are responsible for signaling your own thread. A mutex is used to serialize access to the Axon NPU hardware. The workqueue, interrupt, semaphore, and mutex are all initialized by the function nrf_axon_platform_init, which is called once at start up. The initialization process is described in the following sections.

Integration steps

Complete the following steps:

  1. Initializing Axon driver

  2. Initializing model

  3. Executing inference

  4. Integrating the model into an application

Initializing driver

Follow these steps to initialize the Axon driver:

  1. Call the platform initialization function

    nrf_axon_platform_init()

    This function is platform-specific, but you must provide the Axon base address (nrf_axon_driver_init(base_address). You can obtain base_address from the device tree on Zephyr.

    During initialization, the driver powers on Axon by calling the nrf_axon_platform_vote_for_power() function. The driver then verifies that Axon NPU exists at the specified base address.

    Note

    Do not create or manage a driver handle. Axon is implemented as a singleton, and the driver serializes access internally.

  2. Before starting a new inference session on a streaming-style model (where intermediate results are fed forward), initialize the model’s persistent variables:

    nrf_axon_nn_model_init_vars(&my_model_wrapper);

    This sets all persistent variables to their quantized zero-point values.

  3. Refer to further instructions on integrating the driver into your application.

Initializing model

The following section describes a compiled Axon model for synchronous or asynchronous inference, and explains how to initialize each type.

Synchronous model inference

Synchronous inference means that the inference call waits for completion before returning. The synchronous call will wait for the Axon hardware to be available and then claim its exclusive use.

Asynchnronous requests can be made while the Axon is in synchronous mode. These requests will be serviced upon exiting of synchronous mode, regardless of any pending synchronous requests.

To initialize the model, invoke the function nrf_axon_nn_model_validate(&nrf_model_<model_name>) one time at start-up to do basic model validation. This will confirm that the global buffers are large enough to handle the model.

Asynchronous model inference

Compiled models need to be initialized prior to asynchronous inference. The initialization binds the static, compiled model stored in non-volatile memory (NVM) to a RAM wrapper struct that the driver then manages. First, you must declare a static (not on the stack) instance of nrf_axon_nn_model_async_inference_wrapper_s (included in include/drivers/nrf_axon_infer.h), and then invoke the model by initializing the nrf_axon_nn_model_async_init() function:

static nrf_axon_nn_model_async_inference_wrapper_s my_model_wrapper;
nrf_axon_nn_model_async_init(&my_model_wrapper, nrf_axon_model_<model_name>);

This function also verifies if the interlayer buffer is large enough to accommodate the model’s needs.

Executing inference

This section describes the general steps for executing inference using a compiled Axon model. The same high‑level flow applies to both synchronous and asynchronous inference, with differences in how inference is submitted and how completion is handled.

Follow these steps to execute inference with a compiled Axon model:

  1. Ensure the Axon driver has been initialized and the model has been initialized for the selected execution mode.

  2. Prepare the input data in packed format:

    input_vector[input_channel_cnt][input_height][input_width]

    The input to populate externally is identified by nrf_axon_model_<model_name>.external_input_ndx.

    The input element size is defined by inputs[external_input_ndx].byte_width:

    • 1 for int8

    • 2 for int16

    Note

    The data ordering in Axon NPU input and output buffers differs from TFLite. Axon NPU stores data with channels being the outermost dimension, while TFLite stores data with channels as the innermost dimension.

  3. Prepare an output buffer outside of the interlayer buffer, sized to hold the packed output:

    output_buffer[output_channel_cnt][output_height][output_width]

    Output dimension information is available in nrf_axon_model_<model_name>.output_dimensions. The output rank ordering is channels, height, width.

    The model also declares an internal output buffer in nrf_axon_model_<model_name>.packed_output_buf.

    To allocate and use this buffer, define the following macro before including the model header file

    #define NRF_AXON_MODEL_ALLOCATE_PACKED_OUTPUT_BUFFER 1
#include "nrf_axon_model_<model_name>_.h"

  4. Submit the inference request using the appropriate API for the selected execution mode.

    1. Call the following function:

      nrf_axon_nn_model_infer_sync(&nrf_axon_model_<model_name>,
                            input_vector,
                            output_buffer)

      The call blocks until inference completes. The Axon hardware is reserved exclusively for the duration of the call.

    2. Observe that when nrf_axon_nn_model_infer_sync() returns, the output_buffer is populated with the inference results.

    In both modes, the driver fills and drains the interlayer buffer in a thread‑safe manner. The dimensions of the output tensor are provided in the nrf_axon_model_<model_name>.output_dimensions field, and the output data is ordered in channels, height, width format. This ordering differs from TensorFlow Lite, which uses height, width, channels.

Optional variations

The following variations apply to both execution modes:

  • input_vector may be NULL if input data is written directly to the model input buffer.

  • output_buffer may be NULL if output data is read directly from the model output buffer.

  • Either or both buffers may be NULL only if there are no other users of the Axon hardware.

When using NULL buffers, the application is responsible for copying input data into the model and extracting output data from the model.

This approach is useful when the input data is not stored in a contiguous memory block and double buffering would otherwise be required.

To provide input data directly to the model buffer, complete the following steps:

  1. Reserve the Axon hardware explicitly:

    nrf_axon_platform_reserve_for_user();

  2. Obtain the model input address:

    const nrf_axon_nn_compiled_model_input_info_s *model_input =
    nrf_axon_nn_model_1st_external_input(&nrf_axon_model_<model_name>);

  3. Copy the input data to the model input address.

  4. Invoke the inference API with input_vector == NULL.

Asynchronous inference requests may be submitted while a synchronous inference is executing. These requests are queued and serviced once the synchronous execution completes.

Integrating the model into an application

It is recommended to first test your compiled model by using the inference test application (/tests/axon/inference). Once you have verified that the model works correctly on the test application, you can integrate it into your application. It is your responsibility to feed data into the model, schedule inference, and respond to the model output.

Ensure you have completed the following:

  1. Placed the nrf_axon_<model_name>_.h file in the application’s include path.

  2. Included the header files:

    #include "drivers/axon/nrf_axon_driver.h"
#include "drivers/axon/nrf_axon_inference.h"

  3. Included the model header file nrf_axon_model_<model_name>_.h in exactly one source file. The model symbols are intentionally not declared static, to avoid compiling multiple instances of the model into the application.

  4. Updated the following Kconfig values in the application’s prj.conf file:

    • Enable the NRF_AXON Kconfig option.

    • Set NRF_AXON_INTERLAYER_BUFFER_SIZE to the maximum value needed across all models in the application. This value is printed near the top of the compiled model header file.

  5. Initialized driver one time at start-up.

  6. Initialized the model one-time at start-up for the desired mode of execution, synchronous (nrf_axon_nn_model_validate) or asynchronous (nrf_axon_nn_model_async_init).