1 files changed, 90 insertions, 0 deletions
diff --git a/mm/tlb.cc b/mm/tlb.cc
new file mode 100644
index 0000000..b7a8ba7
--- /dev/null
+++ b/mm/tlb.cc
@@ -0,0 +1,90 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+/*
+ * tlb.cc — TLB shootdown implementation
+ *
+ * Uses RISC-V SBI Remote Fence extension to perform cross-hart
+ * TLB invalidation. SBI firmware handles the IPI delivery and
+ * synchronization — the call blocks until all target harts have
+ * completed their local sfence.vma.
+ *
+ * When a VmSpace is provided, we use its cpu_mask to target only
+ * the CPUs that have this address space loaded in satp. This avoids
+ * unnecessary IPI storms on unrelated cores.
+ *
+ * Reference: docs/specs/linux_abi_mm_plan.md §SMP §2.1
+ */
+
+#include <lume/tlb.h>
+#include <lume/vmm.h>
+#include <lume/config.h>
+#include <arch/mmu.h>
+#include <arch/tlb.h>
+#include <arch/cpu.h>
+
+/*
+ * If a shootdown range exceeds this many pages, do a full flush
+ * instead of per-page invalidation.
+ */
+static constexpr uint64 kShootdownPageThreshold = 32;
+
+/*
+ * Build a cpu mask for all online CPUs (fallback when mm is null).
+ */
+namespace {
+uint64 all_cpus_mask()
+{
+    uint64 mask = 0;
+    for (uint32 i = 0; i < kMaxCpus; i++) {
+        mask |= (1ULL << i);
+    }
+    return mask;
+}
+
+/*
+ * Get the target cpu mask: use mm->cpu_mask when available,
+ * fall back to all_cpus_mask for kernel address space (mm == nullptr).
+ */
+uint64 target_cpu_mask(VmSpace* mm)
+{
+    if (mm) {
+        uint64 mask = mm->cpu_mask.load(__ATOMIC_RELAXED);
+        /* Always include our own CPU — we may still have stale TLB */
+        mask |= (1ULL << arch::cpu::id());
+        return mask;
+    }
+    return all_cpus_mask();
+}
+} // namespace
+
+void tlb_shootdown_range(VmSpace* mm, uint64 start, uint64 end)
+{
+    if (start >= end) {
+        return;
+    }
+
+    uint64 num_pages = (end - start) / kPageSize;
+    uint64 cpu_mask = target_cpu_mask(mm);
+
+    if (num_pages > kShootdownPageThreshold) {
+        /* Full flush on target CPUs via arch backend */
+        arch::tlb::flush_remote(cpu_mask, 0, ~0ULL);
+        return;
+    }
+
+    /* Range flush on target CPUs via arch backend */
+    arch::tlb::flush_remote(cpu_mask, start, end - start);
+}
+
+void tlb_shootdown_page(VmSpace* mm, uint64 va)
+{
+    uint64 cpu_mask = target_cpu_mask(mm);
+    arch::tlb::flush_remote(cpu_mask, va, kPageSize);
+}
+
+void tlb_shootdown_all()
+{
+    uint64 cpu_mask = all_cpus_mask();
+    arch::tlb::flush_remote(cpu_mask, 0, ~0ULL);
+}