derek/pim-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Crude implementation of matrices using nalgebra

Browse Source
This commit is contained in:
Derek Christ
2023-12-19 15:37:32 +01:00
parent 2e44890c53
commit 8a2c675d71
7 changed files with 287 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
pim-os/src/main.rs
View File

@@ -1,20 +1,27 @@
#![feature(generic_const_exprs)]
#![no_std]
#![no_main]
use aarch64_cpu::asm::barrier;
use core::{
cell::RefCell,
fmt::Write,
panic::PanicInfo,
sync::atomic::{compiler_fence, Ordering},
};
use half::f16;
use nalgebra::{Const, Matrix, Matrix2, SMatrixView};
use pim::{
array::{BankArray, ComputeArray},
array::{PimMatrixArena, PimStorage},
// array::PimMatrix,
// array::{BankArray, ComputeArray},
kernel::TEST_KERNEL,
matrix::{F16x1, F16x16},
state::PimState,
};
use pim_isa::BankMode;
use uart::Uart0;
mod boot;
mod m5ops;
mod pim;
@@ -22,53 +29,74 @@ mod uart;
#[no_mangle]
pub extern "C" fn entry() -> ! {
let mut pim_state = PimState::new(&TEST_KERNEL);
let mut compute_array: ComputeArray<3> = ComputeArray([
BankArray([f16::from_f32(0.1); 512]),
BankArray([f16::from_f32(0.2); 512]),
BankArray([f16::from_f32(0.3); 512]),
]);
let dummy_array = BankArray::default();
let mut uart = Uart0;
let mut pim_state = PimState::new(&TEST_KERNEL);
writeln!(
&mut uart,
"PIM array is at {:x?}",
core::ptr::addr_of!(compute_array)
)
.unwrap();
let mut arena = RefCell::new(PimMatrixArena([[F16x16::default(); 8]; 8]));
let pim_storage0 = PimStorage {
arena: &arena,
index: 0,
};
let pim_storage1 = PimStorage {
arena: &arena,
index: 1,
};
let pim_storage2 = PimStorage {
arena: &arena,
index: 2,
};
let mut matrix0 = Matrix::from_data(pim_storage0);
let mut matrix1 = Matrix::from_data(pim_storage1);
matrix0.fill_lower_triangle(F16x1(f16::ONE), 0);
matrix1.fill_upper_triangle(F16x1(f16::from_f32(2.0)), 0);
writeln!(&mut uart, "{}", matrix0).unwrap();
writeln!(&mut uart, "{}", matrix1).unwrap();
writeln!(
&mut uart,
"BankArray0: [{:?}, ...]\nBankArray1: [{:?}, ...]\nBankArray2: [{:?}, ...]",
compute_array.0[0].0[0], compute_array.0[1].0[0], compute_array.0[2].0[0]
)
.unwrap();
// let mut compute_array: ComputeArray<3> = ComputeArray([
// BankArray([F16x16([f16::from_f32(0.1); 16]); 32]),
// BankArray([f16::from_f32(0.2); 512]),
// BankArray([f16::from_f32(0.3); 512]),
// ]);
// let dummy_array = BankArray::default();
writeln!(&mut uart, "MAC: BankArray2 += BankArray0 * BankArray1",).unwrap();
// writeln!(
// &mut uart,
// "PIM array is at {:x?}",
// core::ptr::addr_of!(compute_array)
// )
// .unwrap();
// Invalidate and flush array just in case
compute_array.invalidate_flush();
dummy_array.invalidate_flush();
barrier::dsb(barrier::SY);
// writeln!(
// &mut uart,
// "BankArray0: [{:?}, ...]\nBankArray1: [{:?}, ...]\nBankArray2: [{:?}, ...]",
// compute_array.0[0].0[0], compute_array.0[1].0[0], compute_array.0[2].0[0]
// )
// .unwrap();
pim_state.set_bank_mode(BankMode::PimAllBank);
compute_array.0[1].execute_instruction_read_dual_bank();
compute_array.0[2].execute_instruction_read_dual_bank();
compute_array.0[0].execute_instruction_read_dual_bank();
compute_array.0[2].execute_instruction_write_dual_bank();
dummy_array.execute_instruction_read_single_bank();
pim_state.set_bank_mode(BankMode::SingleBank);
// writeln!(&mut uart, "MAC: BankArray2 += BankArray0 * BankArray1",).unwrap();
compute_array.invalidate();
barrier::dsb(barrier::SY);
// // Invalidate and flush array just in case
// compute_array.invalidate_flush();
// dummy_array.invalidate_flush();
// barrier::dsb(barrier::SY);
writeln!(
&mut uart,
"BankArray2: [{:?}, ...]",
compute_array.0[2].0[0]
)
.unwrap();
// pim_state.set_bank_mode(BankMode::PimAllBank);
// compute_array.0[1].execute_instruction_read_dual_bank();
// compute_array.0[2].execute_instruction_read_dual_bank();
// compute_array.0[0].execute_instruction_read_dual_bank();
// compute_array.0[2].execute_instruction_write_dual_bank();
// dummy_array.execute_instruction_read_single_bank();
// pim_state.set_bank_mode(BankMode::SingleBank);
// compute_array.invalidate();
// barrier::dsb(barrier::SY);
// writeln!(
// &mut uart,
// "BankArray2: [{:?}, ...]",
// compute_array.0[2].0[0]
// )
// .unwrap();
// writeln!(&mut uart, "ComputeArray:\n{:#?}", compute_array.0[2]).unwrap();