[PATCH v1 12/16] gpu: nova-core: mm: Add virtual address range tracking to VMM

[Joel Fernandes <joelagnelf@nvidia.com>]

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

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

diff --git a/drivers/gpu/nova-core/mm/vmm.rs b/drivers/gpu/nova-core/mm/vmm.rs
index 3e18adc23b68..05ff77c5f888 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,
+    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,18 +44,74 @@ 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> {
+        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