1 files changed, 335 insertions, 0 deletions

diff --git a/kernel/memory.c.iommu b/kernel/memory.c.iommu
new file mode 100644
index 0000000..335446b
--- /dev/null
+++ b/kernel/memory.c.iommu
@@ -0,0 +1,335 @@
+#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"
+
+int nv_vgpu_memory_init(struct nv_vgpu_device *vgpu)
+{
+	struct nv_vgpu_mm_ctx *ctx;
+	int ret;
+
+	if (vgpu->dry_run) {
+		nv_vgpu_info(vgpu, "DRY-RUN: Memory init validation");
+		/* Check IOMMU support */
+		if (iommu_present(&pci_bus_type)) {
+			nv_vgpu_info(vgpu, "DRY-RUN: IOMMU present");
+			ctx = kzalloc(sizeof(*ctx), GFP_KERNEL);
+			if (!ctx)
+				return -ENOMEM;
+			ctx->iommu_domain = iommu_domain_alloc(&pci_bus_type);
+			if (ctx->iommu_domain) {
+				nv_vgpu_info(vgpu, "DRY-RUN: IOMMU domain allocatable");
+				iommu_domain_free(ctx->iommu_domain);
+			} else {
+				nv_vgpu_warn(vgpu, "DRY-RUN: IOMMU domain allocation failed");
+			}
+			kfree(ctx);
+		} else {
+			nv_vgpu_warn(vgpu, "DRY-RUN: No IOMMU present");
+		}
+		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 = iommu_domain_alloc(&pci_bus_type);
+	if (!ctx->iommu_domain) {
+		ret = -ENODEV;
+		goto err_free;
+	}
+
+	ret = iommu_attach_device(ctx->iommu_domain, &vgpu->pdev->dev);
+	if (ret) {
+		nv_vgpu_warn(vgpu, "Failed to attach IOMMU domain: %d", ret);
+	}
+
+	vgpu->mm_ctx = ctx;
+	return 0;
+
+err_free:
+	kfree(ctx);
+	return ret;
+}
+
+void nv_vgpu_memory_cleanup(struct nv_vgpu_device *vgpu)
+{
+	struct nv_vgpu_mm_ctx *ctx = vgpu->mm_ctx;
+	struct nv_vgpu_mapping *mapping, *tmp;
+
+	if (!ctx)
+		return;
+
+	if (vgpu->dry_run) {
+		nv_vgpu_info(vgpu, "DRY-RUN: Memory cleanup (no hardware resources to release)");
+		if (ctx->iommu_domain)
+			iommu_domain_free(ctx->iommu_domain);
+		kfree(ctx);
+		vgpu->mm_ctx = NULL;
+		return;
+	}
+
+	mutex_lock(&ctx->lock);
+	list_for_each_entry_safe(mapping, tmp, &ctx->mappings, list) {
+		if (mapping->mapped) {
+			iommu_unmap(ctx->iommu_domain, mapping->iova, mapping->size);
+			mapping->mapped = false;
+		}
+		if (mapping->pinned) {
+			int i;
+			for (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);
+		ctx->mapping_count--;
+	}
+	mutex_unlock(&ctx->lock);
+
+	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;
+}
+
+static struct nv_vgpu_mapping *nv_vgpu_find_mapping(struct nv_vgpu_mm_ctx *ctx, u64 iova, u64 size)
+{
+	struct nv_vgpu_mapping *mapping;
+
+	list_for_each_entry(mapping, &ctx->mappings, list) {
+		if (mapping->iova == iova && mapping->size == size)
+			return mapping;
+	}
+	return 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_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;
+
+	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;
+	}
+
+	mutex_lock(&ctx->lock);
+
+	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;
+	list_add_tail(&mapping->list, &ctx->mappings);
+	ctx->mapping_count++;
+
+	mutex_unlock(&ctx->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);
+	mutex_unlock(&ctx->lock);
+	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_mapping *mapping;
+	u64 iova = 0;
+	int ret = 0;
+
+	if (!ctx)
+		return -ENODEV;
+
+	mutex_lock(&ctx->lock);
+	list_for_each_entry(mapping, &ctx->mappings, list) {
+		if (mapping->gpa == gpa && mapping->size == size) {
+			if (mapping->mapped) {
+				iova = mapping->iova;
+				ret = -EBUSY;
+			} else if (mapping->pinned) {
+				int i;
+				for (i = 0; i < mapping->page_count; i++) {
+					if (mapping->pages[i])
+						put_page(mapping->pages[i]);
+				}
+				mapping->pinned = false;
+			}
+			list_del(&mapping->list);
+			ctx->mapping_count--;
+			vfree(mapping->pages);
+			kfree(mapping);
+			break;
+		}
+	}
+	mutex_unlock(&ctx->lock);
+
+	return ret;
+}
+
+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_mapping *mapping;
+	u64 map_iova;
+	int ret;
+
+	if (!ctx || !ctx->iommu_domain)
+		return -ENODEV;
+
+	mutex_lock(&ctx->lock);
+	mapping = nv_vgpu_find_mapping(ctx, 0, 0);
+	if (!mapping || mapping->gpa != gpa || mapping->size != size) {
+		list_for_each_entry(mapping, &ctx->mappings, list) {
+			if (mapping->gpa == gpa && mapping->size == size)
+				break;
+		}
+		if (&mapping->list == &ctx->mappings) {
+			mutex_unlock(&ctx->lock);
+			return -ENOENT;
+		}
+	}
+
+	if (!mapping->pinned) {
+		mutex_unlock(&ctx->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(&ctx->lock);
+		return -ENOMEM;
+	}
+
+	mapping->iova = map_iova;
+	mapping->mapped = true;
+	*iova = map_iova;
+
+	ret = iommu_map(ctx->iommu_domain, map_iova, gpa, size,
+			IOMMU_READ | IOMMU_WRITE | IOMMU_CACHE, GFP_KERNEL);
+	if (ret) {
+		dma_unmap_page(&vgpu->pdev->dev, map_iova, size, DMA_BIDIRECTIONAL);
+		mapping->mapped = false;
+		mapping->iova = 0;
+	}
+
+	mutex_unlock(&ctx->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_mapping *mapping;
+	int ret = 0;
+
+	if (!ctx || !ctx->iommu_domain)
+		return -ENODEV;
+
+	mutex_lock(&ctx->lock);
+	mapping = nv_vgpu_find_mapping(ctx, iova, size);
+	if (!mapping || !mapping->mapped) {
+		mutex_unlock(&ctx->lock);
+		return -ENOENT;
+	}
+
+	ret = iommu_unmap(ctx->iommu_domain, iova, size);
+	if (ret == 0) {
+		dma_unmap_page(&vgpu->pdev->dev, iova, size, DMA_BIDIRECTIONAL);
+		mapping->mapped = false;
+		mapping->iova = 0;
+	}
+
+	mutex_unlock(&ctx->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_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;
+
+	mutex_lock(&ctx->lock);
+	list_for_each_entry(mapping, &ctx->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(&ctx->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;
+}
\ No newline at end of file