/rk3588_npu

rust wrapper around @mtx512's reverse engineering of the rk3588 npu https://github.com/mtx512/rk3588-npu

Primary LanguageRust

rk3588_npu

Important

rk3588_npu is a work in progress and is not ready for production use.

This is a small rust wrapper around the awesome work by mtx512 that reverse engineered the Rockchip3588 NPU.

Example Usage

use half::f16;
use rk3588_npu::multiply;

fn main() {
    let m: usize = 32;
    let k: usize = 64;
    let n: usize = 512;

    // Initialize test matrices
    let mut matrix_a = vec![f16::ZERO; m * k];
    let mut matrix_b = vec![f16::ZERO; n * k];

    for i in 0..m * k {
        matrix_a[i] = f16::from_f32((i % 8 + 1) as f32);
    }
    for i in 0..n * k {
        matrix_b[i] = f16::from_f32((i % 4 + 1) as f32);
    }

    // multiply expects matrices in row-major order (i.e. C-style)
    // matrix_a: m x k
    // matrix_b: k x n
    let out = multiply(matrix_a, matrix_b, m, k, n).expect("Failed to multiply matrices");

    // Print the first 10 elements of the output matrix
    for i in 0..10 {
        println!("out[{}] = {}", i, out[i]);
    }
}

and run it with:

cargo run --example basic
#    Compiling rk3588_npu v0.1.0 (/home/orangepi/rk3588_npu)
#     Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.01s
#      Running `target/debug/examples/basic`
# drm name is rknpu - 20240322 - RKNPU driver
# out[0] = 800
# out[1] = 800
# out[2] = 800
# out[3] = 800
# out[4] = 800
# out[5] = 800
# out[6] = 800
# out[7] = 800
# out[8] = 800
# out[9] = 800