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.
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