From b729c823b0fc7dc904f37acf8e4ceee5de37d73d Mon Sep 17 00:00:00 2001
From: Derek Christ <christ.derek@gmail.com>
Date: Tue, 7 Nov 2023 20:32:01 +0100
Subject: [PATCH] Define pim_config region and add some m5ops

---
 Cargo.toml      |  2 ++
 aarch64-gem5.ld |  2 ++
 src/m5ops.rs    | 29 +++++++++++++++++++++++++++++
 src/main.rs     | 21 +++++++++++++++++----
 src/pim.rs      | 40 ++++++++++++++++++++++++++++++++++++++++
 5 files changed, 90 insertions(+), 4 deletions(-)
 create mode 100644 src/m5ops.rs
 create mode 100644 src/pim.rs

diff --git a/Cargo.toml b/Cargo.toml
index eeabd12..dd40a7d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,6 +6,8 @@ edition = "2021"
 [dependencies]
 aarch64-cpu = "9.4.0"
 half = { version = "2.3.1", default-features = false }
+serde = { version = "1.0", default-features = false, features = ["derive"] }
+serde-json-core = "0.5.1"
 
 [profile.dev]
 panic = "abort"
diff --git a/aarch64-gem5.ld b/aarch64-gem5.ld
index b20f756..126be6d 100644
--- a/aarch64-gem5.ld
+++ b/aarch64-gem5.ld
@@ -8,6 +8,8 @@ ENTRY(_start)
 SECTIONS
 {
     .init : { *(.init) } > bootmem
+    .pim_config : { KEEP(*(.pim_config)) } > dram
+    # . = . + 0x4000;
     .text : { KEEP(*(.text)) } > dram
     .data : { *(.data) } > dram
     .rodata : { *(.rodata) } > dram
diff --git a/src/m5ops.rs b/src/m5ops.rs
new file mode 100644
index 0000000..1435428
--- /dev/null
+++ b/src/m5ops.rs
@@ -0,0 +1,29 @@
+pub fn exit() {
+    unsafe {
+        core::ptr::read_volatile((0x10010000 + (0x21 << 8)) as *mut u64);
+    }
+}
+
+pub fn checkpoint() {
+    unsafe {
+        core::ptr::read_volatile((0x10010000 + (0x43 << 8)) as *mut u64);
+    }
+}
+
+// pub fn dump_stats() {
+//     unsafe {
+//         core::ptr::read_volatile((0x10010000 + (0x40 << 8)) as *mut u64);
+//     }
+// }
+
+// pub fn reset_stats() {
+//     unsafe {
+//         core::ptr::read_volatile((0x10010000 + (0x41 << 8)) as *mut u64);
+//     }
+// }
+
+// pub fn dump_reset_stats() {
+//     unsafe {
+//         core::ptr::read_volatile((0x10010000 + (0x42 << 8)) as *mut u64);
+//     }
+// }
diff --git a/src/main.rs b/src/main.rs
index 752badf..96b9af1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,14 +6,27 @@ use core::arch::asm;
 use core::sync::atomic::{self, Ordering};
 use core::{arch::global_asm, fmt::Write, panic::PanicInfo};
 
+mod m5ops;
+mod pim;
 mod uart;
 
 global_asm!(include_str!("start.s"));
 
 #[no_mangle]
 pub extern "C" fn entry() -> ! {
-    let mut uart = Uart0 {};
+    let pim_config = pim::PimConfig {
+        bank_mode: pim::BankMode::SingleBank,
+    };
+    let mut pim_writer = pim::PimWriter;
 
+    write!(
+        &mut pim_writer,
+        "{}",
+        serde_json_core::to_string::<_, 256>(&pim_config).unwrap()
+    )
+    .unwrap();
+
+    let mut uart = Uart0 {};
     for i in 0..3 {
         writeln!(&mut uart, "Hello from Rust {i}!").unwrap();
     }
@@ -25,9 +38,9 @@ pub extern "C" fn entry() -> ! {
         }
     }
 
-    loop {
-        atomic::compiler_fence(Ordering::SeqCst);
-    }
+    m5ops::checkpoint();
+    m5ops::exit();
+    unreachable!();
 }
 
 #[panic_handler]
diff --git a/src/pim.rs b/src/pim.rs
new file mode 100644
index 0000000..ea013e9
--- /dev/null
+++ b/src/pim.rs
@@ -0,0 +1,40 @@
+use core::arch::asm;
+use core::fmt::Write;
+use serde::{Deserialize, Serialize};
+
+#[link_section = ".pim_config"]
+static mut PIM_CONFIG_REGION: [u8; 0x4000] = [0; 0x4000];
+
+pub struct PimWriter;
+
+impl Write for PimWriter {
+    fn write_str(&mut self, s: &str) -> core::fmt::Result {
+        unsafe {
+            PIM_CONFIG_REGION[..s.len()].copy_from_slice(s.as_bytes());
+            PIM_CONFIG_REGION[s.len()] = b'\0';
+
+            // Flush all cache lines that were affected by write operation
+            // A cache line is 64 bytes so we only need to flush every 64th virtual address
+            for element in PIM_CONFIG_REGION[..s.len()].iter().step_by(64) {
+                asm!("dc cvac, {val}", val = in(reg) element);
+            }
+
+            // Wait on all flushes to complete
+            asm!("dsb sy");
+        }
+
+        Ok(())
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct PimConfig {
+    pub bank_mode: BankMode,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub enum BankMode {
+    SingleBank,
+    AllBank,
+    PimAllBank,
+}