78 lines
1.8 KiB
Rust
78 lines
1.8 KiB
Rust
#![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,
|
|
};
|
|
|
|
#[cfg(feature = "X1")]
|
|
const ROWS: usize = 256;
|
|
|
|
#[cfg(feature = "X2")]
|
|
const ROWS: usize = 512;
|
|
|
|
#[cfg(feature = "X3")]
|
|
const ROWS: usize = 1024;
|
|
|
|
#[cfg(feature = "X4")]
|
|
const ROWS: usize = 2048;
|
|
|
|
const ELEMENTS_PER_BANK: usize = 16;
|
|
const BANKS: usize = 16;
|
|
const BLOCKS: usize = ROWS / (ELEMENTS_PER_BANK * BANKS);
|
|
|
|
#[no_mangle]
|
|
pub extern "C" fn main() {
|
|
#[cfg(feature = "X1")]
|
|
pim::state::set_kernel(&vadd::KERNEL_X1);
|
|
|
|
#[cfg(feature = "X2")]
|
|
pim::state::set_kernel(&vadd::KERNEL_X2);
|
|
|
|
#[cfg(feature = "X3")]
|
|
pim::state::set_kernel(&vadd::KERNEL_X3);
|
|
|
|
#[cfg(feature = "X4")]
|
|
pim::state::set_kernel(&vadd::KERNEL_X4);
|
|
|
|
let a = Box::new(pim::continuous_array::Vector(
|
|
SVector::<F16x1, ROWS>::from_fn(|i, _| F16x1(f16::from_f32(i as _))),
|
|
));
|
|
let b = Box::new(pim::continuous_array::Vector(
|
|
SVector::<F16x1, ROWS>::from_fn(|i, _| F16x1(f16::from_f32((ROWS - i) 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::<ROWS, BLOCKS>(&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();
|
|
}
|