1 files changed, 411 insertions, 3 deletions

diff --git a/mmu.c b/mmu.c
index ababef5..e2b9a91 100644
--- a/mmu.c
+++ b/mmu.c
@@ -1,9 +1,13 @@
 /* Copyright 2024 Joshua Bakita
  * Helpers to deal with NVIDIA's MMU and associated page tables
  */
+#include <linux/dma-mapping.h>  // dma_map_page() and dma_unmap_page()
 #include <linux/err.h>  // ERR_PTR() etc.
+#include <linux/gfp.h>  // alloc_pages()
 #include <linux/iommu.h>  // iommu_get_domain_for_dev() and iommu_iova_to_phys()
 #include <linux/kernel.h>  // Kernel types
+#include <linux/list.h>  // struct list_head and associated functions
+#include <linux/mm.h>  // put_page()
 
 #include "nvdebug.h"
 
@@ -15,6 +19,11 @@ int g_verbose = 0;
 #define printk_debug if (g_verbose >= 2) printk
 #define printk_info  if (g_verbose >= 1) printk
 
+// At least map_page_directory() assumes that pages are 4 KiB
+#if PAGE_SIZE != 4096
+#error nvdebug assumes and requires a 4 KiB page size.
+#endif
+
 /* Convert a page directory (PD) pointer and aperture to be kernel-accessible
 
   I/O MMU handling inspired by amdgpu_iomem_read() in amdgpu_ttm.c of the
@@ -22,7 +31,8 @@ int g_verbose = 0;
 
   @param addr  Pointer from page directory entry (PDE)
   @param pd_ap PD-type aperture (target address space) for `addr`
-  @return A dereferencable kernel address, or an ERR_PTR-wrapped error
+  @return A dereferencable kernel address, 0 if an I/O MMU is in use and has
+          no available mapping for the bus address, or an ERR_PTR-wrapped error
  */
 static void __iomem *pd_deref(struct nvdebug_state *g, uintptr_t addr,
                               enum PD_TARGET pd_ap) {
@@ -56,7 +66,7 @@ static void __iomem *pd_deref(struct nvdebug_state *g, uintptr_t addr,
         // Check for, and translate through, the I/O MMU (if any)
         if ((dom = iommu_get_domain_for_dev(g->dev))) {
                 phys = iommu_iova_to_phys(dom, addr);
-                printk_debug(KERN_DEBUG "[nvdebug] I/O MMU translated SYS_MEM I/O VA %#lx to physical address %#llx.\n", addr, phys);
+                printk_debug(KERN_DEBUG "[nvdebug] %s: I/O MMU translated SYS_MEM I/O VA %#lx to physical address %#llx.\n", __func__, addr, phys);
         } else
                 phys = addr;
 
@@ -143,6 +153,327 @@ uint64_t search_page_directory(struct nvdebug_state *g,
         return 0;
 }
 
+/* GPU Virtual address -> Physical address ("forward" translation) for V2 tables
+  Index the page directories and tables used by the GPU MMU to determine which
+  physical address a given GPU virtual address has been mapped to.
+
+  The page directory and tables may be located in VID_MEM, SYS_MEM, or spread
+  across multiple apertures.
+
+  @param pd_config      Page Directory configuration, containing pointer and
+                        aperture for the start of the PDE3 entries
+  @param addr_to_find   Virtual address to translate to a physical address
+  @param found_addr     Where to store found physical address (0 if unfound)
+  @param found_aperture Where to store aperture of found physical address
+  @return 0 on success, -ENXIO if not found, and -errno on error.
+*/
+int translate_page_directory(struct nvdebug_state *g,
+                             page_dir_config_t pd_config,
+                             uint64_t addr_to_find,
+                             uint64_t *found_addr /* out */,
+                             enum INST_TARGET *found_aperture /* out */) {
+        page_dir_entry_t entry;
+        void __iomem *next_kva;
+        unsigned int level, pde_idx;
+        uintptr_t next = (uintptr_t)pd_config.page_dir << 12;
+        enum PD_TARGET next_target = INST2PD_TARGET(pd_config.target);
+
+        *found_addr = 0;
+        *found_aperture = TARGET_INVALID;
+
+        // Make sure that the query is page-aligned (likely mistake otherwise)
+        if (addr_to_find & 0xfff) {
+                printk(KERN_WARNING "[nvdebug] Attempting to translate unaligned address %#llx in translate_page_directory()!\n", addr_to_find);
+                return -EINVAL;
+        }
+
+        printk_info(KERN_INFO "[nvdebug] Translating addr %#018llx in V2 page table with base %#018llx\n", (u64)addr_to_find, (u64)next);
+
+        // Step through each PDE level and the PTE level
+        for (level = 0; level < 5; level++) {
+                // Index into this level
+                pde_idx = (addr_to_find >> NV_MMU_PT_V2_LSB[level]) & (NV_MMU_PT_V2_SZ[level] - 1);
+                printk_debug(KERN_DEBUG "[nvdebug] Using index %u in lvl %d\n", pde_idx, level);
+                // Hack to workaround PDE0 being double-size and strangely formatted
+                if (NV_MMU_PT_V2_ENTRY_SZ[level] == 16)
+                        next += 8;
+                // Obtain a kernel-dereferencable address
+                next_kva = pd_deref(g, next, next_target);
+                if (IS_ERR_OR_NULL(next_kva)) {
+                        printk(KERN_ERR "[nvdebug] %s: Unable to resolve %#lx in GPU %s to a kernel-accessible address. Error %ld.\n", __func__, next, pd_target_to_text(next_target), PTR_ERR(next_kva));
+                        return PTR_ERR(next_kva);
+                }
+                // Obtain entry at this level
+                entry.raw_w = readq(next_kva + NV_MMU_PT_V2_ENTRY_SZ[level] * pde_idx);
+                if (entry.target == PD_AND_TARGET_INVALID)
+                        return -ENXIO;
+                printk_debug(KERN_DEBUG "[nvdebug] Found %s pointing to %#018llx in ap '%s' at lvl %d (raw: %#018llx)\n", entry.is_pte ? "PTE" : "PDE", ((u64)entry.addr) << 12, pd_target_to_text(entry.target), level, entry.raw_w);
+                // Just return the physical address if this is the PTE level
+                if (entry.is_pte) { // level == 4 for 4 KiB pages, == 3 for 2 MiB
+                        *found_addr = ((uint64_t)entry.addr) << 12;
+                        *found_aperture = entry.aperture;
+                        return 0;
+                }
+                // Otherwise step to the next table level
+                // TODO: Use addr_w as appropriate
+                next = (uint64_t)entry.addr << 12;
+                next_target = entry.target;
+        }
+
+        return 0;
+}
+
+// This struct is very special. We will never directly allocate this struct;
+// its sole purpose is to provide more intuitive names to the offsets at which
+// we store data in Linux's struct page. Such (ab)use of struct page is
+// explictly permitted (see linux/mm_types.h). This struct is thus used by
+// casting a pointer of struct page to a pointer of struct nvdebug_pd_page,
+// then accessing the associated fields. This pointer may also be freely cast
+// back to a sturct page pointer.
+// We have 24 (32-bit) or 44 (64-bit) bytes available in the page struct
+// (according to the documentation on struct page). Our comments indicate what
+// available parts of struct page we repurpose for our own needs.
+struct nvdebug_pd_page {
+        unsigned long __flags; // From struct page; do not touch!
+        // Overlaps struct page.lru
+        struct list_head list; // 4/8 bytes
+        // Overlaps struct page.mapping (and page.share on 32-bit)
+        uintptr_t parent_addr; // 8 bytes
+        // Overlaps struct page.share (page.private on 32-bit)
+        enum PD_TARGET parent_aperture; // 4 bytes
+        // Overlaps page.private (page.page_type on 32-bit)
+        dma_addr_t dma_addr; // 4/8 bytes
+};
+
+/* Collect and free any now-unused page directory/table allocations
+
+  @param force Deallocate all page directories/tables created by this module,
+               no matter if they appear to be in-use or not.
+  @returns Number of freed pages on success, -errno on error.
+*/
+int gc_page_directory(struct nvdebug_state *g, bool force) {
+        struct nvdebug_pd_page  *page, *_page;
+        void __iomem *parent_kva;
+        page_dir_entry_t parent_entry;
+        int freed_pages = 0;
+
+        // Depth-first traversal (from perspective of each page table) of page
+        // allocations.
+        // (This is depth-first because map_page_directory() always allocates and
+        // pushes page directory allocations before page table allocations.)
+        list_for_each_entry_safe_reverse(page, _page, &g->pd_allocs, list) {
+                printk_debug(KERN_DEBUG "[nvdebug] %s: Checking if page directory/table at %llx (SYS_MEM_?) with parent at %lx (%s) is unused...\n", __func__, page->dma_addr, page->parent_addr, pd_target_to_text(page->parent_aperture));
+                // Try to determine if we're still in-use. We consider ourselves
+                // potentially in-use if our parent still points to us.
+                parent_kva = pd_deref(g, page->parent_addr, page->parent_aperture);
+                if (IS_ERR(parent_kva)) {
+                        printk(KERN_ERR "[nvdebug] %s: Error resolving %#lx in GPU %s to a kernel-accessible address. Error %ld.\n", __func__, page->parent_addr, pd_target_to_text(page->parent_aperture), PTR_ERR(parent_kva));
+                        return -ENOTRECOVERABLE;
+                }
+                // A NULL kva indicates parent no longer exists
+                parent_entry.raw_w = parent_kva ? readq(parent_kva) : 0;
+                // Page directory/table still in-use; do not free unless forced
+                if (parent_entry.addr_w == (page->dma_addr >> 12) && !force)
+                        continue;
+                // Free this page table/directory and delete our parent's pointer to us
+                if (parent_entry.addr_w == (page->dma_addr >> 12)) {
+                        printk(KERN_WARNING "[nvdebug] %s: Deleting page table/directory at %llx (SYS_MEM_?) with parent at %lx (%s) that may still be in-use!\n", __func__, page->dma_addr, page->parent_addr, pd_target_to_text(page->parent_aperture));
+                        writeq(0, parent_kva);
+                }
+                // Unmap, zero, free, and remove from tracking (these all return void)
+                dma_unmap_page(g->dev, page->dma_addr, PAGE_SIZE, DMA_TO_DEVICE);
+                memset(page_to_virt((struct page*)page), 0, PAGE_SIZE);
+                // Necessary to reset mapcount as we (ab)use its state for other things
+                page_mapcount_reset((struct page*)page);
+                // Same reset needed for mapping
+                ((struct page*)page)->mapping = NULL;
+                // Remove this page from our list of allocated pages
+                list_del(&page->list);
+                // Free the page
+                put_page((struct page*)page);
+                freed_pages++;
+        }
+        printk_debug(KERN_DEBUG "[nvdebug] %s: Freed %d pages.", __func__, freed_pages);
+        return freed_pages;
+}
+
+/* Map a GPU virtual address to a physical address in a GPU page table
+  Search for a mapping for specified GPU virtual address, and create a new one
+  if none is found. Automatically creates page directories and page table
+  entries as necessary.
+
+  The page directory and tables may be located in VID_MEM, SYS_MEM, or spread
+  across multiple apertures.
+
+  @param pd_config      Page Directory configuration, containing pointer and
+                        aperture for the start of the PDE3 entries
+  @param vaddr_to_find  Virtual address to check, and map to a physical address
+                        if nothing is already mapped (up to 49 bits long)
+  @param paddr_to_map   Physical address to use (up to 36 bits long if VID_MEM,
+                        and up to 58 bits if SYS_MEM)
+  @param paddr_target   Which space does the physical address refer to?
+  @param huge_page      Set to map a 2 MiB, rather than 4 KiB, page
+  @return 0 on success, 1 if mapping already exists, -EADDRINUSE if virtual
+          address is already mapped to something else, and -errno on error
+*/
+int map_page_directory(struct nvdebug_state *g,
+                       page_dir_config_t pd_config,