[PATCH v2] gpu: nova-core: mm: reserve room for every prepared page table page

[PATCH v2] gpu: nova-core: mm: reserve room for every prepared page table page

[Vladislav Zaharov]

install_mappings() moves every page in pt_pages into page_table_allocs
with push_within_capacity(), but prepare_map() reserves room only for an
estimate of the pages of the mapping it prepares. The estimate is too
low for a range that crosses a page table boundary, and it leaves out
pages that earlier calls put in pt_pages. Preparing several mappings
before executing them is then enough to make the push fail. The page is
freed although its PDE has been written, and pages still in pt_pages
can have their PDEs installed into it later.

Reserve room for each page when it is added to pt_pages, and drop the
estimate. Also move a page into page_table_allocs before writing its
PDE, so that a failing PDE write does not free it either.

The problem was found by an LLM-assisted code review.

Fixes: 4e4182282d4e ("gpu: nova-core: mm: Add multi-page mapping API to VMM")
Signed-off-by: Vladislav Zaharov <vladazaharova2018@gmail.com>
---
v2: reserve room for each page as it is added to pt_pages, so that pages
left by earlier or failed calls are covered too. The Sashiko review
pointed out that v1 alone still lets a failing push free a parent page
while its children stay in pt_pages.

Tested on a GB203 with a temporary self-test that prepares eight
single-page mappings 1 TiB apart in a fresh Vmm and then executes them:
without this patch every execute_map() fails with -ENOMEM, with it they
all succeed. The MM self-tests still pass.

v1: https://lore.kernel.org/nova-gpu/20261113075422.943779-1-vladazaharova2018@gmail.com/

 drivers/gpu/nova-core/mm/pagetable.rs     | 24 ---------------
 drivers/gpu/nova-core/mm/pagetable/map.rs | 37 ++++++++++++++---------
 2 files changed, 22 insertions(+), 39 deletions(-)

diff --git a/drivers/gpu/nova-core/mm/pagetable.rs b/drivers/gpu/nova-core/mm/pagetable.rs
index 63a7e1855caa..03a132406a15 100644
--- a/drivers/gpu/nova-core/mm/pagetable.rs
+++ b/drivers/gpu/nova-core/mm/pagetable.rs
@@ -248,30 +248,6 @@ fn entry_size(level: PageTableLevel) -> usize {
             8 // 64-bit PDE/PTE
         }
     }
-
-    /// Compute upper bound on page table pages needed for `num_virt_pages`.
-    ///
-    /// Walks from PTE level up through PDE levels, accumulating the tree.
-    fn pt_pages_upper_bound(num_virt_pages: usize) -> usize {
-        let mut total = 0;
-
-        // PTE pages at the leaf level.
-        let pte_epp = Self::entries_per_page(Self::PTE_LEVEL);
-        let mut pages_at_level = num_virt_pages.div_ceil(pte_epp);
-        total += pages_at_level;
-
-        // Walk PDE levels bottom-up (reverse of PDE_LEVELS).
-        for &level in Self::PDE_LEVELS.iter().rev() {
-            let epp = Self::entries_per_page(level);
-
-            // How many pages at this level do we need to point to
-            // the previous pages_at_level?
-            pages_at_level = pages_at_level.div_ceil(epp);
-            total += pages_at_level;
-        }
-
-        total
-    }
 }
 
 /// Marker struct for MMU v2 (Turing/Ampere/Ada).
diff --git a/drivers/gpu/nova-core/mm/pagetable/map.rs b/drivers/gpu/nova-core/mm/pagetable/map.rs
index f7774b9cc972..2517fdcc5217 100644
--- a/drivers/gpu/nova-core/mm/pagetable/map.rs
+++ b/drivers/gpu/nova-core/mm/pagetable/map.rs
@@ -116,6 +116,7 @@ fn ensure_single_pte_path(
         &self,
         mm: &mut GpuMm<'_>,
         vfn: Vfn,
+        page_table_allocs: &mut KVec<Pin<KBox<AllocatedBlocks>>>,
         pt_pages: &mut RBTree<VramAddress, PreparedPtPage>,
     ) -> Result {
         let max_iter = 2 * M::PDE_LEVELS.len();
@@ -140,6 +141,10 @@ fn ensure_single_pte_path(
                     install_addr,
                     level,
                 } => {
+                    // Reserve so install_mappings() can use push_within_capacity (no alloc in
+                    // fence signalling critical path). It installs every page in `pt_pages`,
+                    // including pages left by earlier calls, so keep room for all of them.
+                    page_table_allocs.reserve(pt_pages.iter().count() + 1, GFP_KERNEL)?;
                     let page = Self::alloc_and_zero_page(mm, level)?;
                     let node = RBTreeNode::new(install_addr, page, GFP_KERNEL)?;
                     let old = pt_pages.insert(node);
@@ -163,8 +168,8 @@ fn ensure_single_pte_path(
 
     /// 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.
+    /// Walks the hierarchy per-VFN to prepare pages for all missing PDEs, and
+    /// reserves capacity in `page_table_allocs` for them.
     pub(super) fn prepare_map(
         &self,
         mm: &mut GpuMm<'_>,
@@ -173,16 +178,11 @@ pub(super) fn prepare_map(
         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)?;
+            self.ensure_single_pte_path(mm, vfn, page_table_allocs, pt_pages)?;
         }
         Ok(())
     }
@@ -207,10 +207,21 @@ pub(super) fn install_mappings(
             while let Some(c) = cursor {
                 let (next, node) = c.remove_current();
                 let (install_addr, page) = node.to_key_value();
-                let page_vram =
-                    VramAddress::from_raw(page.alloc.iter().next().ok_or(ENOMEM)?.offset());
+                let page_vram = page
+                    .alloc
+                    .iter()
+                    .next()
+                    .map(|b| VramAddress::from_raw(b.offset()));
+                let level = page.level;
+
+                // Keep the page before pointing a PDE at it, so that no failure below can free a
+                // page that a PDE points at. ensure_single_pte_path() reserved room for it.
+                page_table_allocs
+                    .push_within_capacity(page.alloc)
+                    .map_err(|_| ENOMEM)?;
+                let page_vram = page_vram.ok_or(ENOMEM)?;
 
-                if page.level == M::DUAL_PDE_LEVEL {
+                if level == M::DUAL_PDE_LEVEL {
                     let new_dpde = M::DualPde::new_small(Pfn::from(page_vram));
                     new_dpde.write(pramin, install_addr)?;
                 } else {
@@ -218,10 +229,6 @@ pub(super) fn install_mappings(
                     new_pde.write(pramin, install_addr)?;
                 }
 
-                page_table_allocs
-                    .push_within_capacity(page.alloc)
-                    .map_err(|_| ENOMEM)?;
-
                 cursor = next;
             }
 

base-commit: 0ff9e5b9ff6736c37593cd7db89f4593ec4f8e86
-- 
2.56.0