[PATCH 0/4] mm: Reduce struct page overhead for FS-DAX pmem

[PATCH 0/4] mm: Reduce struct page overhead for FS-DAX pmem

[Muchun Song]

AI and code-execution services increasingly run user workloads in
short-lived sandbox VMs. These VMs are usually small, densely packed, and
backed by sparse disk images, so both memory footprint and VM startup time
matter directly for consolidation density.

Using virtio-pmem with FS-DAX for the guest filesystem is one way to
reduce that footprint. It avoids keeping a second copy of file data in the
guest page cache, and turns that cache into host-side file cache that can
be accounted and reclaimed with the host's global memory view. Host-side
page-cache reclaim is also more direct than asking the guest to drop its
page cache, waiting for the freed pages to be reported back, and then
reclaiming the memory from the host side. This makes higher overcommit
more practical for dense sandbox deployments.

The remaining problem is that virtio-pmem with FS-DAX still pays the
guest-side ZONE_DEVICE metadata cost up front. The guest registers the
whole pmem aperture and allocates and initializes struct page metadata for
every advertised PFN. The vmemmap overhead is about 1.56% of the pmem
device size, and initializing all of those struct pages can noticeably
slow down startup for lightweight VMs.

Much of that private metadata is unnecessary in common sandbox setups.
Holes in a sparse rootfs image have no host storage allocated, but the
guest still allocates struct page metadata for the corresponding pmem PFNs.
Also, many filesystem workloads access files through read(2) and write(2)
rather than mmap(2); those DAX blocks are copied through the kernel and do
not need to be inserted into userspace page tables, so they do not need
private per-PFN struct page state either.

The key observation is that when sizeof(struct page) is a power of two,
each PAGE_SIZE vmemmap page contains a naturally aligned, repeatable group
of struct page slots. For an FS-DAX pmem range at device registration time,
those slots only need the same ZONE_DEVICE and dev_pagemap state before a
PFN is exposed to userspace. The kernel mainly needs the vmemmap to resolve
the PFN back to a valid device page and its dev_pagemap. It does not need
independent writable per-PFN state for ranges that are only accessed
through the DAX direct-access path, or for ranges that are never accessed
at all.

This series takes advantage of that split by separating vmemmap population
from private metadata allocation. Device registration still gives every
advertised PFN a valid struct page representation, but the vmemmap mappings
initially point at a shared read-only vmemmap page containing the common
ZONE_DEVICE state. This avoids allocating and initializing private metadata
for PFNs that may never need it.

Private writable metadata is materialized only when it becomes necessary:
before a DAX fault inserts the PFN into a userspace mapping. At that point
the PFN can participate in the normal page-based MM paths, so the shared
vmemmap page is replaced with a private writable copy. PTE faults
materialize the corresponding vmemmap page, and PMD faults materialize the
whole PMD-sized metadata range. PFNs that are never faulted continue to use
the shared vmemmap page, avoiding both the memory cost and the struct page
initialization work.

A natural follow-up is to make this optimization reversible. Once a DAX
entry is removed from the address_space, and after all mappings, references
and pins that require private metadata are gone, the corresponding vmemmap
page could be remapped back to the shared read-only vmemmap page and the
private metadata page could be freed. With that, the guest-side struct page
overhead would track the live DAX working set rather than the full advertised
pmem device size.

Patch 1 renames the architecture opt-in for runtime vmemmap remapping so
it describes the generic capability instead of the HugeTLB user.

Patch 2 avoids touching PG_hwpoison state for clean pmem pages. This keeps
the normal clean-I/O path compatible with read-only shared metadata.

Patch 3 adds the shared read-only FS-DAX vmemmap infrastructure and the
helper that materializes a private metadata page on demand.

Patch 4 opts pmem FS-DAX mappings into the new mode and materializes the
metadata before DAX inserts a PFN into a userspace mapping.

This does not change FS-DAX data path semantics. It only changes when
private struct page metadata is allocated for pmem FS-DAX PFNs.

