[PATCH 02/16] gpu: nova-core: mm: Add buddy allocator and TLB to GpuMm

[Eliot Courtney <ecourtney@nvidia.com>]

From: Joel Fernandes <joelagnelf@nvidia.com>

Extend GpuMm with the remaining two memory-management components:

- Buddy allocator for VRAM allocation.
- TLB manager for translation buffer operations.

PRAMIN was added in an earlier commit; this completes the centralized
ownership model with accessor methods for each component.

Signed-off-by: Joel Fernandes <joelagnelf@nvidia.com>
[ecourtney: rebase GpuMm for borrowed BAR0, pin Tlb, write regs directly]
[ecourtney: size the buddy from the first usable FB region, drop dev_info]
[ecourtney: update for the VramAddress raw API and typed register base]
Signed-off-by: Eliot Courtney <ecourtney@nvidia.com>
---
 drivers/gpu/nova-core/Kconfig   |   1 +
 drivers/gpu/nova-core/gpu.rs    |  27 +++++++--
 drivers/gpu/nova-core/mm.rs     |  24 ++++++++
 drivers/gpu/nova-core/mm/tlb.rs | 120 ++++++++++++++++++++++++++++++++++++++++
 drivers/gpu/nova-core/regs.rs   |  68 +++++++++++++++++++++++
 5 files changed, 234 insertions(+), 6 deletions(-)

diff --git a/drivers/gpu/nova-core/Kconfig b/drivers/gpu/nova-core/Kconfig
index cb7f0b00f796..1934f17baa8b 100644
--- a/drivers/gpu/nova-core/Kconfig
+++ b/drivers/gpu/nova-core/Kconfig
@@ -5,6 +5,7 @@ config NOVA_CORE
 	depends on RUST
 	depends on !CPU_BIG_ENDIAN
 	select AUXILIARY_BUS
+	select GPU_BUDDY
 	select RUST_FW_LOADER_ABSTRACTIONS
 	default n
 	help
diff --git a/drivers/gpu/nova-core/gpu.rs b/drivers/gpu/nova-core/gpu.rs
index a52a3d4d86af..b797472279d3 100644
--- a/drivers/gpu/nova-core/gpu.rs
+++ b/drivers/gpu/nova-core/gpu.rs
@@ -6,11 +6,16 @@
     device,
     dma::Device,
     fmt,
+    gpu::buddy::GpuBuddyParams,
     io::Io,
     num::Bounded,
     pci,
     prelude::*,
-    sizes::SizeConstants, //
+    ptr::Alignment,
+    sizes::{
+        SizeConstants,
+        SZ_4K, //
+    },
 };
 
 use crate::{
@@ -422,11 +427,21 @@ pub(crate) fn new<'a>(
             },
 
             // Create GPU memory manager owning memory management resources.
-            mm: GpuMm::new(
-                bar,
-                gsp_resources.spec.chipset,
-                VramAddress::from_raw(gsp_static_info.total_fb_end),
-            )?,
+            mm: {
+                let usable_vram = gsp_static_info.usable_fb_regions.first().ok_or(ENODEV)?;
+                let buddy_params = GpuBuddyParams {
+                    base_offset: usable_vram.start,
+                    size: usable_vram.end - usable_vram.start,
+                    chunk_size: Alignment::new::<SZ_4K>(),
+                };
+
+                GpuMm::new(
+                    bar,
+                    gsp_resources.spec.chipset,
+                    buddy_params,
+                    VramAddress::from_raw(gsp_static_info.total_fb_end),
+                )?
+            },
         })
     }
 
