feat(tlb): add sync mechanism for ipi and support lazy flush tlb

Author: li041

modules/axipi/Cargo.toml

@@ -11,6 +11,7 @@ documentation = "https://arceos-org.github.io/arceos/axipi/index.html"
 
 [features]
 default = []
+smp = []
 
 [dependencies]
 axconfig.workspace = true
@@ -19,3 +20,7 @@ kspin.workspace = true
 lazyinit.workspace = true
 log.workspace = true
 percpu.workspace = true
+crate_interface = "0.1.4"
+axcpu = { git = "https://github.com/arceos-org/axcpu.git", tag = "dev-v03" }
+axtask = { workspace = true, features = ["task-ext"] }
+page_table_multiarch.workspace = true

modules/axipi/src/event.rs

@@ -1,4 +1,5 @@
 use alloc::{boxed::Box, sync::Arc};
+use core::sync::atomic::AtomicBool;
 
 /// A callback function that will be called when an [`IpiEvent`] is received and handled.
 pub struct Callback(Box<dyn FnOnce()>);
@@ -50,8 +51,10 @@ impl<T: Fn() + 'static> From<T> for MulticastCallback {
 
 /// An IPI event that is sent from a source CPU to the target CPU.
 pub struct IpiEvent {
+    pub name: &'static str,
     /// The source CPU ID that sent the IPI event.
     pub src_cpu_id: usize,
     /// The callback function that will be called when the IPI event is handled.
     pub callback: Callback,
+    pub done: Option<Arc<AtomicBool>>,
 }

modules/axipi/src/lib.rs

@@ -6,16 +6,34 @@
 extern crate log;
 extern crate alloc;
 
-use axhal::irq::{IPI_IRQ, IpiTarget};
-use axhal::percpu::this_cpu_id;
+use alloc::{sync::Arc, vec::Vec};
+use core::sync::atomic::{AtomicBool, Ordering};
+
+use axhal::{
+    irq::{IPI_IRQ, IpiTarget},
+    percpu::this_cpu_id,
+};
+use axtask::AxCpuMask;
 use kspin::SpinNoIrq;
 use lazyinit::LazyInit;
+use queue::IpiEventQueue;
+
+use crate::event::{Callback, MulticastCallback};
 
 mod event;
 mod queue;
+#[cfg(feature = "smp")]
+mod tlb;
 
-pub use event::{Callback, MulticastCallback};
-use queue::IpiEventQueue;
+static SECONDARY_CPUS_STARTED: AtomicBool = AtomicBool::new(false);
+
+pub fn start_secondary_cpus_done() {
+    SECONDARY_CPUS_STARTED.store(true, Ordering::Release);
+}
+
+pub fn secondary_cpus_ready() -> bool {
+    SECONDARY_CPUS_STARTED.load(Ordering::Acquire)
+}
 
 #[percpu::def_percpu]
 static IPI_EVENT_QUEUE: LazyInit<SpinNoIrq<IpiEventQueue>> = LazyInit::new();
@@ -28,34 +46,164 @@ pub fn init() {
 }
 
 /// Executes a callback on the specified destination CPU via IPI.
