feat(smp): support cross-core TLB flush

Author: li041

Cargo.lock

@@ -16,11 +16,13 @@ rust-version.workspace = true
 default = []
 axerrno = ["dep:axerrno"]
 copy-from = ["dep:bitmaps"]
+smp = []
 
 [dependencies]
 axerrno = { version = "0.1", optional = true }
 arrayvec = { version = "0.7", default-features = false }
 bitmaps = { version = "3.2", default-features = false, optional = true }
+crate_interface = "0.1.4"
 log = "0.4"
 memory_addr.workspace = true
 page_table_entry.workspace = true

page_table_multiarch/Cargo.toml

@@ -27,10 +27,17 @@ impl PagingMetaData for A64PagingMetaData {
             if let Some(vaddr) = vaddr {
                 // TLB Invalidate by VA, All ASID, EL1, Inner Shareable
                 const VA_MASK: usize = (1 << 44) - 1; // VA[55:12] => bits[43:0]
-                asm!("tlbi vaae1is, {}; dsb sy; isb", in(reg) ((vaddr.as_usize() >> 12) & VA_MASK))
+                let va = (vaddr.as_usize() >> 12) & VA_MASK;
+                #[cfg(feature = "smp")]
+                asm!("tlbi vaae1is, {}; dsb sy; isb", in(reg) va);
+                #[cfg(not(feature = "smp"))]
+                asm!("tlbi vaae1, {}; dsb sy; isb", in(reg) va);
             } else {
+                #[cfg(feature = "smp")]
                 // TLB Invalidate by VMID, All at stage 1, EL1
-                asm!("tlbi vmalle1; dsb sy; isb")
+                asm!("tlbi vmalle1is; dsb sy; isb");
+                #[cfg(not(feature = "smp"))]
+                asm!("tlbi vmalle1; dsb sy; isb");
             }
         }
     }

page_table_multiarch/src/arch/aarch64.rs

@@ -7,6 +7,34 @@ use page_table_entry::loongarch64::LA64PTE;
 
 use crate::{PageTable64, PageTable64Mut, PagingMetaData};
 
+#[inline]
+fn local_flush_tlb(vaddr: Option<memory_addr::VirtAddr>) {
+    unsafe {
+        if let Some(vaddr) = vaddr {
+            // <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#_dbar>
+            //
+            // Only after all previous load/store access operations are completely
+            // executed, the DBAR 0 instruction can be executed; and only after the
+            // execution of DBAR 0 is completed, all subsequent load/store access
+            // operations can be executed.
+            //
+            // <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#_invtlb>
+            //
+            // formats: invtlb op, asid, addr
+            //
+            // op 0x5: Clear all page table entries with G=0 and ASID equal to the
+            // register specified ASID, and VA equal to the register specified VA.
+            //
+            // When the operation indicated by op does not require an ASID, the
+            // general register rj should be set to r0.
+            asm!("dbar 0; invtlb 0x05, $r0, {reg}", reg = in(reg) vaddr.as_usize());
+        } else {
+            // op 0x0: Clear all page table entries
+            asm!("dbar 0; invtlb 0x00, $r0, $r0");
+        }
+    }
+}
+
 /// Metadata of LoongArch64 page tables.
 #[derive(Copy, Clone, Debug)]
 pub struct LA64MetaData;
