[RFC PATCH 4/4] KVM: guest_memfd: Enforce NUMA mempolicy if available

[Shivank Garg <shivankg@amd.com>]

Enforce memory policy on guest-memfd to provide proper NUMA support.
Previously, guest-memfd allocations were following local NUMA node id
in absence of process mempolicy, resulting in random memory allocation.
Moreover, it cannot use mbind() since memory isn't mapped to userspace.

To support NUMA policies, call fbind() syscall from VMM to store
mempolicy as f_policy in struct kvm_gmem of guest_memfd. The f_policy
is retrieved to pass in filemap_grab_folio_mpol() to ensure that
allocations follow the specified memory policy.

Signed-off-by: Shivank Garg <shivankg@amd.com>
---
 mm/mempolicy.c         |  2 ++
 virt/kvm/guest_memfd.c | 49 ++++++++++++++++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/mm/mempolicy.c b/mm/mempolicy.c
index 3a697080ecad..af2e1ef4dae7 100644
--- a/mm/mempolicy.c
+++ b/mm/mempolicy.c
@@ -347,6 +347,7 @@ void __mpol_put(struct mempolicy *pol)
 		return;
 	kmem_cache_free(policy_cache, pol);
 }
+EXPORT_SYMBOL(__mpol_put);
 
 static void mpol_rebind_default(struct mempolicy *pol, const nodemask_t *nodes)
 {
@@ -2599,6 +2600,7 @@ struct mempolicy *__mpol_dup(struct mempolicy *old)
 	atomic_set(&new->refcnt, 1);
 	return new;
 }
+EXPORT_SYMBOL(__mpol_dup);
 
 /* Slow path of a mempolicy comparison */
 bool __mpol_equal(struct mempolicy *a, struct mempolicy *b)
diff --git a/virt/kvm/guest_memfd.c b/virt/kvm/guest_memfd.c
index 2c6fcf7c3ec9..0237bda4382c 100644
--- a/virt/kvm/guest_memfd.c
+++ b/virt/kvm/guest_memfd.c
@@ -4,6 +4,7 @@
 #include <linux/kvm_host.h>
 #include <linux/pagemap.h>
 #include <linux/anon_inodes.h>
+#include <linux/mempolicy.h>
 
 #include "kvm_mm.h"
 
@@ -11,6 +12,7 @@ struct kvm_gmem {
 	struct kvm *kvm;
 	struct xarray bindings;
 	struct list_head entry;
+	struct mempolicy *f_policy;
 };
 
 /**
@@ -87,7 +89,8 @@ static int kvm_gmem_prepare_folio(struct kvm *kvm, struct kvm_memory_slot *slot,
 }
 
 static struct folio *kvm_gmem_get_huge_folio(struct inode *inode, pgoff_t index,
-					     unsigned int order)
+					     unsigned int order,
+					     struct mempolicy *policy)
 {
 	pgoff_t npages = 1UL << order;
 	pgoff_t huge_index = round_down(index, npages);
@@ -104,7 +107,7 @@ static struct folio *kvm_gmem_get_huge_folio(struct inode *inode, pgoff_t index,
 				   (loff_t)(huge_index + npages - 1) << PAGE_SHIFT))
 		return NULL;
 
-	folio = filemap_alloc_folio(gfp, order);
+	folio = filemap_alloc_folio_mpol(gfp, order, policy);
 	if (!folio)
 		return NULL;
 
@@ -129,12 +132,26 @@ static struct folio *__kvm_gmem_get_folio(struct file *file, pgoff_t index,
 					  bool allow_huge)
 {
 	struct folio *folio = NULL;
+	struct kvm_gmem *gmem = file->private_data;
+	struct mempolicy *policy = NULL;
+
+	/*
+	 * RCU lock is required to prevent any race condition with set_policy().
+	 */
+	if (IS_ENABLED(CONFIG_NUMA)) {
+		rcu_read_lock();
+		policy = READ_ONCE(gmem->f_policy);
+		mpol_get(policy);
+		rcu_read_unlock();
+	}
 
 	if (gmem_2m_enabled && allow_huge)
-		folio = kvm_gmem_get_huge_folio(file_inode(file), index, PMD_ORDER);
+		folio = kvm_gmem_get_huge_folio(file_inode(file), index, PMD_ORDER, policy);
 
 	if (!folio)
-		folio = filemap_grab_folio(file_inode(file)->i_mapping, index);
+		folio = filemap_grab_folio_mpol(file_inode(file)->i_mapping, index, policy);
+
+	mpol_put(policy);
 
 	pr_debug("%s: allocate folio with PFN %lx order %d\n",
 		 __func__, folio_pfn(folio), folio_order(folio));
@@ -338,6 +355,7 @@ static int kvm_gmem_release(struct inode *inode, struct file *file)
 	mutex_unlock(&kvm->slots_lock);
 
 	xa_destroy(&gmem->bindings);
+	mpol_put(gmem->f_policy);
 	kfree(gmem);
 
 	kvm_put_kvm(kvm);
@@ -356,10 +374,32 @@ static inline struct file *kvm_gmem_get_file(struct kvm_memory_slot *slot)
 	return get_file_active(&slot->gmem.file);
 }
 
+#ifdef CONFIG_NUMA
+static int kvm_gmem_set_policy(struct file *file, struct mempolicy *mpol)
+{
+	struct mempolicy *old, *new;
+	struct kvm_gmem *gmem = file->private_data;
+
+	new = mpol_dup(mpol);
+	if (IS_ERR(new))
+		return PTR_ERR(new);
+
+	old = gmem->f_policy;
+	WRITE_ONCE(gmem->f_policy, new);
+	synchronize_rcu();
+	mpol_put(old);
+
+	return 0;
+}
+#endif
+
 static struct file_operations kvm_gmem_fops = {
 	.open		= generic_file_open,
 	.release	= kvm_gmem_release,
 	.fallocate	= kvm_gmem_fallocate,
+#ifdef CONFIG_NUMA
+	.set_policy	= kvm_gmem_set_policy,
+#endif
 };
 
 void kvm_gmem_init(struct module *module)
@@ -489,6 +529,7 @@ static int __kvm_gmem_create(struct kvm *kvm, loff_t size, u64 flags)
 
 	kvm_get_kvm(kvm);
 	gmem->kvm = kvm;
+	gmem->f_policy = NULL;
 	xa_init(&gmem->bindings);
 	list_add(&gmem->entry, &inode->i_mapping->i_private_list);
 
-- 
2.34.1