nrf_axon_nn_infer.h

Go to the documentation of this file.

/*
 * Copyright (c) 2025 Nordic Semiconductor ASA
 *
 * SPDX-License-Identifier: LicenseRef-Nordic-5-Clause
 */
 
#pragma once
 
#ifdef __cplusplus
extern "C" {
#endif
#include <stdarg.h>
#include <stdint.h>
#include "nrf_axon_driver.h"
 
#ifndef NRF_AXON_MODEL_APP_STORAGE
# define NRF_AXON_MODEL_APP_STORAGE
#endif
typedef struct {
        uint16_t height;
        uint16_t width;
        uint16_t channel_cnt;
        int16_t batch_cnt;
        uint8_t byte_width;
} nrf_axon_nn_model_layer_dimensions_s;
 
typedef struct nrf_axon_nn_compiled_model_input_tag_s {
        int8_t *ptr;
        nrf_axon_nn_model_layer_dimensions_s dimensions;
        uint32_t quant_mult;
        int16_t node_id;
        uint16_t stride;
        uint8_t quant_round;
        int8_t quant_zp;
} nrf_axon_nn_compiled_model_input_s;
 
typedef struct {
        int8_t *buf_ptr;
        uint32_t buf_size;
        int32_t initial_value;
        uint8_t byte_width;
} nrf_axon_nn_model_persistent_var_s;
 
typedef struct nrf_axon_compiled_model_output_tag_s {
        int8_t *ptr;
        /* offset into the packed buffer of this output */
        uint32_t packed_buffer_offset;
        /* packed size of this output */
        uint32_t packed_size;
        nrf_axon_nn_model_layer_dimensions_s dimensions;
        /* dequantization multiplier. */
        uint32_t dequant_mult;
        /* node id of the output */
        int16_t node_id;
        uint8_t dequant_round;
        int8_t dequant_zp;
        /*
         * length in bytes of the distance between the start of rows in the unpacked
         * output in the interlayer byffer.
         */
        uint16_t stride;
} nrf_axon_compiled_model_output_s;
 
typedef struct nrf_axon_nn_compiled_model_tag_s  {
        /*
         * version of the compiler that generated the model. bits 23:16 => major,
         * bits 15:8 => minor, bits 8:0 => patch
         */
        uint32_t compiler_version;
        /*  name of the model provided by user at compilation time.*/
        const char *model_name;
        /*
         * optional list of text labels that correspond to classification indices.
         * Applies to single dimension classification models only.
         */
        const char **labels;
        const nrf_axon_nn_compiled_model_input_s *inputs;
        /*
         * list of ptrs to model input vectors that is sized to input_cnt that
         * is provisioned by the model.
         * Each index is the input vector for the corresponding entry in inputs[].
         * User is required to populate this list before calling inference APIs.
         */
        const int8_t **input_vector_list;
        int16_t input_cnt;
        uint32_t interlayer_buffer_needed;
        uint32_t psum_buffer_needed;
        const NRF_AXON_PLATFORM_BITWIDTH_UNSIGNED_TYPE *cmd_buffer_ptr;
        const void *model_const_ptr;
        uint32_t model_const_size;
        uint32_t cmd_buffer_len;
        struct {
                int8_t *buf_ptr;
                uint32_t buf_size;
                const nrf_axon_nn_model_persistent_var_s *vars;
                uint16_t count;
        } persistent_vars;
        /* number of model outputs pointed to by outputs */
        uint16_t output_cnt;
        /* pointer to the outputs of the model */
        const nrf_axon_compiled_model_output_s *outputs;
        int8_t *packed_output_buf;
        /* model requires this version of the driver or later to execute properly. */
        uint32_t min_driver_version_required;
        /*
         * layer model is a superset of the full model. If true, this can be treated as a
         * nrf_axon_nn_compiled_model_layer_s.
         */
        bool is_layer_model;
} nrf_axon_nn_compiled_model_s;
 
static inline void nrf_axon_nn_set_input_vector(const nrf_axon_nn_compiled_model_s *compiled_model,
        uint16_t input_no, const int8_t *input_vector)
{
        if (input_no < compiled_model->input_cnt) {
                compiled_model->input_vector_list[input_no] = input_vector;
        }
}
 
nrf_axon_result_e nrf_axon_nn_model_validate(
        const nrf_axon_nn_compiled_model_s *compiled_model);
 
 
int nrf_axon_nn_model_init_vars(
        const nrf_axon_nn_compiled_model_s *compiled_model);
 
nrf_axon_result_e nrf_axon_nn_model_infer_sync(
        const nrf_axon_nn_compiled_model_s *compiled_model,
        const int8_t *input_vector,
        int8_t *output_buffer);
 
nrf_axon_result_e nrf_axon_nn_model_infer_sync_multi_inputs(
        const nrf_axon_nn_compiled_model_s *compiled_model,
        int8_t *output_buffer);
 
typedef enum {
        NRF_AXON_NN_ASYNC_INFERENCE_STATUS_IDLE,
        NRF_AXON_NN_ASYNC_INFERENCE_STATUS_ACTIVE,
        NRF_AXON_NN_ASYNC_INFERENCE_STATUS_COMPLETE,
} nrf_axon_nn_async_inference_status_e;
 
typedef struct {
        const nrf_axon_nn_compiled_model_s *compiled_model;
        nrf_axon_cmd_buffer_info_s cmd_buf_info;
        nrf_axon_queued_cmd_info_wrapper_s queued_cmd_buf_wrapper;
        void (*inference_callback)(nrf_axon_result_e result, void *callback_context);
        void *callback_context;
        /*
         * populated and managed by the driver. dedicated buffers outside the interlayer buffer
         * for storing model outputs.
         */
        int8_t *output_buffer;
        nrf_axon_nn_async_inference_status_e infer_status;
} nrf_axon_nn_model_async_inference_wrapper_s;
 
nrf_axon_result_e nrf_axon_nn_model_async_init(
        nrf_axon_nn_model_async_inference_wrapper_s *model_wrapper,
        const nrf_axon_nn_compiled_model_s *compiled_model);
 
nrf_axon_nn_async_inference_status_e nrf_axon_nn_get_model_async_infer_status(
        const nrf_axon_nn_model_async_inference_wrapper_s *model_wrapper);
 
/*
 * @brief Starts an asynchronous inference on the provided model.
 
 * In asynchonous mode, models are inferred in a separate thread, one after another.
 * User provides input vector and output buffer information as the interlayer buffer
 * is used by all models, so only when it is this model's turn to execute can its
 * input be populated from the input_vector.
 *
 * Upon completion of inference the next job is queued before the user callback is invoked,
 * so the results have to be copied by the driver to the output_buffer before invoking the
 * user callback.
 *
 * @param[in] model_wrapper Model to run inference on, initialized via a one-time call to
 *            nrf_axon_nn_model_async_init.
 * @param[in] input_vector Input to run inference on. It is not consumed immediately so has to
 *            be in memory that is valid as long as inference is occurring.
 * @param[in] output_buffer buffer to copy inference results to.
 * @param[in] inference_callback Function to invoke when inference has completed.
 * @param[in] callback_context Opaque pointer provided to inference_callback.
 * @retval[0] Inference successfully queued.
 * @retval[NRF_AXON_RESULT_NOT_FINISHED] Model is still busy with an ealier inference
 * @retval[<0] Error code.
 */
nrf_axon_result_e nrf_axon_nn_model_infer_async(
        nrf_axon_nn_model_async_inference_wrapper_s *model_wrapper,
        const int8_t *input_vector,
        int8_t *output_buffer,
        void (*inference_callback)(nrf_axon_result_e result, void *callback_context),
        void *callback_context);
 
nrf_axon_result_e nrf_axon_nn_model_infer_async_multi_inputs(
        nrf_axon_nn_model_async_inference_wrapper_s *model_wrapper,
        const int8_t **input_vector_list,
        int8_t *output_buffer,
        void (*inference_callback)(nrf_axon_result_e result, void *callback_context),
        void *callback_context);
int16_t nrf_axon_nn_get_classification(
        const nrf_axon_nn_compiled_model_s *compiled_model,
        const int8_t *packed_output,
        const char **label, int32_t *score);
 
nrf_axon_result_e nrf_axon_nn_populate_input_vectors(
        const nrf_axon_nn_compiled_model_s *compiled_model);
 
int nrf_axon_nn_offset_to_output_ndx(
        const nrf_axon_nn_compiled_model_s *compiled_model,
        uint8_t output_ndx);
 
void nrf_axon_nn_copy_output_to_packed_buffer(
        const nrf_axon_nn_compiled_model_s *compiled_model,
        int8_t *to_buffer);
 
#ifdef __cplusplus
} /* extern "C" { */
#endif