Add conveniance functions for a PIM array

pim-os/src/pim/array.rs

@@ -1,8 +1,105 @@
+use core::arch::asm;
 use half::f16;
 
-#[derive(Debug)]
-#[repr(C, align(1024))]
-pub struct BankArray(pub [f16; 512]);
+const NUMBER_OF_BANKS: usize = 32;
+const ELEMENTS_PER_CACHE_LINE: usize = 16;
+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()))
+    }
+}