[PATCH v11 16/20] gpu: nova-core: mm: Add multi-page mapping API to VMM

Joel Fernandes
2026-04-15T21:06:46Z

Add the page table mapping and unmapping API to the Virtual Memory
Manager, implementing a two-phase prepare/execute model suitable for
use both inside and outside the DMA fence signalling critical path.

Cc: Nikola Djukic <ndjukic@nvidia.com>
Signed-off-by: Joel Fernandes <joelagnelf@nvidia.com>
---
drivers/gpu/nova-core/mm/pagetable.rs | 1 +
drivers/gpu/nova-core/mm/pagetable/map.rs | 338 ++++++++++++++++++++++
drivers/gpu/nova-core/mm/vmm.rs | 217 ++++++++++++--
3 files changed, 537 insertions(+), 19 deletions(-)
create mode 100644 drivers/gpu/nova-core/mm/pagetable/map.rs

diff --git a/drivers/gpu/nova-core/mm/pagetable.rs b/drivers/gpu/nova-core/mm/pagetable.rs
index b7e0e8e02905..4070070922a4 100644
--- a/drivers/gpu/nova-core/mm/pagetable.rs
+++ b/drivers/gpu/nova-core/mm/pagetable.rs
@@ -8,6 +8,7 @@

#![expect(dead_code)]

