【爱芯派 Pro 开发板试用体验】用rust写模型runner

2023-12-2 22:00:47

475 爱芯派 Rust

之前学习用的都是官网写好的模型runner，这里尝试用rust自己写一个runner，主要是为了熟悉Axera的一些API以便后期扩展。这里主要用到AxEngine，他是神经网络模型芯片侧推理计算库，能够完成模型加载到执行的全部推理任务。这次主要以获取版本为例，介绍一下如何用rust接入以及相关API，完整的runner待后期完善，本例中仅仅涉及三个API，AX_ENGINE_Init，AX_ENGINE_Deinit和AX_ENGINE_GetVersion。

模型runner主要API以及调用流程，请参考文档《51 - AX ENGINE API 使用说明.pdf》，特别是调用流程，也可以参考一些C实现的example。主要用到了libax_engine.so动态库，提供的函数不多。

主要参考代码如下：

extern crate libc;

use libc::c_char;
use std::ffi::CStr;

// 指定动态库文件名
#[link(name = "ax_engine")]
extern "C" {
    // 声明需要调用的函数签名
    fn AX_ENGINE_Init(npu_attr: &AxEngineNpuAttrT) -> i32;
    fn AX_ENGINE_Deinit() -> i32;
    fn AX_ENGINE_GetVersion() -> *const c_char;
}

#[repr(C)]
pub struct AxEngineNpuAttrT {
    e_hard_mode: AxEngineNpuModeT,
    reserve: [u32; 8],
}

#[repr(u8)]  // Use u8 so it matches with C enum size
pub enum AxEngineNpuModeT {
    AxEngineVirtualNpuDisable = 0, // virtual npu disable
    AxEngineVirtualNpuStd=1,
    AxEngineVirtualNpuBigLittle = 2,   
    AxEngineVirtualNpuButt = 3,   
}

fn main() {
    let init_param = AxEngineNpuAttrT {
        e_hard_mode: AxEngineNpuModeT::AxEngineVirtualNpuDisable,
        reserve: Default::default(),
    };

    // 初始化推理引擎资源
    let ret = unsafe { AX_ENGINE_Init(&init_param) };
    println!("init engine: {}", ret);

    // 获取Engine版本号
    let c_buf: *const c_char = unsafe {   AX_ENGINE_GetVersion()  };
    let c_str: &CStr = unsafe {  CStr::from_ptr(c_buf) };
    let ver: &str = c_str.to_str().unwrap();
    println!("{}", ver);

    // 释放推理引擎资源
    let ret = unsafe { AX_ENGINE_Deinit() };
    println!("deinit engine:{}", ret);
}