[PATCH 35/49] mm/sparse-vmemmap: introduce section zone to struct mem_section

[Muchun Song]

Currently, HugeTLB obtains zone information for vmemmap optimization
through early pfn_to_zone(). However, ZONE_DEVICE cannot utilize this
approach because its zone information is updated after vmemmap population.

To pave the way for unifying DAX and HugeTLB vmemmap optimization,
this patch introduces the 'zone' member to struct mem_section. This
allows both DAX and HugeTLB to reliably obtain zone information
directly from the memory section.

Signed-off-by: Muchun Song <songmuchun at bytedance.com>
---
 include/linux/mmzone.h | 31 +++++++++++++++++++++++++++----
 mm/hugetlb.c           |  2 +-
 mm/hugetlb_vmemmap.c   |  4 +++-
 mm/sparse-vmemmap.c    | 19 +++++++++++++------
 4 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/include/linux/mmzone.h b/include/linux/mmzone.h
index 6edcb0cc46c4..846a7ee1334f 100644
--- a/include/linux/mmzone.h
+++ b/include/linux/mmzone.h
@@ -2022,6 +2022,7 @@ struct mem_section {
 	 * multiple sections.
 	 */
 	unsigned int order;
+	enum zone_type zone;
 #endif
 };
 
@@ -2214,32 +2215,54 @@ static inline void section_set_order(struct mem_section *section, unsigned int o
 	section->order = order;
 }
 
+static inline void section_set_zone(struct mem_section *section, enum zone_type zone)
+{
+	section->zone = zone;
+}
+
 static inline unsigned int section_order(const struct mem_section *section)
 {
 	return section->order;
 }
+
+static inline enum zone_type section_zone(const struct mem_section *section)
+{
+	return section->zone;
+}
 #else
 static inline void section_set_order(struct mem_section *section, unsigned int order)
 {
 }
 
+static inline void section_set_zone(struct mem_section *section, enum zone_type zone)
+{
+}
+
 static inline unsigned int section_order(const struct mem_section *section)
 {
 	return 0;
 }
+
+static inline enum zone_type section_zone(const struct mem_section *section)
+{
+	return 0;
+}
 #endif
 
-static inline void section_set_order_pfn_range(unsigned long pfn,
-					       unsigned long nr_pages,
-					       unsigned int order)
+static inline void section_set_compound_range(unsigned long pfn,
+					      unsigned long nr_pages,
+					      unsigned int order,
+					      enum zone_type zone)
 {
 	unsigned long section_nr = pfn_to_section_nr(pfn);
 
 	if (!IS_ALIGNED(pfn | nr_pages, PAGES_PER_SECTION))
 		return;
 
-	for (int i = 0; i < nr_pages / PAGES_PER_SECTION; i++)
+	for (int i = 0; i < nr_pages / PAGES_PER_SECTION; i++) {
 		section_set_order(__nr_to_section(section_nr + i), order);
+		section_set_zone(__nr_to_section(section_nr + i), zone);
+	}
 }
 
 static inline bool section_vmemmap_optimizable(const struct mem_section *section)
diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index 59728e942384..ce5a58aab5c3 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -3281,7 +3281,7 @@ static void __init gather_bootmem_prealloc_node(unsigned long nid)
 
 		if (section_vmemmap_optimizable(__pfn_to_section(folio_pfn(folio))))
 			folio_set_hugetlb_vmemmap_optimized(folio);
-		section_set_order_pfn_range(folio_pfn(folio), folio_nr_pages(folio), 0);
+		section_set_compound_range(folio_pfn(folio), folio_nr_pages(folio), 0, 0);
 
 		if (hugetlb_bootmem_page_earlycma(m))
 			folio_set_hugetlb_cma(folio);
diff --git a/mm/hugetlb_vmemmap.c b/mm/hugetlb_vmemmap.c
index a7ea98fcc18e..92c95ebdbb9a 100644
--- a/mm/hugetlb_vmemmap.c
+++ b/mm/hugetlb_vmemmap.c
@@ -681,11 +681,13 @@ void __init hugetlb_vmemmap_optimize_bootmem_page(struct huge_bootmem_page *m)
 {
 	struct hstate *h = m->hstate;
 	unsigned long pfn = PHYS_PFN(virt_to_phys(m));
+	int nid = early_pfn_to_nid(PHYS_PFN(__pa(m)));
 
 	if (!READ_ONCE(vmemmap_optimize_enabled))
 		return;
 
-	section_set_order_pfn_range(pfn, pages_per_huge_page(h), huge_page_order(h));
+	section_set_compound_range(pfn, pages_per_huge_page(h), huge_page_order(h),
+				   zone_idx(pfn_to_zone(pfn, nid)));
 }
 
 static const struct ctl_table hugetlb_vmemmap_sysctls[] = {
diff --git a/mm/sparse-vmemmap.c b/mm/sparse-vmemmap.c
index 6f959a999d5b..1867b5dcc73c 100644
--- a/mm/sparse-vmemmap.c
+++ b/mm/sparse-vmemmap.c
@@ -143,6 +143,11 @@ void __meminit vmemmap_verify(pte_t *pte, int node,
 			start, end - 1);
 }
 
+static inline struct zone *section_to_zone(const struct mem_section *ms, int nid)
+{
+	return &NODE_DATA(nid)->node_zones[section_zone(ms)];
+}
+
 static pte_t * __meminit vmemmap_pte_populate(pmd_t *pmd, unsigned long addr, int node,
 					      struct vmem_altmap *altmap,
 					      unsigned long ptpfn)
@@ -159,7 +164,7 @@ static pte_t * __meminit vmemmap_pte_populate(pmd_t *pmd, unsigned long addr, in
 			const struct mem_section *ms = __pfn_to_section(pfn);
 
 			page = vmemmap_shared_tail_page(section_order(ms),
-							pfn_to_zone(pfn, node));
+							section_to_zone(ms, node));
 			if (!page)
 				return NULL;
 			ptpfn = page_to_pfn(page);
@@ -471,16 +476,14 @@ static int __meminit vmemmap_populate_compound_pages(unsigned long start,
 	int rc;
 	unsigned long start_pfn = page_to_pfn((struct page *)start);
 	const struct mem_section *ms = __pfn_to_section(start_pfn);
-	struct page *tail = NULL;
+	struct page *tail;
 
 	/* This may occur in sub-section scenarios. */
 	if (!section_vmemmap_optimizable(ms))
 		return vmemmap_populate_range(start, end, node, NULL, -1);
 
-#ifdef CONFIG_ZONE_DEVICE
 	tail = vmemmap_shared_tail_page(section_order(ms),
-					&NODE_DATA(node)->node_zones[ZONE_DEVICE]);
-#endif
+					section_to_zone(ms, node));
 	if (!tail)
 		return -ENOMEM;
 
@@ -834,8 +837,12 @@ int __meminit sparse_add_section(int nid, unsigned long start_pfn,
 		return ret;
 
 	ms = __nr_to_section(section_nr);
-	if (vmemmap_can_optimize(altmap, pgmap) && nr_pages == PAGES_PER_SECTION)
+	if (vmemmap_can_optimize(altmap, pgmap) && nr_pages == PAGES_PER_SECTION) {
 		section_set_order(ms, pgmap->vmemmap_shift);
+#ifdef CONFIG_ZONE_DEVICE
+		section_set_zone(ms, ZONE_DEVICE);
+#endif
+	}
 	memmap = section_activate(nid, start_pfn, nr_pages, altmap, pgmap);
 	if (IS_ERR(memmap))
 		return PTR_ERR(memmap);
-- 
2.20.1