+pub(super) mod map;
pub(super) mod ver2;
pub(super) mod ver3;
pub(super) mod walk;
diff --git a/drivers/gpu/nova-core/mm/pagetable/map.rs b/drivers/gpu/nova-core/mm/pagetable/map.rs
new file mode 100644
index 000000000000..a9719580143e
--- /dev/null
+++ b/drivers/gpu/nova-core/mm/pagetable/map.rs
@@ -0,0 +1,338 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Page table mapping operations for NVIDIA GPUs.
+
+use core::marker::PhantomData;
+
+use kernel::{
+ gpu::buddy::{
+ AllocatedBlocks,
+ GpuBuddyAllocFlags,
+ GpuBuddyAllocMode, //
+ },
+ prelude::*,
+ ptr::Alignment,
+ rbtree::{RBTree, RBTreeNode},
+ sizes::SZ_4K, //
+};
+
+use super::{
+ walk::{
+ PtWalkInner,
+ WalkPdeResult,
+ WalkResult, //
+ },
+ DualPdeOps,
+ MmuConfig,
+ MmuV2,
+ MmuV3,
+ MmuVersion,
+ PageTableLevel,
+ PdeOps,
+ PteOps, //
+};
+use crate::{
+ mm::{
+ GpuMm,
+ Pfn,
+ Vfn,
+ VramAddress,
+ PAGE_SIZE, //
+ },
+ num::{
+ IntoSafeCast, //
+ },
+};
+
+/// A pre-allocated and zeroed page table page.
+///
+/// Created during the mapping prepare phase and consumed during the execute phase.
+/// Stored in an [`RBTree`] keyed by the PDE slot address (`install_addr`).
+pub(in crate::mm) struct PreparedPtPage {
+ /// The allocated and zeroed page table page.
+ pub(in crate::mm) alloc: Pin<KBox<AllocatedBlocks>>,
+ /// Page table level -- needed to determine if this PT page is for a dual PDE.
+ pub(in crate::mm) level: PageTableLevel,
+}
+
+/// Page table mapper.
+pub(in crate::mm) struct PtMapInner<M: MmuConfig> {
+ walker: PtWalkInner<M>,
+ pdb_addr: VramAddress,
+ _phantom: PhantomData<M>,
+}
+
+impl<M: MmuConfig> PtMapInner<M> {
+ /// Create a new [`PtMapInner`].
+ pub(super) fn new(pdb_addr: VramAddress) -> Self {
+ Self {
+ walker: PtWalkInner::<M>::new(pdb_addr),
+ pdb_addr,
+ _phantom: PhantomData,
+ }
+ }
+
+ /// Allocate and zero a physical page table page.
+ fn alloc_and_zero_page(mm: &GpuMm, level: PageTableLevel) -> Result<PreparedPtPage> {
+ let blocks = KBox::pin_init(
+ mm.buddy().alloc_blocks(
+ GpuBuddyAllocMode::Simple,
+ SZ_4K.into_safe_cast(),
+ Alignment::new::<SZ_4K>(),
+ GpuBuddyAllocFlags::default(),
+ ),
+ GFP_KERNEL,
+ )?;
+
+ let page_vram = VramAddress::new(blocks.iter().next().ok_or(ENOMEM)?.offset());
+
+ // Zero via PRAMIN.
+ let mut window = mm.pramin().get_window()?;
+ let base = page_vram.raw();
+ for off in (0..PAGE_SIZE).step_by(8) {
+ window.try_write64(base + off, 0)?;
+ }
+
+ Ok(PreparedPtPage {
+ alloc: blocks,
+ level,
+ })
+ }
+
+ /// Ensure all intermediate page table pages exist for a single VFN.
+ ///
+ /// PRAMIN is released before each allocation and re-acquired after. Memory
+ /// allocations are done outside of holding this lock to prevent deadlocks with
+ /// the fence signalling critical path.
+ fn ensure_single_pte_path(
+ &self,
+ mm: &GpuMm,
+ vfn: Vfn,
+ pt_pages: &mut RBTree<VramAddress, PreparedPtPage>,
+ ) -> Result {
+ let max_iter = 2 * M::PDE_LEVELS.len();
+
+ for _ in 0..max_iter {
+ let mut window = mm.pramin().get_window()?;
+
+ let result = self
+ .walker
+ .walk_pde_levels(&mut window, vfn, |install_addr| {
+ pt_pages
+ .get(&install_addr)
+ .and_then(|p| p.alloc.iter().next().map(|b| VramAddress::new(b.offset())))
+ })?;
+
+ match result {
+ WalkPdeResult::Complete { .. } => {
+ return Ok(());
+ }
+ WalkPdeResult::Missing {
+ install_addr,
+ level,
+ } => {
+ // Drop PRAMIN before allocation.
+ drop(window);
+ let page = Self::alloc_and_zero_page(mm, level)?;
+ let node = RBTreeNode::new(install_addr, page, GFP_KERNEL)?;
+ let old = pt_pages.insert(node);
+ if old.is_some() {
+ kernel::pr_warn_once!(
+ "VMM: duplicate install_addr in pt_pages (internal consistency error)\n"
+ );
+ return Err(EIO);
+ }
+ }
+ }
+ }
+
+ kernel::pr_warn!(
+ "VMM: ensure_pte_path: loop exhausted after {} iters (VFN {:?})\n",
+ max_iter,
+ vfn
+ );
+ Err(EIO)
+ }
+
+ /// Prepare page table resources for mapping `num_pages` pages starting at `vfn_start`.
+ ///
+ /// Reserves capacity in `page_table_allocs`, then walks the hierarchy
+ /// per-VFN to prepare pages for all missing PDEs.
+ pub(super) fn prepare_map(
+ &self,
+ mm: &GpuMm,
+ vfn_start: Vfn,
+ num_pages: usize,
+ page_table_allocs: &mut KVec<Pin<KBox<AllocatedBlocks>>>,
+ pt_pages: &mut RBTree<VramAddress, PreparedPtPage>,
+ ) -> Result {
+ // Pre-reserve so install_mappings() can use push_within_capacity (no alloc
+ // in fence signalling critical path).
+ let pt_upper_bound = M::pt_pages_upper_bound(num_pages);
+ page_table_allocs.reserve(pt_upper_bound, GFP_KERNEL)?;
+
+ // Walk the hierarchy per-VFN to prepare pages for all missing PDEs.
+ for i in 0..num_pages {
+ let i_u64: u64 = i.into_safe_cast();
+ let vfn = Vfn::new(vfn_start.raw() + i_u64);
+ self.ensure_single_pte_path(mm, vfn, pt_pages)?;
+ }
+ Ok(())
+ }
+
+ /// Install prepared PDEs and write PTEs, then flush TLB.
+ ///
+ /// Drains `pt_pages` and moves allocations into `page_table_allocs`.
+ pub(super) fn install_mappings(
+ &self,
+ mm: &GpuMm,
+ pt_pages: &mut RBTree<VramAddress, PreparedPtPage>,
+ page_table_allocs: &mut KVec<Pin<KBox<AllocatedBlocks>>>,
+ vfn_start: Vfn,
+ pfns: &[Pfn],
+ writable: bool,
+ ) -> Result {
+ let mut window = mm.pramin().get_window()?;
+
+ // Drain prepared PT pages, install all pending PDEs.
+ let mut cursor = pt_pages.cursor_front_mut();
+ while let Some(c) = cursor {
+ let (next, node) = c.remove_current();
+ let (install_addr, page) = node.to_key_value();
+ let page_vram = VramAddress::new(page.alloc.iter().next().ok_or(ENOMEM)?.offset());
+
+ if page.level == M::DUAL_PDE_LEVEL {
+ let new_dpde = M::DualPde::new_small(Pfn::from(page_vram));
+ new_dpde.write(&mut window, install_addr)?;
+ } else {
+ let new_pde = M::Pde::new_vram(Pfn::from(page_vram));
+ new_pde.write(&mut window, install_addr)?;
+ }
+
+ page_table_allocs
+ .push_within_capacity(page.alloc)
+ .map_err(|_| ENOMEM)?;
+
+ cursor = next;
+ }
+
+ // Write PTEs (all PDEs now installed in HW).
+ for (i, &pfn) in pfns.iter().enumerate() {
+ let i_u64: u64 = i.into_safe_cast();
+ let vfn = Vfn::new(vfn_start.raw() + i_u64);
+ let result = self
+ .walker
+ .walk_to_pte_lookup_with_window(&mut window, vfn)?;
+
+ match result {
+ WalkResult::Unmapped { pte_addr } | WalkResult::Mapped { pte_addr, .. } => {
+ let pte = M::Pte::new_vram(pfn, writable);
+ pte.write(&mut window, pte_addr)?;
+ }
+ WalkResult::PageTableMissing => {
+ kernel::pr_warn_once!("VMM: page table missing for VFN {vfn:?}\n");
+ return Err(EIO);
+ }
+ }
+ }
+
+ drop(window);
+
+ // Flush TLB.
+ mm.tlb().flush(self.pdb_addr)
+ }
+
+ /// Invalidate PTEs for a range and flush TLB.
+ pub(super) fn invalidate_ptes(&self, mm: &GpuMm, vfn_start: Vfn, num_pages: usize) -> Result {
+ let invalid_pte = M::Pte::invalid();
+
+ let mut window = mm.pramin().get_window()?;
+ for i in 0..num_pages {
+ let i_u64: u64 = i.into_safe_cast();
+ let vfn = Vfn::new(vfn_start.raw() + i_u64);
+ let result = self
+ .walker
+ .walk_to_pte_lookup_with_window(&mut window, vfn)?;
+
+ match result {
+ WalkResult::Mapped { pte_addr, .. } | WalkResult::Unmapped { pte_addr } => {
+ invalid_pte.write(&mut window, pte_addr)?;
+ }
+ WalkResult::PageTableMissing => {
+ continue;
+ }
+ }
+ }
+ drop(window);
+
+ mm.tlb().flush(self.pdb_addr)
+ }
+}
+
+macro_rules! pt_map_dispatch {
+ ($self:expr, $method:ident ( $($arg:expr),* $(,)? )) => {
+ match $self {
+ PtMap::V2(inner) => inner.$method($($arg),*),
+ PtMap::V3(inner) => inner.$method($($arg),*),
+ }
+ };
+}
+
+/// Page table mapper dispatch.
+pub(in crate::mm) enum PtMap {
+ /// MMU v2 (Turing/Ampere/Ada).
+ V2(PtMapInner<MmuV2>),
+ /// MMU v3 (Hopper+).
+ V3(PtMapInner<MmuV3>),
+}
+
+impl PtMap {
+ /// Create a new page table mapper for the given MMU version.
+ pub(in crate::mm) fn new(pdb_addr: VramAddress, version: MmuVersion) -> Self {
+ match version {
+ MmuVersion::V2 => Self::V2(PtMapInner::<MmuV2>::new(pdb_addr)),
+ MmuVersion::V3 => Self::V3(PtMapInner::<MmuV3>::new(pdb_addr)),
+ }
+ }
+
+ /// Prepare page table resources for a mapping.
+ pub(in crate::mm) fn prepare_map(
+ &self,
+ mm: &GpuMm,
+ vfn_start: Vfn,
+ num_pages: usize,
+ page_table_allocs: &mut KVec<Pin<KBox<AllocatedBlocks>>>,
+ pt_pages: &mut RBTree<VramAddress, PreparedPtPage>,
+ ) -> Result {
+ pt_map_dispatch!(
+ self,
+ prepare_map(mm, vfn_start, num_pages, page_table_allocs, pt_pages)
+ )
+ }
+
+ /// Install prepared PDEs and write PTEs, then flush TLB.
+ pub(in crate::mm) fn install_mappings(
+ &self,
+ mm: &GpuMm,
+ pt_pages: &mut RBTree<VramAddress, PreparedPtPage>,
+ page_table_allocs: &mut KVec<Pin<KBox<AllocatedBlocks>>>,
+ vfn_start: Vfn,
+ pfns: &[Pfn],
+ writable: bool,
+ ) -> Result {
+ pt_map_dispatch!(
+ self,
+ install_mappings(mm, pt_pages, page_table_allocs, vfn_start, pfns, writable)
+ )
+ }
+
+ /// Invalidate PTEs for a range and flush TLB.
+ pub(in crate::mm) fn invalidate_ptes(
+ &self,
+ mm: &GpuMm,
+ vfn_start: Vfn,
+ num_pages: usize,
+ ) -> Result {
+ pt_map_dispatch!(self, invalidate_ptes(mm, vfn_start, num_pages))
+ }
+}
diff --git a/drivers/gpu/nova-core/mm/vmm.rs b/drivers/gpu/nova-core/mm/vmm.rs
index 0ff71119708d..4109d413e1b7 100644
--- a/drivers/gpu/nova-core/mm/vmm.rs
+++ b/drivers/gpu/nova-core/mm/vmm.rs
@@ -3,8 +3,7 @@
//! Virtual Memory Manager for NVIDIA GPU page table management.
//!
//! The [`Vmm`] provides high-level page mapping and unmapping operations for GPU
-//! virtual address spaces (Channels, BAR1, BAR2). It wraps the page table walker
-//! and handles TLB flushing after modifications.
+//! virtual address spaces (Channels, BAR1, BAR2).

