* Store addresses as pointers

* Remove lots of unsafe code by initialising in the closure passed to
  `call_once`

* Remove `NUMBER_OF_DOMAINS` as size can always be inferred from the
  hashmap's len
This commit is contained in:
R Aadarsh
2026-06-21 14:01:46 +05:30
parent 4106dcbbfa
commit 2ad76496c4
4 changed files with 64 additions and 80 deletions
+14 -12
View File
@@ -1,7 +1,7 @@
use crate::{
acpi::sdt::Sdt,
find_one_sdt,
numa::{self, NumaNode, NUMA_NODES, NUMBER_OF_DOMAINS},
numa::{self, NumaNode, NUMA_NODES},
};
use core::ops::Add;
use hashbrown::HashMap;
@@ -10,7 +10,7 @@ use hashbrown::HashMap;
pub struct Slit {
sdt: &'static Sdt,
no: u64,
address: usize,
address: *const u8,
}
impl Slit {
@@ -18,28 +18,30 @@ impl Slit {
Self {
sdt,
no: unsafe { *(sdt.data_address() as *const u64) },
address: sdt.data_address() + 8,
address: (sdt.data_address() + 8) as *const u8,
}
}
pub fn init(&self) {
let ndom = *NUMBER_OF_DOMAINS.get().unwrap();
let address = self.address as *const u8;
pub fn init(&self, numa_nodes: &mut HashMap<u32, NumaNode>) {
let address = self.address;
let ndom = NUMA_NODES.get().unwrap().len() as u32;
for i in 0..ndom {
for j in i..ndom {
// ignore distances from a domain to itself, since it is always 10
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) });
}
numa::set_distance(numa_nodes, i, j, unsafe {
*address.add((i + ndom * j) as usize)
});
numa::set_distance(numa_nodes, j, i, unsafe {
*address.add((i + ndom * j) as usize)
});
}
}
}
}
}
pub fn init() {
pub fn init(numa_nodes: &mut HashMap<u32, NumaNode>) {
let slit = Slit::new(find_one_sdt!("SLIT"));
slit.init();
slit.init(numa_nodes);
}
+9 -8
View File
@@ -1,9 +1,11 @@
//! See <https://uefi.org/htmlspecs/ACPI_Spec_6_4_html/05_ACPI_Software_Programming_Model/ACPI_Software_Programming_Model.html#system-resource-affinity-table-srat>
use hashbrown::HashMap;
use crate::{
acpi::{find_sdt, sdt::Sdt, srat},
find_one_sdt,
numa::{NUMA_NODES, NUMBER_OF_DOMAINS},
numa::{NumaNode, NUMA_NODES},
};
#[cfg(target_arch = "aarch64")]
@@ -17,20 +19,19 @@ mod arch;
#[repr(C, packed)]
pub struct Srat {
sdt: &'static Sdt,
entries: usize,
entries: *const u8,
}
pub fn init() {
pub fn init(numa_nodes: &mut HashMap<u32, NumaNode>) {
let srat = Srat::new(find_one_sdt!("SRAT"));
arch::init_srat(&srat);
NUMBER_OF_DOMAINS.call_once(|| NUMA_NODES.get().unwrap().len() as u32);
arch::init_srat(numa_nodes, &srat);
}
impl Srat {
pub fn new(sdt: &'static Sdt) -> Self {
Self {
sdt,
entries: sdt.data_address() + 12,
entries: (sdt.data_address() + 12) as *const u8,
}
}
}
@@ -55,8 +56,8 @@ impl<'a> Iterator for SratIter<'a> {
fn next(&mut self) -> Option<Self::Item> {
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) };
let entry = unsafe { self.srat.entries.add(self.i as usize) };
let entry_len = unsafe { *self.srat.entries.add(self.i as usize + 1) };
let entry = Some(match unsafe { *entry } {
0 => SratEntry::LegacyProcessorLocalAffinity(unsafe {
+15 -12
View File
@@ -1,8 +1,10 @@
use core::iter;
use hashbrown::HashMap;
use crate::{
acpi::srat::{to_usize, Srat, SratEntry},
numa::{self, NUMA_NODES},
numa::{self, NumaNode, NUMA_NODES},
};
#[inline(always)]
@@ -13,23 +15,25 @@ fn to_single_int(high: &[u8; 3], low: u8) -> u32 {
u32::from_le_bytes(high_and_low)
}
pub fn init_srat(srat: &Srat) {
pub fn init_srat(numa_nodes: &mut HashMap<u32, NumaNode>, srat: &Srat) {
for affinity in srat {
match affinity {
SratEntry::LegacyProcessorLocalAffinity(legacy_processor_local_affinity) => unsafe {
SratEntry::LegacyProcessorLocalAffinity(legacy_processor_local_affinity) => {
numa::add_cpu(
numa_nodes,
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 {
}
SratEntry::MemoryAffinity(memory_affinity) => {
if memory_affinity.length_low == 0 {
continue;
}
numa::add_memory(
numa_nodes,
memory_affinity.proximity_domain,
to_usize(
memory_affinity.base_address_low,
@@ -37,13 +41,12 @@ pub fn init_srat(srat: &Srat) {
),
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,
)
},
}
SratEntry::ProcessorLocalAffinity(processor_local_affinity) => numa::add_cpu(
numa_nodes,
processor_local_affinity.x2apic_id,
processor_local_affinity.proximity_domain,
),
_ => continue,
}
}
+26 -48
View File
@@ -8,7 +8,6 @@ use hashbrown::HashMap;
use spin::once::Once;
pub static NUMA_NODES: Once<HashMap<u32, NumaNode>> = Once::new();
pub static NUMBER_OF_DOMAINS: Once<u32> = Once::new();
#[derive(Debug)]
pub struct NumaMemory {
@@ -34,39 +33,28 @@ pub struct NumaNode {
}
pub fn init() {
NUMA_NODES.call_once(|| HashMap::new());
let mut flag = false;
NUMA_NODES.call_once(|| {
let mut numa_nodes = HashMap::new();
#[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))]
{
acpi::srat::init(&mut numa_nodes);
acpi::slit::init(&mut numa_nodes);
#[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))]
{
acpi::srat::init();
acpi::slit::init();
flag = true;
}
#[cfg(any(target_arch = "riscv64", target_arch = "aarch64"))]
{
if !flag {
// todo!()
return;
sort_by_distances(&mut numa_nodes);
reorganise(&mut numa_nodes);
shrink(&mut numa_nodes);
}
}
unsafe {
sort_by_distances();
reorganise();
shrink();
}
#[cfg(any(target_arch = "riscv64", target_arch = "aarch64"))]
{
// todo!()
}
// 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.
numa_nodes
});
}
pub unsafe fn add_cpu(id: u32, node_id: u32) {
let numa_nodes = NUMA_NODES.get().unwrap();
let numa_nodes = unsafe { &mut *(&raw const *numa_nodes as *mut HashMap<u32, NumaNode>) };
pub fn add_cpu(numa_nodes: &mut HashMap<u32, NumaNode>, id: u32, node_id: u32) {
if let Some(node) = numa_nodes.get_mut(&id) {
node.cpus.push(NumaCpu { id });
} else {
@@ -83,11 +71,12 @@ pub unsafe fn add_cpu(id: u32, node_id: u32) {
}
}
pub unsafe fn add_memory(node_id: u32, start: usize, length: usize) {
let numa_nodes = NUMA_NODES.get().unwrap();
let numa_nodes = unsafe { &mut *(&raw const *numa_nodes as *mut HashMap<u32, NumaNode>) };
pub fn add_memory(
numa_nodes: &mut HashMap<u32, NumaNode>,
node_id: u32,
start: usize,
length: usize,
) {
if let Some(node) = numa_nodes.get_mut(&node_id) {
node.memory.push(NumaMemory { start, length });
} else {
@@ -104,18 +93,12 @@ pub unsafe fn add_memory(node_id: u32, start: usize, length: usize) {
}
}
pub unsafe fn set_distance(src: u32, target: u32, distance: u8) {
let nodes =
unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
pub fn set_distance(nodes: &mut HashMap<u32, NumaNode>, src: u32, target: u32, distance: u8) {
let src = nodes.get_mut(&src).unwrap();
src.distances.push((target, distance));
}
unsafe fn shrink() {
let nodes =
unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
fn shrink(nodes: &mut HashMap<u32, NumaNode>) {
nodes.shrink_to_fit();
for (id, node) in nodes {
@@ -130,9 +113,7 @@ unsafe fn shrink() {
/// 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() {
let nodes =
unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
fn reorganise(nodes: &mut HashMap<u32, NumaNode>) {
let ids = nodes.keys().map(|e| *e).collect::<Vec<u32>>();
for id in ids {
@@ -199,10 +180,7 @@ fn put_for_adoption(
}
}
unsafe fn sort_by_distances() {
let nodes =
unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap<u32, NumaNode>) };
fn sort_by_distances(nodes: &mut HashMap<u32, NumaNode>) {
for (id, node) in nodes {
node.distances.sort_by_key(|(_, e)| *e);
}