feat(prototyper): add system suspend extension support

Signed-off-by: guttatus <xingzhiang2020@gmail.com>

prototyper/prototyper/src/main.rs

@@ -108,7 +108,7 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
     // Clear all pending IPIs.
     ipi::clear_all();
 
-    // Configure CSRs and trap handling.
+    // Configure CSRs
     unsafe {
         // Delegate all interrupts and exceptions to supervisor mode.
         asm!("csrw mideleg,    {}", in(reg) !0);

prototyper/prototyper/src/platform/mod.rs

@@ -29,6 +29,7 @@ use crate::sbi::logger;
 use crate::sbi::pmu::{EventToCounterMap, RawEventToCounterMap};
 use crate::sbi::reset::SbiReset;
 use crate::sbi::rfence::SbiRFence;
+use crate::sbi::suspend::SbiSuspend;
 
 mod clint;
 mod console;
@@ -87,7 +88,7 @@ impl Platform {
 
         // Get console device, init sbi console and logger.
         self.sbi_find_and_init_console(&root);
-        // Get clint and reset device, init sbi ipi, reset, hsm and rfence.
+        // Get clint and reset device, init sbi ipi, reset, hsm, rfence and susp extension.
         self.sbi_init_ipi_reset_hsm_rfence(&root);
         // Initialize pmu extension
         self.sbi_init_pmu(&root);
@@ -158,6 +159,7 @@ impl Platform {
         self.sbi_hsm_init();
         self.sbi_reset_init();
         self.sbi_rfence_init();
+        self.sbi_susp_init();
     }
 
     fn sbi_init_pmu(&mut self, root: &serde_device_tree::buildin::Node) {
@@ -255,6 +257,12 @@ impl Platform {
             let hart_id = cpu.reg.iter().next().unwrap().0.start;
             if let Some(x) = cpu_list.get_mut(hart_id) {
                 *x = true;
+            } else {
+                error!(
+                    "The maximum supported hart id is {}, but the hart id {} was obtained. Please check the config!",
+                    NUM_HART_MAX - 1,
+                    hart_id
+                );
             }
         }
         self.info.cpu_enabled = Some(cpu_list);
@@ -326,6 +334,14 @@ impl Platform {
         }
     }
 
+    fn sbi_susp_init(&mut self) {
+        if self.sbi.hsm.is_some() {
+            self.sbi.susp = Some(SbiSuspend);
+        } else {
+            self.sbi.susp = None;
+        }
+    }
+
     pub fn print_board_info(&self) {
         info!("RustSBI version {}", rustsbi::VERSION);
         rustsbi::LOGO.lines().for_each(|line| info!("{}", line));
@@ -372,6 +388,7 @@ impl Platform {
         self.print_reset_info();
         self.print_hsm_info();
         self.print_rfence_info();
+        self.print_susp_info();
         self.print_pmu_info();
     }
 
@@ -415,28 +432,29 @@ impl Platform {
 
     #[inline]
     fn print_hsm_info(&self) {
-        info!(
-            "{:<30}: {}",
-            "Platform HSM Extension",
-            if self.have_hsm() {
-                "Available"
-            } else {
-                "Not Available"
-            }
-        );
+        if self.have_hsm() {
+            info!("{:<30}: {}", "Platform HSM Extension", "Available");
+        } else {
+            warn!("{:<30}: {}", "Platform HSM Extension", "Not Available");
+        }
     }
 
     #[inline]
     fn print_rfence_info(&self) {
-        info!(
-            "{:<30}: {}",
-            "Platform RFence Extension",
-            if self.have_rfence() {
-                "Available"
-            } else {
-                "Not Available"
-            }
-        );
+        if self.have_rfence() {
+            info!("{:<30}: {}", "Platform RFence Extension", "Available");
+        } else {
+            warn!("{:<30}: {}", "Platform RFence Extension", "Not Available");
+        }
+    }
+
+    #[inline]
+    fn print_susp_info(&self) {
+        if self.have_susp() {
+            info!("{:<30}: {}", "Platform SUSP Extension", "Available");
+        } else {
+            warn!("{:<30}: {}", "Platform SUSP Extension", "Not Available");
+        }
     }
 
     #[inline]
@@ -498,6 +516,10 @@ impl Platform {
         self.sbi.rfence.is_some()
     }
 
+    pub fn have_susp(&self) -> bool {
+        self.sbi.susp.is_some()
+    }
+
     pub fn have_pmu(&self) -> bool {
         self.sbi.pmu.is_some()
     }

prototyper/prototyper/src/sbi/features.rs

@@ -15,6 +15,12 @@ pub struct HartFeatures {
     mhpm_bits: u32,
 }
 
+impl HartFeatures {
+    pub fn privileged_version(&self) -> PrivilegedVersion {
+        self.privileged_version
+    }
+}
+
 #[derive(Copy, Clone)]
 pub enum Extension {
     Sstc = 0,

prototyper/prototyper/src/sbi/hart_context.rs

@@ -1,8 +1,10 @@
 use crate::sbi::features::HartFeatures;
+use crate::sbi::features::PrivilegedVersion;
 use crate::sbi::hsm::HsmCell;
 use crate::sbi::rfence::RFenceCell;
 use core::ptr::NonNull;
 use core::sync::atomic::AtomicU8;
+use core::sync::atomic::Ordering;
 use fast_trap::FlowContext;
 use riscv::register::mstatus;
 
@@ -38,6 +40,36 @@ impl HartContext {
     pub fn context_ptr(&mut self) -> NonNull<FlowContext> {
         unsafe { NonNull::new_unchecked(&mut self.trap) }
     }
+
+    #[inline]
+    pub fn reset(&mut self) {
+        self.ipi_reset();
+        self.rfence_reset();
+        self.pmu_state_reset();
+    }
+
+    #[inline]
+    fn rfence_reset(&mut self) {
+        self.rfence = RFenceCell::new();
+    }
+
+    #[inline]
+    fn ipi_reset(&mut self) {
+        self.ipi_type.store(0, Ordering::Relaxed);
+    }
+
+    #[inline]
+    fn pmu_state_reset(&mut self) {
+        // stop all hardware pmu event
+        let hart_priv_version = self.features.privileged_version();
+        if hart_priv_version >= PrivilegedVersion::Version1_11 {
+            unsafe {
+                core::arch::asm!("csrw mcountinhibit, {}", in(reg) !0b10);
+            }
+        }
+        // reset hart pmu state
+        self.pmu_state = PmuState::new();
+    }
 }
 
 /// Information needed to boot into the next execution stage.

prototyper/prototyper/src/sbi/hsm.rs

@@ -10,8 +10,9 @@ use crate::platform::PLATFORM;
 use crate::riscv::current_hartid;
 use crate::sbi::hart_context::NextStage;
 use crate::sbi::trap_stack::ROOT_STACK;
+use crate::trap_stack::hart_context_mut;
 
-use super::trap_stack::hart_context;
+use super::{trap::boot::boot, trap_stack::hart_context};
 
 /// Special state indicating a hart is in the process of starting.
 const HART_STATE_START_PENDING_EXT: usize = usize::MAX;
@@ -131,10 +132,36 @@ impl<T: core::fmt::Debug> RemoteHsmCell<'_, T> {
         }
     }
 
+    /// Attempts to resume a suspended hart by providing resume data.
+    ///
+    /// Returns true if successful, false if hart was not in SUSPENDED state.
+    #[inline]
+    pub fn resume(&self, t: T) -> bool {
+        if self
+            .0
+            .status
+            .compare_exchange(
+                hart_state::SUSPENDED,
+                HART_STATE_START_PENDING_EXT,
+                Ordering::Acquire,
+                Ordering::Relaxed,
+            )
+            .is_ok()
+        {
+            unsafe { *self.0.inner.get() = Some(t) };
+            self.0
+                .status
+                .store(hart_state::START_PENDING, Ordering::Release);
+            true
+        } else {
+            false
+        }
+    }
+
     /// Gets the current state of the hart.
     #[allow(unused)]
     #[inline]
-    pub fn sbi_get_status(&self) -> usize {
+    pub fn get_status(&self) -> usize {
         match self.0.status.load(Ordering::Relaxed) {
             HART_STATE_START_PENDING_EXT => hart_state::START_PENDING,
             normal => normal,
@@ -212,33 +239,67 @@ impl rustsbi::Hsm for SbiHsm {
     #[inline]
     fn hart_get_status(&self, hartid: usize) -> SbiRet {
         match remote_hsm(hartid) {
-            Some(remote) => SbiRet::success(remote.sbi_get_status()),
+            Some(remote) => SbiRet::success(remote.get_status()),
             None => SbiRet::invalid_param(),
         }
     }
 
     /// Suspends execution on the current hart.
-    fn hart_suspend(&self, suspend_type: u32, _resume_addr: usize, _opaque: usize) -> SbiRet {
+    fn hart_suspend(&self, suspend_type: u32, resume_addr: usize, opaque: usize) -> SbiRet {
         use rustsbi::spec::hsm::suspend_type::{NON_RETENTIVE, RETENTIVE};
-        if matches!(suspend_type, NON_RETENTIVE | RETENTIVE) {
-            unsafe {
-                PLATFORM
-                    .sbi
-                    .ipi
-                    .as_ref()
-                    .unwrap()
-                    .clear_msip(current_hartid());
+
+        if !matches!(suspend_type, NON_RETENTIVE | RETENTIVE) {
+            return SbiRet::invalid_param();
+        }
+
+        crate::sbi::trap::handler::msoft_ipi_handler();
+        unsafe {
+            PLATFORM
+                .sbi
+                .ipi
+                .as_ref()
+                .unwrap()
+                .clear_msip(current_hartid());
+        }
+        unsafe {
+            riscv::register::mie::set_msoft();
+        }
+        local_hsm().suspend();
+        riscv::asm::wfi();
+        crate::sbi::trap::handler::msoft_ipi_handler();
+
+        match suspend_type {
+            RETENTIVE => {
+                local_hsm().resume();
+                return SbiRet::success(0);
             }
-            unsafe {
-                riscv::register::mie::set_msoft();
+            NON_RETENTIVE => return self.hart_resume(current_hartid(), resume_addr, opaque),
+            _ => return SbiRet::invalid_param(),
+        }
+    }
+}
+
+impl SbiHsm {
+    // non retentive resume
+    fn hart_resume(&self, hartid: usize, resume_addr: usize, opaque: usize) -> SbiRet {
+        match remote_hsm(hartid) {
+            Some(remote) => {
+                if remote.resume(NextStage {
+                    start_addr: resume_addr,
+                    opaque,
+                    next_mode: MPP::Supervisor,
+                }) {
+                    // reset the hart local context to prevent the hart context from being polluted
+                    hart_context_mut(hartid).reset();
+                    // boot resume hart from resume addr
+                    unsafe {
+                        boot();
+                    }
+                } else {
+                    SbiRet::failed()
+                }
             }
-            local_hsm().suspend();
-            riscv::asm::wfi();
-            crate::sbi::trap::handler::msoft_ipi_handler();
-            local_hsm().resume();
-            SbiRet::success(0)
-        } else {
-            SbiRet::not_supported()
+            None => SbiRet::failed(),
         }
     }
 }

prototyper/prototyper/src/sbi/mod.rs

@@ -6,6 +6,7 @@ pub mod ipi;
 pub mod pmu;
 pub mod reset;
 pub mod rfence;
+pub mod suspend;
 
 pub mod early_trap;
 pub mod features;
@@ -22,6 +23,7 @@ use ipi::SbiIpi;
 use pmu::SbiPmu;
 use reset::SbiReset;
 use rfence::SbiRFence;
+use suspend::SbiSuspend;
 
 #[derive(RustSBI, Default)]
 #[rustsbi(dynamic)]
@@ -39,6 +41,8 @@ pub struct SBI {
     pub rfence: Option<SbiRFence>,
     #[rustsbi(pmu)]
     pub pmu: Option<SbiPmu>,
+    #[rustsbi(susp)]
+    pub susp: Option<SbiSuspend>,
 }
 
 impl SBI {
@@ -50,6 +54,7 @@ impl SBI {
             reset: None,
             rfence: None,
             pmu: None,
+            susp: None,
         }
     }
 }

prototyper/prototyper/src/sbi/suspend.rs

@@ -0,0 +1,55 @@
+use riscv::register::mstatus;
+use rustsbi::{Hsm, SbiRet};
+use sbi_spec::hsm::{hart_state::STOPPED, suspend_type::NON_RETENTIVE};
+
+use crate::{platform::PLATFORM, riscv::current_hartid};
+
+use super::hsm::remote_hsm;
+
+const SUSPEND_TO_RAM: u32 = 0x0;
+
+/// Implementation of SBI System Suspend Extension extension.
+pub(crate) struct SbiSuspend;
+
+impl rustsbi::Susp for SbiSuspend {
+    fn system_suspend(&self, sleep_type: u32, resume_addr: usize, opaque: usize) -> SbiRet {
+        if sleep_type != SUSPEND_TO_RAM {
+            return SbiRet::invalid_param();
+        }
+
+        let prev_mode = mstatus::read().mpp();
+        if prev_mode != mstatus::MPP::Supervisor || prev_mode != mstatus::MPP::User {
+            return SbiRet::failed();
+        }
+
+        // Check if all harts except the current hart are stopped
+        let hart_enable_map = if let Some(hart_enable_map) = unsafe { PLATFORM.info.cpu_enabled } {
+            hart_enable_map
+        } else {
+            return SbiRet::failed();
+        };
+        for (hartid, hart_enable) in hart_enable_map.iter().enumerate() {
+            if *hart_enable && hartid != current_hartid() {
+                match remote_hsm(hartid) {
+                    Some(remote) => {
+                        if remote.get_status() != STOPPED {
+                            return SbiRet::denied();
+                        }
+                    }
+                    None => return SbiRet::failed(),
+                }
+            }
+        }
+
+        // TODO: The validity of `resume_addr` should be checked.
+        // If it is invalid, `SBI_ERR_INVALID_ADDRESS` should be returned.
+
+        if let Some(hsm) = unsafe { &PLATFORM.sbi.hsm } {
+            hsm.hart_suspend(NON_RETENTIVE, resume_addr, opaque);
+        } else {
+            return SbiRet::not_supported();
+        }
+
+        unreachable!();
+    }
+}

prototyper/prototyper/src/sbi/trap/boot.rs

@@ -1,6 +1,7 @@
 use crate::riscv::current_hartid;
 use crate::sbi::hsm::local_hsm;
 use crate::sbi::ipi;
+use crate::sbi::trap_stack;
 use core::arch::naked_asm;
 use riscv::register::{mie, mstatus, satp, sstatus};
 
@@ -12,8 +13,9 @@ pub unsafe extern "C" fn boot() -> ! {
     unsafe {
         naked_asm!(
             ".align 2",
-            // Switch stacks
-            "csrrw  sp, mscratch, sp",
+            // Reset hart local stack
+            "call    {locate_stack}",
+            "csrw    mscratch, sp",
             // Allocate stack space
             "addi   sp, sp, -3*8",
             // Call handler with context pointer
@@ -23,15 +25,15 @@ pub unsafe extern "C" fn boot() -> ! {
             "ld     t0, 0*8(sp)
             csrw    mepc, t0",
             // Restore registers
-            "
-        ld      a0, 1*8(sp)
-        ld      a1, 2*8(sp)",
+            "ld      a0, 1*8(sp)",
+            "ld      a1, 2*8(sp)",
             // Restore stack pointer
-            "addi   sp, sp, 3*8",
+            "add     sp, sp, 3*8",
             // Switch stacks back
             "csrrw  sp, mscratch, sp",
             // Return from machine mode
             "mret",
+            locate_stack = sym trap_stack::locate,
             boot_handler = sym boot_handler,
         );
     }

prototyper/test-kernel/src/main.rs

@@ -116,7 +116,19 @@ extern "C" fn rust_main(hartid: usize, dtb_pa: usize) -> ! {
     };
     let test_result = testing.test();
 
-    // PMU test, only available in qemu-system-riscv64 single core
+    pmu_test();
+
+    if test_result {
+        sbi::system_reset(sbi::Shutdown, sbi::NoReason);
+    } else {
+        sbi::system_reset(sbi::Shutdown, sbi::SystemFailure);
+    }
+    unreachable!()
+}
+
+#[inline]
+// PMU test, only available in qemu-system-riscv64 single core
+fn pmu_test() {
     let counters_num = sbi::pmu_num_counters();
     println!("[pmu] counters number: {}", counters_num);
     for idx in 0..counters_num {
@@ -278,13 +290,6 @@ extern "C" fn rust_main(hartid: usize, dtb_pa: usize) -> ! {
     let ipi_num = sbi::pmu_counter_fw_read(ipi_counter_idx);
     assert!(ipi_num.is_ok());
     assert_eq!(ipi_num.value, 27);
-
-    if test_result {
-        sbi::system_reset(sbi::Shutdown, sbi::NoReason);
-    } else {
-        sbi::system_reset(sbi::Shutdown, sbi::SystemFailure);
-    }
-    unreachable!()
 }
 
 #[cfg_attr(not(test), panic_handler)]