Muchun Song (4):
  mm: generalize vmemmap remap architecture support
  nvdimm/pmem: avoid HWPoison flag updates for clean pages
  mm: add shared read-only vmemmap support for FS-DAX
  fsdax: materialize pmem vmemmap metadata on faults

 arch/loongarch/Kconfig   |  2 +-
 arch/riscv/Kconfig       |  2 +-
 arch/x86/Kconfig         |  2 +-
 drivers/nvdimm/pmem.c    |  3 +-
 fs/Kconfig               |  2 +-
 fs/dax.c                 |  5 ++++
 include/linux/memremap.h | 11 ++++++-
 mm/Kconfig               |  6 ++--
 mm/memremap.c            | 38 ++++++++++++++++++++++--
 mm/mm_init.c             | 11 +++++++
 mm/sparse-vmemmap.c      | 64 +++++++++++++++++++++++++++++++++++++---
 11 files changed, 131 insertions(+), 15 deletions(-)


base-commit: 32b6ef9a5d0eca44f9cd91f52f4faa89f145a0de
-- 
2.54.0

[PATCH 1/4] mm: generalize vmemmap remap architecture support

[Muchun Song]

ARCH_WANT_OPTIMIZE_HUGETLB_VMEMMAP is named after its first user, but
the underlying requirement is not HugeTLB-specific. Generic MM only needs
to know whether an architecture can safely remap populated vmemmap PTEs at
run time.

Rename the architecture opt-in to ARCH_SUPPORTS_VMEMMAP_REMAP and make
HugeTLB vmemmap optimization depend on that generic capability. This also
prepares for FS-DAX vmemmap optimization, which relies on the same runtime
remap support.

Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 arch/loongarch/Kconfig | 2 +-
 arch/riscv/Kconfig     | 2 +-
 arch/x86/Kconfig       | 2 +-
 fs/Kconfig             | 2 +-
 mm/Kconfig             | 6 ++++--
 5 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/arch/loongarch/Kconfig b/arch/loongarch/Kconfig
index 9c5def706222..e9bf19172e86 100644
--- a/arch/loongarch/Kconfig
+++ b/arch/loongarch/Kconfig
@@ -72,6 +72,7 @@ config LOONGARCH
 	select ARCH_SUPPORTS_RT
 	select ARCH_SUPPORTS_SCHED_SMT if SMP
 	select ARCH_SUPPORTS_SCHED_MC  if SMP
+	select ARCH_SUPPORTS_VMEMMAP_REMAP if 64BIT
 	select ARCH_USE_BUILTIN_BSWAP
 	select ARCH_USE_CMPXCHG_LOCKREF
 	select ARCH_USE_MEMTEST
@@ -80,7 +81,6 @@ config LOONGARCH
 	select ARCH_WANT_DEFAULT_BPF_JIT if HAVE_EBPF_JIT
 	select ARCH_WANT_DEFAULT_TOPDOWN_MMAP_LAYOUT
 	select ARCH_WANT_LD_ORPHAN_WARN
-	select ARCH_WANT_OPTIMIZE_HUGETLB_VMEMMAP if 64BIT
 	select ARCH_WANTS_NO_INSTR
 	select ARCH_WANTS_THP_SWAP if HAVE_ARCH_TRANSPARENT_HUGEPAGE
 	select BUILDTIME_TABLE_SORT
diff --git a/arch/riscv/Kconfig b/arch/riscv/Kconfig
index 5965666194b0..3614001cafdb 100644
--- a/arch/riscv/Kconfig
+++ b/arch/riscv/Kconfig
@@ -76,6 +76,7 @@ config RISCV
 	select ARCH_SUPPORTS_RT
 	select ARCH_SUPPORTS_SHADOW_CALL_STACK if HAVE_SHADOW_CALL_STACK
 	select ARCH_SUPPORTS_SCHED_MC if SMP
