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"
 
 
typedef struct {
        uint16_t height;
        uint16_t width;
        uint16_t channel_cnt;
        uint8_t byte_width;
} nrf_axon_nn_model_layer_dimensions_s;
 
typedef struct {
        int8_t *ptr;
        nrf_axon_nn_model_layer_dimensions_s dimensions;
        uint32_t quant_mult;
        uint16_t stride;
        uint8_t quant_round;
        int8_t quant_zp;
        bool is_external;
} 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;
        nrf_axon_nn_model_layer_dimensions_s dimensions;
        uint32_t dequant_mult;
        uint8_t dequant_round;
        int8_t dequant_zp;
        uint16_t stride;
} nrf_axon_compiled_model_output_s;
 
#define NRF_AXON_NN_MAX_MODEL_INPUTS 2
typedef struct nrf_axon_nn_compiled_model_tag_s  {
        uint32_t compiler_version;
        const char *model_name;
        const char **labels;
        nrf_axon_nn_compiled_model_input_s inputs[NRF_AXON_NN_MAX_MODEL_INPUTS];
        uint8_t input_cnt;
        int8_t external_input_ndx;
        int8_t *output_ptr;
        int8_t *packed_output_buf;
        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;
        nrf_axon_nn_model_layer_dimensions_s output_dimensions;
 
        uint32_t output_dequant_mult;
        uint8_t output_dequant_round;
        int8_t output_dequant_zp;
        uint16_t output_stride;
        bool is_layer_model;
        uint16_t extra_output_cnt;
        const nrf_axon_compiled_model_output_s *extra_outputs;
        uint32_t min_driver_version_required;
} nrf_axon_nn_compiled_model_s;
 
#define NRF_AXON_COMPILED_MODEL_EXTRA_OUTPUT_CNT(compiled_model) \
        (compiled_model->compiler_version >= 0x10200 ? compiled_model->extra_output_cnt : 0)
 
int8_t nrf_axon_nn_model_1st_external_input_ndx(
        const nrf_axon_nn_compiled_model_s *the_model);
 
const nrf_axon_nn_compiled_model_input_s *nrf_axon_nn_model_1st_external_input(
        const nrf_axon_nn_compiled_model_s *the_model);
 
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);
 
 
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;
        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);
 
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_vector(
        const nrf_axon_nn_compiled_model_s *compiled_model,
        const int8_t *input_vector);
 
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,
        void *to_buffer);
 
#ifdef __cplusplus
} /* extern "C" { */
#endif