[PATCH 2/4] platform/nvidia: Implement mmap and memory scrubbing for EGM chardev

[Ankit Agrawal]

Setup the file operations (open, close and mmap) for the EGM char device.

Memory scrubbing:
- The EGM region is invisible to the host Linux kernel and is not
  managed by the page allocator. So the driver is responsible for
  zeroing it before handing it to a VM.
- Clear the entire EGM region on the first open() call using memremap()
  and memset() in 1 GiB chunks with cond_resched() between iterations
  to avoid RCU stalls on large regions.
- The region is handed to a single VM at a time and must be scrubbed
  before each handover. Serialise openers with a mutex guarding open_count
  and refuse a second opener with -EBUSY while the region is still held.

mmap:
- Implement nvgrace_egm_mmap() using remap_pfn_range(). The EGM region
  is not part of the host EFI map and has no struct pages making
  remap_pfn_range() the correct interface for mapping it directly into
  QEMU's virtual address space.
- Validate that the requested range fits within the EGM region. Reject a
  vm_pgoff that already lies outside the region before it is shifted into
  a byte count so that the page offset cannot overflow and slip past the
  range check.

Suggested-by: Aniket Agashe <aniketa@xxxxxxxxxx>
Suggested-by: Vikram Sethi <vsethi@xxxxxxxxxx>
Assisted-by: Claude:claude-opus-4.8
Signed-off-by: Ankit Agrawal <ankita@xxxxxxxxxx>
---
 drivers/platform/nvidia/egm.c | 125 ++++++++++++++++++++++++++++++++--
 1 file changed, 121 insertions(+), 4 deletions(-)

diff --git a/drivers/platform/nvidia/egm.c b/drivers/platform/nvidia/egm.c
index a340c4fcbc7a..09ab00213de2 100644
--- a/drivers/platform/nvidia/egm.c
+++ b/drivers/platform/nvidia/egm.c
@@ -3,6 +3,9 @@
  * Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved
  */
 
+#include <linux/mutex.h>
+#include <linux/sched/signal.h>
+#include <linux/sizes.h>
 #include <linux/vfio_pci_core.h>
 
 #define NVGRACE_EGM_DEV_NAME "egm"
