00_ndarray/
00_ndarray.rs

1use std::time::{Instant};
2use mlinrust::ndarray::NdArray;
3
4// numpy with 4 threads, 7.7 seconds
5// rust version 30 seconds~
6fn main() {
7    const PSIZE: usize = 512 * 512 * 512;
8    let mut a = NdArray::new((0..PSIZE).map(|i| i as f32).collect::<Vec<f32>>());
9    a.reshape(vec![8, 128, 256, 512]);
10    a += 1.0;
11    a /= 0.5;
12
13    let mut b = NdArray::new((0..PSIZE).map(|i| i as f32).collect::<Vec<f32>>());
14    b.reshape(vec![128, 512, 2048]);
15    b -= 1.5;
16    b = &b + &b;
17
18
19    let start = Instant::now();
20    let _ = a * b; // large matrix multiplication
21    let dur = start.elapsed();
22    println!("execute {dur:?}");
23}