[TEMP] KVM: guest_memfd: Update gmem_prepare hook to handle partially-allocated folios

TODO: move to bitmap-based tracking within gmem

Marking the whole folio uptodate as part of kvm_gmem_populate() doesn't
work as expected for THP because kvm_gmem_populate() might only update a
subset of the pages, leaving the other pages still needing the
gmem_prepare hooks to be called when they are faulted in later.

Eventually this will be tracked by introducing a more granular flag than
uptodate that can track whether any particular 4K page has been prepared
or not.

For now however, modify kvm_gmem_populate() to only set the uptodate
flag if the entire 2M range is initialized, and leave it unset afterward
so gmem_prepare hooks still run. The gmem_prepare() hooks themselves
can already track page states at 4K granularity via the RMP table, so
are fully capable of dealing with partially-initialized folios.

As an optimization, also allow kvm_gmem_prepare() to set the uptodate
flag if the whole folio is already in a private/initialized state to
avoid needing to scan the RMP entries for each subpage.

Effectively, this makes the uptodate flag an optimization to skip RMP
checks within a fully-prepared folio, but in the absence of the uptodate
flag the RMP table is ultimately the authority on preparation state of
an individual sub-page within a folio.

Signed-off-by: Michael Roth <michael.roth@amd.com>

Author: mdroth

arch/x86/kvm/svm/sev.c

@@ -2189,6 +2189,9 @@ static int sev_gmem_post_populate(struct kvm *kvm, gfn_t gfn_start, kvm_pfn_t pf
 	int npages = (1 << order);
 	gfn_t gfn;
 
+	pr_debug("%s: GFN start 0x%llx PFN start 0x%llx order %d\n",
+		 __func__, gfn_start, pfn, order);
+
 	if (WARN_ON_ONCE(sev_populate_args->type != KVM_SEV_SNP_PAGE_TYPE_ZERO && !src))
 		return -EINVAL;
 
@@ -4745,6 +4748,33 @@ void sev_handle_rmp_fault(struct kvm_vcpu *vcpu, gpa_t gpa, u64 error_code)
 	kvm_release_page_unused(page);
 }
 