+	select ARCH_SUPPORTS_VMEMMAP_REMAP
 	select ARCH_USE_CMPXCHG_LOCKREF if 64BIT
 	select ARCH_USE_MEMTEST
 	select ARCH_USE_QUEUED_RWLOCKS
@@ -88,7 +89,6 @@ config RISCV
 	select ARCH_WANT_HUGE_PMD_SHARE if 64BIT
 	select ARCH_WANT_LD_ORPHAN_WARN
 	select ARCH_WANT_OPTIMIZE_DAX_VMEMMAP
-	select ARCH_WANT_OPTIMIZE_HUGETLB_VMEMMAP
 	select ARCH_WANTS_NO_INSTR
 	select ARCH_WANTS_THP_SWAP if HAVE_ARCH_TRANSPARENT_HUGEPAGE
 	select ARCH_WEAK_RELEASE_ACQUIRE if ARCH_USE_QUEUED_SPINLOCKS
diff --git a/arch/x86/Kconfig b/arch/x86/Kconfig
index a8c3b3d31a27..53f72d1500c5 100644
--- a/arch/x86/Kconfig
+++ b/arch/x86/Kconfig
@@ -133,6 +133,7 @@ config X86
 	select ARCH_SUPPORTS_LTO_CLANG
 	select ARCH_SUPPORTS_LTO_CLANG_THIN
 	select ARCH_SUPPORTS_RT
+	select ARCH_SUPPORTS_VMEMMAP_REMAP	if X86_64
 	select ARCH_USE_BUILTIN_BSWAP
 	select ARCH_USE_CMPXCHG_LOCKREF
 	select ARCH_USE_MEMTEST
@@ -148,7 +149,6 @@ config X86
 	select ARCH_WANT_HUGE_PMD_SHARE		if X86_64
 	select ARCH_WANT_LD_ORPHAN_WARN
 	select ARCH_WANT_OPTIMIZE_DAX_VMEMMAP	if X86_64
-	select ARCH_WANT_OPTIMIZE_HUGETLB_VMEMMAP	if X86_64
 	select ARCH_WANTS_THP_SWAP		if X86_64
 	select ARCH_HAS_PARANOID_L1D_FLUSH
 	select ARCH_WANT_IRQS_OFF_ACTIVATE_MM
diff --git a/fs/Kconfig b/fs/Kconfig
index d1c210c6508f..655795fbbd1b 100644
--- a/fs/Kconfig
+++ b/fs/Kconfig
@@ -276,7 +276,7 @@ config HUGETLB_PAGE
 
 config HUGETLB_PAGE_OPTIMIZE_VMEMMAP
 	def_bool HUGETLB_PAGE
-	depends on ARCH_WANT_OPTIMIZE_HUGETLB_VMEMMAP
+	depends on ARCH_SUPPORTS_VMEMMAP_REMAP
 	depends on SPARSEMEM_VMEMMAP
 
 config HUGETLB_PMD_PAGE_TABLE_SHARING
diff --git a/mm/Kconfig b/mm/Kconfig
index c1ddf59c0d71..7ec734d89beb 100644
--- a/mm/Kconfig
+++ b/mm/Kconfig
@@ -463,12 +463,14 @@ config SPARSEMEM_VMEMMAP
 
 #
 # Select this config option from the architecture Kconfig, if it is preferred
-# to enable the feature of HugeTLB/dev_dax vmemmap optimization.
+# to enable the feature of dev_dax vmemmap optimization.
 #
 config ARCH_WANT_OPTIMIZE_DAX_VMEMMAP
 	bool
 
-config ARCH_WANT_OPTIMIZE_HUGETLB_VMEMMAP
+# Select this from architectures where generic MM can safely remap populated
+# vmemmap PTEs at run time.
+config ARCH_SUPPORTS_VMEMMAP_REMAP
 	bool
 
 config HAVE_MEMBLOCK_PHYS_MAP
-- 
2.54.0

