feat(prototyper): refactoring the extension probe mechanism

prototyper/src/dt.rs

@@ -28,6 +28,7 @@ pub struct Cpus<'a> {
 pub struct Cpu<'a> {
     #[serde(rename = "riscv,isa-extensions")]
     pub isa: Option<StrSeq<'a>>,
+    pub reg: Reg<'a>,
 }
 
 #[derive(Deserialize, Debug)]

prototyper/src/main.rs

@@ -18,9 +18,12 @@ mod sbi;
 use core::sync::atomic::{AtomicBool, Ordering};
 use core::{arch::asm, mem::MaybeUninit};
 
+use sbi::extensions;
+
 use crate::board::{SBI_IMPL, SIFIVECLINT, SIFIVETEST, UART};
 use crate::riscv_spec::{current_hartid, menvcfg};
 use crate::sbi::console::SbiConsole;
+use crate::sbi::extensions::{hart_extension_probe, Extension};
 use crate::sbi::hart_context::NextStage;
 use crate::sbi::hsm::{local_remote_hsm, SbiHsm};
 use crate::sbi::ipi::{self, SbiIpi};
@@ -28,7 +31,7 @@ use crate::sbi::logger;
 use crate::sbi::reset::SbiReset;
 use crate::sbi::rfence::SbiRFence;
 use crate::sbi::trap::{self, trap_vec};
-use crate::sbi::trap_stack::{self, NUM_HART_MAX};
+use crate::sbi::trap_stack;
 use crate::sbi::SBI;
 
 #[no_mangle]
@@ -37,21 +40,26 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
     static SBI_READY: AtomicBool = AtomicBool::new(false);
 
     let boot_hart_info = platform::get_boot_hart(opaque, nonstandard_a2);
