* Parse SLIT for distance information

* Order by distance, reorganise domains into nodes

src/acpi/mod.rs

@@ -20,9 +20,12 @@ mod rsdp;
 mod rsdt;
 mod rxsdt;
 pub mod sdt;
+#[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))]
+pub mod slit;
 #[cfg(target_arch = "aarch64")]
 mod spcr;
-mod srat;
+#[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))]
+pub mod srat;
 mod xsdt;
 
 unsafe fn map_linearly(addr: PhysicalAddress, len: usize, mapper: &mut crate::memory::PageMapper) {

src/acpi/slit.rs

@@ -0,0 +1,48 @@
+use crate::{
+    acpi::sdt::Sdt,
+    find_one_sdt,
+    numa::{self, NumaNode, NUMA_NODES, NUMBER_OF_DOMAINS},
+};
+use core::ops::Add;
+use hashbrown::HashMap;
+
+#[derive(Debug)]
+pub struct Slit {
+    sdt: &'static Sdt,
+    no: u64,
+    address: usize,
+}
+
+impl Slit {
+    pub fn new(sdt: &'static Sdt) -> Self {
+        Self {
+            sdt,
+            no: unsafe { *(sdt.data_address() as *const u64) },
+            address: sdt.data_address() + 8,
+        }
+    }
+    pub fn init(&self) {
+        let ndom = *NUMBER_OF_DOMAINS.get().unwrap();
+        let address = self.address as *const u8;
+
+        for i in 0..ndom {
+            for j in i..ndom {
+                if i != j {
+                    unsafe {
+                        numa::set_distance(i, j, unsafe { *address.add((i + ndom * j) as usize) });
+                        numa::set_distance(j, i, unsafe { *address.add((i + ndom * j) as usize) });
+                    }
+                } else {
+                    unsafe {
+                        numa::set_distance(i, j, 10);
+                    }
+                }
+            }
+        }
+    }
+}
+
+pub fn init() {
+    let slit = Slit::new(find_one_sdt!("SLIT"));
+    slit.init();
+}

src/acpi/srat/mod.rs

@@ -3,6 +3,7 @@
 use crate::{
     acpi::{find_sdt, sdt::Sdt, srat},
     find_one_sdt,
+    numa::{NUMA_NODES, NUMBER_OF_DOMAINS},
 };
 
 #[cfg(target_arch = "aarch64")]
@@ -19,21 +20,32 @@ pub struct Srat {
     entries: usize,
 }
 
-impl Srat {
-    pub fn init() {
-        let srat = Self::new(find_one_sdt!("SRAT"));
-        arch::init_srat(&srat);
-    }
+pub fn init() {
+    let srat = Srat::new(find_one_sdt!("SRAT"));
+    arch::init_srat(&srat);
+    NUMBER_OF_DOMAINS.call_once(|| NUMA_NODES.get().unwrap().len() as u32);
+}
 
+impl Srat {
     pub fn new(sdt: &'static Sdt) -> Self {
         Self {
             sdt,
-            entries: sdt.data_address() + 16,
+            entries: sdt.data_address() + 12,
         }
     }
 }
 
-struct SratIter<'a> {
+impl<'a> IntoIterator for &'a Srat {
+    type Item = SratEntry;
+
+    type IntoIter = SratIter<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        SratIter { i: 0, srat: self }
+    }
+}
+
+pub struct SratIter<'a> {
     i: u32,
     srat: &'a Srat,
 }
