[PATCH v12 15/22] gpu: nova-core: mm: Add virtual address range tracking to VMM

Joel Fernandes
2026-04-25T21:15:45Z

Add virtual address range tracking to the VMM using a buddy allocator.
This enables contiguous virtual address range allocation for mappings.

Cc: Nikola Djukic <ndjukic@nvidia.com>
Signed-off-by: Joel Fernandes <joelagnelf@nvidia.com>
---
drivers/gpu/nova-core/mm/vmm.rs | 89 +++++++++++++++++++++++++++++----
1 file changed, 79 insertions(+), 10 deletions(-)

diff --git a/drivers/gpu/nova-core/mm/vmm.rs b/drivers/gpu/nova-core/mm/vmm.rs
index 7a2ab2acaa9b..81ae08ce1cd7 100644
--- a/drivers/gpu/nova-core/mm/vmm.rs
+++ b/drivers/gpu/nova-core/mm/vmm.rs
@@ -9,18 +9,27 @@
use kernel::{
device,
gpu::buddy::AllocatedBlocks,
- prelude::*, //
+ maple_tree::MapleTreeAlloc,
+ prelude::*,
};

-use crate::mm::{
- pagetable::{
- walk::{PtWalk, WalkResult},
- MmuVersion, //
+use core::ops::Range;
+
+use crate::{
+ mm::{
+ pagetable::{
+ walk::{PtWalk, WalkResult},
+ MmuVersion, //
+ },
+ GpuMm,
+ Pfn,
+ Vfn,
+ VramAddress,
+ PAGE_SIZE, //
+ },
+ num::{
+ IntoSafeCast, //
},
- GpuMm,
- Pfn,
- Vfn,
- VramAddress, //
};

/// Virtual Memory Manager for a GPU address space.
@@ -35,23 +44,83 @@ pub(crate) struct Vmm {
mmu_version: MmuVersion,
/// Page table allocations required for mappings.
page_table_allocs: KVec<Pin<KBox<AllocatedBlocks>>>,
+ /// Maple tree allocator for virtual address range tracking.
+ virt_alloc: Pin<KBox<MapleTreeAlloc<()>>>,
+ /// Total number of pages in the virtual address space.
+ va_pages: usize,
}

impl Vmm {
/// Create a new [`Vmm`] for the given Page Directory Base address.
- pub(crate) fn new(pdb_addr: VramAddress, mmu_version: MmuVersion) -> Result<Self> {
+ ///
+ /// The [`Vmm`] will manage a virtual address space of `va_size` bytes.
+ pub(crate) fn new(
+ pdb_addr: VramAddress,
+ mmu_version: MmuVersion,
+ va_size: u64,
+ ) -> Result<Self> {
// Only MMU v2 is supported for now.
if mmu_version != MmuVersion::V2 {
return Err(ENOTSUPP);
}

+ let page_size: u64 = PAGE_SIZE.into_safe_cast();
+ let va_pages: usize = (va_size / page_size).into_safe_cast();
+ let virt_alloc = KBox::pin_init(MapleTreeAlloc::<()>::new(), GFP_KERNEL)?;
+
Ok(Self {
pdb_addr,
mmu_version,
page_table_allocs: KVec::new(),
+ virt_alloc,
+ va_pages,
})
}

+ /// 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,
+ va_range: Option<Range<u64>>,
+ ) -> Result<Vfn> {
+ let page_size: u64 = PAGE_SIZE.into_safe_cast();
+
+ let start_vfn = match va_range {
+ Some(r) => {
+ let num_pages_u64: u64 = num_pages.into_safe_cast();
+ let size = num_pages_u64.checked_mul(page_size).ok_or(EOVERFLOW)?;
+ let range_size = r.end.checked_sub(r.start).ok_or(EOVERFLOW)?;
+ if range_size != size {
+ return Err(EINVAL);
+ }
+ let start_vfn: usize = (r.start / page_size).into_safe_cast();
+ let end_vfn: usize = (r.end / page_size).into_safe_cast();
+ self.virt_alloc
+ .insert_range(start_vfn..end_vfn, (), GFP_KERNEL)?;
+ start_vfn
+ }
+ None => self
+ .virt_alloc
+ .alloc_range(num_pages, (), ..self.va_pages, GFP_KERNEL)?,
+ };
+
+ Ok(Vfn::new(start_vfn.into_safe_cast()))
+ }
+
+ /// Free a virtual frame number range back to the maple tree.
+ fn free_vfn(&self, vfn: Vfn) {
+ let vfn_index: usize = vfn.raw().into_safe_cast();
+ if self.virt_alloc.erase(vfn_index).is_none() {
+ kernel::pr_warn!("free_vfn: VFN {} not found in maple tree\n", vfn_index);
+ }
+ }
+
/// Read the [`Pfn`] for a mapped [`Vfn`] if one is mapped.
pub(super) fn read_mapping(
&self,
--
2.34.1