+static bool is_pfn_range_private(kvm_pfn_t start, kvm_pfn_t end)
+{
+	kvm_pfn_t pfn = start;
+
+	while (pfn < end) {
+		int ret, rmp_level;
+		bool assigned;
+
+		ret = snp_lookup_rmpentry(pfn, &assigned, &rmp_level);
+		if (ret) {
+			pr_warn_ratelimited("SEV: Failed to retrieve RMP entry: PFN 0x%llx GFN start 0x%llx GFN end 0x%llx RMP level %d error %d\n",
+					    pfn, start, end, rmp_level, ret);
+			return false;
+		}
+
+		if (!assigned) {
+			pr_debug("%s: overlap detected, PFN 0x%llx start 0x%llx end 0x%llx RMP level %d\n",
+				 __func__, pfn, start, end, rmp_level);
+			return false;
+		}
+
+		pfn++;
+	}
+
+	return true;
+}
+
 static bool is_pfn_range_shared(kvm_pfn_t start, kvm_pfn_t end)
 {
 	kvm_pfn_t pfn = start;
@@ -4800,6 +4830,7 @@ int sev_gmem_prepare(struct kvm *kvm, kvm_pfn_t pfn, gfn_t gfn, int max_order)
 {
 	struct kvm_sev_info *sev = to_kvm_sev_info(kvm);
 	kvm_pfn_t pfn_aligned;
+	struct folio *folio;
 	gfn_t gfn_aligned;
 	int level, rc;
 	bool assigned;
@@ -4830,6 +4861,16 @@ int sev_gmem_prepare(struct kvm *kvm, kvm_pfn_t pfn, gfn_t gfn, int max_order)
 		gfn_aligned = gfn;
 	}
 
+	folio = page_folio(pfn_to_page(pfn_aligned));
+	if (!folio_test_uptodate(folio)) {
+		unsigned long nr_pages = level == PG_LEVEL_4K ? 1 : 512;
+		int i;
+
+		pr_debug("%s: folio not up-to-date, clearing folio pages.\n", __func__);
+		for (i = 0; i < nr_pages; i++)
+			clear_highpage(pfn_to_page(pfn_aligned + i));
+	}
+
 	rc = rmp_make_private(pfn_aligned, gfn_to_gpa(gfn_aligned), level, sev->asid, false);
 	if (rc) {
 		pr_err_ratelimited("SEV: Failed to update RMP entry: GFN %llx PFN %llx level %d error %d\n",
@@ -4840,6 +4881,15 @@ int sev_gmem_prepare(struct kvm *kvm, kvm_pfn_t pfn, gfn_t gfn, int max_order)
 	pr_debug("%s: updated: gfn %llx pfn %llx pfn_aligned %llx max_order %d level %d\n",
 		 __func__, gfn, pfn, pfn_aligned, max_order, level);
 
+	if (pfn == pfn_aligned && folio_order(folio) == max_order) {
+		folio_mark_uptodate(folio);
+		pr_debug("%s: setting folio up-to-date (full update)\n", __func__);
+	} else if (is_pfn_range_private(pfn_aligned,
+					pfn_aligned + (1 << folio_order(folio)))) {
+		folio_mark_uptodate(folio);
+		pr_debug("%s: setting folio up-to-date (follow-up update)\n", __func__);
+	}
+
 	return 0;
 }
 

virt/kvm/guest_memfd.c

@@ -56,14 +56,9 @@ static inline void kvm_gmem_mark_prepared(struct folio *folio)
 static int kvm_gmem_prepare_folio(struct kvm *kvm, struct kvm_memory_slot *slot,
 				  gfn_t gfn, struct folio *folio)
 {
-	unsigned long nr_pages, i;
 	pgoff_t index;
 	int r;
 
-	nr_pages = folio_nr_pages(folio);
-	for (i = 0; i < nr_pages; i++)
-		clear_highpage(folio_page(folio, i));
-
 	/*
 	 * Preparing huge folios should always be safe, since it should
 	 * be possible to split them later if needed.
@@ -77,12 +72,16 @@ static int kvm_gmem_prepare_folio(struct kvm *kvm, struct kvm_memory_slot *slot,
 	 * The order will be passed when creating the guest_memfd, and
 	 * checked when creating memslots.
 	 */
-	WARN_ON(!IS_ALIGNED(slot->gmem.pgoff, 1 << folio_order(folio)));
+	if (!IS_ALIGNED(slot->gmem.pgoff, 1 << folio_order(folio)))
+		pr_debug_ratelimited("%s: GFN %llx not aligned (slot gfn start %llx pgoff %lx)",
+				     __func__, gfn, slot->base_gfn, slot->gmem.pgoff);
 	index = gfn - slot->base_gfn + slot->gmem.pgoff;
 	index = ALIGN_DOWN(index, 1 << folio_order(folio));
 	r = __kvm_gmem_prepare_folio(kvm, slot, index, folio);
-	if (!r)
-		kvm_gmem_mark_prepared(folio);
+	if (!r) {
+		pr_debug("%s: marking GFN %llx prepared\n", __func__, gfn);
+		//kvm_gmem_mark_prepared(folio);
+	}
 
 	return r;
 }
@@ -126,19 +125,36 @@ static struct folio *kvm_gmem_get_huge_folio(struct inode *inode, pgoff_t index,
  * Ignore accessed, referenced, and dirty flags.  The memory is
  * unevictable and there is no storage to write back to.
  */
-static struct folio *kvm_gmem_get_folio(struct inode *inode, pgoff_t index)
+static struct folio *__kvm_gmem_get_folio(struct inode *inode, pgoff_t index,
+					  bool allow_huge)
 {
 	struct folio *folio = NULL;
 
-	if (gmem_2m_enabled)
+	if (gmem_2m_enabled && allow_huge)
 		folio = kvm_gmem_get_huge_folio(inode, index, PMD_ORDER);
 
 	if (!folio)
 		folio = filemap_grab_folio(inode->i_mapping, index);
 
+	pr_debug("%s: allocate folio with PFN %lx order %d\n",
+		 __func__, folio_pfn(folio), folio_order(folio));
 	return folio;
 }
 
+/*
+ * Returns a locked folio on success.  The caller is responsible for
+ * setting the up-to-date flag before the memory is mapped into the guest.
+ * There is no backing storage for the memory, so the folio will remain
+ * up-to-date until it's removed.
+ *
+ * Ignore accessed, referenced, and dirty flags.  The memory is
+ * unevictable and there is no storage to write back to.
+ */
+static struct folio *kvm_gmem_get_folio(struct inode *inode, pgoff_t index)
+{
+	return __kvm_gmem_get_folio(inode, index, true);
+}
+
 static void kvm_gmem_invalidate_begin(struct kvm_gmem *gmem, pgoff_t start,
 				      pgoff_t end)
 {
@@ -606,7 +622,8 @@ void kvm_gmem_unbind(struct kvm_memory_slot *slot)
 static struct folio *__kvm_gmem_get_pfn(struct file *file,
 					struct kvm_memory_slot *slot,
 					pgoff_t index, kvm_pfn_t *pfn,
-					bool *is_prepared, int *max_order)
+					bool *is_prepared, int *max_order,
+					bool allow_huge)
 {
 	struct file *gmem_file = READ_ONCE(slot->gmem.file);
 	struct kvm_gmem *gmem = file->private_data;
@@ -624,7 +641,7 @@ static struct folio *__kvm_gmem_get_pfn(struct file *file,
 		return ERR_PTR(-EIO);
 	}
 
-	folio = kvm_gmem_get_folio(file_inode(file), index);
+	folio = __kvm_gmem_get_folio(file_inode(file), index, allow_huge);
 	if (IS_ERR(folio))
 		return folio;
 
@@ -636,11 +653,11 @@ static struct folio *__kvm_gmem_get_pfn(struct file *file,
 
 	*pfn = folio_file_pfn(folio, index);
 	if (!max_order)
-		goto success;
+		return folio;
 
-	*max_order = compound_order(compound_head(page));
+	*max_order = folio_order(folio);
 	if (!*max_order)
-		goto success;
+		return folio;
 
 	/*
 	 * The folio can be mapped with a hugepage if and only if the folio is
@@ -670,7 +687,7 @@ int kvm_gmem_get_pfn(struct kvm *kvm, struct kvm_memory_slot *slot,
 	if (!file)
 		return -EFAULT;
 
-	folio = __kvm_gmem_get_pfn(file, slot, index, pfn, &is_prepared, max_order);
+	folio = __kvm_gmem_get_pfn(file, slot, index, pfn, &is_prepared, max_order, true);
 	if (IS_ERR(folio)) {
 		r = PTR_ERR(folio);
 		goto out;
@@ -730,12 +747,16 @@ long kvm_gmem_populate(struct kvm *kvm, gfn_t start_gfn, void __user *src, long
 			break;
 		}
 
-		folio = __kvm_gmem_get_pfn(file, slot, index, &pfn, &is_prepared, &max_order);
+		folio = __kvm_gmem_get_pfn(file, slot, index, &pfn, &is_prepared, &max_order, false);
 		if (IS_ERR(folio)) {
 			ret = PTR_ERR(folio);
 			break;
 		}
 
+		pr_debug("%s: GFN start 0x%llx PFN start 0x%llx max_order %d folio_order %d uptodate %d\n",
+			 __func__, gfn, pfn, max_order, folio_order(folio),
+			 folio_test_uptodate(folio));
+
 		if (is_prepared) {
 			folio_unlock(folio);
 			folio_put(folio);
@@ -758,8 +779,11 @@ long kvm_gmem_populate(struct kvm *kvm, gfn_t start_gfn, void __user *src, long
 
 		p = src ? src + i * PAGE_SIZE : NULL;
 		ret = post_populate(kvm, gfn, pfn, p, max_order, opaque);
-		if (!ret)
+		if (!ret && max_order == folio_order(folio)) {
+			pr_debug("%s: GFN start 0x%llx PFN start 0x%llx max_order %d folio_order %d marking uptodate.\n",
+				 __func__, gfn, pfn, max_order, folio_order(folio));
 			kvm_gmem_mark_prepared(folio);
+		}
 
 put_folio_and_exit:
 		folio_put(folio);