[PATCH v3 17/21] mm/mm_init: don't rely on memblock to get KHO scratch migratetype

[Pratyush Yadav <pratyush@kernel.org>]

From: "Pratyush Yadav (Google)" <pratyush@kernel.org>

Currently struct page init via memmap_init() or deferred_init_memmap()
only queries the migrate type from KHO for each discrete memory range.
That works currently since KHO scratch memory has a different memory
type so it is always it its own region.

An upcoming patch will add support for discovering blocks of memory with
no preservations and it will mark it as MEMBLOCK_KHO_SCRATCH to allow
allocations from them. This can lead to the bootmem KHO scratch areas to
be merged into larger free ranges. This merging breaks the selection of
migrate type.

Get rid of kho_scratch_migratetype(). Instead, let the struct page init
logic set the migratetype to MIGRATE_MOVABLE for all pages, and then
give KHO a chance to update the migrate type of its bootmem scratch
areas to MIGRATE_CMA. Since scratch areas should be a few order of
magnitues smaller than the total system memory, doing it this way is
more efficient than searching the kho_scratch array for every pageblock.

Since kho_scratch_migratetype() no longer exists, drop the migratetype
argument to memmap_init_zone_range() and deferred_init_pages().

Signed-off-by: Pratyush Yadav (Google) <pratyush@kernel.org>
---
 include/linux/kexec_handover.h     |  6 ++++
 include/linux/memblock.h           | 19 ------------
 kernel/liveupdate/kexec_handover.c | 27 +++++++++++++++++
 mm/memblock.c                      | 10 -------
 mm/mm_init.c                       | 47 +++++++++++++++++++-----------
 5 files changed, 63 insertions(+), 46 deletions(-)

diff --git a/include/linux/kexec_handover.h b/include/linux/kexec_handover.h
index 18ba417f1ada..30a71710565a 100644
--- a/include/linux/kexec_handover.h
+++ b/include/linux/kexec_handover.h
@@ -39,6 +39,7 @@ void kho_populate(phys_addr_t fdt_phys, u64 fdt_len, phys_addr_t scratch_phys,
 		  u64 scratch_len);
 
 bool kho_scratch_overlap(phys_addr_t phys, size_t size);
+void kho_init_scratch_migratetype(unsigned long start_pfn, unsigned long end_pfn);
 #else
 static inline bool kho_is_enabled(void)
 {
@@ -122,6 +123,11 @@ static inline bool kho_scratch_overlap(phys_addr_t phys, size_t size)
 {
 	return false;
 }
+
+static inline void kho_init_scratch_migratetype(unsigned long start_pfn,
+						unsigned long end_pfn);
+{
+}
 #endif /* CONFIG_KEXEC_HANDOVER */
 
 #endif /* LINUX_KEXEC_HANDOVER_H */
diff --git a/include/linux/memblock.h b/include/linux/memblock.h
index 5afcd99aa8c1..f142c8a8cd5b 100644
--- a/include/linux/memblock.h
+++ b/include/linux/memblock.h
@@ -613,28 +613,9 @@ static inline void memtest_report_meminfo(struct seq_file *m) { }
 #ifdef CONFIG_MEMBLOCK_KHO_SCRATCH
 void memblock_set_kho_scratch_only(void);
 void memblock_clear_kho_scratch_only(void);
-bool memblock_is_kho_scratch_memory(phys_addr_t addr);
-
-static inline enum migratetype kho_scratch_migratetype(unsigned long pfn,
-						       enum migratetype mt)
-{
-	if (memblock_is_kho_scratch_memory(PFN_PHYS(pfn)))
-		return MIGRATE_CMA;
-	return mt;
-}
 #else
 static inline void memblock_set_kho_scratch_only(void) { }
 static inline void memblock_clear_kho_scratch_only(void) { }
-static inline bool memblock_is_kho_scratch_memory(phys_addr_t addr)
-{
-	return false;
-}
-
-static inline enum migratetype kho_scratch_migratetype(unsigned long pfn,
-						       enum migratetype mt)
-{
-	return mt;
-}
 #endif
 
 #endif /* _LINUX_MEMBLOCK_H */
diff --git a/kernel/liveupdate/kexec_handover.c b/kernel/liveupdate/kexec_handover.c
index 6e8682469bd7..b400976851e7 100644
--- a/kernel/liveupdate/kexec_handover.c
+++ b/kernel/liveupdate/kexec_handover.c
@@ -781,6 +781,33 @@ bool kho_scratch_overlap(phys_addr_t phys, size_t size)
 	return false;
 }
 
