[RFC PATCH 6/7] mm: Add alloc_pages_memdesc family of APIs

["Matthew Wilcox (Oracle)" <willy@infradead.org>]

Convert the alloc_frozen_pages implementations into alloc_pages_memdesc
and add wrappers to keep the frozen pages users working.

This hasn't been widely tested; I bet the build bots will find something
I missed.

Signed-off-by: Matthew Wilcox (Oracle) <willy@infradead.org>
---
 include/linux/gfp.h      | 13 +++++++++++++
 include/linux/mm_types.h | 25 +++++++++++++++++++++++++
 mm/internal.h            | 13 ++++++++++---
 mm/mempolicy.c           | 28 ++++++++++++++++------------
 mm/page_alloc.c          | 12 +++++++++---
 5 files changed, 73 insertions(+), 18 deletions(-)

diff --git a/include/linux/gfp.h b/include/linux/gfp.h
index 0ceb4e09306c..6e13e0b829f8 100644
--- a/include/linux/gfp.h
+++ b/include/linux/gfp.h
@@ -225,6 +225,10 @@ struct page *__alloc_pages_noprof(gfp_t gfp, unsigned int order, int preferred_n
 		nodemask_t *nodemask);
 #define __alloc_pages(...)			alloc_hooks(__alloc_pages_noprof(__VA_ARGS__))
 
+struct page *__alloc_pages_memdesc_noprof(gfp_t gfp, unsigned int order,
+		memdesc_t memdesc, int preferred_nid, nodemask_t *nodemask);
+#define __alloc_pages_memdesc(...)		alloc_hooks(__alloc_pages_memdesc_noprof(__VA_ARGS__))
+
 struct folio *__folio_alloc_noprof(gfp_t gfp, unsigned int order, int preferred_nid,
 		nodemask_t *nodemask);
 #define __folio_alloc(...)			alloc_hooks(__folio_alloc_noprof(__VA_ARGS__))