[PATCH 2/4] nvdimm/pmem: avoid HWPoison flag updates for clean pages

[Muchun Song]

pmem_mkpage_present() is called after persistent-memory poison has been
cleared. It only needs to clear PG_hwpoison and undo the MCE nospec state
for PFNs that are still marked poisoned.

For clean PFNs, test_and_clear_pmem_poison() has no semantic effect, but it
still performs an atomic clear operation against struct page flags. That is
unnecessary today and conflicts with the FS-DAX vmemmap optimization added
later in the series, where clean PFNs may be backed by shared read-only
metadata.

Check PageHWPoison() first so only poisoned PFNs update the page state.
Poison recovery keeps the same behavior, while clean PFNs no longer require
writable per-PFN metadata.

Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 drivers/nvdimm/pmem.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/drivers/nvdimm/pmem.c b/drivers/nvdimm/pmem.c
index 30a51c365ce8..b14f75daeda7 100644
--- a/drivers/nvdimm/pmem.c
+++ b/drivers/nvdimm/pmem.c
@@ -80,7 +80,7 @@ static void pmem_mkpage_present(struct pmem_device *pmem, phys_addr_t offset,
 		 * here since we're in the driver I/O path and
 		 * outstanding I/O requests pin the dev_pagemap.
 		 */
-		if (test_and_clear_pmem_poison(page))
+		if (PageHWPoison(page) && test_and_clear_pmem_poison(page))
 			clear_mce_nospec(pfn);
 	}
 }
-- 
2.54.0

[PATCH 3/4] mm: add shared read-only vmemmap support for FS-DAX

[Muchun Song]

FS-DAX registers persistent-memory ranges as ZONE_DEVICE memory, and the
kernel normally allocates and initializes vmemmap storage for every
advertised PFN up front. Sparse pmem images and workloads that only use the
DAX direct-access path may never need writable per-PFN state for most of
that range, but still pay the memory and initialization cost.

Add an opt-in dev_pagemap mode that populates FS-DAX vmemmap PTEs from a
shared read-only metadata page. The shared page is initialized with the
common ZONE_DEVICE and dev_pagemap state, so every PFN still has a valid
struct page representation while private metadata allocation is deferred.

This relies on sizeof(struct page) being a power of two, so each vmemmap
page contains a naturally aligned and repeatable set of struct page slots.
It also requires architecture support for runtime vmemmap remapping,
because shared mappings must be replaced with private writable pages before
a PFN can enter userspace mappings.

The initial implementation is deliberately limited to a single
memory-block-aligned range. That is not a fundamental requirement, but keeps
the registration and teardown paths simple; support for multiple ranges or
less strict alignment can be added later.

Provide vmemmap_materialize_page() to replace shared mappings in the
requested metadata range with private writable copies. A later patch will
call it from the FS-DAX fault path.

No caller enables the mode yet.

Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 include/linux/memremap.h | 11 ++++++-
 mm/memremap.c            | 38 ++++++++++++++++++++++--
 mm/mm_init.c             | 11 +++++++
 mm/sparse-vmemmap.c      | 64 +++++++++++++++++++++++++++++++++++++---
 4 files changed, 116 insertions(+), 8 deletions(-)

diff --git a/include/linux/memremap.h b/include/linux/memremap.h
index e3c2ccf872a8..21c9b6aeef67 100644
--- a/include/linux/memremap.h
+++ b/include/linux/memremap.h
@@ -9,6 +9,7 @@
 
 struct resource;
 struct device;
