From d600b890242beaca9664b75a41db82ba94d04183 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Gr=C3=A9goire?= <loic@gregoire.dev>
Date: Wed, 15 Apr 2026 20:12:12 -0400
Subject: [PATCH] feat(pc): enter x2APIC mode

---
 src/kern/amd64/init.zig  |  2 ++
 src/platform/pc/apic.zig | 35 +++++++++++++++++++++++++----------
 src/platform/pc/impl.zig | 11 +++++++++++
 3 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/src/kern/amd64/init.zig b/src/kern/amd64/init.zig
index f56919a..63f204d 100644
--- a/src/kern/amd64/init.zig
+++ b/src/kern/amd64/init.zig
@@ -5,6 +5,7 @@ const std = @import("std");
 const impl = @import("impl.zig");
 const ke = b.ke;
 const ki = b.ke.private;
+const pl = b.pl;
 
 var gdt = extern struct {
     entries: [9]u64 align(1),
@@ -85,6 +86,7 @@ fn early_cpu_init() linksection(b.init) void {
 pub fn ap_entry(cpu_id: u32) noreturn {
     early_cpu_init();
     ki.cpu.init_cpu(cpu_id);
+    pl.impl.init_ap();
 
     while (true) {
         std.atomic.spinLoopHint();
diff --git a/src/platform/pc/apic.zig b/src/platform/pc/apic.zig
index a930452..c597d88 100644
--- a/src/platform/pc/apic.zig
+++ b/src/platform/pc/apic.zig
@@ -3,6 +3,7 @@ const amd64 = @import("arch");
 const b = @import("base");
 const acpi = b.pl.acpi;
 const mm = b.mm;
+const ke = b.ke;
 
 const MadtLapic = extern struct {
     header: acpi.MadtEntryHeader,
@@ -72,7 +73,7 @@ pub const ApicRegisters = enum(u32) {
 pub var apics: std.ArrayList(u32) = .empty;
 var xapic_address: ?usize = null;
 pub var xapic_base_physical: usize = 0;
-pub var in_x2apic_mode: bool = false;
+const in_x2apic_mode = ke.CpuLocal(bool, false);
 
 const log = std.log.scoped(.apic);
 
@@ -95,7 +96,7 @@ fn x2apic_read(register: ApicRegisters) u64 {
 }
 
 pub fn write(register: ApicRegisters, value: u32) void {
-    if (in_x2apic_mode) {
+    if (in_x2apic_mode.local().*) {
         x2apic_write(register, value);
     } else {
         xapic_write(register, value);
@@ -103,7 +104,7 @@ pub fn write(register: ApicRegisters, value: u32) void {
 }
 
 pub fn read(register: ApicRegisters) u64 {
-    if (in_x2apic_mode) {
+    if (in_x2apic_mode.local().*) {
         return x2apic_read(register);
     } else {
         return @intCast(xapic_read(register));
@@ -111,7 +112,7 @@ pub fn read(register: ApicRegisters) u64 {
 }
 
 pub fn get_id() u32 {
-    if (in_x2apic_mode) {
+    if (in_x2apic_mode.local().*) {
         return @truncate(x2apic_read(.Id));
     } else {
         return xapic_read(.Id) >> 24;
@@ -119,7 +120,7 @@ pub fn get_id() u32 {
 }
 
 pub fn send_init(apic_id: u32) void {
-    if (in_x2apic_mode) {
+    if (in_x2apic_mode.local().*) {
         x2apic_write(.Icr, (@as(u64, apic_id) << 32) | 0x0000C500);
     } else {
         xapic_write(.Icr1, @as(u32, apic_id) << 24);
@@ -129,7 +130,7 @@ pub fn send_init(apic_id: u32) void {
 }
 
 pub fn send_sipi(apic_id: u32, startup_page: u8) void {
-    if (in_x2apic_mode) {
+    if (in_x2apic_mode.local().*) {
         x2apic_write(.Icr, (@as(u64, apic_id) << 32) | 0x00004600 | startup_page);
     } else {
         xapic_write(.Icr1, @as(u32, apic_id) << 24);
@@ -140,7 +141,7 @@ pub fn send_sipi(apic_id: u32, startup_page: u8) void {
 
 pub fn send_ipi(apic_id: u32, vector: u8, delivery_mode: u8) void {
     const icr_value = (@as(u64, apic_id) << 32) | (@as(u64, delivery_mode) << 8) | vector;
-    if (in_x2apic_mode) {
+    if (in_x2apic_mode.local().*) {
         x2apic_write(.Icr, icr_value);
     } else {
         // Write to ICR1 first, because writing to the low word causes
@@ -198,12 +199,26 @@ pub fn init() linksection(b.init) void {
     if (base & (1 << 10) == 0) {
         xapic_address = mm.p2v(base & 0xFFFFF000);
         xapic_base_physical = base & 0xFFFFF000;
+    } else {
+        in_x2apic_mode.local().* = true;
     }
+
+    init_local();
+}
+
+pub fn init_local() void {
+    write(.Svr, @intCast(read(.Svr) | 0x1FF));
+    write(.Tpr, 0);
+    write(.Eoi, 0);
 }
 
 pub fn enter_x2apic() void {
-    if (in_x2apic_mode) return;
-
     const apic_base = amd64.read_msr(.LapicBase);
-    amd64.wrmsr(.LapicBase, apic_base | (1 << 10) | (1 << 11));
+    in_x2apic_mode.local().* = true;
+    if (apic_base & (1 << 10) != 0) return; // Already in x2apic.
+    if (apic_base & (1 << 11) == 0) {
+        // Ensure xAPIC is enabled first.
+        amd64.write_msr(.LapicBase, apic_base | (1 << 11));
+    }
+    amd64.write_msr(.LapicBase, apic_base | (1 << 10) | (1 << 11));
 }
diff --git a/src/platform/pc/impl.zig b/src/platform/pc/impl.zig
index c535ab7..d0faa01 100644
--- a/src/platform/pc/impl.zig
+++ b/src/platform/pc/impl.zig
@@ -72,6 +72,13 @@ fn init_hypervisor(hv: amd64.hypervisor.Info) void {
     }
 }
 
+pub fn init_ap() void {
+    if (amd64.cpu_features.x2apic) {
+        apic.enter_x2apic();
+    }
+    apic.init_local();
+}
+
 pub fn late_init(boot_info: *pl.BootInfo) linksection(b.init) void {
     pl.acpi.init(boot_info);
 
@@ -85,6 +92,10 @@ pub fn late_init(boot_info: *pl.BootInfo) linksection(b.init) void {
     if (apic.apics.items.len != 0) {
         smp.init();
     }
+
+    if (amd64.cpu_features.x2apic) {
+        apic.enter_x2apic();
+    }
 }
 
 pub fn debug_write(c: u8) void {