1 files changed, 274 insertions, 0 deletions
diff --git a/kernel/memory.c b/kernel/memory.c
new file mode 100644
index 0000000..cd81496
--- /dev/null
+++ b/kernel/memory.c
@@ -0,0 +1,274 @@
+#include <linux/slab.h>
+#include <linux/mutex.h>
+#include <linux/list.h>
+#include <linux/mm.h>
+#include <linux/dma-mapping.h>
+#include <linux/iommu.h>
+#include <linux/vmalloc.h>
+#include <linux/highmem.h>
+#include <linux/pfn.h>
+#include <linux/scatterlist.h>
+
+#include "nvidia-vgpu-vfio.h"
+#include "memory.h"
+
+static LIST_HEAD(simple_mappings);
+static DEFINE_MUTEX(simple_lock);
+
+int nv_vgpu_memory_init(struct nv_vgpu_device *vgpu)
+{
+    struct nv_vgpu_mm_ctx *ctx;
+    
+    if (vgpu->dry_run) {
+        nv_vgpu_info(vgpu, "DRY-RUN: Memory init validation (simple fallback)");
+        return 0;
+    }
+
+    ctx = kzalloc(sizeof(*ctx), GFP_KERNEL);
+    if (!ctx)
+        return -ENOMEM;
+
+    ctx->vgpu = vgpu;
+    mutex_init(&ctx->lock);
+    INIT_LIST_HEAD(&ctx->mappings);
+    ctx->mapping_count = 0;
+    ctx->iommu_domain = NULL;
+    ctx->use_simple = true;
+
+    vgpu->mm_ctx = ctx;
+    nv_vgpu_info(vgpu, "Memory init: using simple fallback (no IOMMU)");
+    return 0;
+}
+
+void nv_vgpu_memory_cleanup(struct nv_vgpu_device *vgpu)
+{
+    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
+    struct nv_vgpu_simple_mapping *mapping, *tmp;
+    
+    if (!ctx)
+        return;
+
+    if (ctx->use_simple) {
+        mutex_lock(&simple_lock);
+        list_for_each_entry_safe(mapping, tmp, &simple_mappings, list) {
+            if (mapping->mapped) {
+                dma_unmap_page(&vgpu->pdev->dev, mapping->iova, mapping->size, DMA_BIDIRECTIONAL);
+                mapping->mapped = false;
+            }
+            if (mapping->pinned) {
+                for (int i = 0; i < mapping->page_count; i++) {
+                    if (mapping->pages[i])
+                        put_page(mapping->pages[i]);
+                }
+                mapping->pinned = false;
+            }
+            list_del(&mapping->list);
+            vfree(mapping->pages);
+            kfree(mapping);
+        }
+        mutex_unlock(&simple_lock);
+    } else if (ctx->iommu_domain) {
+        iommu_detach_device(ctx->iommu_domain, &vgpu->pdev->dev);
+        iommu_domain_free(ctx->iommu_domain);
+    }
+
+    kfree(ctx);
+    vgpu->mm_ctx = NULL;
+}
+
+int nv_vgpu_memory_pin(struct nv_vgpu_device *vgpu, u64 gpa, u64 size)
+{
+    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
+    struct nv_vgpu_simple_mapping *mapping;
+    u64 start_pfn = gpa >> PAGE_SHIFT;
+    u64 end_pfn = (gpa + size + PAGE_SIZE - 1) >> PAGE_SHIFT;
+    int page_count = end_pfn - start_pfn;
+    int i, ret = 0;
+
+    if (!ctx)
+        return -ENODEV;
+
+    if (!ctx->use_simple)
+        return -ENODEV;
+
+    mapping = kzalloc(sizeof(*mapping), GFP_KERNEL);
+    if (!mapping)
+        return -ENOMEM;
+
+    mapping->gpa = gpa;
+    mapping->size = size;
+    mapping->page_count = page_count;
+    mapping->pages = vzalloc(page_count * sizeof(struct page *));
+    if (!mapping->pages) {
+        kfree(mapping);
+        return -ENOMEM;
+    }
+
+    for (i = 0; i < page_count; i++) {
+        u64 pfn = start_pfn + i;
+        struct page *page = pfn_to_page(pfn);
+        if (!page) {
+            ret = -EFAULT;
+            goto err_unlock;
+        }
+        get_page(page);
+        mapping->pages[i] = page;
+    }
+
+    mapping->pinned = true;
+    mutex_lock(&simple_lock);
+    list_add_tail(&mapping->list, &simple_mappings);
+    mutex_unlock(&simple_lock);
+
+    return 0;
+
+err_unlock:
+    for (i = 0; i < page_count; i++) {
+        if (mapping->pages[i])
+            put_page(mapping->pages[i]);
+    }
+    vfree(mapping->pages);
+    kfree(mapping);
+    return ret;
+}
+
+int nv_vgpu_memory_unpin(struct nv_vgpu_device *vgpu, u64 gpa, u64 size)
+{
+    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
+    struct nv_vgpu_simple_mapping *mapping;
+    
+    if (!ctx || !ctx->use_simple)
+        return -ENODEV;
+
+    mutex_lock(&simple_lock);
+    list_for_each_entry(mapping, &simple_mappings, list) {
+        if (mapping->gpa == gpa && mapping->size == size) {
+            if (mapping->mapped) {
+                mutex_unlock(&simple_lock);
+                return -EBUSY;
+            } else if (mapping->pinned) {
+                for (int i = 0; i < mapping->page_count; i++) {
+                    if (mapping->pages[i])
+                        put_page(mapping->pages[i]);
+                }
+                mapping->pinned = false;
+            }
+            list_del(&mapping->list);
+            vfree(mapping->pages);
+            kfree(mapping);
+            break;
+        }
+    }
+    mutex_unlock(&simple_lock);
+
+    return 0;
+}
+
+int nv_vgpu_memory_map(struct nv_vgpu_device *vgpu, u64 gpa, u64 size, u64 *iova)
+{
+    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
+    struct nv_vgpu_simple_mapping *mapping;
+    u64 map_iova;
+    int ret = 0;
+
+    if (!ctx || !ctx->use_simple)
+        return -ENODEV;
+
+    mutex_lock(&simple_lock);
+    list_for_each_entry(mapping, &simple_mappings, list) {
+        if (mapping->gpa == gpa && mapping->size == size)
+            break;
+    }
+    if (&mapping->list == &simple_mappings) {
+        mutex_unlock(&simple_lock);
+        return -ENOENT;
+    }
+
+    if (!mapping->pinned) {
+        mutex_unlock(&simple_lock);
+        return -EINVAL;
+    }
+
+    map_iova = dma_map_page(&vgpu->pdev->dev, mapping->pages[0], 0, size, DMA_BIDIRECTIONAL);
+    if (dma_mapping_error(&vgpu->pdev->dev, map_iova)) {
+        mutex_unlock(&simple_lock);
+        return -ENOMEM;
+    }
+
+    mapping->iova = map_iova;
+    mapping->mapped = true;
+    *iova = map_iova;
+    mutex_unlock(&simple_lock);
+
+    return ret;
+}
+
+int nv_vgpu_memory_unmap(struct nv_vgpu_device *vgpu, u64 iova, u64 size)
+{
+    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
+    struct nv_vgpu_simple_mapping *mapping;
+    int ret = 0;
+
+    if (!ctx || !ctx->use_simple)
+        return -ENODEV;
+
+    mutex_lock(&simple_lock);
+    list_for_each_entry(mapping, &simple_mappings, list) {
+        if (mapping->iova == iova && mapping->size == size) {
+            if (!mapping->mapped) {
+                mutex_unlock(&simple_lock);
+                return -ENOENT;
+            }
+            dma_unmap_page(&vgpu->pdev->dev, iova, size, DMA_BIDIRECTIONAL);
+            mapping->mapped = false;
+            mapping->iova = 0;
+            break;
+        }
+    }
+    mutex_unlock(&simple_lock);
+
+    return ret;
+}
+
+static vm_fault_t nv_vgpu_vm_fault(struct vm_fault *vmf)
+{
+    struct nv_vgpu_device *vgpu = vmf->vma->vm_private_data;
+    struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
+    struct nv_vgpu_simple_mapping *mapping;
+    u64 offset = vmf->address - vmf->vma->vm_start + vmf->vma->vm_pgoff * PAGE_SIZE;
+    u64 pfn;
+    vm_fault_t ret = VM_FAULT_SIGBUS;
+
+    if (!ctx)
+        return VM_FAULT_SIGBUS;
+
+    if (!ctx->use_simple)
+        return VM_FAULT_SIGBUS;
+
+    mutex_lock(&simple_lock);
+    list_for_each_entry(mapping, &simple_mappings, list) {
+        if (mapping->mapped && offset >= mapping->iova &&
+            offset < mapping->iova + mapping->size) {
+            pfn = (mapping->gpa + offset - mapping->iova) >> PAGE_SHIFT;
+            vmf->page = pfn_to_page(pfn);
+            get_page(vmf->page);
+            ret = VM_FAULT_NOPAGE;
+            break;
+        }
+    }
+    mutex_unlock(&simple_lock);
+
+    return ret;
+}
+
+static const struct vm_operations_struct nv_vgpu_vm_ops = {
+    .fault = nv_vgpu_vm_fault,
+};
+
+int nv_vgpu_memory_mmap(struct nv_vgpu_device *vgpu, struct vm_area_struct *vma)
+{
+    vma->vm_ops = &nv_vgpu_vm_ops;
+    vma->vm_private_data = vgpu;
+    vm_flags_set(vma, VM_DONTEXPAND | VM_DONTDUMP);
+    return 0;
+}