Add conveniance functions for a PIM array

pim-os/Cargo.toml

@@ -8,6 +8,7 @@ aarch64-cpu = "9.4.0"
 half = { version = "2.3.1", default-features = false }
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 serde-json-core = "0.5.1"
+
 pim-isa = { path = "../pim-isa", default-features = false }
 
 [profile.dev]

pim-os/aarch64-gem5.ld

@@ -1,6 +1,6 @@
 MEMORY
 {
-    bootmem : ORIGIN = 0x0, LENGTH = 0x10000
+    bootmem : ORIGIN = 0x0, LENGTH = 0x100000
     dram : ORIGIN = 0x80000000, LENGTH = 0x100000000
 }
 
@@ -16,6 +16,6 @@ SECTIONS
     .bss : { *(.bss) } > dram
 
     . = ALIGN(8);
-    . = . + 0x8000;
+    . = . + 0x100000; # 1 MiB Stack
     LD_STACK_PTR = .;
 }

pim-os/src/main.rs

@@ -2,13 +2,12 @@
 #![no_main]
 
 use core::{
-    arch::{asm, global_asm},
+    arch::global_asm,
     fmt::Write,
     panic::PanicInfo,
     sync::atomic::{self, Ordering},
 };
-use half::f16;
+use pim::{array::ComputeArray, kernel::TEST_KERNEL};
-use pim::{array::BankArray, kernel::TEST_KERNEL};
 use pim_isa::{BankMode, PimConfig};
 use uart::Uart0;
 
@@ -18,8 +17,6 @@ mod uart;
 
 global_asm!(include_str!("start.s"));
 
