120 lines
3.0 KiB
Rust
120 lines
3.0 KiB
Rust
#![feature(generic_const_exprs)]
|
|
#![no_std]
|
|
#![no_main]
|
|
|
|
extern crate alloc;
|
|
|
|
use aarch64_cpu::asm::barrier;
|
|
use alloc::{boxed::Box, rc::Rc};
|
|
use core::{
|
|
cell::RefCell,
|
|
fmt::Write,
|
|
mem::MaybeUninit,
|
|
panic::PanicInfo,
|
|
sync::atomic::{compiler_fence, Ordering},
|
|
};
|
|
use embedded_alloc::Heap;
|
|
use half::f16;
|
|
use nalgebra::Matrix;
|
|
use pim::{
|
|
array::{DummyArray, PimMatrixArena, PimStorage, NUMBER_OF_BANKS},
|
|
kernel::{execute_matrix_multiply_rowwise, MATRIX_MUL},
|
|
state::PimState,
|
|
vector::{F16x1, F16x16},
|
|
};
|
|
use uart::Uart0;
|
|
|
|
mod boot;
|
|
mod critical_section;
|
|
mod m5ops;
|
|
mod pim;
|
|
mod uart;
|
|
|
|
#[global_allocator]
|
|
static PIM_ALLOC: Heap = Heap::empty();
|
|
|
|
const PIM_ARENA_SIZE: usize = 0x2000000;
|
|
|
|
#[link_section = ".pim_data"]
|
|
static mut PIM_ARENA: [MaybeUninit<u8>; PIM_ARENA_SIZE] = [MaybeUninit::uninit(); PIM_ARENA_SIZE];
|
|
|
|
#[no_mangle]
|
|
pub extern "C" fn entry() -> ! {
|
|
unsafe {
|
|
PIM_ALLOC.init(PIM_ARENA.as_ptr() as usize, PIM_ARENA_SIZE);
|
|
}
|
|
|
|
let mut uart = Uart0;
|
|
let mut pim_state = PimState::new(&MATRIX_MUL);
|
|
pim_state.set_kernel();
|
|
|
|
let pim_matrix_arena0 = Rc::new(RefCell::new(PimMatrixArena(
|
|
[[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
|
|
)));
|
|
let pim_matrix_arena1 = Rc::new(RefCell::new(PimMatrixArena(
|
|
[[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
|
|
)));
|
|
let pim_matrix_arena2 = Rc::new(RefCell::new(PimMatrixArena(
|
|
[[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
|
|
)));
|
|
let pim_storage0 = PimStorage {
|
|
arena: &pim_matrix_arena0,
|
|
index: 0,
|
|
};
|
|
let pim_storage1 = PimStorage {
|
|
arena: &pim_matrix_arena1,
|
|
index: 0,
|
|
};
|
|
let pim_storage2 = PimStorage {
|
|
arena: &pim_matrix_arena2,
|
|
index: 0,
|
|
};
|
|
|
|
writeln!(
|
|
&mut uart,
|
|
"arena0: {:?}\narena1: {:?}\narena2: {:?}",
|
|
core::ptr::addr_of!(*pim_matrix_arena0.borrow()),
|
|
core::ptr::addr_of!(*pim_matrix_arena1.borrow()),
|
|
core::ptr::addr_of!(*pim_matrix_arena2.borrow()),
|
|
)
|
|
.unwrap();
|
|
|
|
let mut matrix0 = Matrix::from_data(pim_storage0);
|
|
matrix0.fill_lower_triangle(F16x1(f16::ONE), 0);
|
|
|
|
let mut matrix1 = Matrix::from_data(pim_storage1);
|
|
matrix1.fill_lower_triangle(F16x1(f16::ONE), 0);
|
|
|
|
let matrix2 = Matrix::from_data(pim_storage2);
|
|
|
|
writeln!(&mut uart, "{matrix0} * {matrix1}\n=").unwrap();
|
|
|
|
let mut dummy_array = Box::new(DummyArray([F16x16::default(); NUMBER_OF_BANKS]));
|
|
barrier::dsb(barrier::SY);
|
|
|
|
execute_matrix_multiply_rowwise(
|
|
&mut pim_state,
|
|
&mut pim_matrix_arena0.borrow_mut(),
|
|
&mut pim_matrix_arena1.borrow_mut(),
|
|
&mut pim_matrix_arena2.borrow_mut(),
|
|
dummy_array.as_mut(),
|
|
);
|
|
|
|
writeln!(&mut uart, "{matrix2}").unwrap();
|
|
|
|
m5ops::exit();
|
|
|
|
loop {
|
|
compiler_fence(Ordering::SeqCst);
|
|
}
|
|
}
|
|
|
|
#[panic_handler]
|
|
fn panic(info: &PanicInfo) -> ! {
|
|
writeln!(Uart0, "{info}").unwrap();
|
|
|
|
loop {
|
|
compiler_fence(Ordering::SeqCst);
|
|
}
|
|
}
|