-
+    // boot hart task entry
     if boot_hart_info.is_boot_hart {
         let fdt_addr = boot_hart_info.fdt_address;
 
         // 1. Init FDT
         // parse the device tree
+        // TODO: shoule remove `fail:device_tree_format`
         let dtb = dt::parse_device_tree(fdt_addr).unwrap_or_else(fail::device_tree_format);
         let dtb = dtb.share();
+
+        // TODO: should remove `fail:device_tree_deserialize`
         let tree =
             serde_device_tree::from_raw_mut(&dtb).unwrap_or_else(fail::device_tree_deserialize);
 
+
         // 2. Init device
         // TODO: The device base address should be find in a better way
         let console_base = tree.soc.serial.unwrap().iter().next().unwrap();
         let clint_device = tree.soc.clint.unwrap().iter().next().unwrap();
+        let cpu_num = tree.cpus.cpu.len();
         let console_base_address = console_base.at();
         let ipi_base_address = clint_device.at();
 
@@ -64,44 +72,32 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
 
         board::console_dev_init(usize::from_str_radix(console_base_address, 16).unwrap());
         board::ipi_dev_init(usize::from_str_radix(ipi_base_address, 16).unwrap());
-        // Assume sstc is enabled only if all hart has sstc ext
-        let sstc_support = tree
-            .cpus
-            .cpu
-            .iter()
-            .map(|cpu_iter| {
-                use crate::dt::Cpu;
-                let cpu = cpu_iter.deserialize::<Cpu>();
-                let isa = match cpu.isa {
-                    Some(value) => value,
-                    None => return false,
-                };
-                isa.iter().any(|x| x == "sstc")
-            })
-            .all(|x| x);
-        #[cfg(feature = "nemu")]
-        let sstc_support = true;
+
         // 3. Init SBI
         unsafe {
             SBI_IMPL = MaybeUninit::new(SBI {
                 console: Some(SbiConsole::new(&UART)),
-                ipi: Some(SbiIpi::new(&SIFIVECLINT, NUM_HART_MAX, sstc_support)),
+                ipi: Some(SbiIpi::new(&SIFIVECLINT, cpu_num)),
                 hsm: Some(SbiHsm),
                 reset: Some(SbiReset::new(&SIFIVETEST)),
                 rfence: Some(SbiRFence),
             });
         }
+        // 设置陷入栈
+        trap_stack::prepare_for_trap();
+        extensions::init(&tree.cpus.cpu);
         SBI_READY.swap(true, Ordering::AcqRel);
         // 4. Init Logger
         logger::Logger::init();
+
         info!("RustSBI version {}", rustsbi::VERSION);
         rustsbi::LOGO.lines().for_each(|line| info!("{}", line));
         info!("Initializing RustSBI machine-mode environment.");
 
+        info!("Number of CPU: {}", cpu_num);
         if let Some(model) = tree.model {
             info!("Model: {}", model.iter().next().unwrap_or("<unspecified>"));
         }
-        info!("Support sstc: {sstc_support}");
         info!("Clint device: {}", ipi_base_address);
         info!("Console deivce: {}", console_base_address);
         info!(
@@ -122,8 +118,6 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
             pmpaddr1::write(usize::MAX >> 2);
         }
 
-        // 设置陷入栈
-        trap_stack::prepare_for_trap();
 
         let boot_info = platform::get_boot_info(nonstandard_a2);
         let (mpp, next_addr) = (boot_info.mpp, boot_info.next_address);
@@ -167,7 +161,7 @@ extern "C" fn rust_main(_hart_id: usize, opaque: usize, nonstandard_a2: usize) {
         use riscv::register::{medeleg, mtvec};
         medeleg::clear_supervisor_env_call();
         medeleg::clear_illegal_instruction();
-        if ipi::has_sstc() {
+        if hart_extension_probe(current_hartid(),Extension::SSTC) {
             menvcfg::set_bits(
                 menvcfg::STCE | menvcfg::CBIE_INVALIDATE | menvcfg::CBCFE | menvcfg::CBZE,
             );
@@ -230,8 +224,7 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
         "[rustsbi-panic] hart {} {info}",
         riscv::register::mhartid::read()
     );
-    println!(
-        "-----------------------------
+    println!("-----------------------------
 > mcause:  {:?}
 > mepc:    {:#018x}
 > mtval:   {:#018x}

prototyper/src/sbi/extensions.rs

@@ -0,0 +1,61 @@
+use serde_device_tree::buildin::NodeSeq;
+
+use crate::sbi::trap_stack::ROOT_STACK;
+pub struct HartExtensions([bool; Extension::COUNT]);
+
+#[derive(Copy, Clone)]
+pub enum Extension {
+    SSTC = 0,
+}
+
+impl Extension {
+    const COUNT: usize = 1;
+    const ITER: [Self;Extension::COUNT] = [Extension::SSTC];
+
+    pub fn to_str(&self) -> &'static str {
+        match self {
+            Extension::SSTC => "sstc",
+        }
+    }
+
+    #[inline]
+    pub fn index(&self) -> usize {
+        *self as usize
+    }
+}
+
+pub fn hart_extension_probe(hart_id: usize, ext: Extension) -> bool {
+    unsafe {
+        ROOT_STACK
+            .get_mut(hart_id)
+            .map(|x| x.hart_context().extensions.0[ext.index()]).unwrap()
+    }
+}
+
+pub fn init(cpus: &NodeSeq) {
+    use crate::dt::Cpu;
+    for cpu_iter in cpus.iter() {
+        let cpu = cpu_iter.deserialize::<Cpu>();
+        let hart_id = cpu.reg.iter().next().unwrap().0.start;
+        let mut hart_exts = [false;Extension::COUNT];
+        let isa = cpu.isa.unwrap();
+        Extension::ITER.iter().for_each(|ext| {
+            if isa.iter().any(|e| e == ext.to_str()) {
+                hart_exts[ext.index()] = true;
+            } else {
+                hart_exts[ext.index()] = false;
+            }
+        });
+
+        #[cfg(feature = "nemu")] 
+        {
+            hart_exts[Extension::SSTC.index()] = true;
+        }
+
+        unsafe {
+            ROOT_STACK
+                .get_mut(hart_id)
+                .map(|stack| stack.hart_context().extensions = HartExtensions(hart_exts)).unwrap()
+        }
+    }
+}

prototyper/src/sbi/hart_context.rs

@@ -1,10 +1,10 @@
 use core::ptr::NonNull;
+use core::sync::atomic::AtomicU8;
 use fast_trap::FlowContext;
 use riscv::register::mstatus;
-
 use crate::sbi::hsm::HsmCell;
 use crate::sbi::rfence::RFenceCell;
-use core::sync::atomic::AtomicU8;
+use crate::sbi::extensions::HartExtensions;
 
 pub(crate) struct HartContext {
     /// trap context
@@ -12,6 +12,7 @@ pub(crate) struct HartContext {
     pub hsm: HsmCell<NextStage>,
     pub rfence: RFenceCell,
     pub ipi_type: AtomicU8,
+    pub extensions: HartExtensions,
 }
 
 impl HartContext {

prototyper/src/sbi/ipi.rs

@@ -7,6 +7,7 @@ use crate::sbi::hsm::remote_hsm;
 use crate::sbi::rfence;
 use crate::sbi::trap;
 use crate::sbi::trap_stack::ROOT_STACK;
+use crate::sbi::extensions::{hart_extension_probe, Extension};
 
 pub(crate) const IPI_TYPE_SSOFT: u8 = 1 << 0;
 pub(crate) const IPI_TYPE_FENCE: u8 = 1 << 1;
@@ -25,13 +26,12 @@ pub trait IpiDevice {
 pub struct SbiIpi<'a, T: IpiDevice> {
     pub ipi_dev: &'a AtomicPtr<T>,
     pub max_hart_id: usize,
-    pub sstc_support: bool,
 }
 
 impl<'a, T: IpiDevice> rustsbi::Timer for SbiIpi<'a, T> {
     #[inline]
     fn set_timer(&self, stime_value: u64) {
-        if self.sstc_support {
+        if hart_extension_probe(current_hartid(),Extension::SSTC) {
             stimecmp::set(stime_value);
             unsafe {
                 riscv::register::mie::set_mtimer();
@@ -64,11 +64,10 @@ impl<'a, T: IpiDevice> rustsbi::Ipi for SbiIpi<'a, T> {
 }
 
 impl<'a, T: IpiDevice> SbiIpi<'a, T> {
-    pub fn new(ipi_dev: &'a AtomicPtr<T>, max_hart_id: usize, sstc_support: bool) -> Self {
+    pub fn new(ipi_dev: &'a AtomicPtr<T>, max_hart_id: usize) -> Self {
         Self {
             ipi_dev,
             max_hart_id,
-            sstc_support,
         }
     }
 
@@ -175,12 +174,4 @@ pub fn clear_all() {
         .as_ref()
         .unwrap()
         .clear();
-}
-
-pub fn has_sstc() -> bool {
-    unsafe { SBI_IMPL.assume_init_ref() }
-        .ipi
-        .as_ref()
-        .unwrap()
-        .sstc_support
-}
+}

prototyper/src/sbi/mod.rs

@@ -11,6 +11,7 @@ pub mod hart_context;
 pub mod logger;
 pub mod trap;
 pub mod trap_stack;
+pub mod extensions;
 
 use console::{ConsoleDevice, SbiConsole};
 use hsm::SbiHsm;

prototyper/src/sbi/trap_stack.rs

@@ -1,6 +1,5 @@
 use core::mem::forget;
 use fast_trap::FreeTrapStack;
-
 use crate::riscv_spec::current_hartid;
 use crate::sbi::hart_context::HartContext;
 use crate::sbi::trap::fast_handler;