[RFC PATCH 19/47] mm: asi: Support for locally nonsensitive page allocations

[Junaid Shahid <junaids@google.com>]

A new GFP flag, __GFP_LOCAL_NONSENSITIVE, is added to allocate pages
that are considered non-sensitive within the context of the current
process, but sensitive in the context of other processes.

For these allocations, page->asi_mm is set to the current mm during
allocation. It must be set to the same value when the page is freed.
Though it can potentially be overwritten and used for some other
purpose in the meantime, as long as it is restored before freeing.

Signed-off-by: Junaid Shahid <junaids@google.com>


---
 include/linux/gfp.h            |  5 +++-
 include/linux/mm_types.h       | 17 ++++++++++--
 include/trace/events/mmflags.h |  1 +
 mm/page_alloc.c                | 47 ++++++++++++++++++++++++++++------
 tools/perf/builtin-kmem.c      |  1 +
 5 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/include/linux/gfp.h b/include/linux/gfp.h
index 07a99a463a34..2ab394adbda3 100644
--- a/include/linux/gfp.h
+++ b/include/linux/gfp.h
@@ -62,8 +62,10 @@ struct vm_area_struct;
 #endif
 #ifdef CONFIG_ADDRESS_SPACE_ISOLATION
 #define ___GFP_GLOBAL_NONSENSITIVE 0x4000000u
+#define ___GFP_LOCAL_NONSENSITIVE  0x8000000u
 #else
 #define ___GFP_GLOBAL_NONSENSITIVE 0
+#define ___GFP_LOCAL_NONSENSITIVE 0
 #endif
 /* If the above are modified, __GFP_BITS_SHIFT may need updating */
 
@@ -255,9 +257,10 @@ struct vm_area_struct;
 
 /* Allocate non-sensitive memory */
 #define __GFP_GLOBAL_NONSENSITIVE ((__force gfp_t)___GFP_GLOBAL_NONSENSITIVE)
+#define __GFP_LOCAL_NONSENSITIVE ((__force gfp_t)___GFP_LOCAL_NONSENSITIVE)
 
 /* Room for N __GFP_FOO bits */
-#define __GFP_BITS_SHIFT 27
+#define __GFP_BITS_SHIFT 28
 #define __GFP_BITS_MASK ((__force gfp_t)((1 << __GFP_BITS_SHIFT) - 1))
 
 /**
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 8624d2783661..f9702d070975 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -193,8 +193,21 @@ struct page {
 		struct rcu_head rcu_head;
 
 #ifdef CONFIG_ADDRESS_SPACE_ISOLATION
-		/* Links the pages_to_free_async list */
-		struct llist_node async_free_node;
+		struct {
+			/* Links the pages_to_free_async list */
+			struct llist_node async_free_node;
+
+			unsigned long _asi_pad_1;
+			unsigned long _asi_pad_2;
+
+			/*
+			 * Upon allocation of a locally non-sensitive page, set
+			 * to the allocating mm. Must be set to the same mm when
+			 * the page is freed. May potentially be overwritten in
+			 * the meantime, as long as it is restored before free.
+			 */
+			struct mm_struct *asi_mm;
+		};
 #endif
 	};
 
diff --git a/include/trace/events/mmflags.h b/include/trace/events/mmflags.h
index 96e61d838bec..c00b8a4e1968 100644
--- a/include/trace/events/mmflags.h
+++ b/include/trace/events/mmflags.h
@@ -51,6 +51,7 @@
 	{(unsigned long)__GFP_KSWAPD_RECLAIM,	"__GFP_KSWAPD_RECLAIM"},\
 	{(unsigned long)__GFP_ZEROTAGS,		"__GFP_ZEROTAGS"},	\
 	{(unsigned long)__GFP_SKIP_KASAN_POISON,"__GFP_SKIP_KASAN_POISON"},\
+	{(unsigned long)__GFP_LOCAL_NONSENSITIVE, "__GFP_LOCAL_NONSENSITIVE"},\
 	{(unsigned long)__GFP_GLOBAL_NONSENSITIVE, "__GFP_GLOBAL_NONSENSITIVE"}\
 
 #define show_gfp_flags(flags)						\
