From 63cd1046168e5385975301ee2afb807bf7e4b67e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lo=C3=AFc=20Gr=C3=A9goire?= <loic@gregoire.dev>
Date: Sun, 9 Aug 2026 16:44:41 -0400
Subject: [PATCH] ke: optimize TLB shootdowns

---
 src/kern/ipi.zig       |  3 +-
 src/kern/shootdown.zig | 68 +++++++++++++++++++++++++++++++-----------
 src/mm/tlb.zig         |  9 ++++++
 3 files changed, 61 insertions(+), 19 deletions(-)

diff --git a/src/kern/ipi.zig b/src/kern/ipi.zig
index 9e45170..97b8b70 100644
--- a/src/kern/ipi.zig
+++ b/src/kern/ipi.zig
@@ -53,7 +53,6 @@ pub fn unicast(
     const state = remote.state.swap(.Pending, .release);
 
     if (state == .Idle) {
-        std.log.info("Sent ipi", .{});
         pl.send_ipi(cpu);
     }
 
@@ -65,7 +64,7 @@ pub fn unicast(
 }
 
 pub fn broadcast(func: *const fn (u32, ?*anyopaque) void, arg: ?*anyopaque) void {
-    const ipl = ke.ipl.raise(.High);
+    const ipl = ke.ipl.raise(.Dispatch);
     const curcpu = percpu.local();
 
     curcpu.func = func;
diff --git a/src/kern/shootdown.zig b/src/kern/shootdown.zig
index dea329c..ac7e7b4 100644
--- a/src/kern/shootdown.zig
+++ b/src/kern/shootdown.zig
@@ -38,6 +38,7 @@ const PerCpu = struct {
     /// CPUs that have sent this CPU a shootdown.
     senders: ke.AtomicCpuMask,
 
+    npages: std.atomic.Value(u32),
     sync_addr: r.VAddr,
     sync_count: usize,
 };
@@ -118,41 +119,67 @@ pub fn process_shootdowns() void {
     const cpu = percpu.local();
     const curcpu = ke.cpu.current();
 
+    rtl.barrier.fence(.acquire);
+
     if (cpu.senders.is_all(false, .monotonic)) {
         // Nothing to do.
         return;
     }
 
+    var claimed: [config.ncpus]u64 = @splat(0);
+    var nreq: usize = 0;
+
     var iter = cpu.senders.iter(.monotonic);
 
-    // Go through every sender and get their states.
     while (iter.next()) |i| {
         if (i == curcpu) continue;
 
         cpu.senders.clear(i, .monotonic);
 
+        const states = percpu.remote(@truncate(i)).valid_states[curcpu].swap(
+            0,
+            .acquire,
+        );
+
+        claimed[i] = states;
+        nreq += @popCount(states);
+    }
+
+    if (nreq == 0) return;
+
+    // npages is just a hint, we don't need it to be accurate.
+    const full = cpu.npages.load(.monotonic) >= ki.impl.tlb_max_pages;
+    cpu.npages.store(0, .monotonic);
+
+    if (full) {
+        // Just flush the entire thing.
+        ki.impl.flush_full_tlb();
+    }
+
+    // Go through every sender and release their states.
+    // Additionally flush the ranges that didn't get flushed previously.
+    for (0.., &claimed) |i, states| {
+        if (states == 0) continue;
+
         const sender = percpu.remote(@truncate(i));
+        var bits = states;
 
-        while (true) {
-            var states = sender.valid_states[curcpu].swap(0, .acquire);
-            if (states == 0) break;
+        while (bits != 0) : (bits &= bits - 1) {
+            const state = &sender.states[@ctz(bits)];
 
-            while (states != 0) {
-                // Go through each set bit and invalidate the slot.
-                const bit = @ctz(states);
-                states &= ~(@as(u64, 1) << @intCast(bit));
+            if (!full) flush_range(state.base, state.npages);
 
-                const state = &sender.states[bit];
-                flush_range(state.base, state.npages);
-
-                if (state.state.fetchSub(1, .release) == 1) {
-                    shootdowns.insert(&state.link);
-                }
+            if (state.state.fetchSub(1, .release) == 1) {
+                shootdowns.insert(&state.link);
             }
         }
     }
 }
 
+fn process_shootdowns_kick(_: u32, _: ?*anyopaque) void {
+    process_shootdowns();
+}
+
 /// Submit a shootdown to occur asynchronously on `target_mask`.
 /// Returns an error if a synchronous shootdown was done instead.
 pub fn submit(state: ShootdownState, target_mask: ke.CpuMask) !void {
@@ -171,7 +198,10 @@ pub fn submit(state: ShootdownState, target_mask: ke.CpuMask) !void {
     };
 
     cpu.states[slot] = state;
-    cpu.states[slot].state.store(@truncate(target_mask.count()), .release);
+    cpu.states[slot].state.store(
+        @truncate(target_mask.count()),
+        .release,
+    );
 
     var iter = target_mask.iter();
 
@@ -183,9 +213,13 @@ pub fn submit(state: ShootdownState, target_mask: ke.CpuMask) !void {
             .release,
         );
 
-        ki.ipl.set_softint_pending(@truncate(bit), .Dispatch);
+        const remote = percpu.remote(@truncate(bit));
 
         // Tell him we have something for him.
-        percpu.remote(@truncate(bit)).senders.set(curcpu, .release);
+        _ = remote.npages.fetchAdd(state.npages, .monotonic);
+        remote.senders.set(curcpu, .release);
+
+        rtl.barrier.fence(.release);
+        ki.ipl.set_softint_pending(@truncate(bit), .Dispatch);
     }
 }
diff --git a/src/mm/tlb.zig b/src/mm/tlb.zig
index 849c6c5..9a9102e 100644
--- a/src/mm/tlb.zig
+++ b/src/mm/tlb.zig
@@ -74,6 +74,9 @@ pub fn reclaim_range(space: *mm.Space, va: r.VAddr, size: usize) void {
         },
     };
 
+    space.lock.release();
+
+    const ipl = ke.ipl.raise(.Dispatch);
     var mask: ke.CpuMask = .init(false);
 
     for (0..ke.ncpus) |i| {
@@ -84,6 +87,9 @@ pub fn reclaim_range(space: *mm.Space, va: r.VAddr, size: usize) void {
 
     ke.shootdown.submit(state, mask) catch {
         _ = sync_shootdowns.fetchAdd(1, .monotonic);
+
+        ke.ipl.lower(ipl);
+        space.lock.acquire();
         space.arena.free(va, size) catch unreachable;
 
         // Free the physical pages
@@ -94,6 +100,9 @@ pub fn reclaim_range(space: *mm.Space, va: r.VAddr, size: usize) void {
         );
         return;
     };
+
+    ke.ipl.lower(ipl);
+    space.lock.acquire();
 }
 
 fn activation(_: *rtl.HandoffList) void {