[PATCH RFC v6 18/26] nova-core: mm: Add virtual address range tracking to VMM

[Joel Fernandes <joelagnelf@nvidia.com>]

Extend the Virtual Memory Manager with optional virtual address range
tracking using a buddy allocator. This enables BarUser to allocate
contiguous virtual ranges for BAR1 mappings.

Signed-off-by: Joel Fernandes <joelagnelf@nvidia.com>
---
 drivers/gpu/nova-core/mm/vmm.rs | 49 +++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 3 deletions(-)

diff --git a/drivers/gpu/nova-core/mm/vmm.rs b/drivers/gpu/nova-core/mm/vmm.rs
index a5b4af9053a0..0ab80b84e55a 100644
--- a/drivers/gpu/nova-core/mm/vmm.rs
+++ b/drivers/gpu/nova-core/mm/vmm.rs
@@ -32,7 +32,9 @@
     gpu::buddy::{
         AllocatedBlocks,
         BuddyFlags,
-        GpuBuddyAllocParams, //
+        GpuBuddy,
+        GpuBuddyAllocParams,
+        GpuBuddyParams, //
     },
     prelude::*,
     sizes::SZ_4K,
@@ -60,29 +62,48 @@
 /// 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 BAR1 and BAR2 mappings.
+/// Directory Base (`PDB`) address. The [`Vmm`] is used for Channel, BAR1 and BAR2 mappings.
 ///
 /// The [`Vmm`] tracks all page table allocations made during mapping operations
 /// to ensure they remain valid for the lifetime of the address space.
+///
+/// It tracks virtual address allocations via a buddy allocator.
 pub(crate) struct Vmm {
     pdb_addr: VramAddress,
     mmu_version: MmuVersion,
     /// Page table allocations that must persist for the lifetime of mappings.
     page_table_allocs: KVec<Arc<AllocatedBlocks>>,
+    /// Buddy allocator for virtual address range tracking.
+    virt_buddy: GpuBuddy,
 }
 
 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 using
+    /// a buddy allocator. This enables [`Vmm::alloc_vfn_range()`] for allocating
+    /// contiguous virtual ranges.
+    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 virt_buddy = GpuBuddy::new(GpuBuddyParams {
+            base_offset_bytes: 0,
+            physical_memory_size_bytes: va_size,
+            chunk_size_bytes: SZ_4K as u64,
+        })?;
+
         Ok(Self {
             pdb_addr,
             mmu_version,
             page_table_allocs: KVec::new(),
+            virt_buddy,
         })
     }
 
@@ -96,6 +117,28 @@ pub(crate) fn mmu_version(&self) -> MmuVersion {
         self.mmu_version
     }
 
+    /// Allocate a contiguous virtual frame number range.
+    ///
+    /// Returns an [`Arc<AllocatedBlocks>`] representing the allocated range.
+    /// The allocation is automatically freed when the [`Arc`] is dropped.
+    pub(crate) fn alloc_vfn_range(&self, num_pages: usize) -> Result<(Vfn, Arc<AllocatedBlocks>)> {
+        let params = GpuBuddyAllocParams {
+            start_range_address: 0,
+            end_range_address: 0,
+            size_bytes: num_pages.checked_mul(PAGE_SIZE).ok_or(EOVERFLOW)? as u64,
+            min_block_size_bytes: SZ_4K as u64,
+            buddy_flags: BuddyFlags::try_new(BuddyFlags::CONTIGUOUS_ALLOCATION)?,
+        };
+
+        let alloc = self.virt_buddy.alloc_blocks(params)?;
+
+        // Get the starting offset from the first (and only, due to CONTIGUOUS) block.
+        let offset = alloc.iter().next().ok_or(ENOMEM)?.offset();
+        let vfn = Vfn::new(offset / PAGE_SIZE as u64);
+
+        Ok((vfn, alloc))
+    }
+
     /// Allocate a new page table, zero it, and track the allocation.
     ///
     /// This method ensures page table allocations persist for the lifetime of
-- 
2.34.1