+/*
+ * Mark all KHO scratch pageblocks between start_pfn and end_pfn as
+ * MIGRATE_CMA.
+ */
+void kho_init_scratch_migratetype(unsigned long start_pfn, unsigned long end_pfn)
+{
+	for (unsigned int i = 0; i < kho_scratch_cnt; i++) {
+		unsigned long scratch_start, scratch_end, pfn;
+
+		scratch_start = PHYS_PFN(kho_scratch[i].addr);
+		scratch_end = PHYS_PFN(kho_scratch[i].addr + kho_scratch[i].size);
+
+		/* Target range doesn't overlap with scratch. */
+		if (start_pfn >= scratch_end || end_pfn <= scratch_start)
+			continue;
+
+		pfn = max_t(unsigned long, start_pfn, scratch_start);
+		pfn = pageblock_start_pfn(pfn);
+
+		while (pfn < end_pfn && pfn < scratch_end) {
+			init_pageblock_migratetype(pfn_to_page(pfn),
+						   MIGRATE_CMA, false);
+			pfn += pageblock_nr_pages;
+		}
+	}
+}
+
 /**
  * kho_reserve_scratch - Reserve a contiguous chunk of memory for kexec
  *
diff --git a/mm/memblock.c b/mm/memblock.c
index 6349c48154f4..8b2e551435ae 100644
--- a/mm/memblock.c
+++ b/mm/memblock.c
@@ -2511,16 +2511,6 @@ __init void memblock_clear_kho_scratch_only(void)
 {
 	kho_scratch_only = false;
 }
-
-bool __init_memblock memblock_is_kho_scratch_memory(phys_addr_t addr)
-{
-	int i = memblock_search(&memblock.memory, addr);
-
-	if (i == -1)
-		return false;
-
-	return memblock_is_kho_scratch(&memblock.memory.regions[i]);
-}
 #endif
 
 #ifdef CONFIG_KEXEC_HANDOVER
diff --git a/mm/mm_init.c b/mm/mm_init.c
index bd9dc1829005..1d2f302d8cd9 100644
--- a/mm/mm_init.c
+++ b/mm/mm_init.c
@@ -707,9 +707,10 @@ void __meminit __init_page_from_nid(unsigned long pfn, int nid)
 	__init_single_page(pfn_to_page(pfn), pfn, zid, nid);
 
 	if (pageblock_aligned(pfn)) {
-		enum migratetype mt =
-			kho_scratch_migratetype(pfn, MIGRATE_MOVABLE);
-		init_pageblock_migratetype(pfn_to_page(pfn), mt, false);
+		init_pageblock_migratetype(pfn_to_page(pfn), MIGRATE_MOVABLE,
+					   false);
+		/* Update the migrate type in case it is KHO scratch. */
+		kho_init_scratch_migratetype(pfn, pfn + 1);
 	}
 }
 
@@ -938,13 +939,22 @@ void __meminit memmap_init_range(unsigned long size, int nid, unsigned long zone
 		}
 		pfn++;
 	}
