Example: add two matrices

pim-os/src/main.rs

@@ -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,
+        arena: &pim_matrix_arena1,
-        index: 1,
+        index: 0,
-    };
-    let pim_storage2 = PimStorage {
-        arena: &pim_matrix_arena,
-        index: 2,
     };
+
     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!(
+    let mut matrix1 = Matrix::from_data(pim_storage1);
-        &mut uart,
+    matrix1.fill_lower_triangle(F16x1(f16::ONE), 0);
-        "Cache Lines: {}\nRows: {}",
+
-        PimMatrixArena::<8, 8>::OCCUPIED_CACHE_LINES,
+    writeln!(&mut uart, "{matrix0} + {matrix1}\n=").unwrap();
-        PimMatrixArena::<8, 8>::OCCUPIED_ROWS
-    )
-    .unwrap();
-    writeln!(&mut uart, "{matrix0} * 2\n=").unwrap();
 
     // Invalidate and flush array just in case
-    pim_matrix_arena.borrow_mut().invalidate_flush();
+    pim_matrix_arena0.borrow_mut().invalidate_flush();
-    // dummy_array.invalidate_flush();
+    pim_matrix_arena1.borrow_mut().invalidate_flush();
     barrier::dsb(barrier::SY);
 
     pim_state.set_bank_mode(BankMode::PimAllBank);
-    pim_matrix_arena
+    pim_matrix_arena0
-        .borrow_mut()
+        .borrow()
         .execute_instruction_read_dual_bank();
-    pim_matrix_arena
+    pim_matrix_arena1
-        .borrow_mut()
+        .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();

pim-os/src/pim.rs

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

pim-os/src/pim/array.rs

@@ -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 {
+        unsafe { (&self.arena.borrow().0[0][0] as *const F16x16 as *const F16x1).add(self.index) }
-            (&self.arena.borrow().0[0][0] as *const F16x16 as *const F16x1).offset(self.index as _)
-        }
     }
 
     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)
+            (&mut self.arena.borrow_mut().0[0][0] as *mut F16x16 as *mut F16x1).add(self.index)
-                .offset(self.index as _)
         }
     }
 
@@ -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") {
+        for i in (0..Self::OCCUPIED_ROWS).map(|i| i * NUMBER_OF_BANKS) {
-            self.invalidate_bank(EVEN_BANK_INDEX);
+            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") {
+        for i in (0..Self::OCCUPIED_ROWS).map(|i| i * NUMBER_OF_BANKS) {
-            self.preload_zero();
+            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());
         }

pim-os/src/pim/kernel.rs

@@ -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,

pim-os/src/pim/state.rs

@@ -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,
             })

pim-os/src/pim/vector.rs

@@ -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)
     }
 }