+struct page;
 
 /**
  * struct vmem_altmap - pre-allocated storage for vmemmap_populate
@@ -108,7 +109,8 @@ struct dev_pagemap_ops {
 	void (*folio_split)(struct folio *head, struct folio *tail);
 };
 
-#define PGMAP_ALTMAP_VALID	(1 << 0)
+#define PGMAP_ALTMAP_VALID				BIT(0)
+#define PGMAP_VMEMMAP_OPTIMIZATION			BIT(1)
 
 /**
  * struct dev_pagemap - metadata for ZONE_DEVICE mappings
@@ -122,6 +124,7 @@ struct dev_pagemap_ops {
  *	A zero value (default) uses base pages as the vmemmap metadata
  *	representation. A bigger value will set up compound struct pages
  *	of the requested order value.
+ * @vmemmap_shared_page: shared read-only vmemmap page for optimized FS-DAX
  * @ops: method table
  * @owner: an opaque pointer identifying the entity that manages this
  *	instance.  Used by various helpers to make sure that no
@@ -137,6 +140,7 @@ struct dev_pagemap {
 	enum memory_type type;
 	unsigned int flags;
 	unsigned long vmemmap_shift;
+	struct page *vmemmap_shared_page;
 	const struct dev_pagemap_ops *ops;
 	void *owner;
 	int nr_range;
@@ -232,6 +236,7 @@ void *devm_memremap_pages(struct device *dev, struct dev_pagemap *pgmap);
 void devm_memunmap_pages(struct device *dev, struct dev_pagemap *pgmap);
 struct dev_pagemap *get_dev_pagemap(unsigned long pfn);
 bool pgmap_pfn_valid(struct dev_pagemap *pgmap, unsigned long pfn);
+int vmemmap_materialize_page(struct page *page, unsigned int order);
 
 unsigned long memremap_compat_align(void);
 
@@ -307,4 +312,8 @@ static inline void put_dev_pagemap(struct dev_pagemap *pgmap)
 		percpu_ref_put(&pgmap->ref);
 }
 
+static inline bool pgmap_vmemmap_optimizable(struct dev_pagemap *pgmap)
+{
+	return pgmap && pgmap->vmemmap_shared_page != NULL;
+}
 #endif /* _LINUX_MEMREMAP_H_ */
diff --git a/mm/memremap.c b/mm/memremap.c
index accba23aef28..a53d09b84eaa 100644
--- a/mm/memremap.c
+++ b/mm/memremap.c
@@ -3,6 +3,7 @@
 #include <linux/device.h>
 #include <linux/io.h>
 #include <linux/kasan.h>
+#include <linux/memory.h>
 #include <linux/memory_hotplug.h>
 #include <linux/memremap.h>
 #include <linux/swap.h>
@@ -83,6 +84,30 @@ static unsigned long pfn_len(struct dev_pagemap *pgmap, unsigned long range_id)
 		pfn_first(pgmap, range_id)) >> pgmap->vmemmap_shift;
 }
 