@@ -51,30 +79,14 @@ impl PagingMetaData for LA64MetaData {
 
     #[inline]
     fn flush_tlb(vaddr: Option<VirtAddr>) {
-        unsafe {
-            if let Some(vaddr) = vaddr {
-                // <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#_dbar>
-                //
-                // Only after all previous load/store access operations are completely
-                // executed, the DBAR 0 instruction can be executed; and only after the
-                // execution of DBAR 0 is completed, all subsequent load/store access
-                // operations can be executed.
-                //
-                // <https://loongson.github.io/LoongArch-Documentation/LoongArch-Vol1-EN.html#_invtlb>
-                //
-                // formats: invtlb op, asid, addr
-                //
-                // op 0x5: Clear all page table entries with G=0 and ASID equal to the
-                // register specified ASID, and VA equal to the register specified VA.
-                //
-                // When the operation indicated by op does not require an ASID, the
-                // general register rj should be set to r0.
-                asm!("dbar 0; invtlb 0x05, $r0, {reg}", reg = in(reg) vaddr.as_usize());
-            } else {
-                // op 0x0: Clear all page table entries
-                asm!("dbar 0; invtlb 0x00, $r0, $r0");
-            }
+        #[cfg(feature = "smp")]
+        {
+            use crate::__TlbFlushIf_mod;
+            use crate_interface::call_interface;
+
+            call_interface!(TlbFlushIf::flush_all(vaddr));
         }
+        local_flush_tlb(vaddr);
     }
 }
 

page_table_multiarch/src/arch/loongarch64.rs

@@ -5,6 +5,14 @@ use page_table_entry::riscv::Rv64PTE;
 
 use crate::{PageTable64, PageTable64Mut, PagingMetaData};
 
+fn local_flush_tlb(vaddr: Option<memory_addr::VirtAddr>) {
+    if let Some(vaddr) = vaddr {
+        riscv::asm::sfence_vma(0, vaddr.as_usize())
+    } else {
+        riscv::asm::sfence_vma_all();
+    }
+}
+
 /// A virtual address that can be used in RISC-V Sv39 and Sv48 page tables.
 pub trait SvVirtAddr: memory_addr::MemoryAddr + Send + Sync {
     /// Flush the TLB.
@@ -14,11 +22,14 @@ pub trait SvVirtAddr: memory_addr::MemoryAddr + Send + Sync {
 impl SvVirtAddr for VirtAddr {
     #[inline]
     fn flush_tlb(vaddr: Option<Self>) {
-        if let Some(vaddr) = vaddr {
-            riscv::asm::sfence_vma(0, vaddr.as_usize())
-        } else {
-            riscv::asm::sfence_vma_all();
+        #[cfg(feature = "smp")]
+        {
+            use crate::__TlbFlushIf_mod;
+            use crate_interface::call_interface;
+
+            call_interface!(TlbFlushIf::flush_all(vaddr));
         }
+        local_flush_tlb(vaddr);
     }
 }
 

page_table_multiarch/src/arch/riscv.rs

@@ -5,6 +5,17 @@ use page_table_entry::x86_64::X64PTE;
 
 use crate::{PageTable64, PageTable64Mut, PagingMetaData};
 
+#[inline]
+fn local_flush_tlb(vaddr: Option<memory_addr::VirtAddr>) {
+    unsafe {
+        if let Some(vaddr) = vaddr {
+            x86::tlb::flush(vaddr.into());
+        } else {
+            x86::tlb::flush_all();
+        }
+    }
+}
+
 /// metadata of x86_64 page tables.
 pub struct X64PagingMetaData;
 
@@ -16,13 +27,14 @@ impl PagingMetaData for X64PagingMetaData {
 
     #[inline]
     fn flush_tlb(vaddr: Option<VirtAddr>) {
-        unsafe {
-            if let Some(vaddr) = vaddr {
-                x86::tlb::flush(vaddr.into());
-            } else {
-                x86::tlb::flush_all();
-            }
+        #[cfg(feature = "smp")]
+        {
+            use crate::__TlbFlushIf_mod;
+            use crate_interface::call_interface;
+
+            call_interface!(TlbFlushIf::flush_all(vaddr));
         }
+        local_flush_tlb(vaddr);
     }
 }
 

page_table_multiarch/src/arch/x86_64.rs

@@ -138,3 +138,18 @@ impl From<PageSize> for usize {
         size as usize
     }
 }
+
+/// Interface for performing TLB flush operations on SMP systems.
+///
+/// This trait is used by the page table subsystem to request
+/// architecture- or platform-specific TLB invalidation.
+#[cfg(feature = "smp")]
+#[crate_interface::def_interface]
+pub trait TlbFlushIf {
+    /// Flush TLB entries.
+    ///
+    /// If `vaddr` is `None`, flush all TLB entries.
+    /// If `Some(vaddr)` is provided, flush the TLB entry corresponding
+    /// to the given virtual address.
+    fn flush_all(vaddr: Option<VirtAddr>);
+}