Implement ADD and JUMP instructions

pim-isa/src/lib.rs

@@ -9,6 +9,17 @@ pub enum Instruction {
     JUMP { offset: i16, count: u16 },
     MOV { src: File, dst: File },
     FILL { src: File, dst: File },
+    ADD { src0: File, src1: File, dst: File },
+}
+
+impl Instruction {
+    pub fn supported_source(&self, src: File) -> bool {
+        true
+    }
+
+    pub fn supported_destination(&self, src: File) -> bool {
+        true
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -19,7 +30,7 @@ pub enum File {
     Bank,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Kernel(pub [Instruction; 32]);
 
 impl Kernel {
@@ -32,7 +43,7 @@ pub struct PimConfig {
     pub kernel: Kernel,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 pub enum BankMode {
     SingleBank,
     AllBank,

pim-os/src/main.rs

@@ -5,10 +5,15 @@ use core::{
     arch::global_asm,
     fmt::Write,
     panic::PanicInfo,
-    sync::atomic::{self, Ordering},
+    sync::atomic::{compiler_fence, Ordering},
 };
-use pim::{array::ComputeArray, kernel::TEST_KERNEL};
+use half::f16;
-use pim_isa::{BankMode, PimConfig};
+use pim::{
+    array::{BankArray, ComputeArray},
+    kernel::TEST_KERNEL,
+    state::PimState,
+};
+use pim_isa::BankMode;
 use uart::Uart0;
 
 mod m5ops;
@@ -19,61 +24,46 @@ global_asm!(include_str!("start.s"));
 
 #[no_mangle]
 pub extern "C" fn entry() -> ! {
-    let mut pim_config = PimConfig {
+    let mut pim_state = PimState::new(&TEST_KERNEL);
-        bank_mode: BankMode::PimAllBank,
+    let mut compute_array: ComputeArray<3> = ComputeArray([
-        kernel: TEST_KERNEL,
+        BankArray([f16::from_f32(0.1); 512]),
-    };
+        BankArray([f16::from_f32(0.2); 512]),
-    let mut pim_writer = pim::config::PimWriter;
+        BankArray([f16::ZERO; 512]),
-
+    ]);
-    let mut test_array: ComputeArray<2> = ComputeArray::default();
+    let dummy_array = BankArray::default();
-
+    let mut uart = Uart0;
-    let mut uart = Uart0 {};
 
     writeln!(
         &mut uart,
         "PIM array is at {:x?}",
-        core::ptr::addr_of!(test_array)
+        core::ptr::addr_of!(compute_array)
     )
     .unwrap();
-    writeln!(&mut uart, "Read from all banks").unwrap();
 
     // Invalidate and flush array just in case
-    test_array.invalidate_flush();
+    compute_array.invalidate_flush();
+    dummy_array.invalidate_flush();
 
-    // Zero array to prevent fetch
+    pim_state.set_bank_mode(BankMode::PimAllBank);
-    let bank_array = &mut test_array.0[0];
+    compute_array.0[0].execute_instruction_read();
-    bank_array.preload_zero();
+    compute_array.0[1].execute_instruction_read();
+    compute_array.0[2].execute_instruction_write();
+    dummy_array.execute_instruction_read();
+    pim_state.set_bank_mode(BankMode::SingleBank);
 
-    pim_writer.write(
+    compute_array.invalidate();
-        serde_json_core::to_string::<PimConfig, 1024>(&pim_config)
-            .unwrap()
-            .as_str(),
-    );
-
-    bank_array.write_data();
-
-    pim_config.bank_mode = BankMode::SingleBank;
-    pim_writer.write(
-        serde_json_core::to_string::<PimConfig, 1024>(&pim_config)
-            .unwrap()
-            .as_str(),
-    );
-
-    // Invalidate and flush
-    test_array.invalidate();
 
     writeln!(
         &mut uart,
-        "{:?}: {:x?}",
+        "BankArray 0: [{:?}, ...]\nBankArray 1: [{:?}, ...]\nBankArray 2: [{:?}, ...]",
-        core::ptr::addr_of!(test_array),
+        compute_array.0[0].0[0], compute_array.0[1].0[0], compute_array.0[2].0[0]
-        test_array
     )
     .unwrap();
 
     m5ops::exit();
 
     loop {
-        atomic::compiler_fence(Ordering::SeqCst);
+        compiler_fence(Ordering::SeqCst);
     }
 }
 
@@ -82,6 +72,6 @@ fn panic(info: &PanicInfo) -> ! {
     writeln!(Uart0, "{info}").unwrap();
 
     loop {
-        atomic::compiler_fence(Ordering::SeqCst);
+        compiler_fence(Ordering::SeqCst);
     }
 }

pim-os/src/pim.rs

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

pim-os/src/pim/array.rs

@@ -16,13 +16,9 @@ impl Default for BankArray {
 }
 
 impl BankArray {
-    pub fn write_data(&mut self) {
+    pub fn execute_instruction_read(&self) {
-        unsafe {
+        self.invalidate_single_bank(0);
-            // Write to first bank
+        self.read_data();
-            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) {
@@ -30,7 +26,20 @@ impl BankArray {
             // Read from first bank
             let first_bank = &self.0[0];
             core::ptr::read_volatile(first_bank);
-            self.invalidate_flush_single_bank(0);
+        }
+    }
+
+    pub fn execute_instruction_write(&mut self) {
+        self.preload_zero();
+        self.write_data();
+        self.invalidate_flush_single_bank(0);
+    }
+
+    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);
         }
     }
 
@@ -90,12 +99,6 @@ impl<const N: usize> ComputeArray<N> {
     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> {

pim-os/src/pim/config.rs

@@ -1,11 +1,11 @@
 use core::arch::asm;
 
-
 #[link_section = ".pim_config"]
 static mut PIM_CONFIG_REGION: [u8; 0x4000] = [0; 0x4000];
 
 const CACHE_LINE_SIZE: usize = 32;
 
+#[derive(Debug)]
 pub struct PimWriter;
 
 impl PimWriter {

pim-os/src/pim/kernel.rs

@@ -1,25 +1,23 @@
 use pim_isa::{File, Instruction, Kernel};
 
 pub const TEST_KERNEL: Kernel = Kernel([
-    Instruction::FILL {
-        src: File::Grf { index: 0 },
-        dst: File::Bank,
-    },
     Instruction::MOV {
         src: File::Bank,
-        dst: File::Grf { index: 1 },
+        dst: File::Grf { index: 0 },
     },
-    Instruction::JUMP {
+    Instruction::ADD {
-        offset: 1,
+        src0: File::Bank,
-        count: 12,
+        src1: File::Grf { index: 0 },
+        dst: File::Grf { index: 0 },
     },
+    Instruction::FILL { src: File::Grf { index: 0 }, dst: File::Bank },
     Instruction::EXIT,
     Instruction::NOP,
-    Instruction::EXIT,
     Instruction::NOP,
-    Instruction::EXIT,
     Instruction::NOP,
-    Instruction::EXIT,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
     Instruction::NOP,
     Instruction::NOP,
     Instruction::NOP,

pim-os/src/pim/operation.rs

@@ -1,18 +0,0 @@
-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() {}

pim-os/src/pim/state.rs

@@ -0,0 +1,32 @@
+use super::config::PimWriter;
+use pim_isa::{BankMode, Kernel, PimConfig};
+
+#[derive(Debug)]
+pub struct PimState {
+    kernel: Kernel,
+    bank_mode: BankMode,
+    writer: PimWriter,
+}
+
+impl PimState {
+    pub fn new(kernel: &Kernel) -> Self {
+        Self {
+            kernel: kernel.clone(),
+            bank_mode: BankMode::SingleBank,
+            writer: PimWriter,
+        }
+    }
+
+    // TODO return token and return to singlebank when dropped
+    pub fn set_bank_mode(&mut self, bank_mode: BankMode) {
+        self.bank_mode = bank_mode;
+        self.writer.write(
+            serde_json_core::to_string::<PimConfig, 1024>(&PimConfig {
+                kernel: self.kernel.clone(),
+                bank_mode: self.bank_mode,
+            })
+            .unwrap()
+            .as_str(),
+        );
+    }
+}

pim-os/src/uart.rs

@@ -2,6 +2,7 @@ use core::{fmt::Write, ptr::write_volatile};
 
 const UART0_ADDR: *mut u32 = 0x1c090000 as _;
 
+#[derive(Debug)]
 pub struct Uart0;
 
 impl Write for Uart0 {

pim-vm/src/lib.rs

@@ -1,6 +1,5 @@
 use half::f16;
 use pim_isa::{BankMode, File, Instruction, Kernel, PimConfig};
-use std::io::Cursor;
 
 #[cxx::bridge(namespace = "pim_vm")]
 mod ffi {
@@ -28,28 +27,32 @@ fn init_logger() {
     env_logger::init();
 }
 
+const BURST_LENGTH: usize = 32;
+
 const GRF_NUM_REGISTERS: usize = 16;
 const SRF_A_NUM_REGISTERS: usize = 8;
 const SRF_M_NUM_REGISTERS: usize = 8;
 
-const REG_NUM_ENTRIES: usize = 16;
+const FP_UNITS: usize = 16;
-type Register = [f16; REG_NUM_ENTRIES];
+type GrfRegister = [f16; FP_UNITS];
 
 #[derive(Clone, Debug)]
 struct PimUnit {
-    grf: [Register; GRF_NUM_REGISTERS],
+    grf: [GrfRegister; GRF_NUM_REGISTERS],
-    srf_a: [Register; SRF_A_NUM_REGISTERS],
+    srf_a: [f16; SRF_A_NUM_REGISTERS],
-    srf_m: [Register; SRF_A_NUM_REGISTERS],
+    srf_m: [f16; SRF_A_NUM_REGISTERS],
     pc: u8,
+    jump_counter: Option<u16>,
 }
 
 impl Default for PimUnit {
     fn default() -> Self {
         Self {
-            grf: [[f16::PI; REG_NUM_ENTRIES]; GRF_NUM_REGISTERS],
+            grf: [[f16::ZERO; FP_UNITS]; GRF_NUM_REGISTERS],
-            srf_a: [[f16::ZERO; REG_NUM_ENTRIES]; SRF_A_NUM_REGISTERS],
+            srf_a: [f16::ZERO; SRF_A_NUM_REGISTERS],
-            srf_m: [[f16::ZERO; REG_NUM_ENTRIES]; SRF_M_NUM_REGISTERS],
+            srf_m: [f16::ZERO; SRF_M_NUM_REGISTERS],
             pc: 0,
+            jump_counter: None,
         }
     }
 }
@@ -64,6 +67,7 @@ impl PimVM {
     fn reset(&mut self) {
         for unit in self.pim_units.iter_mut() {
             unit.pc = 0;
+            unit.jump_counter = None;
         }
     }
 
@@ -96,11 +100,11 @@ fn new_pim_vm(num_pim_units: u32) -> Box<PimVM> {
 }
 
 #[repr(C)]
-struct BankData([f16; REG_NUM_ENTRIES]);
+struct BankData([f16; FP_UNITS]);
 
 impl PimVM {
     pub fn execute_read(&mut self, bank_index: u32, bank_data: &[u8]) {
-        assert_eq!(bank_data.len(), 32);
+        assert_eq!(bank_data.len(), BURST_LENGTH);
 
         let pim_unit = &mut self.pim_units[bank_index as usize];
 
@@ -113,42 +117,48 @@ impl PimVM {
 
         match inst {
             Instruction::NOP => (),
-            Instruction::EXIT => (),
+            Instruction::EXIT => {
-            Instruction::JUMP { offset, count } => todo!(),
+                pim_unit.jump_counter = None;
-            Instruction::MOV { src, dst } => {
+                pim_unit.pc = 0;
-                let data: [f16; REG_NUM_ENTRIES] = match src {
+            }
-                    File::Grf { index } => pim_unit.grf[*index as usize],
+            Instruction::JUMP { offset, count } => {
-                    File::Bank => unsafe {
+                pim_unit.jump_counter = match pim_unit.jump_counter {
-                        std::ptr::read(bank_data.as_ptr() as *const BankData).0
+                    Some(jump_counter) => jump_counter.checked_sub(1),
-                    },
+                    None => Some(*count),
-                    _ => panic!("Unsupported src operand: {src:?}"),
                 };
 
-                match dst {
+                if pim_unit.jump_counter != None {
-                    File::Grf { index } => pim_unit.grf[*index as usize] = data,
+                    let new_pc = pim_unit.pc as i32 + *offset as i32;
-                    File::SrfM { index } => pim_unit.srf_m[*index as usize] = data,
+
-                    File::SrfA { index } => pim_unit.srf_a[*index as usize] = data,
+                    if new_pc < 0 || new_pc >= 32 {
-                    _ => panic!("Unsupported dst operand: {dst:?}"),
+                        panic!("Invalid PC {new_pc} after JUMP: {inst:?}");
+                    }
+
+                    pim_unit.pc = new_pc as _;
                 }
             }
-            Instruction::FILL { src, dst } => {
+            Instruction::MOV { src, dst } | Instruction::FILL { src, dst } => {
-                let data: [f16; REG_NUM_ENTRIES] = match src {
+                let data = PimVM::load(*src, pim_unit, &bank_data);
-                    File::Grf { index } => pim_unit.grf[*index as usize],
+                PimVM::store(*dst, pim_unit, &data);
-                    File::Bank => unsafe {
+            }
-                        std::ptr::read(bank_data.as_ptr() as *const BankData).0
+            Instruction::ADD { src0, src1, dst } => {
-                    },
+                let data0 = PimVM::load(*src0, pim_unit, &bank_data);
-                    _ => panic!("Unsupported src operand: {src:?}"),
+                let data1 = PimVM::load(*src1, pim_unit, &bank_data);
-                };
 
-                match dst {
+                let sum: [f16; FP_UNITS] = data0
-                    File::Grf { index } => pim_unit.grf[*index as usize] = data,
+                    .into_iter()
-                    _ => panic!("Unsupported dst operand: {dst:?}"),
+                    .zip(data1)
-                }
+                    .map(|(src0, src1)| src0 + src1)
+                    .collect::<Vec<_>>()
+                    .try_into()
+                    .unwrap();
+
+                PimVM::store(*dst, pim_unit, &sum);
             }
         }
     }
 
-    pub fn execute_write(&mut self, bank_index: u32) -> [u8; 32] {
+    pub fn execute_write(&mut self, bank_index: u32) -> [u8; BURST_LENGTH] {
         let pim_unit = &mut self.pim_units[bank_index as usize];
 
         let current_pc = pim_unit.pc;
@@ -160,7 +170,7 @@ impl PimVM {
 
         let data = match inst {
             Instruction::FILL { src, dst } => {
-                let data: [f16; REG_NUM_ENTRIES] = match src {
+                let data: [f16; FP_UNITS] = match src {
                     File::Grf { index } => pim_unit.grf[*index as usize],
                     _ => panic!("Unsupported src operand: {src:?}"),
                 };
@@ -176,4 +186,22 @@ impl PimVM {
 
         unsafe { std::mem::transmute(data) }
     }
+
+    fn load(src: File, pim_unit: &PimUnit, bank_data: &[u8]) -> [f16; FP_UNITS] {
+        match src {
+            File::Grf { index } => pim_unit.grf[index as usize],
+            File::SrfM { index } => [pim_unit.srf_m[index as usize]; FP_UNITS],
+            File::SrfA { index } => [pim_unit.srf_a[index as usize]; FP_UNITS],
+            File::Bank => unsafe { std::ptr::read(bank_data.as_ptr() as *const BankData).0 },
+        }
+    }
+
+    fn store(dst: File, pim_unit: &mut PimUnit, data: &[f16; FP_UNITS]) {
+        match dst {
+            File::Grf { index } => pim_unit.grf[index as usize] = data.clone(),
+            File::SrfM { index } => pim_unit.srf_m[index as usize] = data[0],
+            File::SrfA { index } => pim_unit.srf_a[index as usize] = data[0],
+            File::Bank => panic!("Unsupported dst operand: {dst:?}"),
+        }
+    }
 }