@@ -315,6 +319,8 @@ static inline struct page *alloc_pages_node_noprof(int nid, gfp_t gfp_mask,
 
 #ifdef CONFIG_NUMA
 struct page *alloc_pages_noprof(gfp_t gfp, unsigned int order);
+struct page *alloc_pages_memdesc_noprof(gfp_t gfp, unsigned int order,
+		memdesc_t memdesc);
 struct folio *folio_alloc_noprof(gfp_t gfp, unsigned int order);
 struct folio *folio_alloc_mpol_noprof(gfp_t gfp, unsigned int order,
 		struct mempolicy *mpol, pgoff_t ilx, int nid);
@@ -325,6 +331,12 @@ static inline struct page *alloc_pages_noprof(gfp_t gfp_mask, unsigned int order
 {
 	return alloc_pages_node_noprof(numa_node_id(), gfp_mask, order);
 }
+static inline struct page *alloc_pages_memdesc_noprof(gfp_t gfp,
+		unsigned int order, memdesc_t memdesc)
+{
+	return __alloc_pages_memdesc_noprof(gfp, order, memdesc,
+			numa_node_id(), NULL);
+}
 static inline struct folio *folio_alloc_noprof(gfp_t gfp, unsigned int order)
 {
 	return __folio_alloc_node_noprof(gfp, order, numa_node_id());
@@ -339,6 +351,7 @@ static inline struct folio *folio_alloc_mpol_noprof(gfp_t gfp, unsigned int orde
 #endif
 
 #define alloc_pages(...)			alloc_hooks(alloc_pages_noprof(__VA_ARGS__))
+#define alloc_pages_memdesc(...)		alloc_hooks(alloc_pages_memdesc_noprof(__VA_ARGS__))
 #define folio_alloc(...)			alloc_hooks(folio_alloc_noprof(__VA_ARGS__))
 #define folio_alloc_mpol(...)			alloc_hooks(folio_alloc_mpol_noprof(__VA_ARGS__))
 #define vma_alloc_folio(...)			alloc_hooks(vma_alloc_folio_noprof(__VA_ARGS__))
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 90e5790c318f..f5d9e0afe0fa 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -38,6 +38,30 @@ typedef struct {
 	unsigned long f;
 } memdesc_flags_t;
 
+/**
+ * typedef memdesc_t - A typed memory descriptor.
+ *
+ * The bottom few bits of this encoded pointer determine the type
+ * of the memdesc.
+ */
+typedef struct {
+	unsigned long v;
+} memdesc_t;
+
+#define MEMDESC_TYPE_PAGE_TABLE		15
+
+static inline memdesc_t memdesc_create(void *p, unsigned long type)
+{
+	VM_BUG_ON((unsigned long)p & 15);
+	VM_BUG_ON(type > 15);
+	return (memdesc_t) { .v = type | (unsigned long)p };
+}
+
+static inline unsigned long memdesc_type(memdesc_t memdesc)
+{
+	return memdesc.v & 15;
+}
+
 /*
  * Each physical page in the system has a struct page associated with
  * it to keep track of whatever it is we are using the page for at the
@@ -126,6 +150,7 @@ struct page {
 		};
 		struct {	/* Tail pages of compound page */
 			unsigned long compound_head;	/* Bit zero is set */
+			memdesc_t memdesc;	/* All pages, not just tail */
 		};
 		struct {	/* ZONE_DEVICE pages */
 			/*
diff --git a/mm/internal.h b/mm/internal.h
index 1561fc2ff5b8..15d64601289b 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -824,15 +824,22 @@ extern bool free_pages_prepare(struct page *page, unsigned int order);
 
 extern int user_min_free_kbytes;
 
-struct page *__alloc_frozen_pages_noprof(gfp_t, unsigned int order, int nid,
-		nodemask_t *);
+static inline struct page *__alloc_frozen_pages_noprof(gfp_t gfp,
+		unsigned int order, int nid, nodemask_t *mask)
+{
+	return __alloc_pages_memdesc_noprof(gfp, order,
+			memdesc_create(NULL, 0), nid, mask);
+}
 #define __alloc_frozen_pages(...) \
 	alloc_hooks(__alloc_frozen_pages_noprof(__VA_ARGS__))
 void free_frozen_pages(struct page *page, unsigned int order);
 void free_unref_folios(struct folio_batch *fbatch);
 
 #ifdef CONFIG_NUMA
-struct page *alloc_frozen_pages_noprof(gfp_t, unsigned int order);
+static inline struct page *alloc_frozen_pages_noprof(gfp_t gfp, unsigned int order)
+{
+	return alloc_pages_memdesc_noprof(gfp, order, memdesc_create(NULL, 0));
+}
 #else
 static inline struct page *alloc_frozen_pages_noprof(gfp_t gfp, unsigned int order)
 {
diff --git a/mm/mempolicy.c b/mm/mempolicy.c
index eb83cff7db8c..866d6609a758 100644
--- a/mm/mempolicy.c
+++ b/mm/mempolicy.c
@@ -2338,7 +2338,7 @@ bool mempolicy_in_oom_domain(struct task_struct *tsk,
 }
 
 static struct page *alloc_pages_preferred_many(gfp_t gfp, unsigned int order,
-						int nid, nodemask_t *nodemask)
+		memdesc_t memdesc, int nid, nodemask_t *nodemask)
 {
 	struct page *page;
 	gfp_t preferred_gfp;
@@ -2351,9 +2351,11 @@ static struct page *alloc_pages_preferred_many(gfp_t gfp, unsigned int order,
 	 */
 	preferred_gfp = gfp | __GFP_NOWARN;
 	preferred_gfp &= ~(__GFP_DIRECT_RECLAIM | __GFP_NOFAIL);
-	page = __alloc_frozen_pages_noprof(preferred_gfp, order, nid, nodemask);
+	page = __alloc_pages_memdesc_noprof(preferred_gfp, order, memdesc,
+			nid, nodemask);
 	if (!page)
-		page = __alloc_frozen_pages_noprof(gfp, order, nid, NULL);
+		page = __alloc_pages_memdesc_noprof(gfp, order, memdesc,
+				nid, NULL);
 
 	return page;
 }
@@ -2362,6 +2364,7 @@ static struct page *alloc_pages_preferred_many(gfp_t gfp, unsigned int order,
  * alloc_pages_mpol - Allocate pages according to NUMA mempolicy.
  * @gfp: GFP flags.
  * @order: Order of the page allocation.
+ * @memdesc: Memory descriptor.
  * @pol: Pointer to the NUMA mempolicy.
  * @ilx: Index for interleave mempolicy (also distinguishes alloc_pages()).
  * @nid: Preferred node (usually numa_node_id() but @mpol may override it).
@@ -2369,7 +2372,7 @@ static struct page *alloc_pages_preferred_many(gfp_t gfp, unsigned int order,
  * Return: The page on success or NULL if allocation fails.
  */
 static struct page *alloc_pages_mpol(gfp_t gfp, unsigned int order,
-		struct mempolicy *pol, pgoff_t ilx, int nid)
+		memdesc_t memdesc, struct mempolicy *pol, pgoff_t ilx, int nid)
 {
 	nodemask_t *nodemask;
 	struct page *page;
@@ -2377,7 +2380,7 @@ static struct page *alloc_pages_mpol(gfp_t gfp, unsigned int order,
 	nodemask = policy_nodemask(gfp, pol, ilx, &nid);
 
 	if (pol->mode == MPOL_PREFERRED_MANY)
-		return alloc_pages_preferred_many(gfp, order, nid, nodemask);
+		return alloc_pages_preferred_many(gfp, order, memdesc, nid, nodemask);
 
 	if (IS_ENABLED(CONFIG_TRANSPARENT_HUGEPAGE) &&
 	    /* filter "hugepage" allocation, unless from alloc_pages() */
@@ -2399,9 +2402,9 @@ static struct page *alloc_pages_mpol(gfp_t gfp, unsigned int order,
 			 * First, try to allocate THP only on local node, but
 			 * don't reclaim unnecessarily, just compact.
 			 */
-			page = __alloc_frozen_pages_noprof(
+			page = __alloc_pages_memdesc_noprof(
 				gfp | __GFP_THISNODE | __GFP_NORETRY, order,
-				nid, NULL);
+				memdesc, nid, NULL);
 			if (page || !(gfp & __GFP_DIRECT_RECLAIM))
 				return page;
 			/*
@@ -2413,7 +2416,7 @@ static struct page *alloc_pages_mpol(gfp_t gfp, unsigned int order,
 		}
 	}
 
-	page = __alloc_frozen_pages_noprof(gfp, order, nid, nodemask);
+	page = __alloc_pages_memdesc_noprof(gfp, order, memdesc, nid, nodemask);
 
 	if (unlikely(pol->mode == MPOL_INTERLEAVE ||
 		     pol->mode == MPOL_WEIGHTED_INTERLEAVE) && page) {
@@ -2432,8 +2435,8 @@ static struct page *alloc_pages_mpol(gfp_t gfp, unsigned int order,
 struct folio *folio_alloc_mpol_noprof(gfp_t gfp, unsigned int order,
 		struct mempolicy *pol, pgoff_t ilx, int nid)
 {
-	struct page *page = alloc_pages_mpol(gfp | __GFP_COMP, order, pol,
-			ilx, nid);
+	struct page *page = alloc_pages_mpol(gfp | __GFP_COMP, order,
+			memdesc_create(NULL, 0), pol, ilx, nid);
 	if (!page)
 		return NULL;
 
@@ -2473,7 +2476,8 @@ struct folio *vma_alloc_folio_noprof(gfp_t gfp, int order, struct vm_area_struct
 }
 EXPORT_SYMBOL(vma_alloc_folio_noprof);
 
-struct page *alloc_frozen_pages_noprof(gfp_t gfp, unsigned order)
+struct page *alloc_pages_memdesc_noprof(gfp_t gfp, unsigned order,
+		memdesc_t memdesc)
 {
 	struct mempolicy *pol = &default_policy;
 
@@ -2484,7 +2488,7 @@ struct page *alloc_frozen_pages_noprof(gfp_t gfp, unsigned order)
 	if (!in_interrupt() && !(gfp & __GFP_THISNODE))
 		pol = get_task_policy(current);
 
-	return alloc_pages_mpol(gfp, order, pol, NO_INTERLEAVE_INDEX,
+	return alloc_pages_mpol(gfp, order, memdesc, pol, NO_INTERLEAVE_INDEX,
 				       numa_node_id());
 }
 
diff --git a/mm/page_alloc.c b/mm/page_alloc.c
index 600d9e981c23..c1451ca0acc1 100644
--- a/mm/page_alloc.c
+++ b/mm/page_alloc.c
@@ -5144,8 +5144,8 @@ EXPORT_SYMBOL_GPL(alloc_pages_bulk_noprof);
 /*
  * This is the 'heart' of the zoned buddy allocator.
  */
-struct page *__alloc_frozen_pages_noprof(gfp_t gfp, unsigned int order,
-		int preferred_nid, nodemask_t *nodemask)
+struct page *__alloc_pages_memdesc_noprof(gfp_t gfp, unsigned int order,
+		memdesc_t memdesc, int preferred_nid, nodemask_t *nodemask)
 {
 	struct page *page;
 	unsigned int alloc_flags = ALLOC_WMARK_LOW;
@@ -5205,9 +5205,15 @@ struct page *__alloc_frozen_pages_noprof(gfp_t gfp, unsigned int order,
 	trace_mm_page_alloc(page, order, alloc_gfp, ac.migratetype);
 	kmsan_alloc_page(page, order, alloc_gfp);
 
+	if (page && memdesc.v) {
+		unsigned long i, max = 1UL << order;
+
+		for (i = 0; i < max; i++)
+			page->memdesc = memdesc;
+	}
 	return page;
 }
-EXPORT_SYMBOL(__alloc_frozen_pages_noprof);
+EXPORT_SYMBOL(__alloc_pages_memdesc_noprof);
 
 struct page *__alloc_pages_noprof(gfp_t gfp, unsigned int order,
 		int preferred_nid, nodemask_t *nodemask)
-- 
2.47.2