-static mut TEST_ARRAY: BankArray = BankArray([f16::ZERO; 512]);
-
 #[no_mangle]
 pub extern "C" fn entry() -> ! {
     let mut pim_config = PimConfig {
@@ -28,67 +25,50 @@ pub extern "C" fn entry() -> ! {
     };
     let mut pim_writer = pim::config::PimWriter;
 
+    let mut test_array: ComputeArray<2> = ComputeArray::default();
+
     let mut uart = Uart0 {};
-    unsafe {
-        writeln!(
-            &mut uart,
-            "PIM array is at {:x?}",
-            core::ptr::addr_of!(TEST_ARRAY)
-        )
-        .unwrap();
-        writeln!(&mut uart, "Read from all banks").unwrap();
 
-        // Invalidate and flush array just in case
+    writeln!(
-        for element in TEST_ARRAY.0.iter().step_by(8) {
+        &mut uart,
-            asm!("dc civac, {val}", val = in(reg) element);
+        "PIM array is at {:x?}",
-        }
+        core::ptr::addr_of!(test_array)
-        asm!("dsb sy");
+    )
+    .unwrap();
+    writeln!(&mut uart, "Read from all banks").unwrap();
 
-        
+    // Invalidate and flush array just in case
-        // Fetch single cache line in array
+    test_array.invalidate_flush();
-        // core::ptr::read_volatile(&TEST_ARRAY.0[0]);
 
-        // Zero array to prevent fetch
+    // Zero array to prevent fetch
-        // for element in TEST_ARRAY.0.iter().step_by(8) {
+    let bank_array = &mut test_array.0[0];
-        //     asm!("dc zva, {val}", val = in(reg) element);
+    bank_array.preload_zero();
-        // }
-        asm!("dc zva, {val}", val = in(reg) &TEST_ARRAY.0[0]);
-        asm!("dsb sy");
 
-        pim_writer.write(
+    pim_writer.write(
-            serde_json_core::to_string::<PimConfig, 1024>(&pim_config)
+        serde_json_core::to_string::<PimConfig, 1024>(&pim_config)
-                .unwrap()
+            .unwrap()
-                .as_str(),
+            .as_str(),
-        );
+    );
 
-        core::ptr::write_volatile(&mut TEST_ARRAY.0[0], f16::ZERO);
+    bank_array.write_data();
 
-        // Invalidate and flush
+    pim_config.bank_mode = BankMode::SingleBank;
-        // for element in TEST_ARRAY.0.iter().step_by(8) {
+    pim_writer.write(
-        asm!("dc civac, {val}", val = in(reg) &TEST_ARRAY.0[0]);
+        serde_json_core::to_string::<PimConfig, 1024>(&pim_config)
-        // }
+            .unwrap()
-        asm!("dsb sy");
+            .as_str(),
+    );
 
-        pim_config.bank_mode = BankMode::SingleBank;
+    // Invalidate and flush
-        pim_writer.write(
+    test_array.invalidate();
-            serde_json_core::to_string::<PimConfig, 1024>(&pim_config)
-                .unwrap()
-                .as_str(),
-        );
 
-        // Invalidate array
+    writeln!(
-        for element in TEST_ARRAY.0.iter().step_by(8) {
+        &mut uart,
-            asm!("dc ivac, {val}", val = in(reg) element);
+        "{:?}: {:x?}",
-        }
+        core::ptr::addr_of!(test_array),
-
+        test_array
-        writeln!(
+    )
-            &mut uart,
+    .unwrap();
-            "{:?}: {:x?}",
-            core::ptr::addr_of!(TEST_ARRAY),
-            TEST_ARRAY
-        )
-        .unwrap();
-    }
 
     m5ops::exit();
 

pim-os/src/pim.rs

@@ -1,3 +1,4 @@
 pub mod array;
 pub mod config;
 pub mod kernel;
+pub mod operation;

pim-os/src/pim/array.rs

@@ -1,8 +1,105 @@
+use core::arch::asm;
 use half::f16;
 
-#[derive(Debug)]
+const NUMBER_OF_BANKS: usize = 32;
-#[repr(C, align(1024))]
+const ELEMENTS_PER_CACHE_LINE: usize = 16;
-pub struct BankArray(pub [f16; 512]);
+const ELEMENTS_PER_BANK_ARRAY: usize = NUMBER_OF_BANKS * ELEMENTS_PER_CACHE_LINE;
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
+#[repr(C, align(1024))]
+pub struct BankArray(pub [f16; ELEMENTS_PER_BANK_ARRAY]);
+
+impl Default for BankArray {
+    fn default() -> Self {
+        Self([f16::ZERO; ELEMENTS_PER_BANK_ARRAY])
+    }
+}
+
+impl BankArray {
+    pub fn write_data(&mut self) {
+        unsafe {
+            // Write to first bank
+            let first_bank = &mut self.0[0];
+            core::ptr::write_volatile(first_bank, f16::ZERO);
+            self.invalidate_flush_single_bank(0);
+        }
+    }
+
+    pub fn read_data(&self) {
+        unsafe {
+            // Read from first bank
+            let first_bank = &self.0[0];
+            core::ptr::read_volatile(first_bank);
+            self.invalidate_flush_single_bank(0);
+        }
+    }
+
+    pub fn invalidate(&self) {
+        (0..ELEMENTS_PER_BANK_ARRAY)
+            .step_by(ELEMENTS_PER_CACHE_LINE)
+            .for_each(|idx| self.invalidate_single_bank(idx));
+    }
+
+    pub fn invalidate_single_bank(&self, idx: usize) {
+        unsafe {
+            // Invalidate and flush first bank
+            asm!("dc ivac, {val}", val = in(reg) &self.0[idx]);
+            asm!("dsb sy");
+        }
+    }
+
+    pub fn invalidate_flush(&self) {
+        (0..ELEMENTS_PER_BANK_ARRAY)
+            .step_by(ELEMENTS_PER_CACHE_LINE)
+            .for_each(|idx| self.invalidate_flush_single_bank(idx));
+    }
+
+    pub fn invalidate_flush_single_bank(&self, idx: usize) {
+        unsafe {
+            // Invalidate and flush first bank
+            asm!("dc civac, {val}", val = in(reg) &self.0[idx]);
+            asm!("dsb sy");
+        }
+    }
+
+    pub fn preload_zero(&self) {
+        (0..ELEMENTS_PER_BANK_ARRAY)
+            .step_by(ELEMENTS_PER_CACHE_LINE)
+            .for_each(|idx| self.preload_zero_single_bank(idx));
+    }
+
+    pub fn preload_zero_single_bank(&self, idx: usize) {
+        unsafe {
+            // Preload first bank
+            asm!("dc zva, {val}", val = in(reg) &self.0[idx]);
+            asm!("dsb sy");
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
 pub struct ComputeArray<const N: usize>(pub [BankArray; N]);
+
+impl<const N: usize> ComputeArray<N> {
+    pub fn invalidate_flush(&self) {
+        self.0
+            .iter()
+            .for_each(|bank_array| bank_array.invalidate_flush());
+    }
+
+    pub fn invalidate(&self) {
+        self.0.iter().for_each(|bank_array| bank_array.invalidate());
+    }
+
+    pub fn preload_zero(&self) {
+        self.0
+            .iter()
+            .for_each(|bank_array| bank_array.preload_zero());
+    }
+}
+
+impl<const N: usize> Default for ComputeArray<N> {
+    fn default() -> Self {
+        Self(core::array::from_fn(|_| Default::default()))
+    }
+}

pim-os/src/pim/operation.rs

@@ -0,0 +1,18 @@
+use super::array::BankArray;
+use core::arch::asm;
+
+fn execute_bank_write(bank_array: &BankArray) {
+    unsafe {
+        // Invalidate and flush cache line
+        asm!("dc civac, {val}", val = in(reg) bank_array);
+        asm!("dsb sy");
+
+        // Zero cache line
+        asm!("dc zva, {val}", val = in(reg) bank_array);
+        asm!("dsb sy");
+    }
+}
+
+fn execute_bank_read() {}
+
+fn execute_generic_inst() {}