use kernel::{
gpu::buddy::{
@@ -16,15 +15,25 @@
},
prelude::*,
ptr::Alignment,
+ rbtree::RBTree,
sizes::SZ_4K, //
};

-use core::ops::Range;
+use core::{
+ cell::Cell,
+ ops::Range, //
+};

use crate::{
mm::{
pagetable::{
- walk::{PtWalk, WalkResult},
+ map::{
+ PtMap, //
+ },
+ walk::{
+ PtWalk,
+ WalkResult, //
+ },
MmuVersion, //
},
GpuMm,
@@ -38,20 +47,77 @@
},
};

+/// Multi-page prepared mapping -- VA range allocated, ready for execute.
+///
+/// Produced by [`Vmm::prepare_map()`], consumed by [`Vmm::execute_map()`].
+/// The struct owns the VA space allocation between prepare and execute phases.
+pub(crate) struct PreparedMapping {
+ vfn_start: Vfn,
+ num_pages: usize,
+ vfn_alloc: Pin<KBox<AllocatedBlocks>>,
+}
+
+/// Result of a mapping operation -- tracks the active mapped range.
+///
+/// Returned by [`Vmm::execute_map()`] and [`Vmm::map_pages()`].
+/// Owns the VA allocation; the VA range is freed when this is dropped.
+/// Callers must call [`Vmm::unmap_pages()`] before dropping to invalidate
+/// PTEs (dropping only frees the VA range, not the PTE entries).
+pub(crate) struct MappedRange {
+ pub(super) vfn_start: Vfn,
+ pub(super) num_pages: usize,
+ /// VA allocation -- freed when [`MappedRange`] is dropped.
+ _vfn_alloc: Pin<KBox<AllocatedBlocks>>,
+ /// Logs a warning if dropped without unmapping.
+ _drop_guard: MustUnmapGuard,
+}
+
+/// Guard that logs a warning once if a [`MappedRange`] is dropped without
+/// calling [`Vmm::unmap_pages()`].
+struct MustUnmapGuard {
+ armed: Cell<bool>,
+}
+
+impl MustUnmapGuard {
+ const fn new() -> Self {
+ Self {
+ armed: Cell::new(true),
+ }
+ }
+
+ fn disarm(&self) {
+ self.armed.set(false);
+ }
+}
+
+impl Drop for MustUnmapGuard {
+ fn drop(&mut self) {
+ if self.armed.get() {
+ kernel::pr_warn!("MappedRange dropped without calling unmap_pages()\n");
+ }
+ }
+}
+
/// Virtual Memory Manager for a GPU address space.
///
/// Each [`Vmm`] instance manages a single address space identified by its Page
-/// Directory Base (`PDB`) address. The [`Vmm`] is used for Channel, BAR1 and
-/// BAR2 mappings.
+/// Directory Base (`PDB`) address. Used for Channel, BAR1 and BAR2 mappings.
pub(crate) struct Vmm {
/// Page Directory Base address for this address space.
pdb_addr: VramAddress,
- /// MMU version used for page table layout.
- mmu_version: MmuVersion,
+ /// Page table walker for reading existing mappings.
+ pt_walk: PtWalk,
+ /// Page table mapper for prepare/execute operations.
+ pt_map: PtMap,
/// Page table allocations required for mappings.
page_table_allocs: KVec<Pin<KBox<AllocatedBlocks>>>,
/// Buddy allocator for virtual address range tracking.
virt_buddy: GpuBuddy,
+ /// Prepared PT pages pending PDE installation, keyed by `install_addr`.
+ ///
+ /// Populated during prepare phase and drained in execute phase. Shared by all
+ /// pending maps, preventing races on the same PDE slot.
+ pt_pages: RBTree<VramAddress, super::pagetable::map::PreparedPtPage>,
}