+static int pgmap_vmemmap_shared_page_alloc(struct dev_pagemap *pgmap, int nid)
+{
+	const struct range *range = &pgmap->range;
+
+	if (!is_power_of_2(sizeof(struct page)) ||
+	    !IS_ENABLED(CONFIG_ARCH_SUPPORTS_VMEMMAP_REMAP) ||
+	    !(pgmap->flags & PGMAP_VMEMMAP_OPTIMIZATION))
+		return 0;
+
+	if (pgmap->nr_range != 1 ||
+	    !IS_ALIGNED(range->start | range_len(range), MIN_MEMORY_BLOCK_SIZE))
+		return 0;
+
+	pgmap->vmemmap_shared_page = alloc_pages_node(nid, GFP_KERNEL, 0);
+
+	return pgmap->vmemmap_shared_page ? 0 : -ENOMEM;
+}
+
+static inline void pgmap_vmemmap_shared_page_free(struct dev_pagemap *pgmap)
+{
+	if (pgmap->vmemmap_shared_page)
+		put_page(pgmap->vmemmap_shared_page);
+}
+
 static void pageunmap_range(struct dev_pagemap *pgmap, int range_id)
 {
 	struct range *range = &pgmap->ranges[range_id];
@@ -93,8 +118,9 @@ static void pageunmap_range(struct dev_pagemap *pgmap, int range_id)
 
 	/* pages are dead and unused, undo the arch mapping */
 	mem_hotplug_begin();
-	remove_pfn_range_from_zone(page_zone(first_page), PHYS_PFN(range->start),
-				   PHYS_PFN(range_len(range)));
+	if (!pgmap_vmemmap_optimizable(pgmap))
+		remove_pfn_range_from_zone(page_zone(first_page), PHYS_PFN(range->start),
+					   PHYS_PFN(range_len(range)));
 	if (pgmap->type == MEMORY_DEVICE_PRIVATE) {
 		__remove_pages(PHYS_PFN(range->start),
 			       PHYS_PFN(range_len(range)), NULL, pgmap);
@@ -123,6 +149,7 @@ void memunmap_pages(struct dev_pagemap *pgmap)
 
 	for (i = 0; i < pgmap->nr_range; i++)
 		pageunmap_range(pgmap, i);
+	pgmap_vmemmap_shared_page_free(pgmap);
 	percpu_ref_exit(&pgmap->ref);
 
 	WARN_ONCE(pgmap->altmap.alloc, "failed to free all reserved pages\n");
@@ -310,6 +337,9 @@ void *memremap_pages(struct dev_pagemap *pgmap, int nid)
 		break;
 	case MEMORY_DEVICE_FS_DAX:
 		params.pgprot = pgprot_decrypted(params.pgprot);
+		error = pgmap_vmemmap_shared_page_alloc(pgmap, nid);
+		if (error)
+			return ERR_PTR(error);
 		break;
 	case MEMORY_DEVICE_GENERIC:
 		break;
@@ -324,8 +354,10 @@ void *memremap_pages(struct dev_pagemap *pgmap, int nid)
 	init_completion(&pgmap->done);
 	error = percpu_ref_init(&pgmap->ref, dev_pagemap_percpu_release, 0,
 				GFP_KERNEL);
-	if (error)
+	if (error) {
+		pgmap_vmemmap_shared_page_free(pgmap);
 		return ERR_PTR(error);
+	}
 
 	/*
 	 * Clear the pgmap nr_range as it will be incremented for each
diff --git a/mm/mm_init.c b/mm/mm_init.c
index 2ed17cc707ed..7dd03b8a8d28 100644
--- a/mm/mm_init.c
+++ b/mm/mm_init.c
@@ -33,6 +33,7 @@
 #include <linux/vmstat.h>
 #include <linux/kexec_handover.h>
 #include <linux/hugetlb.h>
+#include <linux/memremap.h>
 #include "internal.h"
 #include "mm_init.h"
 #include "page_alloc.h"
@@ -1133,6 +1134,15 @@ void __ref memmap_init_zone_device(struct zone *zone,
 	if (!nr_pages)
 		return;
 
+	if (pgmap_vmemmap_optimizable(pgmap)) {
+		struct page *page = page_address(pgmap->vmemmap_shared_page);
+
+		for (int i = 0; i < PAGE_SIZE / sizeof(struct page); i++)
+			__init_zone_device_page(page + i, start_pfn + i,
+						ZONE_DEVICE, nid, pgmap);
+		goto pageblock_init;
+	}
+
 	/*
 	 * Seed the reusable head-page template from the first real struct
 	 * page. The normal page-init and refcount helpers must operate on
@@ -1164,6 +1174,7 @@ void __ref memmap_init_zone_device(struct zone *zone,
 				     compound_nr_pages(pfn, altmap, pgmap));
 	}
 
+pageblock_init:
 	pageblock_migratetype_init_range(start_pfn, nr_pages, MIGRATE_MOVABLE,
 					 /* isolate */ false, /* atomic */ false);
 
diff --git a/mm/sparse-vmemmap.c b/mm/sparse-vmemmap.c
index e62e6aa07f12..4dc7f020ed10 100644
--- a/mm/sparse-vmemmap.c
+++ b/mm/sparse-vmemmap.c
@@ -37,6 +37,8 @@
  */
 /* Get a ref on the head page struct page, for ZONE_DEVICE compound pages */
 #define VMEMMAP_POPULATE_PAGEREF	0x0001
+/* Read-only shared vmemmap mappings for FS-DAX base pages */
+#define VMEMMAP_POPULATE_FSDAX_SHARED	0x0002
 
 #include "internal.h"
 #include "mm_init.h"
@@ -262,10 +264,11 @@ static pte_t * __meminit vmemmap_pte_populate(pmd_t *pmd, unsigned long addr, in
 			 * and through vmemmap_populate_compound_pages() when
 			 * slab is available.
 			 */
-			if (flags & VMEMMAP_POPULATE_PAGEREF)
+			if (flags & (VMEMMAP_POPULATE_PAGEREF | VMEMMAP_POPULATE_FSDAX_SHARED))
 				get_page(pfn_to_page(ptpfn));
 		}
-		entry = pfn_pte(ptpfn, PAGE_KERNEL);
+		entry = pfn_pte(ptpfn, flags & VMEMMAP_POPULATE_FSDAX_SHARED ?
+				PAGE_KERNEL_RO : PAGE_KERNEL);
 		set_pte_at(&init_mm, addr, pte, entry);
 	} else if (WARN_ON_ONCE(vmemmap_optimizable_pfn(pfn)))
 		return NULL;
