From a13de610816da45474aac5835871ad4bfe92fe5b Mon Sep 17 00:00:00 2001
From: Derek Christ <christ.derek@gmail.com>
Date: Thu, 22 Feb 2024 22:28:43 +0100
Subject: [PATCH] Add a GEMV layer bin

---
 pim-os/src/bin/classic_gemv_layers.rs | 24 +++++++++
 pim-os/src/bin/gemv.rs                |  7 ++-
 pim-os/src/bin/gemv_layers.rs         | 72 +++++++++++++++++++++++++++
 3 files changed, 101 insertions(+), 2 deletions(-)
 create mode 100644 pim-os/src/bin/classic_gemv_layers.rs
 create mode 100644 pim-os/src/bin/gemv_layers.rs

diff --git a/pim-os/src/bin/classic_gemv_layers.rs b/pim-os/src/bin/classic_gemv_layers.rs
new file mode 100644
index 0000000..d587be4
--- /dev/null
+++ b/pim-os/src/bin/classic_gemv_layers.rs
@@ -0,0 +1,24 @@
+#![no_std]
+#![no_main]
+
+extern crate alloc;
+
+use core::fmt::Write;
+use nalgebra::{SMatrix, SVector};
+use pim_os::{pim::vector::F16x1, uart::Uart0};
+
+const ROWS: usize = 128;
+const COLUMNS: usize = 128;
+
+#[no_mangle]
+pub extern "C" fn main() {
+    let matrix = SMatrix::<F16x1, ROWS, COLUMNS>::zeros();
+    let mut input_vector = SVector::<F16x1, COLUMNS>::zeros();
+
+    for _ in 0..5 {
+        input_vector = matrix * input_vector;
+        core::hint::black_box(input_vector);
+    }
+
+    writeln!(Uart0, "Done").unwrap();
+}
diff --git a/pim-os/src/bin/gemv.rs b/pim-os/src/bin/gemv.rs
index 9a51d9b..4938071 100644
--- a/pim-os/src/bin/gemv.rs
+++ b/pim-os/src/bin/gemv.rs
@@ -77,7 +77,7 @@ pub extern "C" fn main() {
         pim::state::set_bank_mode(BankMode::SingleBank);
     }
 
-    writeln!(Uart0, "{output_partial_sum_vector}").unwrap();
+    // writeln!(Uart0, "{output_partial_sum_vector}").unwrap();
 
     let output_vector = SVector::<F16x1, ROWS>::from_fn(|r, _| {
         output_partial_sum_vector[r]
@@ -86,5 +86,8 @@ pub extern "C" fn main() {
             .fold(F16x1::zero(), |acc, val| acc + *val)
     });
 
-    writeln!(Uart0, "{output_vector}").unwrap();
+    core::hint::black_box(output_vector);
+
+    // writeln!(Uart0, "{output_vector}").unwrap();
+    writeln!(Uart0, "Done").unwrap();
 }
diff --git a/pim-os/src/bin/gemv_layers.rs b/pim-os/src/bin/gemv_layers.rs
new file mode 100644
index 0000000..10960e6
--- /dev/null
+++ b/pim-os/src/bin/gemv_layers.rs
@@ -0,0 +1,72 @@
+#![no_std]
+#![no_main]
+
+extern crate alloc;
+
+use aarch64_cpu::asm::barrier;
+use alloc::boxed::Box;
+use core::fmt::Write;
+use nalgebra::{SMatrix, SVector};
+use num_traits::{One, Zero};
+use pim_isa::BankMode;
+use pim_os::{
+    kernel::gemv,
+    pim::{
+        self, interleaved_array,
+        vector::{F16x1, F16x16},
+    },
+    uart::Uart0,
+};
+
+const ROWS: usize = 128;
+const COLUMNS: usize = 128;
+const X16_ROWS: usize = ROWS / 16;
+const X16_COLUMNS: usize = COLUMNS / 16;
+
+#[no_mangle]
+pub extern "C" fn main() {
+    pim::state::set_kernel(&gemv::KERNEL_X4);
+
+    let mut matrix = SMatrix::<_, ROWS, COLUMNS>::zeros();
+    matrix.fill_lower_triangle(F16x1::one(), 0);
+
+    let pim_matrix = Box::new(pim::continuous_array::Matrix::<X16_ROWS, X16_COLUMNS>::from(matrix));
+
+    let input_vector = SVector::<_, X16_COLUMNS>::from_element(F16x16::one());
+
+    let mut output_partial_sum_vector = Box::new(SVector::<F16x16, ROWS>::zeros());
+
+    let dummy = Box::new(0);
+
+    for _ in 0..5 {
+        let interleaved_input_vector =
+            Box::new(interleaved_array::Vector::from(input_vector.clone()));
+
+        // Verify everything is correctly initialized before PIM operation
+        barrier::dsb(barrier::SY);
+
+        // Execute kernel
+        {
+            pim::state::set_bank_mode(BankMode::PimAllBank);
+
+            gemv::execute(
+                pim_matrix.as_ref(),
+                interleaved_input_vector.as_ref(),
+                output_partial_sum_vector.as_mut(),
+                dummy.as_ref(),
+            );
+
+            pim::state::set_bank_mode(BankMode::SingleBank);
+        }
+
+        let output_vector = SVector::<F16x1, ROWS>::from_fn(|r, _| {
+            output_partial_sum_vector[r]
+                .0
+                .iter()
+                .fold(F16x1::zero(), |acc, val| acc + *val)
+        });
+
+        core::hint::black_box(output_vector);
+    }
+    writeln!(Uart0, "Done").unwrap();
+}