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

Example: add two matrices

Browse Source
This commit is contained in:
Derek Christ
2023-12-24 11:27:32 +01:00
parent e49b667b8a
commit 329c0191cc
6 changed files with 65 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -14,8 +14,8 @@ use nalgebra::Matrix;
use pim::{
array::{PimMatrixArena, PimRegion, PimStorage},
kernel::TEST_KERNEL,
matrix::{F16x1, F16x16},
state::PimState,
vector::{F16x1, F16x16},
};
use pim_isa::BankMode;
use uart::Uart0;
@@ -30,21 +30,18 @@ pub extern "C" fn entry() -> ! {
let mut uart = Uart0;
let mut pim_state = PimState::new(&TEST_KERNEL);
let pim_matrix_arena = RefCell::new(PimMatrixArena([[F16x16::default(); 8]; 8]));
let pim_matrix_arena0 = RefCell::new(PimMatrixArena([[F16x16::default(); 8]; 8]));
let pim_matrix_arena1 = RefCell::new(PimMatrixArena([[F16x16::default(); 8]; 8]));
let pim_storage0 = PimStorage {
arena: &pim_matrix_arena,
arena: &pim_matrix_arena0,
index: 0,
};
let pim_storage1 = PimStorage {
arena: &pim_matrix_arena,
index: 1,
};
let pim_storage2 = PimStorage {
arena: &pim_matrix_arena,
index: 2,
arena: &pim_matrix_arena1,
index: 0,
};
let mut matrix0 = Matrix::from_data(pim_storage0);
let matrix1 = Matrix::from_data(pim_storage1);
matrix0.fill_column(0, F16x1(f16::ZERO));
matrix0.fill_column(1, F16x1(f16::ONE));
matrix0.fill_column(2, F16x1(f16::PI));
@@ -54,36 +51,31 @@ pub extern "C" fn entry() -> ! {
matrix0.fill_column(6, F16x1(f16::LN_2));
matrix0.fill_column(7, F16x1(f16::LN_10));
writeln!(
&mut uart,
"Cache Lines: {}\nRows: {}",
PimMatrixArena::<8, 8>::OCCUPIED_CACHE_LINES,
PimMatrixArena::<8, 8>::OCCUPIED_ROWS
)
.unwrap();
writeln!(&mut uart, "{matrix0} * 2\n=").unwrap();
let mut matrix1 = Matrix::from_data(pim_storage1);
matrix1.fill_lower_triangle(F16x1(f16::ONE), 0);
writeln!(&mut uart, "{matrix0} + {matrix1}\n=").unwrap();
// Invalidate and flush array just in case
pim_matrix_arena.borrow_mut().invalidate_flush();
// dummy_array.invalidate_flush();
pim_matrix_arena0.borrow_mut().invalidate_flush();
pim_matrix_arena1.borrow_mut().invalidate_flush();
barrier::dsb(barrier::SY);
pim_state.set_bank_mode(BankMode::PimAllBank);
pim_matrix_arena
.borrow_mut()
pim_matrix_arena0
.borrow()
.execute_instruction_read_dual_bank();
pim_matrix_arena
.borrow_mut()
pim_matrix_arena1
.borrow()
.execute_instruction_read_dual_bank();
pim_matrix_arena
pim_matrix_arena0
.borrow_mut()
.execute_instruction_write_dual_bank();
pim_state.set_bank_mode(BankMode::SingleBank);
pim_matrix_arena.borrow_mut().invalidate();
pim_matrix_arena0.borrow_mut().invalidate();
barrier::dsb(barrier::SY);
// writeln!(&mut uart, "{matrix0}+{matrix1}").unwrap();
writeln!(&mut uart, "{matrix0}").unwrap();
m5ops::exit();

2
pim-os/src/pim.rs
View File

@@ -1,5 +1,5 @@
pub mod array;
pub mod config;
pub mod kernel;
pub mod matrix;
pub mod vector;
pub mod state;

65
pim-os/src/pim/array.rs
View File

@@ -1,4 +1,4 @@
use super::matrix::{F16x1, F16x16};
use super::vector::{F16x1, F16x16};
use aarch64_cpu::asm::barrier;
use core::{arch::asm, cell::RefCell};
use half::f16;
@@ -17,15 +17,15 @@ impl<const R: usize, const C: usize> PimRegion for PimMatrixArena<R, C> {
const OCCUPIED_ROWS: usize = Self::OCCUPIED_CACHE_LINES / NUMBER_OF_BANKS;
fn bank_ptr(&self, bank_index: usize) -> *const f16 {
unsafe { (self.0.as_ptr() as *const F16x16).offset(bank_index as _) as *const f16 }
unsafe { (self.0.as_ptr() as *const F16x16).add(bank_index) as *const f16 }
}
fn bank_ptr_mut(&mut self, bank_index: usize) -> *mut f16 {
unsafe { (self.0.as_mut_ptr() as *mut F16x16).offset(bank_index as _) as *mut f16 }
unsafe { (self.0.as_mut_ptr() as *mut F16x16).add(bank_index) as *mut f16 }
}
}
#[derive(Debug)]
#[derive(Debug, Clone)]
pub struct PimStorage<'a, const R: usize, const C: usize> {
pub arena: &'a RefCell<PimMatrixArena<R, C>>,
pub index: usize,
@@ -38,9 +38,7 @@ unsafe impl<'a, const R: usize, const C: usize> RawStorage<F16x1, Const<R>, Cons
type CStride = Dyn;
fn ptr(&self) -> *const F16x1 {
unsafe {
(&self.arena.borrow().0[0][0] as *const F16x16 as *const F16x1).offset(self.index as _)
}
unsafe { (&self.arena.borrow().0[0][0] as *const F16x16 as *const F16x1).add(self.index) }
}
fn shape(&self) -> (Const<R>, Const<C>) {
@@ -65,8 +63,7 @@ unsafe impl<'a, const R: usize, const C: usize> RawStorageMut<F16x1, Const<R>, C
{
fn ptr_mut(&mut self) -> *mut F16x1 {
unsafe {
(&mut self.arena.borrow_mut().0[0][0] as *mut F16x16 as *mut F16x1)
.offset(self.index as _)
(&mut self.arena.borrow_mut().0[0][0] as *mut F16x16 as *mut F16x1).add(self.index)
}
}
@@ -83,15 +80,17 @@ pub trait PimRegion {
fn bank_ptr_mut(&mut self, bank_index: usize) -> *mut f16;
fn execute_instruction_read_single_bank(&self) {
if !cfg!(feature = "cacheless") {
self.invalidate_bank(EVEN_BANK_INDEX);
for i in (0..Self::OCCUPIED_ROWS).map(|i| i * NUMBER_OF_BANKS) {
if !cfg!(feature = "cacheless") {
self.invalidate_bank(EVEN_BANK_INDEX + i);
barrier::dsb(barrier::SY);
}
// Read from first bank
self.read_data_bank(EVEN_BANK_INDEX + i);
barrier::dsb(barrier::SY);
}
// Read from first bank
self.read_data_bank(EVEN_BANK_INDEX);
barrier::dsb(barrier::SY);
}
fn execute_instruction_read_dual_bank(&self) {
@@ -113,25 +112,32 @@ pub trait PimRegion {
fn read_data_bank(&self, bank_index: usize) {
let bank = self.bank_ptr(bank_index);
// For some reason, this is needed...
use core::fmt::Write;
write!(&mut crate::uart::Uart0 {}, "").unwrap();
unsafe {
core::ptr::read_volatile(bank);
}
}
fn execute_instruction_write_single_bank(&mut self) {
if !cfg!(feature = "cacheless") {
self.preload_zero();
for i in (0..Self::OCCUPIED_ROWS).map(|i| i * NUMBER_OF_BANKS) {
if !cfg!(feature = "cacheless") {
self.preload_zero_bank(EVEN_BANK_INDEX + i);
barrier::dsb(barrier::SY);
}
// Write to first bank
self.write_data_bank(EVEN_BANK_INDEX + i);
if !cfg!(feature = "cacheless") {
self.invalidate_flush_bank(EVEN_BANK_INDEX + i);
}
barrier::dsb(barrier::SY);
}
// Write to first bank
self.write_data_bank(EVEN_BANK_INDEX);
if !cfg!(feature = "cacheless") {
self.invalidate_flush_bank(EVEN_BANK_INDEX);
}
barrier::dsb(barrier::SY);
}
fn execute_instruction_write_dual_bank(&mut self) {
@@ -157,6 +163,11 @@ pub trait PimRegion {
fn write_data_bank(&mut self, bank_index: usize) {
let bank = self.bank_ptr_mut(bank_index);
// For some reason, this is needed...
use core::fmt::Write;
write!(&mut crate::uart::Uart0 {}, "").unwrap();
unsafe {
core::ptr::write_volatile(bank, Default::default());
}

8
pim-os/src/pim/kernel.rs
View File

@@ -18,25 +18,25 @@ pub const TEST_KERNEL: Kernel = Kernel([
dst: File::GrfB { index: 1 },
},
Instruction::ADD {
src0: File::GrfA { index: 0 },
src0: File::Bank,
src1: File::GrfA { index: 0 },
dst: File::GrfA { index: 0 },
aam: false,
},
Instruction::ADD {
src0: File::GrfB { index: 0 },
src0: File::Bank,
src1: File::GrfB { index: 0 },
dst: File::GrfB { index: 0 },
aam: false,
},
Instruction::ADD {
src0: File::GrfA { index: 1 },
src0: File::Bank,
src1: File::GrfA { index: 1 },
dst: File::GrfA { index: 1 },
aam: false,
},
Instruction::ADD {
src0: File::GrfB { index: 1 },
src0: File::Bank,
src1: File::GrfB { index: 1 },
dst: File::GrfB { index: 1 },
aam: false,

2
pim-os/src/pim/state.rs
View File

@@ -28,7 +28,7 @@ impl PimState {
}
self.writer.write(
serde_json_core::to_string::<PimConfig, 1024>(&PimConfig {
serde_json_core::to_string::<PimConfig, 2048>(&PimConfig {
kernel: self.kernel.clone(),
bank_mode,
})

4
pim-os/src/pim/matrix.rs → pim-os/src/pim/vector.rs
View File

@@ -8,13 +8,13 @@ pub struct F16x1(pub f16);
impl core::fmt::Debug for F16x1 {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
Ok(self.0.fmt(f)?)
self.0.fmt(f)
}
}
impl core::fmt::Display for F16x1 {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
Ok(self.0.fmt(f)?)
self.0.fmt(f)
}
}