ANDROID: gunyah: vm_mgr: Bindings can now be CMA backed or GUP backed

Make the bindings generic so that the APIs can now handle the binding's
memory backed either by CMA memory or GUP.

Bug: 399219478
Change-Id: I889fe1fbe5664c7c40938db78f350d62bfe33431
Signed-off-by: Prakruthi Deepak Heragu <quic_pheragu@quicinc.com>

drivers/virt/gunyah/gunyah_vcpu.c

@@ -97,7 +97,7 @@ static bool gunyah_handle_page_fault(
 	bool write = !!vcpu_run_resp->state_data[1];
 	int ret = 0;
 
-	ret = gunyah_gup_demand_page(vcpu->ghvm, addr, write);
+	ret = gunyah_demand_page(vcpu->ghvm, addr, write);
 	if (!ret || ret == -EAGAIN)
 		return true;
 
@@ -120,7 +120,7 @@ gunyah_handle_mmio(struct gunyah_vcpu *vcpu, unsigned long resume_data[3],
 	if (WARN_ON(len > sizeof(u64)))
 		len = sizeof(u64);
 
-	ret = gunyah_gup_demand_page(vcpu->ghvm, addr,
+	ret = gunyah_demand_page(vcpu->ghvm, addr,
 					vcpu->vcpu_run->mmio.is_write);
 	if (!ret || ret == -EAGAIN) {
 		resume_data[1] = GUNYAH_ADDRSPACE_VMMIO_ACTION_RETRY;

drivers/virt/gunyah/vm_mgr.c

@@ -625,47 +625,6 @@ static inline int gunyah_vm_fill_boot_context(struct gunyah_vm *ghvm)
 	return 0;
 }
 
-int gunyah_gup_setup_demand_paging(struct gunyah_vm *ghvm)
-{
-	struct gunyah_rm_mem_entry *entries;
-	struct gunyah_vm_gup_binding *b;
-	unsigned long index = 0;
-	u32 count = 0, i;
-	int ret = 0;
-
-	down_read(&ghvm->bindings_lock);
-	mt_for_each(&ghvm->bindings, b, index, ULONG_MAX)
-		if (b->share_type == VM_MEM_LEND &&
-			(b->guest_phys_addr != ghvm->fw.config.guest_phys_addr))
-			count++;
-
-	if (!count)
-		goto out;
-
-	entries = kcalloc(count, sizeof(*entries), GFP_KERNEL);
-	if (!entries) {
-		ret = -ENOMEM;
-		goto out;
-	}
-
-	index = i = 0;
-	mt_for_each(&ghvm->bindings, b, index, ULONG_MAX) {
-		if (b->share_type != VM_MEM_LEND ||
-			(b->guest_phys_addr == ghvm->fw.config.guest_phys_addr))
-			continue;
-		entries[i].phys_addr = cpu_to_le64(b->guest_phys_addr);
-		entries[i].size = cpu_to_le64(b->size);
-		if (++i == count)
-			break;
-	}
-
-	ret = gunyah_rm_vm_set_demand_paging(ghvm->rm, ghvm->vmid, i, entries);
-	kfree(entries);
-out:
-	up_read(&ghvm->bindings_lock);
-	return ret;
-}
-
 static int gunyah_vm_start(struct gunyah_vm *ghvm)
 {
 	struct gunyah_rm_hyp_resources *resources;
@@ -691,11 +650,11 @@ static int gunyah_vm_start(struct gunyah_vm *ghvm)
 	ghvm->vmid = ret;
 	ghvm->vm_status = GUNYAH_RM_VM_STATUS_LOAD;
 
-	ghvm->dtb.parcel_start = ghvm->dtb.config.guest_phys_addr >> PAGE_SHIFT;
-	ghvm->dtb.parcel_pages = ghvm->dtb.config.size >> PAGE_SHIFT;
-	ret = gunyah_gup_share_parcel(ghvm, &ghvm->dtb.parcel,
-					&ghvm->dtb.parcel_start,
-					&ghvm->dtb.parcel_pages);
+	ghvm->dtb.parcel.start = ghvm->dtb.config.guest_phys_addr >> PAGE_SHIFT;
+	ghvm->dtb.parcel.pages = ghvm->dtb.config.size >> PAGE_SHIFT;
+	ret = gunyah_share_parcel(ghvm, &ghvm->dtb.parcel,
+					&ghvm->dtb.parcel.start,
+					&ghvm->dtb.parcel.pages);
 	if (ret) {
 		dev_warn(ghvm->parent,
 			 "Failed to allocate parcel for DTB: %d\n", ret);
@@ -703,11 +662,11 @@ static int gunyah_vm_start(struct gunyah_vm *ghvm)
 	}
 
 	if (ghvm->auth == GUNYAH_RM_VM_AUTH_QCOM_ANDROID_PVM) {
-		ghvm->fw.parcel_start = ghvm->fw.config.guest_phys_addr >> PAGE_SHIFT;
-		ghvm->fw.parcel_pages = ghvm->fw.config.size >> PAGE_SHIFT;
-		ret = gunyah_gup_share_parcel(ghvm, &ghvm->fw.parcel,
-				&ghvm->fw.parcel_start,
-				&ghvm->fw.parcel_pages);
+		ghvm->fw.parcel.start = ghvm->fw.config.guest_phys_addr >> PAGE_SHIFT;
+		ghvm->fw.parcel.pages = ghvm->fw.config.size >> PAGE_SHIFT;
+		ret = gunyah_share_parcel(ghvm, &ghvm->fw.parcel,
+				&ghvm->fw.parcel.start,
+				&ghvm->fw.parcel.pages);
 		if (ret) {
 			dev_warn(ghvm->parent,
 					"Failed to allocate parcel for FW: %d\n", ret);
@@ -715,9 +674,9 @@ static int gunyah_vm_start(struct gunyah_vm *ghvm)
 		}
 	}
 	ret = gunyah_rm_vm_configure(ghvm->rm, ghvm->vmid, ghvm->auth,
-				     ghvm->dtb.parcel.mem_handle, 0, 0,
+				     ghvm->dtb.parcel.parcel.mem_handle, 0, 0,
 				     ghvm->dtb.config.guest_phys_addr -
-					     (ghvm->dtb.parcel_start
+					     (ghvm->dtb.parcel.start
 					      << PAGE_SHIFT),
 				     ghvm->dtb.config.size);
 	if (ret) {
@@ -726,26 +685,26 @@ static int gunyah_vm_start(struct gunyah_vm *ghvm)
 	}
 
 	if (ghvm->auth == GUNYAH_RM_VM_AUTH_QCOM_ANDROID_PVM) {
-		ret = gunyah_rm_vm_set_firmware_mem(ghvm->rm, ghvm->vmid, &ghvm->fw.parcel,
-				ghvm->fw.config.guest_phys_addr - (ghvm->fw.parcel_start << PAGE_SHIFT),
-				ghvm->fw.config.size);
+		ret = gunyah_rm_vm_set_firmware_mem(ghvm->rm, ghvm->vmid, &ghvm->fw.parcel.parcel,
+			ghvm->fw.config.guest_phys_addr - (ghvm->fw.parcel.start << PAGE_SHIFT),
+			ghvm->fw.config.size);
 		if (ret) {
 			pr_warn("Failed to configure pVM firmware\n");
 			goto err;
 		}
 	}
 
-	ret = gunyah_gup_setup_demand_paging(ghvm);
+	ret = gunyah_setup_demand_paging(ghvm, 0, ULONG_MAX);
 	if (ret) {
 		dev_warn(ghvm->parent,
-			 "Failed to set up gmem demand paging: %d\n", ret);
-		goto err;
+			"Failed to set up demand paging: %d\n", ret);
+		return ret;
 	}
 
 	ret = gunyah_rm_vm_set_address_layout(
 		ghvm->rm, ghvm->vmid, GUNYAH_RM_RANGE_ID_IMAGE,
-		ghvm->dtb.parcel_start << PAGE_SHIFT,
-		ghvm->dtb.parcel_pages << PAGE_SHIFT);
+		ghvm->dtb.parcel.start << PAGE_SHIFT,
+		ghvm->dtb.parcel.pages << PAGE_SHIFT);
 	if (ret) {
 		dev_warn(ghvm->parent,
 			 "Failed to set location of DTB mem parcel: %d\n", ret);
@@ -788,9 +747,9 @@ static int gunyah_vm_start(struct gunyah_vm *ghvm)
 		gunyah_vm_add_resource(ghvm, ghrsc);
 	}
 
-	ret = gunyah_vm_parcel_to_paged(ghvm, &ghvm->dtb.parcel,
-					ghvm->dtb.parcel_start,
-					ghvm->dtb.parcel_pages);
+	ret = gunyah_vm_parcel_to_paged(ghvm, &ghvm->dtb.parcel.parcel,
+					ghvm->dtb.parcel.start,
+					ghvm->dtb.parcel.pages);
 	if (ret)
 		goto err;
 
@@ -800,8 +759,8 @@ static int gunyah_vm_start(struct gunyah_vm *ghvm)
 		 * need to rollback parcel_to_paged because RM is still
 		 * tracking the parcel
 		 */
-		gunyah_vm_mm_erase_range(ghvm, ghvm->dtb.parcel_start,
-					 ghvm->dtb.parcel_pages);
+		gunyah_vm_mm_erase_range(ghvm, ghvm->dtb.parcel.start,
+					 ghvm->dtb.parcel.pages);
 		dev_warn(ghvm->parent, "Failed to start VM: %d\n", ret);
 		goto err;
 	}
@@ -961,51 +920,10 @@ int __must_check gunyah_vm_get(struct gunyah_vm *ghvm)
 }
 EXPORT_SYMBOL_GPL(gunyah_vm_get);
 
-int gunyah_gup_reclaim_parcel(struct gunyah_vm *ghvm,
-			       struct gunyah_rm_mem_parcel *parcel, u64 gfn,
-			       u64 nr)
-{
-	struct gunyah_rm_mem_entry *entry;
-	struct folio *folio;
-	pgoff_t i;
-	int ret;
-
-	if (parcel->mem_handle != GUNYAH_MEM_HANDLE_INVAL) {
-		ret = gunyah_rm_mem_reclaim(ghvm->rm, parcel);
-		if (ret) {
-			dev_err(ghvm->parent, "Failed to reclaim parcel: %d\n",
-				ret);
-			/* We can't reclaim the pages -- hold onto the pages
-			 * forever because we don't know what state the memory
-			 * is in
-			 */
-			return ret;
-		}
-		parcel->mem_handle = GUNYAH_MEM_HANDLE_INVAL;
-
-		for (i = 0; i < parcel->n_mem_entries; i++) {
-			entry = &parcel->mem_entries[i];
-
-			folio = pfn_folio(PHYS_PFN(le64_to_cpu(entry->phys_addr)));
-
-			if (folio_test_private(folio))
-				gunyah_folio_host_reclaim(folio);
-
-			unpin_user_page(folio_page(folio, 0));
-			account_locked_vm(ghvm->mm_s, 1, false);
-		}
-
-		kfree(parcel->mem_entries);
-		kfree(parcel->acl_entries);
-	}
-
-	return 0;
-}
-
 static void _gunyah_vm_put(struct kref *kref)
 {
 	struct gunyah_vm *ghvm = container_of(kref, struct gunyah_vm, kref);
-	struct gunyah_vm_gup_binding *b;
+	struct gunyah_vm_binding *b;
 	unsigned long index = 0;
 	void *entry;
 	int ret;
@@ -1017,28 +935,8 @@ static void _gunyah_vm_put(struct kref *kref)
 	if (ghvm->vm_status == GUNYAH_RM_VM_STATUS_RUNNING)
 		gunyah_vm_stop(ghvm);
 
-	if (ghvm->vm_status == GUNYAH_RM_VM_STATUS_LOAD ||
-	    ghvm->vm_status == GUNYAH_RM_VM_STATUS_READY ||
-	    ghvm->vm_status == GUNYAH_RM_VM_STATUS_INIT_FAILED) {
-		ret = gunyah_gup_reclaim_parcel(ghvm, &ghvm->dtb.parcel,
-						 ghvm->dtb.parcel_start,
-						 ghvm->dtb.parcel_pages);
-		if (ret)
-			dev_err(ghvm->parent,
-				"Failed to reclaim DTB parcel: %d\n", ret);
-	}
-
 	gunyah_vm_remove_functions(ghvm);
 
-	down_write(&ghvm->bindings_lock);
-	mt_for_each(&ghvm->bindings, b, index, ULONG_MAX) {
-		mtree_erase(&ghvm->bindings, gunyah_gpa_to_gfn(b->guest_phys_addr));
-		kfree(b);
-	}
-	up_write(&ghvm->bindings_lock);
-	WARN_ON(!mtree_empty(&ghvm->bindings));
-	mtree_destroy(&ghvm->bindings);
-
 	/**
 	 * If this fails, we're going to lose the memory for good and is
 	 * BUG_ON-worthy, but not unrecoverable (we just lose memory).
@@ -1046,6 +944,8 @@ static void _gunyah_vm_put(struct kref *kref)
 	 * running and RM will let us reclaim all the memory.
 	 */
 	WARN_ON(gunyah_vm_reclaim_range(ghvm, 0, U64_MAX));
+	WARN_ON(!mtree_empty(&ghvm->mm));
+	mtree_destroy(&ghvm->mm);
 
 	/* clang-format off */
 	gunyah_vm_remove_resource_ticket(ghvm, &ghvm->addrspace_ticket);
@@ -1070,17 +970,15 @@ static void _gunyah_vm_put(struct kref *kref)
 		/* clang-format on */
 	}
 
-	WARN_ON(!mtree_empty(&ghvm->mm));
-	mtree_destroy(&ghvm->mm);
-
-	if (ghvm->auth == GUNYAH_RM_VM_AUTH_QCOM_ANDROID_PVM) {
-		ret = gunyah_gup_reclaim_parcel(ghvm, &ghvm->fw.parcel,
-				ghvm->fw.parcel_start,
-				ghvm->fw.parcel_pages);
-		if (ret)
-			dev_err(ghvm->parent,
-					"Failed to reclaim firmware parcel: %d\n", ret);
+	WARN_ON(gunyah_reclaim_parcels(ghvm, 0, ULONG_MAX));
+	down_write(&ghvm->bindings_lock);
+	mt_for_each(&ghvm->bindings, b, index, ULONG_MAX) {
+		mtree_erase(&ghvm->bindings, gunyah_gpa_to_gfn(b->guest_phys_addr));
+		kfree(b);
 	}
+	up_write(&ghvm->bindings_lock);
+	WARN_ON(!mtree_empty(&ghvm->bindings));
+	mtree_destroy(&ghvm->bindings);
 
 	if (ghvm->vm_status > GUNYAH_RM_VM_STATUS_NO_STATE) {
 		gunyah_rm_notifier_unregister(ghvm->rm, &ghvm->nb);

drivers/virt/gunyah/vm_mgr.h

@@ -7,6 +7,7 @@
 #define _GUNYAH_VM_MGR_PRIV_H
 
 #include <linux/device.h>
+#include <linux/file.h>
 #include <linux/kref.h>
 #include <linux/maple_tree.h>
 #include <linux/mutex.h>
@@ -20,17 +21,20 @@
 
 #include "rsc_mgr.h"
 
-enum gunyah_vm_mem_share_type {
-	VM_MEM_SHARE,
-	VM_MEM_LEND,
+struct gunyah_vm_parcel {
+	u64 start, pages;
+	struct gunyah_rm_mem_parcel parcel;
 };
 
-struct gunyah_vm_gup_binding {
-	enum gunyah_vm_mem_share_type share_type;
-	u64 guest_phys_addr;
+struct gunyah_vm_binding {
+	enum { VM_MEM_USER } mem_type;
 	u64 userspace_addr;
+	struct gunyah_vm_parcel *vm_parcel;
+	enum { VM_MEM_SHARE, VM_MEM_LEND } share_type;
+	u64 guest_phys_addr;
 	u64 size;
 	u32 flags;
+	u32 label;
 };
 
 static inline u64 gunyah_gpa_to_gfn(u64 gpa)
@@ -134,13 +138,11 @@ struct gunyah_vm {
 	enum gunyah_rm_vm_auth_mechanism auth;
 	struct {
 		struct gunyah_vm_dtb_config config;
-		u64 parcel_start, parcel_pages;
-		struct gunyah_rm_mem_parcel parcel;
+		struct gunyah_vm_parcel parcel;
 	} dtb;
 	struct {
 		struct gunyah_vm_firmware_config config;
-		u64 parcel_start, parcel_pages;
-		struct gunyah_rm_mem_parcel parcel;
+		struct gunyah_vm_parcel parcel;
 	} fw;
 	struct xarray boot_context;
 };
@@ -210,22 +212,24 @@ int gunyah_vm_parcel_to_paged(struct gunyah_vm *ghvm,
 			      struct gunyah_rm_mem_parcel *parcel, u64 gfn,
 			      u64 nr);
 void gunyah_vm_mm_erase_range(struct gunyah_vm *ghvm, u64 gfn, u64 nr);
-int gunyah_vm_reclaim_parcel(struct gunyah_vm *ghvm,
-			     struct gunyah_rm_mem_parcel *parcel, u64 gfn);
 int gunyah_vm_provide_folio(struct gunyah_vm *ghvm, struct folio *folio,
 			    u64 gfn, bool share, bool write);
 int gunyah_vm_reclaim_folio(struct gunyah_vm *ghvm, u64 gfn, struct folio *folio);
 int gunyah_vm_reclaim_range(struct gunyah_vm *ghvm, u64 gfn, u64 nr);
-
 int gunyah_vm_binding_alloc(struct gunyah_vm *ghvm,
 			    struct gunyah_userspace_memory_region *region,
 			    bool lend);
-int gunyah_gup_setup_demand_paging(struct gunyah_vm *ghvm);
-int gunyah_gup_share_parcel(struct gunyah_vm *ghvm,
-			      struct gunyah_rm_mem_parcel *parcel,
+int gunyah_share_parcel(struct gunyah_vm *ghvm,
+			      struct gunyah_vm_parcel *parcel,
 			      u64 *gfn, u64 *nr);
-int gunyah_gup_demand_page(struct gunyah_vm *ghvm, u64 gpa, bool write);
-
+int gunyah_share_range_as_parcels(struct gunyah_vm *ghvm, u64 start_addr,
+				  u64 end_addr,
+				  struct gunyah_vm_parcel **parcels);
+int gunyah_reclaim_parcels(struct gunyah_vm *ghvm,
+				u64 start_gfn, u64 end_gfn);
+int gunyah_demand_page(struct gunyah_vm *ghvm, u64 gpa, bool write);
+int gunyah_setup_demand_paging(struct gunyah_vm *ghvm, u64 start_gfn,
+				u64 end_gfn);
 #ifdef CONFIG_DMA_CMA
 int gunyah_cma_mem_init(void);
 void gunyah_cma_mem_exit(void);

drivers/virt/gunyah/vm_mgr_mem.c

@@ -22,16 +22,37 @@ __first_resource(struct gunyah_vm_resource_ticket *ticket)
 					struct gunyah_resource, list);
 }
 
+/*
+ * Once the parcel is converted to paged, vm_mgr only tracks
+ * the pages. The parcel needs to be reclaimed by the caller.
+ */
 int gunyah_vm_parcel_to_paged(struct gunyah_vm *ghvm,
 			      struct gunyah_rm_mem_parcel *parcel, u64 gfn,
 			      u64 nr)
 {
+	struct gunyah_vm_parcel *vm_parcel;
 	struct gunyah_rm_mem_entry *entry;
+	struct gunyah_vm_binding *b;
 	unsigned long i, tag = 0;
 	struct folio *folio;
 	pgoff_t off = 0;
 	u64 size;
-	int ret;
+	int ret = 0;
+
+	down_write(&ghvm->bindings_lock);
+	b = mtree_load(&ghvm->bindings, gfn);
+	if (!b) {
+		WARN_ON("No backing binding for the parcel being paged");
+		ret = -ENOENT;
+		goto unlock;
+	}
+
+	vm_parcel = b->vm_parcel;
+	if (!vm_parcel) {
+		WARN_ON("No parcel found");
+		ret = -ENOENT;
+		goto unlock;
+	}
 
 	if (parcel->n_acl_entries > 1)
 		tag |= SHARE_TAG;
@@ -50,7 +71,7 @@ int gunyah_vm_parcel_to_paged(struct gunyah_vm *ghvm,
 			if (ret) {
 				WARN_ON(ret != -ENOMEM);
 				gunyah_vm_mm_erase_range(ghvm, gfn, off - 1);
-				return ret;
+				goto unlock;
 			}
 			off += folio_nr_pages(folio);
 			size -= folio_size(folio);
@@ -58,8 +79,12 @@ int gunyah_vm_parcel_to_paged(struct gunyah_vm *ghvm,
 		}
 	}
 	BUG_ON(off != nr);
+	vm_parcel->start = 0;
+	b->vm_parcel = NULL;
 
-	return 0;
+unlock:
+	up_write(&ghvm->bindings_lock);
+	return ret;
 }
 
 /**
@@ -74,9 +99,13 @@ void gunyah_vm_mm_erase_range(struct gunyah_vm *ghvm, u64 gfn, u64 nr)
 {
 	struct folio *folio;
 	u64 off = gfn;
+	void *entry;
 
 	while (off < gfn + nr) {
-		folio = xa_untag_pointer(mtree_erase(&ghvm->mm, off));
+		entry = mtree_erase(&ghvm->mm, off);
+		if (!entry)
+			return;
+		folio = xa_untag_pointer(entry);
 		if (!folio)
 			return;
 		off += folio_nr_pages(folio);
@@ -135,9 +164,8 @@ int gunyah_vm_provide_folio(struct gunyah_vm *ghvm, struct folio *folio,
 	/* clang-format on */
 	addrspace = __first_resource(&ghvm->addrspace_ticket);
 
-	if (!addrspace || !guest_extent || !host_extent) {
+	if (!addrspace || !guest_extent || !host_extent)
 		return -ENODEV;
-	}
 
 	if (share) {
 		map_flags |= BIT(GUNYAH_ADDRSPACE_MAP_FLAG_VMMIO);
@@ -356,7 +384,7 @@ int gunyah_vm_binding_alloc(struct gunyah_vm *ghvm,
 			    struct gunyah_userspace_memory_region *region,
 			    bool lend)
 {
-	struct gunyah_vm_gup_binding *binding;
+	struct gunyah_vm_binding *binding;
 	int ret = 0;
 
 	if (!region->memory_size || !PAGE_ALIGNED(region->memory_size) ||
@@ -372,10 +400,13 @@ int gunyah_vm_binding_alloc(struct gunyah_vm *ghvm,
 		return -ENOMEM;
 	}
 
+	binding->mem_type = VM_MEM_USER;
 	binding->userspace_addr = region->userspace_addr;
+	binding->vm_parcel = NULL;
 	binding->guest_phys_addr = region->guest_phys_addr;
 	binding->size = region->memory_size;
 	binding->flags = region->flags;
+	binding->label = region->label;
 
 	if (lend) {
 		binding->share_type = VM_MEM_LEND;
@@ -388,36 +419,25 @@ int gunyah_vm_binding_alloc(struct gunyah_vm *ghvm,
 				 gunyah_gpa_to_gfn(binding->guest_phys_addr + region->memory_size - 1),
 				 binding, GFP_KERNEL);
 
-
 	if(ret != 0)
 		kfree(binding);
-
 	up_write(&ghvm->bindings_lock);
 
 	return ret;
 }
 
-int gunyah_gup_demand_page(struct gunyah_vm *ghvm, u64 gpa, bool write)
+static int gunyah_gup_demand_page(struct gunyah_vm *ghvm, struct gunyah_vm_binding *b,
+								u64 gpa, bool write)
 {
 	unsigned long gfn = gunyah_gpa_to_gfn(gpa);
-	struct gunyah_vm_gup_binding *b;
 	unsigned int gup_flags;
 	u64 offset;
 	int pinned, ret;
 	struct page *page;
 	struct folio *folio;
 
-	down_read(&ghvm->bindings_lock);
-	b = mtree_load(&ghvm->bindings, gfn);
-	if (!b) {
-		ret = -ENOENT;
-		goto unlock;
-	}
-
-	if (write && !(b->flags & GUNYAH_MEM_ALLOW_WRITE)) {
-		ret = -EPERM;
-		goto unlock;
-	}
+	if (write && !(b->flags & GUNYAH_MEM_ALLOW_WRITE))
+		return -EPERM;
 	gup_flags = FOLL_LONGTERM;
 	if (b->flags & GUNYAH_MEM_ALLOW_WRITE)
 		gup_flags |= FOLL_WRITE;
@@ -426,7 +446,7 @@ int gunyah_gup_demand_page(struct gunyah_vm *ghvm, u64 gpa, bool write)
 
 	ret = account_locked_vm(current->mm, 1, true);
 	if (ret)
-		goto unlock;
+		return ret;
 
 	pinned = pin_user_pages_fast(b->userspace_addr + offset, 1,
 					gup_flags, &page);
@@ -455,48 +475,50 @@ int gunyah_gup_demand_page(struct gunyah_vm *ghvm, u64 gpa, bool write)
 				gpa, ret);
 		goto unpin_page;
 	}
-	goto unlock;
+	return ret;
 
 unpin_page:
 	unpin_user_page(page);
 unlock_page:
 	account_locked_vm(current->mm, 1, false);
-unlock:
-	up_read(&ghvm->bindings_lock);
 	return ret;
 }
-EXPORT_SYMBOL_GPL(gunyah_gup_demand_page);
 
-
-int gunyah_gup_share_parcel(struct gunyah_vm *ghvm, struct gunyah_rm_mem_parcel *parcel,
-			     u64 *gfn, u64 *nr)
+int gunyah_demand_page(struct gunyah_vm *ghvm, u64 gpa, bool write)
 {
-	struct gunyah_vm_gup_binding *b;
-	bool lend = false;
-	struct page **pages;
-	int pinned, ret;
-	u16 vmid;
-	struct folio *folio;
-	unsigned int gup_flags;
-	unsigned long i, offset, entries, entry_size;
+	unsigned long gfn = gunyah_gpa_to_gfn(gpa);
+	struct gunyah_vm_binding *b;
+	int ret;
 
-	parcel->mem_handle = GUNYAH_MEM_HANDLE_INVAL;
-
-	if (!*nr)
-		return -EINVAL;
 	down_read(&ghvm->bindings_lock);
-	b = mtree_load(&ghvm->bindings, *gfn);
+	b = mtree_load(&ghvm->bindings, gfn);
 	if (!b) {
 		ret = -ENOENT;
 		goto unlock;
 	}
+	ret = gunyah_gup_demand_page(ghvm, b, gpa, write);
+
+unlock:
+	up_read(&ghvm->bindings_lock);
+	return ret;
+}
+EXPORT_SYMBOL_GPL(gunyah_demand_page);
+
+static int gunyah_gup_share_parcel(struct gunyah_vm *ghvm,
+			struct gunyah_vm_parcel *vm_parcel,
+			struct gunyah_vm_binding *b, u64 *gfn, u64 *nr)
+{
+	struct gunyah_rm_mem_parcel *parcel = &vm_parcel->parcel;
+	struct page **pages;
+	int pinned, ret;
+	struct folio *folio;
+	unsigned int gup_flags;
+	unsigned long i, offset, entries, entry_size;
 
 	offset = gunyah_gfn_to_gpa(*gfn) - b->guest_phys_addr;
 	pages = kcalloc(*nr, sizeof(*pages), GFP_KERNEL_ACCOUNT);
-	if (!pages) {
-		ret = -ENOMEM;
-		goto unlock;
-	}
+	if (!pages)
+		return -ENOMEM;
 
 	gup_flags = FOLL_LONGTERM;
 	if (b->flags & GUNYAH_MEM_ALLOW_WRITE)
@@ -516,49 +538,12 @@ int gunyah_gup_share_parcel(struct gunyah_vm *ghvm, struct gunyah_rm_mem_parcel
 	if (ret)
 		goto unpin_pages;
 
-	if (b->share_type == VM_MEM_LEND) {
-		parcel->n_acl_entries = 1;
-		lend = true;
-	} else {
-		lend = false;
-		parcel->n_acl_entries = 2;
-	}
-	parcel->acl_entries = kcalloc(parcel->n_acl_entries,
-				      sizeof(*parcel->acl_entries), GFP_KERNEL);
-	if (!parcel->acl_entries) {
-		ret = -ENOMEM;
-		goto unaccount_pages;
-	}
-
-	/* acl_entries[0].vmid will be this VM's vmid. We'll fill it when the
-	 * VM is starting and we know the VM's vmid.
-	 */
-	parcel->acl_entries[0].vmid = cpu_to_le16(ghvm->vmid);
-	if (b->flags & GUNYAH_MEM_ALLOW_READ)
-		parcel->acl_entries[0].perms |= GUNYAH_RM_ACL_R;
-	if (b->flags & GUNYAH_MEM_ALLOW_WRITE)
-		parcel->acl_entries[0].perms |= GUNYAH_RM_ACL_W;
-	if (b->flags & GUNYAH_MEM_ALLOW_EXEC)
-		parcel->acl_entries[0].perms |= GUNYAH_RM_ACL_X;
-
-	if (!lend) {
-		ret = gunyah_rm_get_vmid(ghvm->rm, &vmid);
-		if (ret)
-			goto free_acl;
-
-		parcel->acl_entries[1].vmid = cpu_to_le16(vmid);
-		/* Host assumed to have all these permissions. Gunyah will not
-		* grant new permissions if host actually had less than RWX
-		*/
-		parcel->acl_entries[1].perms = GUNYAH_RM_ACL_R | GUNYAH_RM_ACL_W | GUNYAH_RM_ACL_X;
-	}
-
 	/* overallocate & assume no large folios */
 	parcel->mem_entries = kcalloc(pinned, sizeof(parcel->mem_entries[0]),
 					GFP_KERNEL_ACCOUNT);
 	if (!parcel->mem_entries) {
 		ret = -ENOMEM;
-		goto free_acl;
+		goto unaccount_pages;
 	}
 	folio = page_folio(pages[0]);
 	*gfn -= folio_page_idx(folio, pages[0]);
@@ -587,11 +572,15 @@ int gunyah_gup_share_parcel(struct gunyah_vm *ghvm, struct gunyah_rm_mem_parcel
 	parcel->mem_entries[entries].size = entry_size;
 	parcel->n_mem_entries = entries + 1;
 	ret = gunyah_rm_mem_share(ghvm->rm, parcel);
+	if (ret)
+		goto free_mem_entries;
+
+	vm_parcel->start = *gfn;
+	vm_parcel->pages = *nr;
+	b->vm_parcel = vm_parcel;
 	goto free_pages;
 
-free_acl:
-	kfree(parcel->acl_entries);
-	parcel->acl_entries = NULL;
+free_mem_entries:
 	kfree(parcel->mem_entries);
 	parcel->mem_entries = NULL;
 	parcel->n_mem_entries = 0;
@@ -601,7 +590,278 @@ unpin_pages:
 	unpin_user_pages(pages, pinned);
 free_pages:
 	kfree(pages);
+	return ret;
+}
+
+int gunyah_share_parcel(struct gunyah_vm *ghvm, struct gunyah_vm_parcel *vm_parcel,
+			     u64 *gfn, u64 *nr)
+{
+	struct gunyah_rm_mem_parcel *parcel = &vm_parcel->parcel;
+	struct gunyah_vm_binding *b;
+	bool lend;
+	u16 vmid;
+	int ret;
+
+	if (!*nr)
+		return -EINVAL;
+
+	down_write(&ghvm->bindings_lock);
+	b = mtree_load(&ghvm->bindings, *gfn);
+	if (!b) {
+		ret = -ENOENT;
+		goto unlock;
+	}
+
+	parcel->mem_handle = GUNYAH_MEM_HANDLE_INVAL;
+	if (b->share_type == VM_MEM_LEND) {
+		parcel->n_acl_entries = 1;
+		lend = true;
+	} else {
+		lend = false;
+		parcel->n_acl_entries = 2;
+		parcel->label = b->label;
+	}
+	parcel->acl_entries = kcalloc(parcel->n_acl_entries,
+				      sizeof(*parcel->acl_entries), GFP_KERNEL);
+	if (!parcel->acl_entries) {
+		ret = -ENOMEM;
+		goto unlock;
+	}
+
+	/* acl_entries[0].vmid will be this VM's vmid. We'll fill it when the
+	 * VM is starting and we know the VM's vmid.
+	 */
+	parcel->acl_entries[0].vmid = cpu_to_le16(ghvm->vmid);
+	if (b->flags & GUNYAH_MEM_ALLOW_READ)
+		parcel->acl_entries[0].perms |= GUNYAH_RM_ACL_R;
+	if (b->flags & GUNYAH_MEM_ALLOW_WRITE)
+		parcel->acl_entries[0].perms |= GUNYAH_RM_ACL_W;
+	if (b->flags & GUNYAH_MEM_ALLOW_EXEC)
+		parcel->acl_entries[0].perms |= GUNYAH_RM_ACL_X;
+
+	if (!lend) {
+		ret = gunyah_rm_get_vmid(ghvm->rm, &vmid);
+		if (ret)
+			goto free_acl;
+
+		parcel->acl_entries[1].vmid = cpu_to_le16(vmid);
+		/* Host assumed to have all these permissions. Gunyah will not
+		* grant new permissions if host actually had less than RWX
+		*/
+		parcel->acl_entries[1].perms = GUNYAH_RM_ACL_R | GUNYAH_RM_ACL_W | GUNYAH_RM_ACL_X;
+	}
+
+	ret = gunyah_gup_share_parcel(ghvm, vm_parcel, b, gfn, nr);
+	if (ret) {
+		dev_warn(ghvm->parent, "Failed to share GUP memory: %d\n", ret);
+		goto free_acl;
+	}
+	goto unlock;
+
+free_acl:
+	kfree(parcel->acl_entries);
+	parcel->acl_entries = NULL;
 unlock:
+	up_write(&ghvm->bindings_lock);
+	return ret;
+}
+
+/*
+ * This function will provide the number of bindings from
+ * start_addr to end_addr.
+ * Use ULONG_MAX as the end_addr to get all the bindings of the VM.
+ */
+static u32 gunyah_count_bindings(struct gunyah_vm *ghvm, u64 start_addr,
+						u64 end_addr)
+{
+	struct gunyah_vm_binding *b;
+	unsigned long addr = start_addr;
+	u32 count = 0;
+
+	down_read(&ghvm->bindings_lock);
+	mt_for_each(&ghvm->bindings, b, addr, end_addr)
+		count++;
+	up_read(&ghvm->bindings_lock);
+
+	return count;
+}
+
+static int gunyah_gup_reclaim_parcel(struct gunyah_vm *ghvm,
+		struct gunyah_vm_parcel *vm_parcel, struct gunyah_vm_binding *b)
+{
+	struct gunyah_rm_mem_parcel *parcel = &vm_parcel->parcel;
+	struct gunyah_rm_mem_entry *entry;
+	struct folio *folio;
+	pgoff_t i;
+	int ret;
+
+	if (parcel->mem_handle == GUNYAH_MEM_HANDLE_INVAL)
+		return 0;
+
+	ret = gunyah_rm_mem_reclaim(ghvm->rm, parcel);
+	if (ret) {
+		dev_err(ghvm->parent, "Failed to reclaim parcel: %d\n",
+			ret);
+		/* We can't reclaim the pages -- hold onto the pages
+		 * forever because we don't know what state the memory
+		 * is in
+		 */
+		return ret;
+	}
+
+	for (i = 0; i < parcel->n_mem_entries; i++) {
+		entry = &parcel->mem_entries[i];
+
+		folio = pfn_folio(PHYS_PFN(le64_to_cpu(entry->phys_addr)));
+
+		if (folio_test_private(folio))
+			gunyah_folio_host_reclaim(folio);
+
+		unpin_user_page(folio_page(folio, 0));
+		account_locked_vm(ghvm->mm_s, 1, false);
+	}
+
+	parcel->mem_handle = GUNYAH_MEM_HANDLE_INVAL;
+	kfree(parcel->mem_entries);
+	kfree(parcel->acl_entries);
+	vm_parcel->start = 0;
+	vm_parcel->pages = 0;
+	b->vm_parcel = NULL;
+	return ret;
+}
+
+static int gunyah_reclaim_parcel(struct gunyah_vm *ghvm,
+			    struct gunyah_vm_parcel *vm_parcel)
+{
+	struct gunyah_vm_binding *b;
+	int ret;
+
+	down_write(&ghvm->bindings_lock);
+	b = mtree_load(&ghvm->bindings, vm_parcel->start);
+	if (!b) {
+		ret = -ENOENT;
+		goto unlock;
+	}
+	ret = gunyah_gup_reclaim_parcel(ghvm, vm_parcel, b);
+
+unlock:
+	up_write(&ghvm->bindings_lock);
+	return ret;
+}
+
+int gunyah_reclaim_parcels(struct gunyah_vm *ghvm, u64 start_gfn,
+							u64 end_gfn)
+{
+	unsigned long gfn = start_gfn;
+	struct gunyah_vm_binding *b;
+	int ret, ret2 = 0;
+
+	mt_for_each(&ghvm->bindings, b, gfn, end_gfn) {
+		if (b->vm_parcel)
+			ret = gunyah_reclaim_parcel(ghvm, b->vm_parcel);
+		if (ret)
+			ret2 = ret;
+	}
+
+	return ret2;
+}
+
+/*
+ * gunyah_share_range_as_parcels() - Share all bindings as parcels from start_gfn to end_gfn
+ * @ghvm - The gunyah vm
+ * @start_gfn: Start guest page number
+ * @end_gfn: Last guest page number
+ * @parcels: Array of parcels allocated.
+ *
+ * Use ULONG_MAX as the end_gfn to share all the bindings of the VM
+ * provided enough space for parcels is present.
+ * Caller is responsible to free the parcels when parcels are done
+ * being used.
+ */
+int gunyah_share_range_as_parcels(struct gunyah_vm *ghvm, u64 start_gfn,
+				u64 end_gfn, struct gunyah_vm_parcel **parcels)
+{
+	struct gunyah_vm_binding *b;
+	unsigned long gfn = start_gfn;
+	u32 count = 0, n;
+	int ret, ret_err;
+
+	/* Find the number of parcels needed to be created within the requested range*/
+	n = gunyah_count_bindings(ghvm, start_gfn, end_gfn);
+
+	*parcels = kzalloc(sizeof(struct gunyah_vm_parcel) * n, GFP_KERNEL);
+	if (!*parcels)
+		return -ENOMEM;
+
+	mt_for_each(&ghvm->bindings, b, gfn, end_gfn) {
+		u64 parcel_start = b->guest_phys_addr >> PAGE_SHIFT;
+		u64 parcel_pages = b->size >> PAGE_SHIFT;
+
+		ret = gunyah_share_parcel(ghvm, &(*parcels)[count++], &parcel_start, &parcel_pages);
+		if (ret) {
+			dev_err(ghvm->parent, "Failed to share parcel of %llx: %d\n",
+								parcel_start, ret);
+			/* Let's roll back.*/
+			while (count--) {
+				if ((*parcels)[count].parcel.mem_handle !=
+					GUNYAH_MEM_HANDLE_INVAL) {
+					ret_err = gunyah_reclaim_parcel(ghvm, &(*parcels)[count]);
+					if (ret_err)
+						dev_err(ghvm->parent, "Failed to reclaim parcel: %d, memory will leak\n",
+										ret_err);
+				}
+			}
+			goto err;
+		}
+	}
+	return ret;
+
+err:
+	kfree(*parcels);
+	*parcels = NULL;
+	return ret;
+}
+EXPORT_SYMBOL_GPL(gunyah_share_range_as_parcels);
+
+int gunyah_setup_demand_paging(struct gunyah_vm *ghvm, u64 start_gfn,
+				u64 end_gfn)
+{
+	struct gunyah_rm_mem_entry *entries;
+	unsigned long gfn = start_gfn;
+	struct gunyah_vm_binding *b;
+	u32 count = 0, i;
+	int ret = 0;
+
+	down_read(&ghvm->bindings_lock);
+	mt_for_each(&ghvm->bindings, b, gfn, end_gfn)
+		if (b->share_type == VM_MEM_LEND &&
+			(b->guest_phys_addr != ghvm->fw.config.guest_phys_addr))
+			count++;
+
+	if (!count)
+		goto out;
+
+	entries = kcalloc(count, sizeof(*entries), GFP_KERNEL);
+	if (!entries) {
+		ret = -ENOMEM;
+		goto out;
+	}
+
+	gfn = start_gfn;
+	i = 0;
+	mt_for_each(&ghvm->bindings, b, gfn, end_gfn) {
+		if (b->share_type != VM_MEM_LEND ||
+			(b->guest_phys_addr == ghvm->fw.config.guest_phys_addr))
+			continue;
+		entries[i].phys_addr = cpu_to_le64(b->guest_phys_addr);
+		entries[i].size = cpu_to_le64(b->size);
+		if (++i == count)
+			break;
+	}
+
+	ret = gunyah_rm_vm_set_demand_paging(ghvm->rm, ghvm->vmid, i, entries);
+	kfree(entries);
+out:
 	up_read(&ghvm->bindings_lock);
 	return ret;
 }