#![allow(static_mut_refs)]
use alloc::vec::Vec;
use core::sync::atomic::{AtomicU64, Ordering};
use semx_bitops::{bitmap::BitMap, make_64bit_mask};
use semx_cpumask::Cpumask;
use crate::{
println,
processor::{nr_cpus, this_processor_id},
space::mm::{MmStruct, entries::Pgd, tlb::local_flush_tlb_all},
sync::spinlock::Spinlock,
};
pub(crate) struct MmContext(pub(crate) AtomicU64);
impl MmContext {
pub(crate) const fn new() -> Self {
MmContext(AtomicU64::new(0))
}
}
const ID_AA64MMFR0_ASID_SHIFT: usize = 4;
static mut ASID_BITS: u64 = 0;
static CPU_ASID_LOCK: Spinlock<usize> = Spinlock::new(0);
static ASID_GENERATION: AtomicU64 = AtomicU64::new(0);
static mut ASID_MAP: BitMap = BitMap::new();
static mut ACTIVE_ASIDS: Vec<AtomicU64> = Vec::new();
static mut RESERVED_ASIDS: Vec<u64> = Vec::new();
static TLB_FLUSH_PENDING: Cpumask = Cpumask::new();
#[inline(always)]
fn asid_mask() -> u64 {
unsafe { !make_64bit_mask(0, ASID_BITS as usize) }
}
#[inline(always)]
fn asid_first_version() -> usize {
unsafe { 1 << ASID_BITS }
}
#[inline(always)]
fn num_user_asids() -> usize {
asid_first_version()
}
#[inline(always)]
fn asid2idx(asid: u64) -> u64 {
asid & !asid_mask()
}
#[inline(always)]
fn idx2asid(idx: u64) -> u64 {
asid2idx(idx)
}
fn get_cpu_asid_bits() -> u32 {
let value: u64;
unsafe {
core::arch::asm!(
"mrs {0}, ID_AA64MMFR0_EL1",
out(reg) value
);
}
let fld = (value << (64 - 4 - ID_AA64MMFR0_ASID_SHIFT)) >> (64 - 4);
match fld {
0 => 8,
2 => 16,
_ => {
println!("CPU{}: Unknown ASID size ({}); assuming 8-bit", this_processor_id(), fld);
8
},
}
}
fn verify_cpu_asid_bits() {
let asid = u64::from(get_cpu_asid_bits());
let boot_cpu_asid = unsafe { ASID_BITS };
assert!(
asid >= boot_cpu_asid,
"CPU{}: smaller ASID size({}) than boot CPU ({})",
this_processor_id(),
asid,
boot_cpu_asid
);
}
pub(crate) fn context_init() {
unsafe {
for _ in 0..nr_cpus() {
RESERVED_ASIDS.push(0);
ACTIVE_ASIDS.push(AtomicU64::new(0));
}
ASID_BITS = u64::from(get_cpu_asid_bits());
assert!(num_user_asids() - 1 > nr_cpus());
ASID_MAP.init(num_user_asids());
}
ASID_GENERATION.store(asid_first_version() as u64, Ordering::Relaxed);
println!("ASID allocator initialised with {} entries", num_user_asids());
}
pub(crate) fn context_init_cpu() {
verify_cpu_asid_bits();
}
fn flush_context() {
unsafe {
ASID_MAP.clear_range(0, num_user_asids());
for i in 0..nr_cpus() {
let active_asids = &ACTIVE_ASIDS[i];
let mut asid = active_asids.swap(0, Ordering::Relaxed);
let reserved_asids = &mut RESERVED_ASIDS[i];
if asid == 0 {
asid = *reserved_asids;
}
ASID_MAP.set_bit(asid2idx(asid) as usize);
*reserved_asids = asid;
TLB_FLUSH_PENDING.set(i);
}
}
}
fn check_update_reserved_asid(asid: u64, newasid: u64) -> bool {
unsafe {
let mut hit = false;
for reserved_asids in RESERVED_ASIDS.iter_mut().take(nr_cpus()) {
if *reserved_asids == asid {
hit = true;
*reserved_asids = newasid;
}
}
hit
}
}
fn set_aisd(asid: u64, generation: u64) -> u64 {
unsafe {
ASID_MAP.set_bit(asid as usize);
CUR_IDX = asid as usize;
idx2asid(asid) | generation
}
}
static mut CUR_IDX: usize = 1;
impl MmStruct {
fn new_context(&self) -> u64 {
unsafe {
let mut asid = self.context.0.load(Ordering::Relaxed);
let mut generation = ASID_GENERATION.load(Ordering::Relaxed);
if asid != 0 {
let newasid = generation | (asid & !asid_mask());
if check_update_reserved_asid(asid, newasid) {
return newasid;
}
if !ASID_MAP.test_and_set_bit(asid2idx(asid) as usize) {
return newasid;
}
}
asid = ASID_MAP.find_next_zero_bit(CUR_IDX) as u64;
if asid != num_user_asids() as u64 {
return set_aisd(asid, generation);
}
generation = ASID_GENERATION.fetch_add(asid_first_version() as u64, Ordering::Relaxed)
+ asid_first_version() as u64;
flush_context();
asid = ASID_MAP.find_next_zero_bit(1) as u64;
set_aisd(asid, generation)
}
}
pub(crate) fn check_and_switch_context(&self) {
unsafe {
let cpu = this_processor_id();
let mut asid = self.context.0.load(Ordering::Relaxed);
let active_asids = &ACTIVE_ASIDS[cpu];
let old_active_asid = active_asids.load(Ordering::Relaxed);
if old_active_asid != 0
&& (((asid ^ ASID_GENERATION.load(Ordering::Relaxed)) >> ASID_BITS) == 0)
&& active_asids
.compare_exchange(old_active_asid, asid, Ordering::Relaxed, Ordering::Relaxed)
.is_ok()
{
cpu_switch_mm(&self.pgd, self.context.0.load(Ordering::Relaxed));
return;
}
{
let mut lock = CPU_ASID_LOCK.lock_irq_save();
asid = self.context.0.load(Ordering::Relaxed);
if (asid ^ ASID_GENERATION.load(Ordering::Relaxed) >> ASID_BITS) != 0 {
asid = self.new_context();
self.context.0.store(asid, Ordering::Relaxed);
}
if TLB_FLUSH_PENDING.test_and_clear(cpu) {
local_flush_tlb_all();
}
active_asids.store(asid, Ordering::Relaxed);
*lock = 1;
}
cpu_switch_mm(&self.pgd, self.context.0.load(Ordering::Relaxed));
}
}
}
#[doc(hidden)]
#[unsafe(no_mangle)]
#[allow(unreachable_pub)]
pub extern "C" fn post_ttbr_update_workaround() {
unsafe {
core::arch::asm!("nop; nop; nop");
}
}
#[inline(always)]
fn cpu_switch_mm(pgd: &Pgd, asid: u64) {
unsafe extern "C" {
fn cpu_do_switch_mm(phys: usize, asid: u64);
}
unsafe {
cpu_do_switch_mm(pgd.to_phys().to_value(), asid);
}
}
core::arch::global_asm!(
r"
TTBR_CNP_BIT = 1
.global cpu_do_switch_mm; .align 2; cpu_do_switch_mm:
mrs x2, ttbr1_el1
mov x3, x0
cbz x1, 1f // skip CNP for reserved ASID
orr x3, x3, #TTBR_CNP_BIT
1:
bfi x2, x1, #48, #16 // set the ASID
msr ttbr1_el1, x2 // in TTBR1 (since TCR.A1 is set)
isb
msr ttbr0_el1, x3 // now update TTBR0
isb
b post_ttbr_update_workaround // Back to C code...
.type cpu_do_switch_mm, @function; .size cpu_do_switch_mm, .-cpu_do_switch_mm
"
);