Merge pull request #1685 from stlankes/smp

Improve aarch64 support

Author: stlankes

.github/workflows/ci.yml

@@ -241,3 +241,4 @@ jobs:
       - run: UHYVE=$CARGO_HOME/bin/uhyve cargo xtask ci rs --arch ${{ matrix.arch }} --profile ${{ matrix.profile }} --package rusty_demo --smp 4 uhyve --sudo
         if: matrix.arch == 'x86_64'
       - run: cargo xtask ci rs --arch ${{ matrix.arch }} --profile ${{ matrix.profile }} --package thread_test --smp 4 qemu ${{ matrix.flags }}
+        if: matrix.arch == 'x86_64'

src/arch/aarch64/kernel/core_local.rs

@@ -32,7 +32,7 @@ pub(crate) struct CoreLocal {
 
 impl CoreLocal {
 	pub fn install() {
-		let core_id = CPU_ONLINE.load(Ordering::Relaxed);
+		let core_id = CPU_ONLINE.0.load(Ordering::Relaxed);
 
 		let irq_statistics = if core_id == 0 {
 			static FIRST_IRQ_STATISTICS: IrqStatistics = IrqStatistics::new();

src/arch/aarch64/kernel/interrupts.rs

@@ -46,9 +46,11 @@ pub(crate) static GIC: SpinMutex<Option<GicV3<'_>>> = SpinMutex::new(None);
 pub fn enable() {
 	unsafe {
 		asm!(
+			"dmb ish",
 			"msr daifclr, {mask}",
+			"dmb ish",
 			mask = const 0b111,
-			options(nostack, nomem),
+			options(nostack),
 		);
 	}
 }
@@ -61,9 +63,11 @@ pub fn enable() {
 pub fn enable_and_wait() {
 	unsafe {
 		asm!(
+			"dmb ish",
 			"msr daifclr, {mask}; wfi",
+			"dmb ish",
 			mask = const 0b111,
-			options(nostack, nomem),
+			options(nostack),
 		);
 	}
 }
@@ -73,9 +77,11 @@ pub fn enable_and_wait() {
 pub fn disable() {
 	unsafe {
 		asm!(
+			"dmb ish",
 			"msr daifset, {mask}",
+			"dmb ish",
 			mask = const 0b111,
-			options(nostack, nomem),
+			options(nostack),
 		);
 	}
 }
@@ -206,8 +212,14 @@ pub(crate) extern "C" fn do_sync(state: &State) {
 		}
 	} else if ec == 0x3c {
 		error!("Trap to debugger, PC={pc:#x}");
+		loop {
+			core::hint::spin_loop();
+		}
 	} else {
 		error!("Unsupported exception class: {ec:#x}, PC={pc:#x}");
+		loop {
+			core::hint::spin_loop();
+		}
 	}
 }
 

src/arch/aarch64/kernel/mod.rs

@@ -12,15 +12,18 @@ mod start;
 pub mod switch;
 pub mod systemtime;
 
+use alloc::alloc::{Layout, alloc};
 use core::arch::global_asm;
-use core::str;
-use core::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use core::sync::atomic::{AtomicPtr, AtomicU32, Ordering};
 use core::task::Waker;
+use core::{ptr, str};
 
 use memory_addresses::arch::aarch64::{PhysAddr, VirtAddr};
 
 use crate::arch::aarch64::kernel::core_local::*;
 use crate::arch::aarch64::kernel::serial::SerialPort;
+use crate::arch::aarch64::mm::paging::{BasePageSize, PageSize};
+use crate::config::*;
 use crate::env;
 
 const SERIAL_PORT_BAUDRATE: u32 = 115_200;
@@ -69,12 +72,15 @@ impl Default for Console {
 	}
 }
 
+#[repr(align(8))]
+pub(crate) struct AlignedAtomicU32(AtomicU32);
+
 /// `CPU_ONLINE` is the count of CPUs that finished initialization.
 ///
 /// It also synchronizes initialization of CPU cores.
-pub(crate) static CPU_ONLINE: AtomicU32 = AtomicU32::new(0);
+pub(crate) static CPU_ONLINE: AlignedAtomicU32 = AlignedAtomicU32(AtomicU32::new(0));
 
-pub(crate) static CURRENT_STACK_ADDRESS: AtomicU64 = AtomicU64::new(0);
+pub(crate) static CURRENT_STACK_ADDRESS: AtomicPtr<u8> = AtomicPtr::new(ptr::null_mut());
 
 #[cfg(target_os = "none")]
 global_asm!(include_str!("start.s"));
@@ -102,12 +108,25 @@ pub fn get_limit() -> usize {
 
 #[cfg(feature = "smp")]
 pub fn get_possible_cpus() -> u32 {
-	CPU_ONLINE.load(Ordering::Acquire)
+	use hermit_dtb::Dtb;
+
+	let dtb = unsafe {
+		Dtb::from_raw(core::ptr::with_exposed_provenance(
+			env::boot_info().hardware_info.device_tree.unwrap().get() as usize,
+		))
+		.expect(".dtb file has invalid header")
+	};
+
+	dtb.enum_subnodes("/cpus")
+		.filter(|name| name.contains("cpu@"))
+		.count()
+		.try_into()
+		.unwrap()
 }
 
 #[cfg(feature = "smp")]
 pub fn get_processor_count() -> u32 {
-	1
+	CPU_ONLINE.0.load(Ordering::Acquire)
 }
 
 #[cfg(not(feature = "smp"))]
@@ -149,10 +168,16 @@ pub fn application_processor_init() {
 
 fn finish_processor_init() {
 	debug!("Initialized Processor");
+
+	// Allocate stack for the CPU and pass the addresses.
+	let layout = Layout::from_size_align(KERNEL_STACK_SIZE, BasePageSize::SIZE as usize).unwrap();
+	let stack = unsafe { alloc(layout) };
+	assert!(!stack.is_null());
+	CURRENT_STACK_ADDRESS.store(stack, Ordering::Relaxed);
 }
 
 pub fn boot_next_processor() {
-	CPU_ONLINE.fetch_add(1, Ordering::Release);
+	CPU_ONLINE.0.fetch_add(1, Ordering::Release);
 }
 
 pub fn print_statistics() {

src/arch/aarch64/kernel/scheduler.rs

@@ -180,7 +180,7 @@ impl TaskStacks {
 	}
 
 	pub fn from_boot_stacks() -> TaskStacks {
-		let stack = VirtAddr::new(CURRENT_STACK_ADDRESS.load(Ordering::Relaxed));
+		let stack = VirtAddr::new(CURRENT_STACK_ADDRESS.load(Ordering::Relaxed) as u64);
 		debug!("Using boot stack {stack:p}");
 
 		TaskStacks::Boot(BootStack { stack })

src/arch/aarch64/kernel/start.rs

@@ -29,9 +29,36 @@ pub unsafe extern "C" fn _start(boot_info: Option<&'static RawBootInfo>, cpu_id:
 
 	unsafe {
 		naked_asm!(
-			"msr spsel, {l1}", // we want to use sp_el1
+			// use core::sync::atomic::{AtomicU32, Ordering};
+			//
+			// pub static CPU_ONLINE: AtomicU32 = AtomicU32::new(0);
+			//
+			// while CPU_ONLINE.load(Ordering::Acquire) != this {
+			//     core::hint::spin_loop();
+			// }
+			"mrs x4, mpidr_el1",
+			"and x4, x4, #0xff",
+			"1:",
+			"adrp x8, {cpu_online}",
+			"ldr x5, [x8, #:lo12:{cpu_online}]",
+			"cmp x4, x5",
+			"b.eq 2f",
+			"b 1b",
+			"2:",
+
+			// we want to use sp_el1
+			"msr spsel, #1",
+
+			// Overwrite RSP if `CURRENT_STACK_ADDRESS != 0`
 			"adrp x8, {current_stack_address}",
+			"ldr x4, [x8, #:lo12:{current_stack_address}]",
+			"cmp x4, 0",
+			"b.eq 3f",
+			"mov sp, x4",
+			"b 4f",
+			"3:",
 			"mov x4, sp",
+			"4:",
 			"str x4, [x8, #:lo12:{current_stack_address}]",
 
 			// Add stack top offset
@@ -41,7 +68,7 @@ pub unsafe extern "C" fn _start(boot_info: Option<&'static RawBootInfo>, cpu_id:
 			// Jump to Rust code
 			"b {pre_init}",
 
-			l1 = const 1,
+			cpu_online = sym super::CPU_ONLINE,
 			stack_top_offset = const KERNEL_STACK_SIZE - TaskStacks::MARKER_SIZE,
 			current_stack_address = sym super::CURRENT_STACK_ADDRESS,
 			pre_init = sym pre_init,