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

["Gupta, Pankaj" <pankaj.gupta@amd.com>]

> 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)

I liked the overall idea.

Minor naming suggestion I have:

'PGMAP_VMEMMAP_OPTIMIZATION' feel too generic to me.

Maybe something on the lines to reflect the actual optimization:

|PGMAP_SHARED_VMEMMAP or PGMAP_VMEMMAP_ON_DEMAND or some_other_name?|

|Thanks,|

|Pankaj|

>   
>   /**
>    * 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);