derek/pim-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Trivial matrix multiplication kernel

Browse Source
This commit is contained in:
Derek Christ
2024-01-01 23:08:46 +01:00
parent 05f184d51f
commit 6380385bd0
3 changed files with 60 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
pim-os/src/main.rs
View File

@@ -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);

173
pim-os/src/pim/kernel.rs
View File

@@ -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);
let mut index = 0;
while index < (R * C) {
let left_index = index % R;
let right_index = (index / R) * R;
pim_matrix_arena0
.borrow()
.execute_instruction_read_single_bank(0);
.execute_instruction_read_single_bank(left_index);
pim_matrix_arena0
.borrow()
.execute_instruction_read_single_bank(1);
.execute_instruction_read_single_bank(left_index + R);
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);
.execute_instruction_read_single_bank(left_index + R * 2);
pim_matrix_arena1
.borrow()
.execute_instruction_read_single_bank(0);
.execute_instruction_read_single_bank(right_index);
pim_matrix_arena1
.borrow()
.execute_instruction_read_single_bank(1);
.execute_instruction_read_single_bank(right_index + 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);
.execute_instruction_read_single_bank(right_index + 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);
.execute_instruction_write_single_bank(index);
dummy_array.execute_instruction_read_single_bank(0);
index += 1;
}
pim_state.set_bank_mode(BankMode::SingleBank);
}

14
pim-vm/src/lib.rs
View File

@@ -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);