yggdrasil/kernel/arch/x86/src/cpuid.rs

use core::ops::{BitAnd, BitOr};

use crate::registers::{CR0, CR4, XCR0};
use kernel_arch_interface::cpu::CpuFeatureSet;
use tock_registers::interfaces::ReadWriteable;

macro_rules! cpuid_features {
    ($vis:vis $name:ident: $ty:ty [ $($feature:ident: $bit:literal),* ]) => {
        bitflags::bitflags! {
            #[derive(Clone, Copy, PartialEq, Eq, Debug)]
            $vis struct $name: $ty {
                $(const $feature = 1 << $bit;)*
            }
        }

        impl $name {
            pub fn as_str(&self) -> &'static str {
                match self {
                    $(&Self::$feature => stringify!($feature),)*
                    _ => "(unknown)"
                }
            }
        }
    };
}

cpuid_features! {
    pub EcxFeatures: u32 [
        SSE3: 0,
        PCLMUL: 1,
        DTES64: 2,
        MONITOR: 3,
        DS_CPL: 4,
        VMX: 5,
        SMX: 6,
        EST: 7,
        TM2: 8,
        SSSE3: 9,
        CID: 10,
        SDBG: 11,
        FMA: 12,
        CX16: 13,
        XTPR: 14,
        PDCM: 15,
        PCID: 17,
        DCA: 18,
        SSE4_1: 19,
        SSE4_2: 20,
        X2APIC: 21,
        MOVBE: 22,
        POPCNT: 23,
        TSC: 24,
        AES: 25,
        XSAVE: 26,
        OSXSAVE: 27,
        AVX: 28,
        F16C: 29,
        RDRAND: 30,
        HYPERVISOR: 31
    ]
}

cpuid_features! {
    pub EdxFeatures: u32 [
        FPU: 0,
        VME: 1,
        DE: 2,
        PSE: 3,
        TSC: 4,
        MSR: 5,
        PAE: 6,
        MCE: 7,
        CX8: 8,
        APIC: 9,
        SEP: 11,
        MTRR: 12,
        PGE: 13,
        MCA: 14,
        CMOV: 15,
        PAT: 16,
        PSE36: 17,
        PSN: 18,
        CLFLUSH: 19,
        DS: 21,
        ACPI: 22,
        MMX: 23,
        FXSR: 24,
        SSE: 25,
        SSE2: 26,
        SS: 27,
        HTT: 28,
        TM: 29,
        IA64: 30,
        PBE: 31
    ]
}

#[derive(Clone, Copy, Debug)]
pub struct CpuFeatures {
    pub ecx: EcxFeatures,
    pub edx: EdxFeatures,
}

impl CpuFeatures {
    pub const fn empty() -> Self {
        Self {
            ecx: EcxFeatures::empty(),
            edx: EdxFeatures::empty(),
        }
    }

    pub fn contains_all(&self, features: &Self) -> bool {
        self.ecx.contains(features.ecx) && self.edx.contains(features.edx)
    }

    pub fn match_all(&self, features: &Self) -> Result<(), Self> {
        if self.contains_all(features) {
            Ok(())
        } else {
            Err(Self {
                ecx: features.ecx & !self.ecx,
                edx: features.edx & !self.edx,
            })
        }
    }
}

impl BitAnd<CpuFeatures> for CpuFeatures {
    type Output = CpuFeatures;

    fn bitand(self, rhs: CpuFeatures) -> Self::Output {
        Self {
            ecx: self.ecx & rhs.ecx,
            edx: self.edx & rhs.edx,
        }
    }
}

impl BitOr<CpuFeatures> for CpuFeatures {
    type Output = CpuFeatures;

    fn bitor(self, rhs: CpuFeatures) -> Self::Output {
        Self {
            ecx: self.ecx | rhs.ecx,
            edx: self.edx | rhs.edx,
        }
    }
}

impl CpuFeatureSet for CpuFeatures {
    fn iter(&self) -> impl Iterator<Item = &'static str> {
        let ecx = self.ecx.iter().map(|e| e.as_str());
        let edx = self.edx.iter().map(|e| e.as_str());

        core::iter::chain(ecx, edx)
    }
}

#[cfg(any(target_arch = "x86_64", rust_analyzer))]
unsafe fn raw_cpuid(eax: u32, result: &mut [u32]) {
    core::arch::asm!(
        r#"
            push %rbx
            cpuid
            mov %ebx, {0:e}
            pop %rbx
        "#,
        out(reg) result[0],
        out("edx") result[1],
        out("ecx") result[2],
        in("eax") eax,
        options(att_syntax)
    );
}

#[cfg(any(target_arch = "x86", rust_analyzer))]
unsafe fn raw_cpuid(eax: u32, result: &mut [u32]) {
    core::arch::asm!(
        r#"
            push %ebx
            cpuid
            mov %ebx, {0:e}
            pop %ebx
        "#,
        out(reg) result[0],
        out("edx") result[1],
        out("ecx") result[2],
        in("eax") eax,
        options(att_syntax)
    );
}

fn cpuid_features() -> (EcxFeatures, EdxFeatures) {
    let mut raw = [0; 3];

    unsafe {
        raw_cpuid(0x1, &mut raw);
    }

    (
        EcxFeatures::from_bits_truncate(raw[2]),
        EdxFeatures::from_bits_truncate(raw[1]),
    )
}

fn enable_features(ecx: EcxFeatures, edx: EdxFeatures) {
    if ecx.contains(EcxFeatures::XSAVE) {
        CR4.modify(CR4::OSXSAVE::SET);
    }
    if edx.contains(EdxFeatures::FXSR) {
        CR4.modify(CR4::OSFXSR::SET);
    }
    if edx.contains(EdxFeatures::FPU) {
        CR0.modify(CR0::EM::CLEAR);
        if ecx.contains(EcxFeatures::XSAVE) {
            XCR0.modify(XCR0::X87::SET);
        }
    }
    if edx.contains(EdxFeatures::SSE) && ecx.contains(EcxFeatures::XSAVE) {
        assert!(edx.contains(EdxFeatures::FPU));
        CR0.modify(CR0::MP::SET);
        CR4.modify(CR4::OSXMMEXCPT::SET);
        XCR0.modify(XCR0::SSE::SET);
    }
    if ecx.contains(EcxFeatures::AVX) {
        // Must also have FPU enabled
        assert!(edx.contains(EdxFeatures::FPU) && ecx.contains(EcxFeatures::XSAVE));
        XCR0.modify(XCR0::AVX::SET);
    }
    if ecx.contains(EcxFeatures::PCID) {
        CR4.modify(CR4::PCIDE::SET);
    }
    if edx.contains(EdxFeatures::PSE) {
        CR4.modify(CR4::PSE::SET);
    }

    CR0.modify(CR0::TS::CLEAR);
}

fn read_features() -> CpuFeatures {
    let (ecx, edx) = cpuid_features();
    CpuFeatures { ecx, edx }
}

pub fn setup_features(
    want_features: CpuFeatures,
    need_features: CpuFeatures,
) -> (CpuFeatures, Result<CpuFeatures, CpuFeatures>) {
    let have_features = read_features();
    let will_features = have_features & (want_features | need_features);

    if let Err(missing_features) = will_features.match_all(&need_features) {
        return (have_features, Err(missing_features));
    }

    enable_features(will_features.ecx, will_features.edx);

    (have_features, Ok(will_features))
}