impl Vmm {
@@ -76,19 +142,15 @@ pub(crate) fn new(

Ok(Self {
pdb_addr,
- mmu_version,
+ pt_walk: PtWalk::new(pdb_addr, mmu_version),
+ pt_map: PtMap::new(pdb_addr, mmu_version),
page_table_allocs: KVec::new(),
virt_buddy,
+ pt_pages: RBTree::new(),
})
}

/// Allocate a contiguous virtual frame number range.
- ///
- /// # Arguments
- ///
- /// - `num_pages`: Number of pages to allocate.
- /// - `va_range`: `None` = allocate anywhere, `Some(range)` = constrain allocation to the given
- /// range.
fn alloc_vfn_range(
&self,
num_pages: usize,
@@ -119,7 +181,6 @@ fn alloc_vfn_range(
GFP_KERNEL,
)?;

- // Get the starting offset of the first block (only block as range is contiguous).
let offset = alloc.iter().next().ok_or(ENOMEM)?.offset();
let vfn = Vfn::new(offset / page_size);

@@ -128,11 +189,129 @@ fn alloc_vfn_range(

/// Read the [`Pfn`] for a mapped [`Vfn`] if one is mapped.
pub(super) fn read_mapping(&self, mm: &GpuMm, vfn: Vfn) -> Result<Option<Pfn>> {
- let walker = PtWalk::new(self.pdb_addr, self.mmu_version);
-
- match walker.walk_to_pte_lookup(mm, vfn)? {
+ match self.pt_walk.walk_to_pte(mm, vfn)? {
WalkResult::Mapped { pfn, .. } => Ok(Some(pfn)),
WalkResult::Unmapped { .. } | WalkResult::PageTableMissing => Ok(None),
}
}
+
+ /// Prepare resources for mapping `num_pages` pages.
+ ///
+ /// Allocates a contiguous VA range, then walks the hierarchy per-VFN to prepare pages
+ /// for all missing PDEs. Returns a [`PreparedMapping`] with the VA allocation.
+ ///
+ /// If `va_range` is not `None`, the VA range is constrained to the given range. Safe
+ /// to call outside the fence signalling critical path.
+ pub(crate) fn prepare_map(
+ &mut self,
+ mm: &GpuMm,
+ num_pages: usize,
+ va_range: Option<Range<u64>>,
+ ) -> Result<PreparedMapping> {
+ if num_pages == 0 {
+ return Err(EINVAL);
+ }
+
+ // Allocate contiguous VA range.
+ let (vfn_start, vfn_alloc) = self.alloc_vfn_range(num_pages, va_range)?;
+
+ self.pt_map.prepare_map(
+ mm,
+ vfn_start,
+ num_pages,
+ &mut self.page_table_allocs,
+ &mut self.pt_pages,
+ )?;
+
+ Ok(PreparedMapping {
+ vfn_start,
+ num_pages,
+ vfn_alloc,
+ })
+ }
+
+ /// Execute a prepared multi-page mapping.
+ ///
+ /// Installs all prepared PDEs and writes PTEs into the page table, then flushes TLB.
+ pub(crate) fn execute_map(
+ &mut self,
+ mm: &GpuMm,
+ prepared: PreparedMapping,
+ pfns: &[Pfn],
+ writable: bool,
+ ) -> Result<MappedRange> {
+ if pfns.len() != prepared.num_pages {
+ return Err(EINVAL);
+ }
+
+ let PreparedMapping {
+ vfn_start,
+ num_pages,
+ vfn_alloc,
+ } = prepared;
+
+ self.pt_map.install_mappings(
+ mm,
+ &mut self.pt_pages,
+ &mut self.page_table_allocs,
+ vfn_start,
+ pfns,
+ writable,
+ )?;
+
+ Ok(MappedRange {
+ vfn_start,
+ num_pages,
+ _vfn_alloc: vfn_alloc,
+ _drop_guard: MustUnmapGuard::new(),
+ })
+ }
+
+ /// Map pages doing prepare and execute in the same call.
+ ///
+ /// This is a convenience wrapper for callers outside the fence signalling critical
+ /// path (e.g., BAR mappings). For DRM usecases, [`Vmm::prepare_map()`] and
+ /// [`Vmm::execute_map()`] will be called separately.
+ pub(crate) fn map_pages(
+ &mut self,
+ mm: &GpuMm,
+ pfns: &[Pfn],
+ va_range: Option<Range<u64>>,
+ writable: bool,
+ ) -> Result<MappedRange> {
+ if pfns.is_empty() {
+ return Err(EINVAL);
+ }
+
+ // Check if provided VA range is sufficient (if provided).
+ if let Some(ref range) = va_range {
+ let required: u64 = pfns
+ .len()
+ .checked_mul(PAGE_SIZE)
+ .ok_or(EOVERFLOW)?
+ .into_safe_cast();
+ let available = range.end.checked_sub(range.start).ok_or(EINVAL)?;
+ if available < required {
+ return Err(EINVAL);
+ }
+ }
+
+ let prepared = self.prepare_map(mm, pfns.len(), va_range)?;
+ self.execute_map(mm, prepared, pfns, writable)
+ }
+
+ /// Unmap all pages in a [`MappedRange`] with a single TLB flush.
+ pub(crate) fn unmap_pages(&mut self, mm: &GpuMm, range: MappedRange) -> Result {
+ self.pt_map
+ .invalidate_ptes(mm, range.vfn_start, range.num_pages)?;
+
+ // TODO: Internal page table pages (PDE, PTE pages) are still kept around.
+ // This is by design as repeated maps/unmaps will be fast. As a future TODO,
+ // we can add a reclaimer here to reclaim if VRAM is short. For now, the PT
+ // pages are dropped once the `Vmm` is dropped.
+
+ // Unmap complete, safe to drop `MappedRange`.
+ range._drop_guard.disarm();
+ Ok(())
+ }
}
--
2.34.1