@@ -379,6 +382,54 @@ int __meminit vmemmap_populate_basepages(unsigned long start, unsigned long end,
 	return vmemmap_populate_range(start, end, node, altmap, -1, 0);
 }
 
+#ifdef CONFIG_ZONE_DEVICE
+static int __vmemmap_materialize_page(struct page *page)
+{
+	unsigned long addr = PAGE_ALIGN_DOWN((unsigned long)page);
+	struct dev_pagemap *pgmap = page_pgmap(page);
+	struct page *candidate, *template = pgmap->vmemmap_shared_page;
+	pte_t *pte = virt_to_kpte(addr);
+
+	if (pte_page(ptep_get(pte)) != template)
+		return 0;
+
+	candidate = alloc_pages_node(page_to_nid(page), GFP_KERNEL, 0);
+	if (!candidate)
+		return -ENOMEM;
+	copy_page(page_address(candidate), page_address(template));
+
+	spin_lock(&init_mm.page_table_lock);
+	if (pte_page(ptep_get(pte)) != template) {
+		__free_page(candidate);
+		goto out;
+	}
+	/* Make the copied struct page contents visible before the PTE update. */
+	smp_wmb();
+	set_pte_at(&init_mm, addr, pte, mk_pte(candidate, PAGE_KERNEL));
+	flush_tlb_kernel_range(addr, addr + PAGE_SIZE);
+	put_page(template);
+out:
+	spin_unlock(&init_mm.page_table_lock);
+
+	return 0;
+}
+
+int vmemmap_materialize_page(struct page *page, unsigned int order)
+{
+	struct dev_pagemap *pgmap = page_pgmap(page);
+	unsigned long end = (unsigned long)(page + (1UL << order));
+
+	if (!pgmap_vmemmap_optimizable(pgmap))
+		return 0;
+
+	for (unsigned long addr = (unsigned long)page; addr < end; addr += PAGE_SIZE)
+		if (__vmemmap_materialize_page((struct page *)addr))
+			return -ENOMEM;
+
+	return 0;
+}
+#endif /* CONFIG_ZONE_DEVICE */
+
 /*
  * Write protect the mirrored tail page structs for HVO. This will be
  * called from the hugetlb code when gathering and initializing the
@@ -581,7 +632,11 @@ struct page * __meminit __populate_section_memmap(unsigned long pfn,
 		!IS_ALIGNED(nr_pages, PAGES_PER_SUBSECTION)))
 		return NULL;
 
-	if (vmemmap_can_optimize(altmap, pgmap))
+	if (pgmap_vmemmap_optimizable(pgmap))
+		r = vmemmap_populate_range(start, end, nid, NULL,
+					   page_to_pfn(pgmap->vmemmap_shared_page),
+					   VMEMMAP_POPULATE_FSDAX_SHARED);
+	else if (vmemmap_can_optimize(altmap, pgmap))
 		r = vmemmap_populate_compound_pages(pfn, start, end, nid, pgmap);
 	else
 		r = vmemmap_populate(start, end, nid, altmap);
@@ -887,7 +942,8 @@ int __meminit sparse_add_section(int nid, unsigned long start_pfn,
 	 * Poison uninitialized struct pages in order to catch invalid flags
 	 * combinations.
 	 */