@@ -42,43 +54,58 @@ impl<'a> Iterator for SratIter<'a> {
     type Item = SratEntry;
 
     fn next(&mut self) -> Option<Self::Item> {
-        while self.i < self.srat.sdt.length {
+        while self.i < self.srat.sdt.data_len() as u32 {
             let entry = (self.srat.entries + self.i as usize) as *const u8;
+            let entry_len = unsafe { *((self.srat.entries + self.i as usize + 1) as *const u8) };
 
-            return Some(match unsafe { *entry } {
+            let entry = Some(match unsafe { *entry } {
                 0 => SratEntry::LegacyProcessorLocalAffinity(unsafe {
+                    assert!(entry_len as usize == size_of::<LegacyProcessorLocalAffinity>() + 2);
                     *(entry.add(2) as *const LegacyProcessorLocalAffinity)
                 }),
 
-                1 => SratEntry::MemoryAffinity(unsafe { *(entry.add(2) as *const MemoryAffinity) }),
+                1 => SratEntry::MemoryAffinity(unsafe {
+                    assert!(entry_len as usize == size_of::<MemoryAffinity>() + 10);
+                    *(entry.add(2) as *const MemoryAffinity)
+                }),
                 2 => SratEntry::ProcessorLocalAffinity(unsafe {
+                    assert!(entry_len as usize == size_of::<ProcessorLocalAffinity>() + 8);
                     *(entry.add(4) as *const ProcessorLocalAffinity)
                 }),
-                3 => SratEntry::GiccAffinity(unsafe { *(entry.add(2) as *const GiccAffinity) }),
-                4 => SratEntry::GicItsAffinity(unsafe { *(entry.add(2) as *const GicItsAffinity) }),
+                3 => SratEntry::GiccAffinity(unsafe {
+                    assert!(entry_len as usize == size_of::<GiccAffinity>() + 2);
+                    *(entry.add(2) as *const GiccAffinity)
+                }),
+                4 => SratEntry::GicItsAffinity(unsafe {
+                    assert!(entry_len as usize == size_of::<GicItsAffinity>() + 2);
+                    *(entry.add(2) as *const GicItsAffinity)
+                }),
                 // ignore Generic Initiator Affinity
                 5 => {
-                    self.i += 1;
+                    self.i += entry_len as u32;
                     continue;
                 }
-                _ => panic!("Unknown value in Srat"),
+                _ => SratEntry::Unknown(unsafe { *entry }),
             });
+            self.i += entry_len as u32;
+            return entry;
         }
         None
     }
 }
 
-enum SratEntry {
+#[derive(Debug, Clone, Copy)]
+pub enum SratEntry {
     LegacyProcessorLocalAffinity(LegacyProcessorLocalAffinity),
     MemoryAffinity(MemoryAffinity),
     ProcessorLocalAffinity(ProcessorLocalAffinity),
     GiccAffinity(GiccAffinity),
     GicItsAffinity(GicItsAffinity),
-    // unimplemented: Generic Initiator Affinity; our current focus is only on memory and cpus
+    Unknown(u8), // unimplemented: Generic Initiator Affinity; our current focus is only on memory and cpus
 }
 
 #[repr(C, packed)]
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 /// For legacy xAPIC systems
 struct LegacyProcessorLocalAffinity {
     proximity_domain_low: u8,
@@ -90,7 +117,7 @@ struct LegacyProcessorLocalAffinity {
 }
 
 #[repr(C, packed)]
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 struct MemoryAffinity {
     proximity_domain: u32,
     _reserved0: u16,
@@ -103,7 +130,7 @@ struct MemoryAffinity {
 }
 
 #[repr(C, packed)]
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 /// For x2APIC systems
 struct ProcessorLocalAffinity {
     proximity_domain: u32,
@@ -113,7 +140,7 @@ struct ProcessorLocalAffinity {
 }
 
 #[repr(C, packed)]
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 struct GiccAffinity {
     proximity_domain: u32,
     processor_uid: u32,
@@ -122,7 +149,7 @@ struct GiccAffinity {
 }
 
 #[repr(C, packed)]
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 struct GicItsAffinity {
     proximity_domain: u32,
     _reserved: u16,

src/acpi/srat/x86.rs

@@ -1,3 +1,65 @@
-use crate::acpi::srat::Srat;
+use core::iter;
 
-pub fn init_srat(srat: &Srat) {}
+use crate::{
+    acpi::srat::{Srat, SratEntry},
+    numa::{self, NUMA_NODES},
+};
+
+#[inline(always)]
+fn to_usize(low: u32, high: u32) -> usize {
+    #[cfg(target_pointer_width = "32")]
+    return low as usize;
+
+    #[cfg(target_pointer_width = "64")]
+    {
+        let mut low_and_high = [0u8; 8];
+        low_and_high[0..=3].copy_from_slice(low.to_ne_bytes().as_slice());
+        low_and_high[4..=7].copy_from_slice(high.to_ne_bytes().as_slice());
+        usize::from_ne_bytes(low_and_high)
+    }
+}
+
+#[inline(always)]
+fn to_single_int(high: &[u8; 3], low: u8) -> u32 {
+    let mut high_and_low = [0u8; 4];
+    high_and_low[0] = low;
+    (high_and_low[1], high_and_low[2], high_and_low[3]) = (high[0], high[1], high[2]);
+    u32::from_ne_bytes(high_and_low)
+}
+
+pub fn init_srat(srat: &Srat) {
+    for affinity in srat {
+        match affinity {
+            SratEntry::LegacyProcessorLocalAffinity(legacy_processor_local_affinity) => unsafe {
+                numa::add_cpu(
+                    legacy_processor_local_affinity.apic_id as u32,
+                    to_single_int(
+                        &legacy_processor_local_affinity.proximity_domain_high,
+                        legacy_processor_local_affinity.proximity_domain_low,
+                    ),
+                )
+            },
+            SratEntry::MemoryAffinity(memory_affinity) => unsafe {
+                if memory_affinity.length_low == 0 {
+                    continue;
+                }
+                numa::add_memory(
+                    memory_affinity.proximity_domain,
+                    to_usize(
+                        memory_affinity.base_address_low,
+                        memory_affinity.base_address_high,
+                    ),
+                    to_usize(memory_affinity.length_low, memory_affinity.length_high),
+                );
+            },
+            SratEntry::ProcessorLocalAffinity(processor_local_affinity) => unsafe {
+                numa::add_cpu(
+                    processor_local_affinity.x2apic_id,
+                    processor_local_affinity.proximity_domain,
+                )
+            },
+            _ => continue,
+        }
+    }
+    println!("{:?}", NUMA_NODES.get().unwrap());
+}

src/numa.rs

@@ -1,14 +1,184 @@
+#[cfg(all(
+    feature = "acpi",
+    any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")
+))]
+use crate::acpi;
 use crate::{
     cpu_set::LogicalCpuId,
-    sync::{Mutex, L0},
+    sync::{CleanLockToken, Mutex, L0},
 };
 use alloc::{sync::Arc, vec::Vec};
+use hashbrown::HashMap;
+use spin::once::Once;
 
-pub static NUMA_NODES: Mutex<L0, Vec<Arc<NumaNode>>> = Mutex::new(Vec::new());
+pub static NUMA_NODES: Once<HashMap<u32, NumaNode>> = Once::new();
+pub static NUMBER_OF_DOMAINS: Once<u32> = Once::new();
 
-pub struct NumaMemory;
-
-pub struct NumaNode {
-    cpus: Vec<LogicalCpuId>,
-    memory: Vec<NumaMemory>,
+#[derive(Debug)]
+pub struct NumaMemory {
+    pub start: usize,
+    pub length: usize,
+}
+
+#[derive(Debug)]
+pub struct NumaCpu {
+    pub id: u32,
+}
+
+#[derive(Default, Debug)]
+/// Represents a single NUMA logical node. A node is different from a domain. NUMA domain
+/// refers to what exists physically. A NUMA node on the other hand is a logical one, with domains having
+/// only CPUs or memory grouped together with other CPUs or memories.
+///
+/// See the function `reorganise` below.
+pub struct NumaNode {
+    cpus: Vec<NumaCpu>,
+    memory: Vec<NumaMemory>,
+    distances: Vec<(u32, u8)>,
+}
+
+pub fn init() {
+    NUMA_NODES.call_once(|| HashMap::new());
+
+    #[cfg(all(
+        feature = "acpi",
+        any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")
+    ))]
+    {
+        acpi::srat::init();
+        acpi::slit::init();
+        return;
+    }
+
+    #[cfg(any(target_arch = "riscv64", target_arch = "aarch64"))]
+    {
+        todo!()
+    }
+
+    unsafe {
+        sort_by_distances();
+        reorganise();
+        shrink();
+    }
+
+    // From this point onwards, the global static `NUMA_NODES` or any of its elements
+    // MUST NOT be mutated by the usual unsafe magic that functions in this file use.
+}
+
+pub unsafe fn add_cpu(id: u32, node_id: u32) {
+    let numa_nodes = NUMA_NODES.get().unwrap();
+
+    /* 💀💀💀 HIGHLY UNSAFE 💀💀💀 */
+    let numa_nodes = unsafe { &mut *(&raw const *numa_nodes as *mut HashMap<u32, NumaNode>) };
+
+    if let Some(node) = numa_nodes.get_mut(&id) {
+        node.cpus.push(NumaCpu { id });
+    } else {
+        let mut cpus = Vec::new();
+        cpus.push(NumaCpu { id });
+        numa_nodes.insert(
+            node_id,
+            NumaNode {
+                cpus,
+                memory: Vec::new(),
+                distances: Vec::new(),
+            },
+        );
+    }
+}
+
+pub unsafe fn add_memory(node_id: u32, start: usize, length: usize) {
+    let numa_nodes = NUMA_NODES.get().unwrap();
+
+    /* 💀💀💀 HIGHLY UNSAFE 💀💀💀 */
+    let numa_nodes = unsafe { &mut *(&raw const *numa_nodes as *mut HashMap<u32, NumaNode>) };
+
+    if let Some(node) = numa_nodes.get_mut(&node_id) {
+        node.memory.push(NumaMemory { start, length });
+    } else {
+        let mut memory = Vec::new();
+        memory.push(NumaMemory { start, length });
+        numa_nodes.insert(
+            node_id,
+            NumaNode {
+                cpus: Vec::new(),
+                memory,
+                distances: Vec::new(),
+            },
+        );
+    }
+}
+
+pub unsafe fn set_distance(src: u32, target: u32, distance: u8) {
+    /* 💀💀💀 HIGHLY UNSAFE 💀💀💀 */
+    let nodes =
+        unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
+
+    let src = nodes.get_mut(&src).unwrap();
+    src.distances.push((target, distance));
+}
+
+unsafe fn shrink() {
+    /* 💀💀💀 HIGHLY UNSAFE 💀💀💀 */
+    let nodes =
+        unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
+
+    nodes.shrink_to_fit();
+
+    for (id, node) in nodes {
+        node.cpus.shrink_to_fit();
+        node.distances.shrink_to_fit();
+        node.memory.shrink_to_fit();
+    }
+}
+
+/// Reorganises CPUs and memories into nodes. If a NUMA domain has only a CPU but no memory, it is
+/// put into a node with a memory that is nearest to it. Similarly, if a NUMA domain has only memory but no
+/// CPUs, the memory is put into a node that has a CPU that is nearest to it.
+///
+/// See the comment above the definition of `NumaNode`.
+unsafe fn reorganise() {
+    /* 💀💀💀 HIGHLY UNSAFE 💀💀💀 */
+    let nodes =
+        unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
+    let ids = nodes.keys().map(|e| *e).collect::<Vec<u32>>();
+
+    for id in ids {
+        let node = nodes.remove(&id).unwrap();
+
+        if node.cpus.len() == 0 {
+            put_for_adoption(nodes, node.distances, Some(node.memory), None);
+        } else {
+            nodes.insert(id, node);
+        }
+    }
+}
+
+fn put_for_adoption(
+    nodes: &mut HashMap<u32, NumaNode>,
+    distances: Vec<(u32, u8)>,
+    memories: Option<Vec<NumaMemory>>,
+    cpus: Option<Vec<NumaCpu>>,
+) {
+    if let Some(memories) = memories {
+        assert!(cpus.is_none());
+        let (nearest_node_id, _) = distances.first().unwrap();
+        let nearest_node = nodes.get_mut(nearest_node_id).unwrap();
+        nearest_node.memory.extend(memories);
+    } else if let Some(cpus) = cpus {
+        assert!(memories.is_none());
+        let (nearest_node_id, _) = distances.first().unwrap();
+        let nearest_node = nodes.get_mut(nearest_node_id).unwrap();
+        nearest_node.cpus.extend(cpus);
+    }
+}
+
+unsafe fn sort_by_distances() {
+    /* 💀💀💀 HIGHLY UNSAFE 💀💀💀 */
+    let nodes =
+        unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
+
+    for (id, node) in nodes {
+        node.distances.sort_by_key(|(_, e)| *e);
+    }
 }

src/startup/mod.rs

@@ -7,10 +7,9 @@ use core::{
 
 use crate::{
     arch::interrupt,
-    context,
-    context::switch::SwitchResult,
+    context::{self, switch::SwitchResult},
     memory::{PhysicalAddress, RmmA, RmmArch},
-    profiling, scheme,
+    numa, profiling, scheme,
     sync::CleanLockToken,
 };
 
@@ -187,6 +186,8 @@ pub(crate) fn kmain(bootstrap: Bootstrap) -> ! {
         }
     }
 
+    numa::init();
+
     run_userspace(&mut token)
 }