From aecb19b4f42ba8c5f840d1e04a33c446efe3daa2 Mon Sep 17 00:00:00 2001
From: Derek Christ <christ.derek@gmail.com>
Date: Tue, 12 Dec 2023 17:15:29 +0100
Subject: [PATCH] Zero-cycle JUMP

---
 pim-os/src/main.rs       |  6 +++--
 pim-os/src/pim/kernel.rs | 10 ++++----
 pim-vm/src/lib.rs        | 53 ++++++++++++++++++++++------------------
 3 files changed, 38 insertions(+), 31 deletions(-)

diff --git a/pim-os/src/main.rs b/pim-os/src/main.rs
index 017e54f..f594167 100644
--- a/pim-os/src/main.rs
+++ b/pim-os/src/main.rs
@@ -53,9 +53,9 @@ pub extern "C" fn entry() -> ! {
     barrier::dsb(barrier::SY);
 
     pim_state.set_bank_mode(BankMode::PimAllBank);
-    compute_array.0[0].execute_instruction_read();
-    compute_array.0[2].execute_instruction_read();
     compute_array.0[1].execute_instruction_read();
+    compute_array.0[2].execute_instruction_read();
+    compute_array.0[0].execute_instruction_read();
     compute_array.0[2].execute_instruction_write();
     dummy_array.execute_instruction_read();
     pim_state.set_bank_mode(BankMode::SingleBank);
@@ -70,6 +70,8 @@ pub extern "C" fn entry() -> ! {
     )
     .unwrap();
 
+    writeln!(&mut uart, "ComputeArray:\n{:?}", compute_array).unwrap();
+
     m5ops::exit();
 
     loop {
diff --git a/pim-os/src/pim/kernel.rs b/pim-os/src/pim/kernel.rs
index 708c15e..bfdf99b 100644
--- a/pim-os/src/pim/kernel.rs
+++ b/pim-os/src/pim/kernel.rs
@@ -7,17 +7,17 @@ pub const TEST_KERNEL: Kernel = Kernel([
     },
     Instruction::MOV {
         src: File::Bank,
-        dst: File::GrfA { index: 1 },
+        dst: File::GrfB { index: 0 },
     },
     Instruction::MAC {
         src0: File::Bank,
         src1: File::GrfA { index: 0 },
-        src2: File::GrfA { index: 1 },
-        dst: File::GrfA { index: 1 },
-        aam: false
+        src2: File::GrfB { index: 0 },
+        dst: File::GrfB { index: 0 },
+        aam: true
     },
     Instruction::FILL {
-        src: File::GrfA { index: 1 },
+        src: File::GrfB { index: 0 },
         dst: File::Bank,
     },
     Instruction::EXIT,
diff --git a/pim-vm/src/lib.rs b/pim-vm/src/lib.rs
index 28bb9d7..2f0dc86 100644
--- a/pim-vm/src/lib.rs
+++ b/pim-vm/src/lib.rs
@@ -27,7 +27,8 @@ fn init_logger() {
     env_logger::init();
 }
 
-const COLUMN_BIT_OFFSET: usize = 10;
+const GRF_A_BIT_OFFSET: usize = 10;
+const GRF_B_BIT_OFFSET: usize = 13;
 const BURST_LENGTH: usize = 32;
 
 const GRF_NUM_REGISTERS: usize = 8;
@@ -111,15 +112,35 @@ impl PimVM {
 
         let pim_unit = &mut self.pim_units[bank_index as usize];
 
-        let current_pc = pim_unit.pc;
+        let mut inst = self.pim_config.kernel.0[pim_unit.pc as usize];
         pim_unit.pc += 1;
 
-        let inst = self.pim_config.kernel.0[current_pc as usize];
+        let aam_grf_a_index = (address >> GRF_A_BIT_OFFSET) & 0b111;
+        let aam_grf_b_index = (address >> GRF_B_BIT_OFFSET) & 0b111;
 
-        let aam_grf_a_index = (address >> COLUMN_BIT_OFFSET) & 0b111;
-        let aam_grf_b_index = (address >> COLUMN_BIT_OFFSET + 3) & 0b111;
+        log::debug!("PimUnit {bank_index} Execute PC {}: {inst:?}", pim_unit.pc);
 
-        log::debug!("PimUnit {bank_index} Execute PC {current_pc}: {inst:?}");
+        // The JUMP instruction is zero-cycle and not actually executed
+        while let Instruction::JUMP { offset, count } = inst {
+            pim_unit.jump_counter = match pim_unit.jump_counter {
+                Some(jump_counter) => jump_counter.checked_sub(1),
+                None => count.checked_sub(1),
+            };
+
+            if pim_unit.jump_counter != None {
+                let new_pc = pim_unit.pc as i32 + offset as i32;
+
+                if new_pc < 0 || new_pc >= 32 {
+                    panic!("Invalid PC {new_pc} after JUMP: {inst:?}");
+                }
+
+                pim_unit.pc = new_pc as _;
+                log::debug!("PimUnit {bank_index} New PC {new_pc}: {inst:?}");
+            }
+
+            inst = self.pim_config.kernel.0[pim_unit.pc as usize];
+            pim_unit.pc += 1;
+        }
 
         match inst {
             Instruction::NOP => (),
@@ -127,23 +148,7 @@ impl PimVM {
                 pim_unit.jump_counter = None;
                 pim_unit.pc = 0;
             }
-            Instruction::JUMP { offset, count } => {
-                pim_unit.jump_counter = match pim_unit.jump_counter {
-                    Some(jump_counter) => jump_counter.checked_sub(1),
-                    None => count.checked_sub(1),
-                };
-
-                if pim_unit.jump_counter != None {
-                    let new_pc = current_pc as i32 + offset as i32;
-
-                    if new_pc < 0 || new_pc >= 32 {
-                        panic!("Invalid PC {new_pc} after JUMP: {inst:?}");
-                    }
-
-                    pim_unit.pc = new_pc as _;
-                    log::debug!("PimUnit {bank_index} New PC {new_pc}: {inst:?}");
-                }
-            }
+            Instruction::JUMP { .. } => unreachable!(),
             Instruction::MOV { src, dst } | Instruction::FILL { src, dst } => {
                 let data = PimVM::load(src, pim_unit, &bank_data);
                 PimVM::store(dst, pim_unit, &data);
@@ -284,7 +289,7 @@ impl PimVM {
                     .try_into()
                     .unwrap();
 
-                log::debug!("{data0:#?}\n{data1:#?}\n{data2:#?}\n{product:#?}\n{sum:#?}");
+                log::debug!("{data0:?}\n{data1:?}\n{data2:?}\n{product:?}\n{sum:?}");
                 PimVM::store(dst, pim_unit, &sum);
             }
         }