diff --git a/drivers/gpu/nova-core/mm.rs b/drivers/gpu/nova-core/mm.rs
index 4234ba4d596a..202bb8f8ebed 100644
--- a/drivers/gpu/nova-core/mm.rs
+++ b/drivers/gpu/nova-core/mm.rs
@@ -40,6 +40,10 @@ macro_rules! impl_pfn_bounded {
 use kernel::{
     bitfield,
     fmt,
+    gpu::buddy::{
+        GpuBuddy,
+        GpuBuddyParams, //
+    },
     num::Bounded,
     prelude::*,
     ptr::{
@@ -54,16 +58,23 @@ macro_rules! impl_pfn_bounded {
     gpu::Chipset, //
 };
 
+pub(crate) use tlb::Tlb;
+
 mod hal;
 mod pramin;
 mod regs;
+pub(super) mod tlb;
 
 /// GPU Memory Manager - owns all core MM components.
 ///
 /// Provides centralized ownership of memory management resources:
+/// - [`GpuBuddy`] allocator for VRAM page table allocation.
 /// - [`pramin::Pramin`] for direct VRAM access.
+/// - [`Tlb`] manager for translation buffer flush operations.
 pub(crate) struct GpuMm<'gpu> {
+    buddy: GpuBuddy,
     pramin: pramin::Pramin<'gpu>,
+    tlb: Pin<KBox<Tlb<'gpu>>>,
 }
 
 impl<'gpu> GpuMm<'gpu> {
@@ -71,6 +82,7 @@ impl<'gpu> GpuMm<'gpu> {
     pub(crate) fn new(
         bar: Bar0<'gpu>,
         chipset: Chipset,
+        buddy_params: GpuBuddyParams,
         total_fb_end: VramAddress,
     ) -> Result<Self> {
         // PRAMIN covers all physical VRAM (including GSP-reserved areas
@@ -78,14 +90,26 @@ pub(crate) fn new(
         let vram_region = VramAddress::ZERO..total_fb_end;
 
         Ok(Self {
+            buddy: GpuBuddy::new(buddy_params)?,
             pramin: pramin::Pramin::new(bar, chipset, vram_region)?,
+            tlb: KBox::pin_init(Tlb::new(bar), GFP_KERNEL)?,
         })
     }
 
+    /// Access the [`GpuBuddy`] allocator.
+    pub(crate) fn buddy(&self) -> &GpuBuddy {
+        &self.buddy
+    }
+
     /// Access the [`pramin::Pramin`].
     fn pramin_mut(&mut self) -> &mut pramin::Pramin<'gpu> {
         &mut self.pramin
     }
+
+    /// Access the [`Tlb`] manager.
+    pub(crate) fn tlb(&self) -> &Tlb<'gpu> {
+        self.tlb.as_ref().get_ref()
+    }
 }
 
 /// Page size in bytes (4 KiB).
diff --git a/drivers/gpu/nova-core/mm/tlb.rs b/drivers/gpu/nova-core/mm/tlb.rs
new file mode 100644
index 000000000000..cc862e8159a1
--- /dev/null
+++ b/drivers/gpu/nova-core/mm/tlb.rs
@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! TLB (Translation Lookaside Buffer) flush support for GPU MMU.
+//!
+//! After modifying page table entries, the GPU's TLB must be flushed to
+//! ensure the new mappings take effect. This module provides TLB flush
+//! functionality for virtual memory managers.
+//!
+//! # Examples
+//!
+//! ```ignore
+//! use crate::mm::tlb::Tlb;
+//!
+//! fn page_table_update(tlb: &Tlb, pdb_addr: VramAddress) -> Result<()> {
+//!     // ... modify page tables ...
+//!
+//!     // Flush TLB to make changes visible (polls for completion).
+//!     tlb.flush(pdb_addr)?;
+//!
+//!     Ok(())
+//! }
+//! ```
+
+use kernel::{
+    io::poll::read_poll_timeout,
+    io::Io,
+    new_mutex,
+    prelude::*,
+    sync::Mutex,
+    time::Delta, //
+};
+
+use crate::{
+    bounded_enum,
+    driver::Bar0,
+    mm::VramAddress,
+    regs, //
+};
+
+bounded_enum! {
+    /// TLB invalidation acknowledgment scope.
+    ///
+    /// Controls how far the hardware waits for the invalidation to propagate
+    /// before clearing the `trigger` bit of `NV_TLB_FLUSH_CTRL`.
+    #[derive(Debug, Copy, Clone, PartialEq, Eq)]
+    pub(crate) enum TlbAckMode with TryFrom<Bounded<u32, 2>> {
+        /// Fire-and-forget: no acknowledgment required.
+        None = 0,
+        /// Wait for acknowledgment from all consumers, including remote GPUs
+        /// reachable over NVLink.
+        ///
+        /// Globally is strictly required only during unmap or permission
+        /// tightening, because the backing memory may be reassigned after the
+        /// flush returns and a stale TLB entry could let the GPU access freed
+        /// memory. For new mapping or relaxing permissions, a stale entry would
+        /// merely cause a redundant fault and retry, so [`TlbAckMode::None`]
+        /// would suffice.
+        Globally = 1,
+        /// Wait for acknowledgment from consumers within the local NVLink
+        /// fabric node only; skip cross-node ack.
+        Intranode = 2,
+    }
+}
+
+/// TLB manager for GPU translation buffer operations.
+#[pin_data]
+pub(crate) struct Tlb<'gpu> {
+    bar: Bar0<'gpu>,
+    /// TLB flush serialization lock: This lock is designed to be acquired during
+    /// the DMA fence signalling critical path. It should NEVER be held across any
+    /// reclaimable CPU memory allocations because the memory reclaim path can
+    /// call `dma_fence_wait()` (when implemented), which would deadlock if lock held.
+    #[pin]
+    lock: Mutex<()>,
+}
+
+impl<'gpu> Tlb<'gpu> {
+    /// Create a new TLB manager.
+    pub(super) fn new(bar: Bar0<'gpu>) -> impl PinInit<Self> {
+        pin_init!(Self {
+            bar,
+            lock <- new_mutex!((), "tlb_flush"),
+        })
+    }
+
+    /// Flush the GPU TLB for a specific page directory base.
+    ///
+    /// This invalidates all TLB entries associated with the given PDB address.
+    /// Must be called after modifying page table entries to ensure the GPU sees
+    /// the updated mappings.
+    pub(super) fn flush(&self, pdb_addr: VramAddress) -> Result {
+        let _guard = self.lock.lock();
+
+        // Write PDB address.
+        self.bar.write_reg(regs::NV_TLB_FLUSH_PDB_LO::from_pdb_addr(
+            pdb_addr.into_raw(),
+        ));
+        self.bar.write_reg(regs::NV_TLB_FLUSH_PDB_HI::from_pdb_addr(
+            pdb_addr.into_raw(),
+        ));
+
+        // Trigger flush.
+        self.bar.write_reg(
+            regs::NV_TLB_FLUSH_CTRL::zeroed()
+                .with_all_va(true)
+                .with_ack(TlbAckMode::None)
+                .with_trigger(true),
+        );
+
+        // Poll for completion.
+        read_poll_timeout(
+            || Ok(self.bar.read(regs::NV_TLB_FLUSH_CTRL)),
+            |ctrl: &regs::NV_TLB_FLUSH_CTRL| !ctrl.trigger(),
+            Delta::ZERO,
+            Delta::from_secs(2),
+        )?;
+
+        Ok(())
+    }
+}
diff --git a/drivers/gpu/nova-core/regs.rs b/drivers/gpu/nova-core/regs.rs
index bf1f3a97c632..9978fb2803b0 100644
--- a/drivers/gpu/nova-core/regs.rs
+++ b/drivers/gpu/nova-core/regs.rs
@@ -10,6 +10,7 @@
     sizes::SizeConstants,
     time, //
 };
