[RFC PATCH v3 1/4] mm: allow smaller large folios to use lru_cache

["Barry Song (Xiaomi)" <baohua@kernel.org>]

For systems that primarily use smaller-order large folios, enabling the
lru_cache can help reduce lock contention.

For higher-order large folios, the number of folios involved is likely
to be smaller, making lock contention less significant.

This patch enables the lru_cache for large folios whose `nr_pages` is
smaller than `FOLIO_BATCH_SIZE`. To avoid holding too many pages in the
lru_cache, which could affect accounting and reclamation, we also limit
the total number of pages in the cache to `FOLIO_BATCH_SIZE`.

To track the number of pages, this patch adds an `unsigned short
nr_pages` field to `struct folio_batch`. It cannot overflow because the
batch contains at most `FOLIO_BATCH_SIZE` folios, each of which has fewer
than `FOLIO_BATCH_SIZE` pages.

For non-LRU caches, `folio_batch` only needs to track the number of
folios, so `nr_pages` is left at zero.

Signed-off-by: Barry Song (Xiaomi) <baohua@kernel.org>
---
 include/linux/folio_batch.h | 25 +++++++++++++++++++++++++
 mm/folio.c                  | 10 +++++++++-
 mm/internal.h               |  4 ++--
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/include/linux/folio_batch.h b/include/linux/folio_batch.h
index b45946adc50b..ffc7de091fa3 100644
--- a/include/linux/folio_batch.h
+++ b/include/linux/folio_batch.h
@@ -10,6 +10,7 @@
 #define _LINUX_FOLIO_BATCH_H
 
 #include <linux/types.h>
+#include <linux/mm.h>
 
 /* 31 pointers + header align the folio_batch structure to a power of two */
 #define FOLIO_BATCH_SIZE	31
@@ -28,6 +29,7 @@ struct folio;
 struct folio_batch {
 	unsigned char nr;
 	unsigned char i;
+	unsigned short nr_pages;
 	bool percpu_pvec_drained;
 	struct folio *folios[FOLIO_BATCH_SIZE];
 };
@@ -42,6 +44,7 @@ static inline void folio_batch_init(struct folio_batch *fbatch)
 {
 	fbatch->nr = 0;
 	fbatch->i = 0;
+	fbatch->nr_pages = 0;
 	fbatch->percpu_pvec_drained = false;
 }
 
@@ -49,6 +52,7 @@ static inline void folio_batch_reinit(struct folio_batch *fbatch)
 {
 	fbatch->nr = 0;
 	fbatch->i = 0;
+	fbatch->nr_pages = 0;
 }
 
 static inline unsigned int folio_batch_count(const struct folio_batch *fbatch)
@@ -78,6 +82,27 @@ static inline unsigned folio_batch_add(struct folio_batch *fbatch,
 	return folio_batch_space(fbatch);
 }
 
+/**
+ * folio_batch_add_lru_cache() - Add a folio to a batch of lru_cache
+ * @fbatch: The folio batch.
+ * @folio: The folio to add.
+ *
+ * The folio is added to the end of the batch.
+ * The batch must have previously been initialised using folio_batch_init().
+ *
+ * Return: 0 if the lru_cache is filled with more than FOLIO_BATCH_SIZE
+ *         pages; otherwise, the number of available slots.
+ */
+static inline unsigned folio_batch_add_lru_cache(struct folio_batch *fbatch,
+		struct folio *folio)
+{
+	fbatch->folios[fbatch->nr++] = folio;
+	fbatch->nr_pages += (unsigned short)folio_nr_pages(folio);
+	if (fbatch->nr_pages > FOLIO_BATCH_SIZE)
+		return 0;
+	return folio_batch_space(fbatch);
+}
+
 /**
  * folio_batch_next - Return the next folio to process.
  * @fbatch: The folio batch being processed.
diff --git a/mm/folio.c b/mm/folio.c
index 59c477120b9a..e5820d7263e8 100644
--- a/mm/folio.c
+++ b/mm/folio.c
@@ -219,7 +219,7 @@ static void __folio_batch_add_and_move(struct folio_batch __percpu *fbatch,
 	else
 		local_lock(&cpu_fbatches.lock);
 
-	if (!folio_batch_add(this_cpu_ptr(fbatch), folio) ||
+	if (!folio_batch_add_lru_cache(this_cpu_ptr(fbatch), folio) ||
 			!folio_may_be_lru_cached(folio) || lru_cache_disabled())
 		folio_batch_move_lru(this_cpu_ptr(fbatch), move_fn);
 
@@ -981,6 +981,7 @@ void folios_put_refs(struct folio_batch *folios, unsigned int *refs)
 	int i, j;
 	struct lruvec *lruvec = NULL;
 	unsigned long flags = 0;
+	unsigned long nr_pages = 0;
 
 	for (i = 0, j = 0; i < folios->nr; i++) {
 		struct folio *folio = folios->folios[i];
@@ -1020,6 +1021,7 @@ void folios_put_refs(struct folio_batch *folios, unsigned int *refs)
 
 		if (j != i)
 			folios->folios[j] = folio;
+		nr_pages += folio_nr_pages(folio);
 		j++;
 	}
 	if (lruvec)
@@ -1030,6 +1032,12 @@ void folios_put_refs(struct folio_batch *folios, unsigned int *refs)
 	}
 
 	folios->nr = j;
+	/*
+	 * For lru_cache, track the number of pages; for non-LRU caches,
+	 * folio_batch->nr_pages is always 0.
+	 */
+	if (folios->nr_pages > 0)
+		folios->nr_pages = nr_pages;
 	mem_cgroup_uncharge_folios(folios);
 	free_unref_folios(folios);
 }
diff --git a/mm/internal.h b/mm/internal.h
index 38b1165212c9..06adf78e13a2 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -48,9 +48,9 @@ static inline bool folio_may_be_lru_cached(const struct folio *folio)
 	/*
 	 * Holding PMD-sized folios in per-CPU LRU cache unbalances accounting.
 	 * Holding small numbers of low-order mTHP folios in per-CPU LRU cache
-	 * will be sensible, but nobody has implemented and tested that yet.
+	 * will be sensible.
 	 */
-	return !folio_test_large(folio);
+	return folio_nr_pages(folio) < FOLIO_BATCH_SIZE;
 }
 
 static inline void lru_cache_enable(void)
-- 
2.34.1