-	page_init_poison(memmap, sizeof(struct page) * nr_pages);
+	if (!pgmap_vmemmap_optimizable(pgmap))
+		page_init_poison(memmap, sizeof(struct page) * nr_pages);
 
 	ms = __nr_to_section(section_nr);
 	__section_mark_present(ms, section_nr);
-- 
2.54.0

[PATCH 4/4] fsdax: materialize pmem vmemmap metadata on faults

[Muchun Song]

Pmem FS-DAX registers the device range as ZONE_DEVICE memory, and the
kernel normally allocates and initializes vmemmap storage for every
advertised PFN up front. Sparse backing storage and workloads that only use
the DAX direct-access path may never need writable per-PFN state for most
of that range, but still pay the memory and initialization cost.

Opt pmem FS-DAX into the shared read-only vmemmap mode and materialize
private metadata before inserting a PFN into a userspace mapping. PMD
faults materialize the whole PMD-sized metadata range. PFNs that are never
faulted continue to use the shared metadata.

This shifts private vmemmap allocation from device registration to the
first DAX fault for each metadata page, so the struct page overhead tracks
the faulted DAX working set instead of the full advertised pmem device
size.

Signed-off-by: Muchun Song <songmuchun@bytedance.com>
---
 drivers/nvdimm/pmem.c | 1 +
 fs/dax.c              | 5 +++++
 2 files changed, 6 insertions(+)

diff --git a/drivers/nvdimm/pmem.c b/drivers/nvdimm/pmem.c
index b14f75daeda7..5530b32163d3 100644
--- a/drivers/nvdimm/pmem.c
+++ b/drivers/nvdimm/pmem.c
@@ -543,6 +543,7 @@ static int pmem_attach_disk(struct device *dev,
 		pmem->pgmap.nr_range = 1;
 		pmem->pgmap.type = MEMORY_DEVICE_FS_DAX;
 		pmem->pgmap.ops = &fsdax_pagemap_ops;
+		pmem->pgmap.flags |= PGMAP_VMEMMAP_OPTIMIZATION;
 		addr = devm_memremap_pages(dev, &pmem->pgmap);
 		bb_range = pmem->pgmap.range;
 	} else {
diff --git a/fs/dax.c b/fs/dax.c
index 1fbba0d21c13..89377301d51b 100644
--- a/fs/dax.c
+++ b/fs/dax.c
@@ -13,6 +13,7 @@
 #include <linux/fs.h>
 #include <linux/highmem.h>
 #include <linux/memcontrol.h>
+#include <linux/memremap.h>
 #include <linux/mm.h>
 #include <linux/mutex.h>
 #include <linux/sched.h>
@@ -1877,6 +1878,10 @@ static vm_fault_t dax_fault_iter(struct vm_fault *vmf,
 	if (err)
 		return pmd ? VM_FAULT_FALLBACK : dax_fault_return(err);
 
+	err = vmemmap_materialize_page(pfn_to_page(pfn), pmd ? PMD_ORDER : 0);
+	if (err)
+		return dax_fault_return(err);
+
 	*entry = dax_insert_entry(xas, vmf, iter, *entry, pfn, entry_flags);
 
 	if (write && iomap->flags & IOMAP_F_SHARED) {
-- 
2.54.0