+use pin_init::Zeroable;
 
 use crate::{
     driver::NovaRegisters,
@@ -26,6 +27,7 @@
         PFalconRegisters,
         PeregrineCoreSelect, //
     },
+    mm::tlb::TlbAckMode, //
 };
 
 // PBUS
@@ -471,3 +473,69 @@ pub(crate) mod gb202 {
         }
     }
 }
+
+// MMU TLB
+
+register! {
+    base: NovaRegisters;
+
+    /// TLB flush register: PDB address lower bits.
+    pub(crate) NV_TLB_FLUSH_PDB_LO(u32) @ 0x00b830a0 {
+        /// PDB address bits [39:8].
+        31:0    pdb_lo => u32;
+    }
+
+    /// TLB flush register: PDB address higher bits.
+    pub(crate) NV_TLB_FLUSH_PDB_HI(u32) @ 0x00b830a4 {
+        /// PDB address bits [47:40].
+        7:0     pdb_hi => u8;
+    }
+
+    /// TLB flush control register.
+    pub(crate) NV_TLB_FLUSH_CTRL(u32) @ 0x00b830b0 {
+        /// Invalidate every VA in the PDB selected by `NV_TLB_FLUSH_PDB_LO/HI`.
+        0:0     all_va => bool;
+        /// Invalidate TLBs for all PDBs (ignores `NV_TLB_FLUSH_PDB_LO/HI`).
+        1:1     all_pdb => bool;
+        /// Restrict the flush to the HUB MMU's TLBs; skip broadcasting to the
+        /// per-GPC L2 TLBs.
+        ///
+        /// The GPU MMU has a two-level TLB hierarchy:
+        /// 1. The *HUB MMU* sits at the top and serves memory requests from
+        ///    "host-side" engines: the host/channel interface, copy engines,
+        ///    display, and BAR1/BAR2 accesses.
+        /// 2. Each GPC (Graphics Processing Cluster — the block that houses
+        ///    shader cores / SMs) has its own L2 TLB that serves requests from
+        ///    the compute and graphics engines inside the cluster.
+        ///
+        /// When set, only the HUB TLBs are invalidated. This is a performance
+        /// optimization for flushes that only affect HUB-side mappings (e.g.
+        /// BAR1/BAR2 windows), where fanning the invalidation out to every
+        /// GPC's L2 TLB would be wasted work. Must be false when flushing
+        /// mappings that may be cached by compute/graphics engines.
+        2:2     hubtlb_only => bool;
+        /// Invalidation acknowledgment scope. See [`TlbAckMode`] for details.
+        8:7     ack ?=> TlbAckMode;
+        /// Write 1 to kick off the flush. Hardware clears this bit when the
+        /// flush completes; reads as 1 while the flush is in progress.
+        31:31   trigger => bool;
+    }
+}
+
+impl NV_TLB_FLUSH_PDB_LO {
+    /// Create a register value from a PDB address.
+    ///
+    /// Extracts bits [39:8] of the address and shifts it right by 8 bits.
+    pub(crate) fn from_pdb_addr(addr: u64) -> Self {
+        Self::zeroed().with_pdb_lo(((addr >> 8) & 0xFFFF_FFFF) as u32)
+    }
+}
+
+impl NV_TLB_FLUSH_PDB_HI {
+    /// Create a register value from a PDB address.
+    ///
+    /// Extracts bits [47:40] of the address and shifts it right by 40 bits.
+    pub(crate) fn from_pdb_addr(addr: u64) -> Self {
+        Self::zeroed().with_pdb_hi(((addr >> 40) & 0xFF) as u8)
+    }
+}

-- 
2.55.0