+
+	/*
+	 * Update the migrate type for any pageblocks that fall in a KHO scratch
+	 * area. Do this here outside the loop because most pageblocks will not
+	 * be KHO scratch, so only search once.
+	 *
+	 * Only pageblocks up to pfn are updated, deferred pageblocks will be
+	 * updated later.
+	 */
+	kho_init_scratch_migratetype(start_pfn, pfn);
 }
 
 static void __init memmap_init_zone_range(struct zone *zone,
 					  unsigned long start_pfn,
 					  unsigned long end_pfn,
-					  unsigned long *hole_pfn,
-					  enum migratetype mt)
+					  unsigned long *hole_pfn)
 {
 	unsigned long zone_start_pfn = zone->zone_start_pfn;
 	unsigned long zone_end_pfn = zone_start_pfn + zone->spanned_pages;
@@ -957,7 +967,8 @@ static void __init memmap_init_zone_range(struct zone *zone,
 		return;
 
 	memmap_init_range(end_pfn - start_pfn, nid, zone_id, start_pfn,
-			  zone_end_pfn, MEMINIT_EARLY, NULL, mt, false);
+			  zone_end_pfn, MEMINIT_EARLY, NULL, MIGRATE_MOVABLE,
+			  false);
 
 	if (*hole_pfn < start_pfn)
 		init_unavailable_range(*hole_pfn, start_pfn, zone_id, nid);
@@ -973,8 +984,6 @@ static void __init memmap_init(void)
 
 	for_each_mem_pfn_range(i, MAX_NUMNODES, &start_pfn, &end_pfn, &nid) {
 		struct pglist_data *node = NODE_DATA(nid);
-		enum migratetype mt =
-			kho_scratch_migratetype(start_pfn, MIGRATE_MOVABLE);
 
 		for (j = 0; j < MAX_NR_ZONES; j++) {
 			struct zone *zone = node->node_zones + j;
@@ -983,7 +992,7 @@ static void __init memmap_init(void)
 				continue;
 
 			memmap_init_zone_range(zone, start_pfn, end_pfn,
-					       &hole_pfn, mt);
+					       &hole_pfn);
 			zone_id = j;
 		}
 	}
@@ -1973,7 +1982,7 @@ unsigned long __init node_map_pfn_alignment(void)
 
 #ifdef CONFIG_DEFERRED_STRUCT_PAGE_INIT
 static void __init deferred_free_pages(unsigned long pfn,
-		unsigned long nr_pages, enum migratetype mt)
+		unsigned long nr_pages)
 {
 	struct page *page;
 	unsigned long i;
@@ -2012,8 +2021,7 @@ static inline void __init pgdat_init_report_one_done(void)
  * Return number of pages initialized.
  */
 static unsigned long __init deferred_init_pages(struct zone *zone,
-		unsigned long start_pfn, unsigned long end_pfn,
-		enum migratetype mt)
+		unsigned long start_pfn, unsigned long end_pfn)
 {
 	int nid = zone_to_nid(zone);
 	unsigned long nr_pages = end_pfn - start_pfn, pfn = start_pfn;
@@ -2027,7 +2035,14 @@ static unsigned long __init deferred_init_pages(struct zone *zone,
 	pfn = pageblock_align(start_pfn);
 	page = pfn_to_page(pfn);
 	for (; pfn < end_pfn; pfn += pageblock_nr_pages, page += pageblock_nr_pages)
-		init_pageblock_migratetype(page, mt, false);
+		init_pageblock_migratetype(page, MIGRATE_MOVABLE, false);
+
+	/*
+	 * Update the migrate type for any pageblocks that fall in a KHO scratch
+	 * area. Since most pageblocks will not be KHO scratch, do it outside
+	 * the loop to only do the search once.
+	 */
+	kho_init_scratch_migratetype(start_pfn, end_pfn);
 
 	return nr_pages;
 }
@@ -2057,8 +2072,6 @@ deferred_init_memmap_chunk(unsigned long start_pfn, unsigned long end_pfn,
 	for_each_free_mem_range(i, nid, 0, &start, &end, NULL) {
 		unsigned long spfn = PFN_UP(start);
 		unsigned long epfn = PFN_DOWN(end);
-		enum migratetype mt =
-			kho_scratch_migratetype(spfn, MIGRATE_MOVABLE);
 
 		if (spfn >= end_pfn)
 			break;
@@ -2070,8 +2083,8 @@ deferred_init_memmap_chunk(unsigned long start_pfn, unsigned long end_pfn,
 			unsigned long mo_pfn = ALIGN(spfn + 1, MAX_ORDER_NR_PAGES);
 			unsigned long chunk_end = min(mo_pfn, epfn);
 
-			nr_pages += deferred_init_pages(zone, spfn, chunk_end, mt);
-			deferred_free_pages(spfn, chunk_end - spfn, mt);
+			nr_pages += deferred_init_pages(zone, spfn, chunk_end);
+			deferred_free_pages(spfn, chunk_end - spfn);
 
 			spfn = chunk_end;
 
-- 
2.55.0.141.g00534a21ce-goog