[RFC PATCH 1/2] mm/secretmem: try to restore large page mappings in direct map

[Lance Yang <lance.yang@linux.dev>]

From: Lance Yang <lance.yang@linux.dev>

secretmem removes the pages backing secretmem mappings from the direct map.

Removing one base page from the direct map can split the covering large
mapping down to PTE mappings. Repeated splits can leave more of the direct
map mapped with PTEs, meaning more TLB entries for the same range and
potentially more TLB pressure.

So let's try to restore large page mappings whenever secretmem restores a
folio to the direct map.

Tested-by: Xueyuan Chen <xueyuan.chen21@gmail.com>
Signed-off-by: Lance Yang <lance.yang@linux.dev>
---
 include/linux/set_memory.h |  6 ++++++
 mm/secretmem.c             | 12 ++++++++++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/include/linux/set_memory.h b/include/linux/set_memory.h
index 3030d9245f5a..ad2fa414a22d 100644
--- a/include/linux/set_memory.h
+++ b/include/linux/set_memory.h
@@ -58,6 +58,12 @@ static inline bool can_set_direct_map(void)
 #endif
 #endif /* CONFIG_ARCH_HAS_SET_DIRECT_MAP */
 
+#ifndef arch_try_collapse_direct_map
+static inline void arch_try_collapse_direct_map(struct page *page)
+{
+}
+#endif
+
 #ifdef CONFIG_X86_64
 int set_mce_nospec(unsigned long pfn);
 int clear_mce_nospec(unsigned long pfn);
diff --git a/mm/secretmem.c b/mm/secretmem.c
index 5f57ac4720d3..82634b99aa2c 100644
--- a/mm/secretmem.c
+++ b/mm/secretmem.c
@@ -47,6 +47,14 @@ bool secretmem_active(void)
 	return !!atomic_read(&secretmem_users);
 }
 
+static void secretmem_restore_direct_map(struct folio *folio)
+{
+	struct page *page = folio_page(folio, 0);
+
+	if (!set_direct_map_default_noflush(page))
+		arch_try_collapse_direct_map(page);
+}
+
 static vm_fault_t secretmem_fault(struct vm_fault *vmf)
 {
 	struct address_space *mapping = vmf->vma->vm_file->f_mapping;
@@ -87,7 +95,7 @@ static vm_fault_t secretmem_fault(struct vm_fault *vmf)
 			 * already happened when we marked the page invalid
 			 * which guarantees that this call won't fail
 			 */
-			set_direct_map_default_noflush(folio_page(folio, 0));
+			secretmem_restore_direct_map(folio);
 			folio_put(folio);
 			if (err == -EEXIST)
 				goto retry;
@@ -151,7 +159,7 @@ static int secretmem_migrate_folio(struct address_space *mapping,
 
 static void secretmem_free_folio(struct folio *folio)
 {
-	set_direct_map_default_noflush(folio_page(folio, 0));
+	secretmem_restore_direct_map(folio);
 	folio_zero_segment(folio, 0, folio_size(folio));
 }
 
-- 
2.49.0