-pub fn run_on_cpu<T: Into<Callback>>(dest_cpu: usize, callback: T) {
-    info!("Send IPI event to CPU {dest_cpu}");
+pub fn run_on_cpu<T: Into<Callback>>(name: &'static str, dest_cpu: usize, callback: T, wait: bool) {
+    info!("Send IPI event to CPU {}", dest_cpu);
     if dest_cpu == this_cpu_id() {
         // Execute callback on current CPU immediately
         callback.into().call();
     } else {
+        let done_flag = if wait {
+            Some(Arc::new(AtomicBool::new(false)))
+        } else {
+            None
+        };
         unsafe { IPI_EVENT_QUEUE.remote_ref_raw(dest_cpu) }
             .lock()
-            .push(this_cpu_id(), callback.into());
+            .push(name, this_cpu_id(), callback.into(), done_flag.clone());
         axhal::irq::send_ipi(IPI_IRQ, IpiTarget::Other { cpu_id: dest_cpu });
+        if wait {
+            if let Some(df) = done_flag {
+                while !df.load(Ordering::Acquire) {
+                    core::hint::spin_loop();
+                }
+            }
+        }
+    }
+}
+
+pub fn run_on_bitmask_except_self<T: Into<MulticastCallback>>(
+    name: &'static str,
+    callback: T,
+    cpu_mask: AxCpuMask,
+    wait: bool,
+) {
+    let current_cpu_id = this_cpu_id();
+    let cpu_num = axconfig::plat::CPU_NUM;
+    let callback = callback.into();
+
+    let mut done_flags: Vec<Arc<AtomicBool>> = Vec::new();
+
+    for cpu_id in 0..cpu_num {
+        if cpu_id != current_cpu_id && cpu_mask.get(cpu_id) {
+            let done_flag = if wait {
+                Some(Arc::new(AtomicBool::new(false)))
+            } else {
+                None
+            };
+            if let Some(df) = &done_flag {
+                done_flags.push(df.clone());
+            }
+
+            unsafe { IPI_EVENT_QUEUE.remote_ref_raw(cpu_id) }
+                .lock()
+                .push(
+                    name,
+                    current_cpu_id,
+                    callback.clone().into_unicast(),
+                    done_flag,
+                );
+        }
+    }
+    if done_flags.is_empty() {
+        return;
+    }
+    for cpu_id in 0..cpu_num {
+        if cpu_id != current_cpu_id && cpu_mask.get(cpu_id) {
+            axhal::irq::send_ipi(IPI_IRQ, IpiTarget::Other { cpu_id });
+        }
+    }
+    if wait {
+        for df in done_flags {
+            while !df.load(Ordering::Acquire) {
+                core::hint::spin_loop();
+            }
+        }
+    }
+}
+
+pub fn run_on_each_cpu_except_self<T: Into<MulticastCallback>>(
+    name: &'static str,
+    callback: T,
+    wait: bool,
+) {
+    let current_cpu_id = this_cpu_id();
+    let cpu_num = axconfig::plat::CPU_NUM;
+    let callback = callback.into();
+
+    let mut done_flags: Vec<Arc<AtomicBool>> = Vec::new();
+
+    // Push the callback to all other CPUs' IPI event queues
+    for cpu_id in 0..cpu_num {
+        if cpu_id != current_cpu_id {
+            let done_flag = if wait {
+                Some(Arc::new(AtomicBool::new(false)))
+            } else {
+                None
+            };
+            if let Some(df) = &done_flag {
+                done_flags.push(df.clone());
+            }
+
+            unsafe { IPI_EVENT_QUEUE.remote_ref_raw(cpu_id) }
+                .lock()
+                .push(
+                    name,
+                    current_cpu_id,
+                    callback.clone().into_unicast(),
+                    done_flag,
+                );
+        }
+    }
+    if done_flags.is_empty() {
+        return;
+    }
+    // Send IPI to all other CPUs to trigger their callbacks
+    axhal::irq::send_ipi(
+        IPI_IRQ,
+        IpiTarget::AllExceptCurrent {
+            cpu_id: current_cpu_id,
+            cpu_num,
+        },
+    );
+    if wait {
+        for df in done_flags {
+            while !df.load(Ordering::Acquire) {
+                core::hint::spin_loop();
+            }
+        }
     }
 }
 
 /// Executes a callback on all other CPUs via IPI.
-pub fn run_on_each_cpu<T: Into<MulticastCallback>>(callback: T) {
+pub fn run_on_each_cpu<T: Into<MulticastCallback>>(name: &'static str, callback: T, wait: bool) {
     info!("Send IPI event to all other CPUs");
     let current_cpu_id = this_cpu_id();
     let cpu_num = axconfig::plat::CPU_NUM;
     let callback = callback.into();
 
     // Execute callback on current CPU immediately
     callback.clone().call();
+
+    let mut done_flags: Vec<Arc<AtomicBool>> = Vec::new();
     // Push the callback to all other CPUs' IPI event queues
     for cpu_id in 0..cpu_num {
         if cpu_id != current_cpu_id {
+            let done_flag = if wait {
+                Some(Arc::new(AtomicBool::new(false)))
+            } else {
+                None
+            };
+            if let Some(df) = &done_flag {
+                done_flags.push(df.clone());
+            }
             unsafe { IPI_EVENT_QUEUE.remote_ref_raw(cpu_id) }
                 .lock()
-                .push(current_cpu_id, callback.clone().into_unicast());
+                .push(
+                    name,
+                    current_cpu_id,
+                    callback.clone().into_unicast(),
+                    done_flag,
+                );
         }
     }
     // Send IPI to all other CPUs to trigger their callbacks
@@ -66,15 +214,27 @@ pub fn run_on_each_cpu<T: Into<MulticastCallback>>(callback: T) {
             cpu_num,
         },
     );
+    if wait {
+        for df in done_flags {
+            while !df.load(Ordering::Acquire) {
+                core::hint::spin_loop();
+            }
+        }
+    }
 }
 
-/// The handler for IPI events. It retrieves the events from the queue and calls the corresponding callbacks.
+/// The handler for IPI events. It retrieves the events from the queue and calls
+/// the corresponding callbacks.
 pub fn ipi_handler() {
-    while let Some((src_cpu_id, callback)) = unsafe { IPI_EVENT_QUEUE.current_ref_mut_raw() }
-        .lock()
-        .pop_one()
+    while let Some((_name, src_cpu_id, callback, done)) =
+        unsafe { IPI_EVENT_QUEUE.current_ref_raw() }
+            .lock()
+            .pop_one()
     {
         debug!("Received IPI event from CPU {src_cpu_id}");
         callback.call();
+        if let Some(done) = done {
+            done.store(true, Ordering::Release);
+        }
     }
 }

modules/axipi/src/queue.rs

@@ -1,4 +1,5 @@
-use alloc::collections::VecDeque;
+use alloc::{collections::VecDeque, sync::Arc};
+use core::sync::atomic::AtomicBool;
 
 use crate::event::{Callback, IpiEvent};
 
@@ -26,20 +27,28 @@ impl IpiEventQueue {
     }
 
     /// Push a new event into the queue.
-    pub fn push(&mut self, src_cpu_id: usize, callback: Callback) {
+    pub fn push(
+        &mut self,
+        name: &'static str,
+        src_cpu_id: usize,
+        callback: Callback,
+        done: Option<Arc<AtomicBool>>,
+    ) {
         self.events.push_back(IpiEvent {
+            name,
             src_cpu_id,
             callback,
+            done,
         });
     }
 
     /// Try to pop the latest event that exists in the queue.
     ///
     /// Return `None` if no event is available.
     #[must_use]
-    pub fn pop_one(&mut self) -> Option<(usize, Callback)> {
+    pub fn pop_one(&mut self) -> Option<(&'static str, usize, Callback, Option<Arc<AtomicBool>>)> {
         if let Some(e) = self.events.pop_front() {
-            Some((e.src_cpu_id, e.callback))
+            Some((e.name, e.src_cpu_id, e.callback, e.done))
         } else {
             None
         }

modules/axipi/src/tlb.rs

@@ -0,0 +1,30 @@
+use axhal::mem::VirtAddr;
+use axtask::{TaskExt, current};
+use page_table_multiarch::TlbFlushIf;
+
+use crate::{
+    MulticastCallback, run_on_bitmask_except_self, run_on_each_cpu_except_self,
+    secondary_cpus_ready, this_cpu_id,
+};
+
+struct TlbFlushImpl;
+
+#[crate_interface::impl_interface]
+impl TlbFlushIf for TlbFlushImpl {
+    fn flush_all(vaddr: Option<VirtAddr>) {
+        if axconfig::plat::CPU_NUM == 1 || !secondary_cpus_ready() {
+            // local
+            axhal::asm::flush_tlb(None);
+        } else {
+            let callback = MulticastCallback::new(move || {
+                axhal::asm::flush_tlb(vaddr);
+            });
+            if let Some(ext) = current().task_ext() {
+                let on_cpu_mask = ext.on_cpu_mask();
+                run_on_bitmask_except_self("flush", callback, on_cpu_mask, true);
+            } else {
+                run_on_each_cpu_except_self("flush", callback, true);
+            }
+        }
+    }
+}

modules/axmm/Cargo.toml

@@ -12,6 +12,7 @@ documentation = "https://arceos-org.github.io/arceos/axmm/index.html"
 [features]
 default = []
 copy = ["page_table_multiarch/copy-from"]
+smp = ["page_table_multiarch/smp"]
 
 [dependencies]
 axalloc = { workspace = true }

modules/axruntime/src/lib.rs

@@ -230,6 +230,8 @@ pub fn rust_main(cpu_id: usize, arg: usize) -> ! {
     while !is_init_ok() {
         core::hint::spin_loop();
     }
+    #[cfg(all(feature = "smp", feature = "ipi"))]
+    axipi::start_secondary_cpus_done();
 
     unsafe { main() };
 
@@ -300,9 +302,12 @@ fn init_interrupt() {
     });
 
     #[cfg(feature = "ipi")]
-    axhal::irq::register(axhal::irq::IPI_IRQ, || {
-        axipi::ipi_handler();
-    });
+    {
+        axipi::init();
+        axhal::irq::register(axhal::irq::IPI_IRQ, || {
+            axipi::ipi_handler();
+        });
+    }
 
     // Enable IRQs before starting app
     axhal::asm::enable_irqs();

modules/axtask/src/task.rs

@@ -54,6 +54,9 @@ pub unsafe trait TaskExt {
     fn on_enter(&self) {}
     /// Called when the task is switched out.
     fn on_leave(&self) {}
+    fn on_cpu_mask(&self) -> AxCpuMask {
+        AxCpuMask::new()
+    }
 }
 
 /// The inner task structure.