diff --git a/mm/page_alloc.c b/mm/page_alloc.c
index a4048fa1868a..01784bff2a80 100644
--- a/mm/page_alloc.c
+++ b/mm/page_alloc.c
@@ -5231,19 +5231,33 @@ early_initcall(asi_page_alloc_init);
 static int asi_map_alloced_pages(struct page *page, uint order, gfp_t gfp_mask)
 {
 	uint i;
+	struct asi *asi;
+
+	VM_BUG_ON((gfp_mask & (__GFP_GLOBAL_NONSENSITIVE |
+			      __GFP_LOCAL_NONSENSITIVE)) ==
+		  (__GFP_GLOBAL_NONSENSITIVE | __GFP_LOCAL_NONSENSITIVE));
 
 	if (!static_asi_enabled())
 		return 0;
 
+	if (!(gfp_mask & (__GFP_GLOBAL_NONSENSITIVE |
+			  __GFP_LOCAL_NONSENSITIVE)))
+		return 0;
+
 	if (gfp_mask & __GFP_GLOBAL_NONSENSITIVE) {
+		asi = ASI_GLOBAL_NONSENSITIVE;
 		for (i = 0; i < (1 << order); i++)
 			__SetPageGlobalNonSensitive(page + i);
-
-		return asi_map_gfp(ASI_GLOBAL_NONSENSITIVE, page_to_virt(page),
-				   PAGE_SIZE * (1 << order), gfp_mask);
+	} else {
+		asi = ASI_LOCAL_NONSENSITIVE;
+		for (i = 0; i < (1 << order); i++) {
+			__SetPageLocalNonSensitive(page + i);
+			page[i].asi_mm = current->mm;
+		}
 	}
 
-	return 0;
+	return asi_map_gfp(asi, page_to_virt(page),
+			   PAGE_SIZE * (1 << order), gfp_mask);
 }
 
 static bool asi_unmap_freed_pages(struct page *page, unsigned int order)
@@ -5251,18 +5265,28 @@ static bool asi_unmap_freed_pages(struct page *page, unsigned int order)
 	void *va;
 	size_t len;
 	bool async_flush_needed;
+	struct asi *asi;
+
+	VM_BUG_ON(PageGlobalNonSensitive(page) && PageLocalNonSensitive(page));
 
 	if (!static_asi_enabled())
 		return true;
 
-	if (!PageGlobalNonSensitive(page))
+	if (PageGlobalNonSensitive(page))
+		asi = ASI_GLOBAL_NONSENSITIVE;
+	else if (PageLocalNonSensitive(page))
+		asi = &page->asi_mm->asi[0];
+	else
 		return true;
 
+	/* Heuristic to check that page->asi_mm is actually an mm_struct */
+	VM_BUG_ON(PageLocalNonSensitive(page) && asi->mm != page->asi_mm);
+
 	va = page_to_virt(page);
 	len = PAGE_SIZE * (1 << order);
 	async_flush_needed = irqs_disabled() || in_interrupt();
 
-	asi_unmap(ASI_GLOBAL_NONSENSITIVE, va, len, !async_flush_needed);
+	asi_unmap(asi, va, len, !async_flush_needed);
 
 	if (!async_flush_needed)
 		return true;
@@ -5476,8 +5500,15 @@ struct page *__alloc_pages(gfp_t gfp, unsigned int order, int preferred_nid,
 		return NULL;
 	}
 
-	if (static_asi_enabled() && (gfp & __GFP_GLOBAL_NONSENSITIVE))
-		gfp |= __GFP_ZERO;
+	if (static_asi_enabled()) {
+		if ((gfp & __GFP_LOCAL_NONSENSITIVE) &&
+		    !mm_asi_enabled(current->mm))
+			gfp &= ~__GFP_LOCAL_NONSENSITIVE;
+
+		if (gfp & (__GFP_GLOBAL_NONSENSITIVE |
+			   __GFP_LOCAL_NONSENSITIVE))
+			gfp |= __GFP_ZERO;
+	}
 
 	gfp &= gfp_allowed_mask;
 	/*
diff --git a/tools/perf/builtin-kmem.c b/tools/perf/builtin-kmem.c
index 5857953cd5c1..a2337fc3404f 100644
--- a/tools/perf/builtin-kmem.c
+++ b/tools/perf/builtin-kmem.c
@@ -661,6 +661,7 @@ static const struct {
 	{ "__GFP_DIRECT_RECLAIM",	"DR" },
 	{ "__GFP_KSWAPD_RECLAIM",	"KR" },
 	{ "__GFP_GLOBAL_NONSENSITIVE",	"GNS" },
+	{ "__GFP_LOCAL_NONSENSITIVE",	"LNS" },
 };
 
 static size_t max_gfp_len;
-- 
2.35.1.473.g83b2b277ed-goog