Trivial matrix multiplication kernel

2024-01-01 23:08:46 +01:00
parent 05f184d51f
commit 6380385bd0
3 changed files with 60 additions and 167 deletions
--- a/pim-os/src/main.rs
+++ b/pim-os/src/main.rs
@@ -32,18 +32,18 @@ pub extern "C" fn entry() -> ! {
    pim_state.set_kernel();

    let pim_matrix_arena0 = RefCell::new(PimMatrixArena(
-        [[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
+        [[[F16x16::default(); NUMBER_OF_BANKS]; 3]; 3],
    ));
    let pim_matrix_arena1 = RefCell::new(PimMatrixArena(
-        [[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
+        [[[F16x16::default(); NUMBER_OF_BANKS]; 3]; 3],
    ));
    let pim_matrix_arena2 = RefCell::new(PimMatrixArena(
-        [[[F16x16::default(); NUMBER_OF_BANKS]; 8]; 8],
+        [[[F16x16::default(); NUMBER_OF_BANKS]; 3]; 3],
    ));
    let pim_storage0 = PimStorage {
        arena: &pim_matrix_arena0,
        index: 0,
-        row_major: true,
+        row_major: false,
    };
    let pim_storage1 = PimStorage {
        arena: &pim_matrix_arena1,
@@ -57,7 +57,7 @@ pub extern "C" fn entry() -> ! {
    };

    let mut matrix0 = Matrix::from_data(pim_storage0);
-    matrix0.fill(F16x1(f16::ONE));
+    matrix0.fill_lower_triangle(F16x1(f16::ONE), 0);

    let mut matrix1 = Matrix::from_data(pim_storage1);
    matrix1.fill_lower_triangle(F16x1(f16::ONE), 0);
--- a/pim-os/src/pim/kernel.rs
+++ b/pim-os/src/pim/kernel.rs
@@ -116,26 +116,6 @@ pub const MATRIX_MUL: Kernel = Kernel([
        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::MAC {
        src0: File::Bank,
        src1: File::GrfA { index: 0 },
@@ -157,73 +137,10 @@ pub const MATRIX_MUL: Kernel = Kernel([
        dst: File::GrfB { index: 0 },
        aam: false,
    },
-    Instruction::MAC {
-        src0: File::Bank,
-        src1: File::GrfA { index: 3 },
-        src2: File::GrfB { index: 0 },
-        dst: File::GrfB { index: 0 },
-        aam: false,
-    },
-    Instruction::MAC {
-        src0: File::Bank,
-        src1: File::GrfA { index: 4 },
-        src2: File::GrfB { index: 0 },
-        dst: File::GrfB { index: 0 },
-        aam: false,
-    },
-    Instruction::MAC {
-        src0: File::Bank,
-        src1: File::GrfA { index: 5 },
-        src2: File::GrfB { index: 0 },
-        dst: File::GrfB { index: 0 },
-        aam: false,
-    },
-    Instruction::MAC {
-        src0: File::Bank,
-        src1: File::GrfA { index: 6 },
-        src2: File::GrfB { index: 0 },
-        dst: File::GrfB { index: 0 },
-        aam: false,
-    },
-    Instruction::MAC {
-        src0: File::Bank,
-        src1: File::GrfA { index: 7 },
-        src2: File::GrfB { index: 0 },
-        dst: File::GrfB { index: 0 },
-        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,
@@ -232,6 +149,23 @@ pub const MATRIX_MUL: Kernel = Kernel([
    Instruction::NOP,
    Instruction::NOP,
    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
+    Instruction::NOP,
 ]);

 pub fn execute_matrix_multiply<const R: usize, const C: usize>(
@@ -243,82 +177,39 @@ pub fn execute_matrix_multiply<const R: usize, const C: usize>(
 ) {
    pim_state.set_bank_mode(BankMode::PimAllBank);

-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(0);
-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(1);
-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(2);
-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(3);
-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(4);
-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(5);
-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(6);
-    pim_matrix_arena0
-        .borrow()
-        .execute_instruction_read_single_bank(7);
+    let mut index = 0;
+    while index < (R * C) {
+        let left_index = index % R;
+        let right_index = (index / R) * R;

-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(0);
-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(1);
-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(2);
-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(3);
-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(4);
-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(5);
-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(6);
-    pim_matrix_arena1
-        .borrow()
-        .execute_instruction_read_single_bank(7);
+        pim_matrix_arena0
+            .borrow()
+            .execute_instruction_read_single_bank(left_index);
+        pim_matrix_arena0
+            .borrow()
+            .execute_instruction_read_single_bank(left_index + R);
+        pim_matrix_arena0
+            .borrow()
+            .execute_instruction_read_single_bank(left_index + R * 2);

-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(0);
-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(1);
-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(2);
-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(3);
-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(4);
-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(5);
-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(6);
-    pim_matrix_arena2
-        .borrow_mut()
-        .execute_instruction_write_single_bank(7);
+        pim_matrix_arena1
+            .borrow()
+            .execute_instruction_read_single_bank(right_index);
+        pim_matrix_arena1
+            .borrow()
+            .execute_instruction_read_single_bank(right_index + 1);
+        pim_matrix_arena1
+            .borrow()
+            .execute_instruction_read_single_bank(right_index + 2);

-    dummy_array.execute_instruction_read_single_bank(0);
+        pim_matrix_arena2
+            .borrow_mut()
+            .execute_instruction_write_single_bank(index);
+
+        dummy_array.execute_instruction_read_single_bank(0);
+
+        index += 1;
+    }

    pim_state.set_bank_mode(BankMode::SingleBank);
 }
--- a/pim-vm/src/lib.rs
+++ b/pim-vm/src/lib.rs
@@ -48,6 +48,12 @@ struct PimUnit {
    jump_counter: Option<u16>,
 }

+impl PimUnit {
+    fn reset(&mut self) {
+        *self = Self::default();
+    }
+}
+
 impl Default for PimUnit {
    fn default() -> Self {
        Self {
@@ -71,8 +77,7 @@ struct PimVM {
 impl PimVM {
    fn reset(&mut self) {
        for unit in self.pim_units.iter_mut() {
-            unit.pc = 0;
-            unit.jump_counter = None;
+            unit.reset();
        }
    }

@@ -166,10 +171,7 @@ impl PimVM {

        match inst {
            Instruction::NOP => (),
-            Instruction::EXIT => {
-                pim_unit.jump_counter = None;
-                pim_unit.pc = 0;
-            }
+            Instruction::EXIT => pim_unit.reset(),
            Instruction::JUMP { .. } => unreachable!(),
            Instruction::MOV { src, dst } | Instruction::FILL { src, dst } => {
                let data = PimVM::load(src, pim_unit, &bank_data);