@@ -19,6 +22,10 @@ struct gpu_node {
 struct nvgrace_egm_dev {
 	struct device device;
 	struct cdev cdev;
+	/* serialises the first-open scrub */
+	struct mutex open_lock;
+	/* protected by open_lock */
+	unsigned int open_count;
 	phys_addr_t egmphys;
 	size_t egmlength;
 	u64 egmpxm;
@@ -39,21 +46,129 @@ static LIST_HEAD(egm_chardevs);
 
 static int nvgrace_egm_open(struct inode *inode, struct file *file)
 {
-	return 0;
+	struct nvgrace_egm_dev *egm_dev =
+		container_of(inode->i_cdev, struct nvgrace_egm_dev, cdev);
+	phys_addr_t phys;
+	size_t remaining, chunk_size;
+	void *chunk_addr;
+	int ret;
+
+	file->private_data = egm_dev;
+
+	/*
+	 * The EGM region is a physical carveout handed to a single VM at a
+	 * time and must be scrubbed before each handover. Take open_lock
+	 * killably around the scrub which can run for long time, so that
+	 * the waiters stay killable.
+	 */
+	if (mutex_lock_killable(&egm_dev->open_lock))
+		return -EINTR;
+
+	/*
+	 * Refuse a second opener while the region is still held. release()
+	 * runs only when the last reference to the struct file drops but an
+	 * mmap keeps that reference (and hence open_count) alive past
+	 * close for the lifetime of the VMA. So open_count returns to zero
+	 * only when every mapping is gone. A concurrent opener cannot be
+	 * handed the region: it cannot be re-scrubbed without destroying the
+	 * current consumer's live data.
+	 */
+	if (egm_dev->open_count) {
+		ret = -EBUSY;
+		goto unlock;
+	}
+
+	/*
+	 * nvgrace-egm module is responsible to manage the EGM memory as
+	 * the host kernel has no knowledge of it. Clear the region before
+	 * handing over to userspace.
+	 *
+	 * The EGM region can be very large (hundreds of GiB). So Map and zero
+	 * one chunk at a time rather than mapping the whole region at once.
+	 */
+	phys = egm_dev->egmphys;
+	remaining = egm_dev->egmlength;
+
+	while (remaining > 0) {
+		/*
+		 * The scrub holds the lock for a long time; let it be
+		 * SIGKILL'd.
+		 */
+		if (fatal_signal_pending(current)) {
+			ret = -EINTR;
+			goto unlock;
+		}
+
+		chunk_size = min(remaining, SZ_1G);
+
+		chunk_addr = memremap(phys, chunk_size, MEMREMAP_WB);
+		if (!chunk_addr) {
+			ret = -ENOMEM;
+			goto unlock;
+		}
+
+		memset(chunk_addr, 0, chunk_size);
+		memunmap(chunk_addr);
+		cond_resched();
+
+		phys += chunk_size;
+		remaining -= chunk_size;
+	}
+
+	/*
+	 * Mark the device open only after the scrub completes so that a
+	 * concurrent opener cannot observe a non-zero count and proceed
+	 * before the region has been cleared.
+	 */
+	egm_dev->open_count = 1;
+	ret = 0;
+
+unlock:
+	mutex_unlock(&egm_dev->open_lock);
+	return ret;
 }
 
 static int nvgrace_egm_release(struct inode *inode, struct file *file)
 {
+	struct nvgrace_egm_dev *egm_dev =
+		container_of(inode->i_cdev, struct nvgrace_egm_dev, cdev);
+
+	guard(mutex)(&egm_dev->open_lock);
+
+	if (!--egm_dev->open_count)
+		file->private_data = NULL;
+
 	return 0;
 }
 
 static int nvgrace_egm_mmap(struct file *file, struct vm_area_struct *vma)
 {
+	struct nvgrace_egm_dev *egm_dev = file->private_data;
+	u64 req_len, pgoff, end;
+	unsigned long start_pfn, num_pages;
+
+	pgoff = vma->vm_pgoff;
+	num_pages = egm_dev->egmlength >> PAGE_SHIFT;
+
+	/* Reject a page offset that already lies outside the EGM region. */
+	if (pgoff >= num_pages)
+		return -EINVAL;
+
+	if (check_sub_overflow(vma->vm_end, vma->vm_start, &req_len) ||
+	    check_add_overflow(PHYS_PFN(egm_dev->egmphys), pgoff, &start_pfn) ||
+	    check_add_overflow(PFN_PHYS(pgoff), req_len, &end))
+		return -EOVERFLOW;
+
+	if (end > egm_dev->egmlength)
+		return -EINVAL;
+
 	/*
-	 * Mapping the EGM region into userspace is implemented by a later
-	 * patch. Until then refuse the mmap.
+	 * EGM memory is invisible to the host kernel and is not managed
+	 * by it. Map the usermode VMA to the EGM region.
 	 */
-	return -EOPNOTSUPP;
+	return remap_pfn_range(vma, vma->vm_start,
+			       start_pfn, req_len,
+			       vma->vm_page_prot);
 }
 
 static const struct file_operations file_ops = {
@@ -122,6 +237,7 @@ static void egm_chardev_release(struct device *dev)
 	struct nvgrace_egm_dev *egm_chardev = container_of(dev, struct nvgrace_egm_dev, device);
 
 	remove_gpus(egm_chardev);
+	mutex_destroy(&egm_chardev->open_lock);
 	kfree(egm_chardev);
 }
 
@@ -155,6 +271,7 @@ static struct nvgrace_egm_dev *setup_egm_chardev(u64 egmphys, u64 egmlength,
 	egm_chardev->egmphys = egmphys;
 	egm_chardev->egmlength = egmlength;
 	egm_chardev->egmpxm = egmpxm;
+	mutex_init(&egm_chardev->open_lock);
 	INIT_LIST_HEAD(&egm_chardev->gpus);
 
 	egm_chardev->device.devt = MKDEV(MAJOR(dev), egm_chardev->egmpxm);
-- 
2.34.1