From 86d0a0d6353319625a590fb22f23b6f05c1d18bc Mon Sep 17 00:00:00 2001
From: Derek Christ <christ.derek@gmail.com>
Date: Thu, 22 Feb 2024 17:30:22 +0100
Subject: [PATCH] Add Vector Mul Kernel

---
 pim-os/src/bin/vmul.rs    |  53 ++++++++++++++
 pim-os/src/kernel/vmul.rs | 145 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 198 insertions(+)
 create mode 100644 pim-os/src/bin/vmul.rs
 create mode 100644 pim-os/src/kernel/vmul.rs

diff --git a/pim-os/src/bin/vmul.rs b/pim-os/src/bin/vmul.rs
new file mode 100644
index 0000000..949b8d7
--- /dev/null
+++ b/pim-os/src/bin/vmul.rs
@@ -0,0 +1,53 @@
+#![no_std]
+#![no_main]
+
+extern crate alloc;
+
+use aarch64_cpu::asm::barrier;
+use alloc::boxed::Box;
+use core::fmt::Write;
+use half::f16;
+use nalgebra::SVector;
+use pim_isa::BankMode;
+use pim_os::{
+    kernel::vadd,
+    pim::{self, vector::F16x1},
+    uart::Uart0,
+};
+
+const ROWS: usize = 2048;
+
+#[no_mangle]
+pub extern "C" fn main() {
+    pim::state::set_kernel(&vadd::KERNEL);
+
+    let a = Box::new(pim::continuous_array::Vector(
+        SVector::<F16x1, ROWS>::from_fn(|_, _| F16x1(f16::from_f32(2 as _))),
+    ));
+    let b = Box::new(pim::continuous_array::Vector(
+        SVector::<F16x1, ROWS>::from_fn(|_, _| F16x1(f16::from_f32(3 as _))),
+    ));
+
+    // writeln!(Uart0, "{}+{}=", a.0, b.0).unwrap();
+
+    let mut c = Box::new(pim::continuous_array::Vector(
+        SVector::<F16x1, ROWS>::zeros(),
+    ));
+
+    let dummy = Box::new(0);
+
+    // Verify everything is correctly initialized before PIM operation
+    barrier::dsb(barrier::SY);
+
+    // Execute kernel
+    {
+        pim::state::set_bank_mode(BankMode::PimAllBank);
+
+        vadd::execute(&a.0, &b.0, &mut c.0, dummy.as_ref());
+
+        pim::state::set_bank_mode(BankMode::SingleBank);
+    }
+
+    // writeln!(Uart0, "{}", c.0).unwrap();
+    writeln!(Uart0, "Done").unwrap();
+}
diff --git a/pim-os/src/kernel/vmul.rs b/pim-os/src/kernel/vmul.rs
new file mode 100644
index 0000000..835878e
--- /dev/null
+++ b/pim-os/src/kernel/vmul.rs
@@ -0,0 +1,145 @@
+use crate::pim::{operation::PimOperand, vector::F16x1};
+use nalgebra::SVector;
+use pim_isa::{File, Instruction, Kernel};
+
+pub const KERNEL: Kernel = Kernel([
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 0 },
+    },
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 1 },
+    },
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 2 },
+    },
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 3 },
+    },
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 4 },
+    },
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 5 },
+    },
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 6 },
+    },
+    Instruction::MOV {
+        src: File::Bank,
+        dst: File::GrfA { index: 7 },
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 0 },
+        dst: File::GrfB { index: 0 },
+        aam: false,
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 1 },
+        dst: File::GrfB { index: 1 },
+        aam: false,
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 2 },
+        dst: File::GrfB { index: 2 },
+        aam: false,
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 3 },
+        dst: File::GrfB { index: 3 },
+        aam: false,
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 4 },
+        dst: File::GrfB { index: 4 },
+        aam: false,
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 5 },
+        dst: File::GrfB { index: 5 },
+        aam: false,
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 6 },
+        dst: File::GrfB { index: 6 },
+        aam: false,
+    },
+    Instruction::MUL {
+        src0: File::Bank,
+        src1: File::GrfA { index: 7 },
+        dst: File::GrfB { index: 7 },
+        aam: false,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 0 },
+        dst: File::Bank,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 1 },
+        dst: File::Bank,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 2 },
+        dst: File::Bank,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 3 },
+        dst: File::Bank,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 4 },
+        dst: File::Bank,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 5 },
+        dst: File::Bank,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 6 },
+        dst: File::Bank,
+    },
+    Instruction::FILL {
+        src: File::GrfB { index: 7 },
+        dst: File::Bank,
+    },
+    Instruction::EXIT,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+]);
+
+pub fn execute<const R: usize>(
+    a: &SVector<F16x1, R>,
+    b: &SVector<F16x1, R>,
+    c: &mut SVector<F16x1, R>,
+    dummy: &impl PimOperand,
+) {
+    a.fixed_rows_with_step::<8>(0, 16 * 16)
+        .iter()
+        .for_each(|entry| entry.execute_read());
+    b.fixed_rows_with_step::<8>(0, 16 * 16)
+        .iter()
+        .for_each(|entry| entry.execute_read());
+    c.fixed_rows_with_step_mut::<8>(0, 16 * 16)
+        .iter_mut()
+        .for_each(|entry| entry.execute_write());
+
+    dummy.execute_read();
+}