From: Tangquan Zheng <zhengtangquan@oppo.com>

This patch introduces support for zsmalloc to store compressed objects
based on multi-pages. Previously, a large folio with nr_pages subpages
would undergo compression one by one, each at the granularity of
PAGE_SIZE. However, by compressing them at a larger granularity, we
can conserve both memory and CPU resources.

We define the granularity with a configuration option called
ZSMALLOC_MULTI_PAGES_ORDER, set to a default value of 2, which matches
the minimum order of anonymous mTHP. As a result, a large folio with
8 subpages will now be split into 2 parts instead of 8.

The introduction of the multi-pages feature necessitates the creation
of new size classes to accommodate it.

Signed-off-by: Tangquan Zheng <zhengtangquan@oppo.com>
Co-developed-by: Barry Song <baohua@kernel.org>
Signed-off-by: Barry Song <baohua@kernel.org>
Co-developed-by: Nanzhe Zhao <zhaonanzhe@xiaomi.com>
Signed-off-by: Nanzhe Zhao <zhaonanzhe@xiaomi.com>
---
 drivers/block/zram/zram_drv.c |   3 +-
 include/linux/zsmalloc.h      |  10 +-
 mm/Kconfig                    |  29 +++-
 mm/zsmalloc.c                 | 245 +++++++++++++++++++++++-----------
 4 files changed, 204 insertions(+), 83 deletions(-)

diff --git a/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index 4f479d905142..24eefcf143de 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -2024,7 +2024,8 @@ static bool zram_meta_alloc(struct zram *zram, u64 disksize)
 	}
 
 	if (!huge_class_size)
-		huge_class_size = zs_huge_class_size(zram->mem_pool);
+		huge_class_size = zs_huge_class_size(zram->mem_pool,
+						ZSMALLOC_TYPE_BASEPAGE);
 
 	lockdep_register_key(&zram->table_lock_key);
 	lockdep_init_map(&zram->table_lock_map, "zram->table[index].lock", &zram->table_lock_key, 0);
diff --git a/include/linux/zsmalloc.h b/include/linux/zsmalloc.h
index 8971eef333ba..738260358a9a 100644
--- a/include/linux/zsmalloc.h
+++ b/include/linux/zsmalloc.h
@@ -17,6 +17,14 @@
 #include <linux/compiler.h>
 #include <linux/types.h>
 
+enum zsmalloc_type {
+	ZSMALLOC_TYPE_BASEPAGE,
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+	ZSMALLOC_TYPE_MULTI_PAGES,
+#endif
+	ZSMALLOC_TYPE_MAX,
+};
+
 struct zs_pool_stats {
 	/* How many pages were migrated (freed) */
 	atomic_long_t pages_compacted;
@@ -55,7 +63,7 @@ void zs_free(struct zs_pool *pool, struct zs_handle *handle);
 void zs_rebind_handle(struct zs_pool *pool, struct zs_handle *dst,
 		      struct zs_handle *src);
 
-size_t zs_huge_class_size(struct zs_pool *pool);
+size_t zs_huge_class_size(struct zs_pool *pool, enum zsmalloc_type type);
 
 unsigned long zs_get_total_pages(struct zs_pool *pool);
 unsigned long zs_compact(struct zs_pool *pool);
diff --git a/mm/Kconfig b/mm/Kconfig
index afff1b25b543..b68720f6fb6b 100644
--- a/mm/Kconfig
+++ b/mm/Kconfig
@@ -140,14 +140,15 @@ config ZSMALLOC_STAT
 	  If unsure, say N.
 
 config ZSMALLOC_CHAIN_SIZE
-	int "Maximum number of physical pages per-zspage"
+	int "Maximum number of chain elements per-zspage"
 	default 8
 	range 4 16
 	help
-	  This option sets the upper limit on the number of physical pages
-	  that a zmalloc page (zspage) can consist of. The optimal zspage
-	  chain size is calculated for each size class during the
-	  initialization of the pool.
+	  This option sets the upper limit on the number of chain elements
+	  that a zsmalloc page (zspage) can consist of. Each element is a
+	  single page, or a multi-page allocation when ZSMALLOC_MULTI_PAGES
+	  is enabled. The optimal zspage chain size is calculated for each
+	  size class during the initialization of the pool.
 
 	  Changing this option can alter the characteristics of size classes,
 	  such as the number of pages per zspage and the number of objects
@@ -157,6 +158,24 @@ config ZSMALLOC_CHAIN_SIZE
 
 	  For more information, see zsmalloc documentation.
 
+config ZSMALLOC_MULTI_PAGES
+	bool "support zsmalloc multiple pages"
+	depends on ZSMALLOC && !HIGHMEM
+	help
+	  This option configures zsmalloc to support allocations larger than
+	  PAGE_SIZE, enabling compression across multiple pages. The size of
+	  these multiple pages is determined by the configured
+	  ZSMALLOC_MULTI_PAGES_ORDER.
+
+config ZSMALLOC_MULTI_PAGES_ORDER
+	int "zsmalloc multiple pages order"
+	default 2
+	range 1 9
+	depends on ZSMALLOC_MULTI_PAGES
+	help
+	  This option is used to configure zsmalloc to support the compression
+	  of multiple pages.
+
 endmenu
 
 menu "Slab allocator options"
diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
index 4a6ec43f418b..2d83fcc8610d 100644
--- a/mm/zsmalloc.c
+++ b/mm/zsmalloc.c
@@ -48,6 +48,12 @@
 
 #define ZSPAGE_MAGIC	0x58
 
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+#define ZSMALLOC_MULTI_PAGES_ORDER	(_AC(CONFIG_ZSMALLOC_MULTI_PAGES_ORDER, UL))
+#define ZSMALLOC_MULTI_PAGES_NR		(1 << ZSMALLOC_MULTI_PAGES_ORDER)
+#define ZSMALLOC_MULTI_PAGES_SIZE	(PAGE_SIZE * ZSMALLOC_MULTI_PAGES_NR)
+#endif
+
 /*
  * This must be power of 2 and greater than or equal to sizeof(link_free).
  * These two conditions ensure that any 'struct link_free' itself doesn't
@@ -114,7 +120,7 @@
 
 #define HUGE_BITS	1
 #define FULLNESS_BITS	4
-#define CLASS_BITS	8
+#define CLASS_BITS	9
 #define MAGIC_VAL_BITS	8
 
 #define ZS_MAX_PAGES_PER_ZSPAGE	(_AC(CONFIG_ZSMALLOC_CHAIN_SIZE, UL))
@@ -172,7 +178,11 @@ static_assert(ZS_OBJ_IDX_BITS >= ZS_PAGES_PER_ZSPAGE_BITS,
 	      "zsmalloc: ZS_MIN_ALLOC_SIZE would exceed ZS_MAX_ALLOC_SIZE");
 
 /* each chunk includes extra space to keep handle */
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+#define ZS_MAX_ALLOC_SIZE	ZSMALLOC_MULTI_PAGES_SIZE
+#else
 #define ZS_MAX_ALLOC_SIZE	PAGE_SIZE
+#endif
 
 /*
  * On systems with 4K page size, this gives 255 size classes! There is a
@@ -187,14 +197,45 @@ static_assert(ZS_OBJ_IDX_BITS >= ZS_PAGES_PER_ZSPAGE_BITS,
  *  ZS_MIN_ALLOC_SIZE and ZS_SIZE_CLASS_DELTA must be multiple of ZS_ALIGN
  *  (reason above)
  */
-#define ZS_SIZE_CLASS_DELTA	(PAGE_SIZE >> CLASS_BITS)
+#define ZS_SIZE_CLASS_DELTA	(PAGE_SIZE >> (CLASS_BITS - 1))
 #define _ZS_MIN_ALLOC_RAW \
 	DIV_ROUND_UP(ZS_MAX_PAGES_PER_ZSPAGE << PAGE_SHIFT, ZS_OBJ_IDX_MASK)
 #define ZS_MIN_ALLOC_SIZE \
 	MIN(ZS_MAX_ALLOC_SIZE, \
 	    MAX(32UL, ALIGN(_ZS_MIN_ALLOC_RAW, ZS_ALIGN)))
-#define ZS_SIZE_CLASSES	(DIV_ROUND_UP(ZS_MAX_ALLOC_SIZE - ZS_MIN_ALLOC_SIZE, \
-				      ZS_SIZE_CLASS_DELTA) + 1)
+
+#define ZS_PAGE_SIZE_CLASSES	(DIV_ROUND_UP(PAGE_SIZE - ZS_MIN_ALLOC_SIZE, \
+					    ZS_SIZE_CLASS_DELTA) + 1)
+
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+#define ZS_MULTI_PAGES_SIZE_CLASS_DELTA	(ZSMALLOC_MULTI_PAGES_SIZE >> (CLASS_BITS - 1))
+#define ZS_MULTI_PAGES_SIZE_CLASSES	(DIV_ROUND_UP(ZS_MAX_ALLOC_SIZE - PAGE_SIZE, \
+					      ZS_MULTI_PAGES_SIZE_CLASS_DELTA) + 1)
+#define ZS_SIZE_CLASSES		(ZS_PAGE_SIZE_CLASSES + ZS_MULTI_PAGES_SIZE_CLASSES)
+#else
+#define ZS_SIZE_CLASSES		(DIV_ROUND_UP(ZS_MAX_ALLOC_SIZE - ZS_MIN_ALLOC_SIZE, \
+					      ZS_SIZE_CLASS_DELTA) + 1)
+#endif
+
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+static inline unsigned int class_size_to_zs_order(unsigned long size)
+{
+	if (size > PAGE_SIZE)
+		return ZSMALLOC_MULTI_PAGES_ORDER;
+
+	return 0;
+}
+#else
+static inline unsigned int class_size_to_zs_order(unsigned long size)
+{
+	return 0;
+}
+#endif
+
+static inline unsigned long class_size_to_zs_size(unsigned long size)
+{
+	return PAGE_SIZE * (1 << class_size_to_zs_order(size));
+}
 
 /*
  * Pages are distinguished by the ratio of used memory (that is the ratio
@@ -230,7 +271,7 @@ struct zs_size_stat {
 static struct dentry *zs_stat_root;
 #endif
 
-static size_t huge_class_size;
+static size_t huge_class_size[ZSMALLOC_TYPE_MAX];
 
 struct size_class {
 	spinlock_t lock;
@@ -241,13 +282,20 @@ struct size_class {
 	 */
 	int size;
 	int objs_per_zspage;
-	/* Number of PAGE_SIZE sized pages to combine to form a 'zspage' */
-	int pages_per_zspage;
+	/* Number of zpdescs to combine to form a 'zspage' */
+	int zpdescs_per_zspage;
 
 	unsigned int index;
 	struct zs_size_stat stats;
 };
 
+/* Number of physical pages in a zspage */
+static inline unsigned long class_pages_per_zspage(const struct size_class *class)
+{
+	return (unsigned long)class->zpdescs_per_zspage <<
+	       class_size_to_zs_order(class->size);
+}
+
 /*
  * Placed within free objects to form a singly linked list.
  * For every zspage, zspage->freeobj gives head of this list.
@@ -295,27 +343,28 @@ struct zs_pool {
 
 static inline void zpdesc_inc_zone_page_state(struct zpdesc *zpdesc)
 {
-	inc_zone_page_state(zpdesc_page(zpdesc), NR_ZSPAGES);
+	__zone_stat_add_folio(zpdesc_folio(zpdesc), NR_ZSPAGES);
 }
 
 static inline void zpdesc_dec_zone_page_state(struct zpdesc *zpdesc)
 {
-	dec_zone_page_state(zpdesc_page(zpdesc), NR_ZSPAGES);
+	__zone_stat_sub_folio(zpdesc_folio(zpdesc), NR_ZSPAGES);
 }
 
-static inline struct zpdesc *alloc_zpdesc(gfp_t gfp, const int nid)
+static inline struct zpdesc *alloc_zpdesc(gfp_t gfp, const int nid,
+					  unsigned int order)
 {
-	struct page *page = alloc_pages_node(nid, gfp, 0);
+	struct page *page = alloc_pages_node(nid, gfp | __GFP_COMP, order);
 
 	return page_zpdesc(page);
 }
 
-static inline void free_zpdesc(struct zpdesc *zpdesc)
+static inline void free_zpdesc(struct zpdesc *zpdesc, unsigned int order)
 {
 	struct page *page = zpdesc_page(zpdesc);
 
 	/* PageZsmalloc is sticky until the page is freed to the buddy. */
-	__free_page(page);
+	__free_pages(page, order);
 }
 
 #define ZS_PAGE_UNLOCKED	0
@@ -420,7 +469,7 @@ static void zspage_write_unlock(struct zspage *zspage)
 	spin_unlock(&zsl->lock);
 }
 
-/* huge object: pages_per_zspage == 1 && maxobj_per_zspage == 1 */
+/* huge object: zpdescs_per_zspage == 1 && maxobj_per_zspage == 1 */
 static void SetZsHugePage(struct zspage *zspage)
 {
 	zspage->huge = 1;
@@ -535,6 +584,15 @@ static int get_size_class_index(int size)
 {
 	int idx = 0;
 
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+	if (size > PAGE_SIZE + ZS_HANDLE_SIZE) {
+		idx = ZS_PAGE_SIZE_CLASSES;
+		idx += DIV_ROUND_UP(size - PAGE_SIZE,
+				    ZS_MULTI_PAGES_SIZE_CLASS_DELTA);
+		return min_t(int, ZS_SIZE_CLASSES - 1, idx);
+	}
+#endif
+
 	if (likely(size > ZS_MIN_ALLOC_SIZE))
 		idx = DIV_ROUND_UP(size - ZS_MIN_ALLOC_SIZE,
 				ZS_SIZE_CLASS_DELTA);
@@ -594,11 +652,11 @@ static int zs_stats_size_show(struct seq_file *s, void *v)
 	unsigned long total_freeable = 0;
 	unsigned long inuse_totals[NR_FULLNESS_GROUPS] = {0, };
 
-	seq_printf(s, " %5s %5s %9s %9s %9s %9s %9s %9s %9s %9s %9s %9s %9s %13s %10s %10s %16s %8s\n",
-			"class", "size", "10%", "20%", "30%", "40%",
+	seq_printf(s, " %5s %5s %5s %9s %9s %9s %9s %9s %9s %9s %9s %9s %9s %9s %13s %10s %10s %16s %16s %8s\n",
+			"class", "size", "order", "10%", "20%", "30%", "40%",
 			"50%", "60%", "70%", "80%", "90%", "99%", "100%",
 			"obj_allocated", "obj_used", "pages_used",
-			"pages_per_zspage", "freeable");
+			"pages_per_zspage", "objs_per_zspage", "freeable");
 
 	for (i = 0; i < ZS_SIZE_CLASSES; i++) {
 
@@ -609,7 +667,8 @@ static int zs_stats_size_show(struct seq_file *s, void *v)
 
 		spin_lock(&class->lock);
 
-		seq_printf(s, " %5u %5u ", i, class->size);
+		seq_printf(s, " %5u %5u %5u", i, class->size,
+			   class_size_to_zs_order(class->size));
 		for (fg = ZS_INUSE_RATIO_10; fg < NR_FULLNESS_GROUPS; fg++) {
 			inuse_totals[fg] += class_stat_read(class, fg);
 			seq_printf(s, "%9lu ", class_stat_read(class, fg));
@@ -622,11 +681,12 @@ static int zs_stats_size_show(struct seq_file *s, void *v)
 
 		objs_per_zspage = class->objs_per_zspage;
 		pages_used = obj_allocated / objs_per_zspage *
-				class->pages_per_zspage;
+				class_pages_per_zspage(class);
 
-		seq_printf(s, "%13lu %10lu %10lu %16d %8lu\n",
+		seq_printf(s, "%13lu %10lu %10lu %16lu %16d %8lu\n",
 			   obj_allocated, obj_used, pages_used,
-			   class->pages_per_zspage, freeable);
+			   class_pages_per_zspage(class), objs_per_zspage,
+			   freeable);
 
 		total_objs += obj_allocated;
 		total_used_objs += obj_used;
@@ -634,13 +694,13 @@ static int zs_stats_size_show(struct seq_file *s, void *v)
 		total_freeable += freeable;
 	}
 
-	seq_printf(s, "\n %5s %5s ", "Total", "");
+	seq_printf(s, "\n %5s %5s %5s ", "Total", "", "");
 
 	for (fg = ZS_INUSE_RATIO_10; fg < NR_FULLNESS_GROUPS; fg++)
 		seq_printf(s, "%9lu ", inuse_totals[fg]);
 
-	seq_printf(s, "%13lu %10lu %10lu %16s %8lu\n",
-		   total_objs, total_used_objs, total_pages, "",
+	seq_printf(s, "%13lu %10lu %10lu %16s %16s %8lu\n",
+		   total_objs, total_used_objs, total_pages, "", "",
 		   total_freeable);
 
 	return 0;
@@ -915,7 +975,7 @@ static void __free_zspage(struct zs_pool *pool, struct size_class *class,
 	assert_spin_locked(&class->lock);
 	__free_zspage_lockless(zspage);
 	class_stat_sub(class, ZS_OBJS_ALLOCATED, class->objs_per_zspage);
-	atomic_long_sub(class->pages_per_zspage, &pool->pages_allocated);
+	atomic_long_sub(class_pages_per_zspage(class), &pool->pages_allocated);
 }
 
 static void free_zspage(struct zs_pool *pool, struct size_class *class,
@@ -943,6 +1003,7 @@ static void init_zspage(struct size_class *class, struct zspage *zspage)
 {
 	unsigned int freeobj = 1;
 	unsigned long off = 0;
+	unsigned long page_size = class_size_to_zs_size(class->size);
 	struct zpdesc *zpdesc = get_first_zpdesc(zspage);
 
 	while (zpdesc) {
@@ -955,7 +1016,7 @@ static void init_zspage(struct size_class *class, struct zspage *zspage)
 		vaddr = kmap_local_zpdesc(zpdesc);
 		link = (struct link_free *)vaddr + off / sizeof(*link);
 
-		while ((off += class->size) < PAGE_SIZE) {
+		while ((off += class->size) < page_size) {
 			link->next = freeobj++ << OBJ_TAG_BITS;
 			link += class->size / sizeof(*link);
 		}
@@ -977,7 +1038,7 @@ static void init_zspage(struct size_class *class, struct zspage *zspage)
 		}
 		kunmap_local(vaddr);
 		zpdesc = next_zpdesc;
-		off %= PAGE_SIZE;
+		off %= page_size;
 	}
 
 	set_freeobj(zspage, 0);
@@ -989,7 +1050,7 @@ static void create_page_chain(struct size_class *class, struct zspage *zspage,
 	int i;
 	struct zpdesc *zpdesc;
 	struct zpdesc *prev_zpdesc = NULL;
-	int nr_zpdescs = class->pages_per_zspage;
+	int nr_zpdescs = class->zpdescs_per_zspage;
 
 	/*
 	 * Allocate individual pages and link them together as:
@@ -1006,7 +1067,7 @@ static void create_page_chain(struct size_class *class, struct zspage *zspage,
 		if (i == 0) {
 			zspage->first_zpdesc = zpdesc;
 			if (unlikely(class->objs_per_zspage == 1 &&
-					class->pages_per_zspage == 1))
+					class->zpdescs_per_zspage == 1))
 				SetZsHugePage(zspage);
 		} else {
 			prev_zpdesc->next = zpdesc;
@@ -1023,6 +1084,7 @@ static struct zspage *alloc_zspage(struct zs_pool *pool,
 				   gfp_t gfp, const int nid)
 {
 	int i;
+	unsigned int order = class_size_to_zs_order(class->size);
 	struct zpdesc *zpdescs[ZS_MAX_PAGES_PER_ZSPAGE];
 	struct zspage *zspage = cache_alloc_zspage(gfp);
 
@@ -1032,19 +1094,24 @@ static struct zspage *alloc_zspage(struct zs_pool *pool,
 	if (!IS_ENABLED(CONFIG_COMPACTION))
 		gfp &= ~__GFP_MOVABLE;
 
+	/*
+	 * Multi-page zspages are not movable, keep them out of ZONE_MOVABLE.
+	 */
+	if (order)
+		gfp &= ~__GFP_MOVABLE;
+
 	zspage->magic = ZSPAGE_MAGIC;
 	zspage->pool = pool;
 	zspage->class = class->index;
 	zspage_lock_init(zspage);
 
-	for (i = 0; i < class->pages_per_zspage; i++) {
-		struct zpdesc *zpdesc;
+	for (i = 0; i < class->zpdescs_per_zspage; i++) {
+		struct zpdesc *zpdesc = alloc_zpdesc(gfp, nid, order);
 
-		zpdesc = alloc_zpdesc(gfp, nid);
 		if (!zpdesc) {
 			while (--i >= 0) {
 				zpdesc_dec_zone_page_state(zpdescs[i]);
-				free_zpdesc(zpdescs[i]);
+				free_zpdesc(zpdescs[i], order);
 			}
 			cache_free_zspage(zspage);
 			return NULL;
@@ -1076,10 +1143,10 @@ static struct zspage *find_get_zspage(struct size_class *class)
 	return zspage;
 }
 
-static bool can_merge(struct size_class *prev, int pages_per_zspage,
+static bool can_merge(struct size_class *prev, int zpdescs_per_zspage,
 					int objs_per_zspage)
 {
-	if (prev->pages_per_zspage == pages_per_zspage &&
+	if (prev->zpdescs_per_zspage == zpdescs_per_zspage &&
 		prev->objs_per_zspage == objs_per_zspage)
 		return true;
 
@@ -1128,7 +1195,7 @@ void zs_obj_read_sg_begin(struct zs_pool *pool, const struct zs_handle *handle,
 {
 	struct zspage *zspage;
 	struct zpdesc *zpdesc;
-	unsigned long obj, off;
+	unsigned long obj, off, page_size;
 	unsigned int obj_idx;
 	struct size_class *class;
 
@@ -1143,12 +1210,13 @@ void zs_obj_read_sg_begin(struct zs_pool *pool, const struct zs_handle *handle,
 	read_unlock(&pool->lock);
 
 	class = zspage_class(pool, zspage);
-	off = offset_in_page(class->size * obj_idx);
+	page_size = class_size_to_zs_size(class->size);
+	off = (class->size * obj_idx) & (page_size - 1);
 
 	if (!ZsHugePage(zspage))
 		off += ZS_HANDLE_SIZE;
 
-	if (off + mem_len <= PAGE_SIZE) {
+	if (off + mem_len <= page_size) {
 		/* this object is contained entirely within a page */
 		sg_init_table(sg, 1);
 		sg_set_page(sg, zpdesc_page(zpdesc), mem_len, off);
@@ -1156,7 +1224,7 @@ void zs_obj_read_sg_begin(struct zs_pool *pool, const struct zs_handle *handle,
 		size_t sizes[2];
 
 		/* this object spans two pages */
-		sizes[0] = PAGE_SIZE - off;
+		sizes[0] = page_size - off;
 		sizes[1] = mem_len - sizes[0];
 
 		sg_init_table(sg, 2);
@@ -1189,16 +1257,17 @@ static void obj_write(struct size_class *class, struct zspage *zspage,
 		      unsigned long obj, const void *src, size_t size)
 {
 	struct zpdesc *zpdesc;
+	unsigned long page_size = class_size_to_zs_size(class->size);
 	unsigned long off;
 	unsigned int obj_idx;
 
 	obj_to_location(obj, &zpdesc, &obj_idx);
-	off = offset_in_page(class->size * obj_idx);
+	off = (class->size * obj_idx) & (page_size - 1);
 
 	if (!ZsHugePage(zspage))
 		off += ZS_HANDLE_SIZE;
 
-	if (off + size <= PAGE_SIZE) {
+	if (off + size <= page_size) {
 		/* this object is contained entirely within a page */
 		void *dst = kmap_local_zpdesc(zpdesc);
 
@@ -1208,7 +1277,7 @@ static void obj_write(struct size_class *class, struct zspage *zspage,
 		/* this object spans two pages */
 		size_t sizes[2];
 
-		sizes[0] = PAGE_SIZE - off;
+		sizes[0] = page_size - off;
 		sizes[1] = size - sizes[0];
 
 		memcpy_to_page(zpdesc_page(zpdesc), off, src, sizes[0]);
@@ -1230,9 +1299,9 @@ static void obj_write(struct size_class *class, struct zspage *zspage,
  *
  * Return: the size (in bytes) of the first huge zsmalloc &size_class.
  */
-size_t zs_huge_class_size(struct zs_pool *pool)
+size_t zs_huge_class_size(struct zs_pool *pool, enum zsmalloc_type type)
 {
-	return huge_class_size;
+	return huge_class_size[type];
 }
 EXPORT_SYMBOL_GPL(zs_huge_class_size);
 
@@ -1245,15 +1314,17 @@ static unsigned long obj_malloc(struct zs_pool *pool,
 	struct size_class *class;
 
 	struct zpdesc *m_zpdesc;
-	unsigned long m_offset;
+	unsigned long m_offset, page_size, page_shift;
 	void *vaddr;
 
 	class = pool->size_class[zspage->class];
 	obj = get_freeobj(zspage);
 
+	page_size = class_size_to_zs_size(class->size);
+	page_shift = PAGE_SHIFT + class_size_to_zs_order(class->size);
 	offset = obj * class->size;
-	nr_zpdesc = offset >> PAGE_SHIFT;
-	m_offset = offset_in_page(offset);
+	nr_zpdesc = offset >> page_shift;
+	m_offset = offset & (page_size - 1);
 	m_zpdesc = get_first_zpdesc(zspage);
 
 	for (i = 0; i < nr_zpdesc; i++)
@@ -1348,12 +1419,13 @@ int zs_store(struct zs_pool *pool, struct zs_handle *handle, const void *src,
 	spin_lock(&class->lock);
 	newfg = get_fullness_group(class, zspage);
 	insert_zspage(class, zspage, newfg);
-	atomic_long_add(class->pages_per_zspage, &pool->pages_allocated);
+	atomic_long_add(class_pages_per_zspage(class), &pool->pages_allocated);
 	class_stat_add(class, ZS_OBJS_ALLOCATED, class->objs_per_zspage);
 	class_stat_add(class, ZS_OBJS_INUSE, 1);
 
 	/* We completely set up zspage so mark them as movable */
-	SetZsPageMovable(pool, zspage);
+	if (class_size_to_zs_order(class->size) == 0)
+		SetZsPageMovable(pool, zspage);
 	spin_unlock(&class->lock);
 
 	return 0;
@@ -1365,13 +1437,15 @@ static void obj_free(int class_size, unsigned long obj)
 	struct link_free *link;
 	struct zspage *zspage;
 	struct zpdesc *f_zpdesc;
-	unsigned long f_offset;
+	unsigned long f_offset, page_size, page_mask;
 	unsigned int f_objidx;
 	void *vaddr;
 
 
 	obj_to_location(obj, &f_zpdesc, &f_objidx);
-	f_offset = offset_in_page(class_size * f_objidx);
+	page_size = class_size_to_zs_size(class_size);
+	page_mask = ~(page_size - 1);
+	f_offset = (class_size * f_objidx) & ~page_mask;
 	zspage = get_zspage(f_zpdesc);
 
 	vaddr = kmap_local_zpdesc(f_zpdesc);
@@ -1482,7 +1556,7 @@ void zs_free(struct zs_pool *pool, struct zs_handle *handle)
 
 	if (zspage_to_free) {
 		__free_zspage_lockless(zspage_to_free);
-		atomic_long_sub(class->pages_per_zspage, &pool->pages_allocated);
+		atomic_long_sub(class_pages_per_zspage(class), &pool->pages_allocated);
 	}
 }
 EXPORT_SYMBOL_GPL(zs_free);
@@ -1540,6 +1614,8 @@ static void zs_object_copy(struct size_class *class, unsigned long dst,
 	struct zpdesc *s_zpdesc, *d_zpdesc;
 	unsigned int s_objidx, d_objidx;
 	unsigned long s_off, d_off;
+	unsigned long page_size = class_size_to_zs_size(class->size);
+	unsigned long page_mask = ~(page_size - 1);
 	void *s_addr, *d_addr;
 	int s_size, d_size, size;
 	int written = 0;
@@ -1549,14 +1625,14 @@ static void zs_object_copy(struct size_class *class, unsigned long dst,
 	obj_to_location(src, &s_zpdesc, &s_objidx);
 	obj_to_location(dst, &d_zpdesc, &d_objidx);
 
-	s_off = offset_in_page(class->size * s_objidx);
-	d_off = offset_in_page(class->size * d_objidx);
+	s_off = (class->size * s_objidx) & ~page_mask;
+	d_off = (class->size * d_objidx) & ~page_mask;
 
-	if (s_off + class->size > PAGE_SIZE)
-		s_size = PAGE_SIZE - s_off;
+	if (s_off + class->size > page_size)
+		s_size = page_size - s_off;
 
-	if (d_off + class->size > PAGE_SIZE)
-		d_size = PAGE_SIZE - d_off;
+	if (d_off + class->size > page_size)
+		d_size = page_size - d_off;
 
 	s_addr = kmap_local_zpdesc(s_zpdesc);
 	d_addr = kmap_local_zpdesc(d_zpdesc);
@@ -1581,7 +1657,7 @@ static void zs_object_copy(struct size_class *class, unsigned long dst,
 		 * kunmap_local(d_addr). For more details see
 		 * Documentation/mm/highmem.rst.
 		 */
-		if (s_off >= PAGE_SIZE) {
+		if (s_off >= page_size) {
 			kunmap_local(d_addr);
 			kunmap_local(s_addr);
 			s_zpdesc = get_next_zpdesc(s_zpdesc);
@@ -1591,7 +1667,7 @@ static void zs_object_copy(struct size_class *class, unsigned long dst,
 			s_off = 0;
 		}
 
-		if (d_off >= PAGE_SIZE) {
+		if (d_off >= page_size) {
 			kunmap_local(d_addr);
 			d_zpdesc = get_next_zpdesc(d_zpdesc);
 			d_addr = kmap_local_zpdesc(d_zpdesc);
@@ -1615,13 +1691,13 @@ static struct zs_handle *find_alloced_obj(struct size_class *class,
 	int index = *obj_idx;
 	struct zs_handle *handle = NULL;
 	void *addr = kmap_local_zpdesc(zpdesc);
+	unsigned long page_size = class_size_to_zs_size(class->size);
 
 	offset = get_first_obj_offset(zpdesc);
 	offset += class->size * index;
 
-	while (offset < PAGE_SIZE) {
-		handle = obj_allocated(zpdesc, addr + offset);
-		if (handle)
+	while (offset < page_size) {
+		if (obj_allocated(zpdesc, addr + offset, &handle))
 			break;
 
 		offset += class->size;
@@ -2002,7 +2078,7 @@ static unsigned long zs_can_compact(struct size_class *class)
 	obj_wasted = obj_allocated - obj_used;
 	obj_wasted /= class->objs_per_zspage;
 
-	return obj_wasted * class->pages_per_zspage;
+	return obj_wasted * class_pages_per_zspage(class);
 }
 
 static unsigned long __zs_compact(struct zs_pool *pool,
@@ -2041,7 +2117,7 @@ static unsigned long __zs_compact(struct zs_pool *pool,
 		fg = putback_zspage(class, src_zspage);
 		if (fg == ZS_INUSE_RATIO_0) {
 			free_zspage(pool, class, src_zspage);
-			pages_freed += class->pages_per_zspage;
+			pages_freed += class_pages_per_zspage(class);
 		}
 		src_zspage = NULL;
 
@@ -2164,6 +2240,7 @@ static int calculate_zspage_chain_size(int class_size)
 {
 	int i, min_waste = INT_MAX;
 	int chain_size = 1;
+	unsigned long page_size = class_size_to_zs_size(class_size);
 
 	if (is_power_of_2(class_size))
 		return chain_size;
@@ -2171,7 +2248,7 @@ static int calculate_zspage_chain_size(int class_size)
 	for (i = 1; i <= ZS_MAX_PAGES_PER_ZSPAGE; i++) {
 		int waste;
 
-		waste = (i * PAGE_SIZE) % class_size;
+		waste = (i * page_size) % class_size;
 		if (waste < min_waste) {
 			min_waste = waste;
 			chain_size = i;
@@ -2215,16 +2292,32 @@ struct zs_pool *zs_create_pool(const char *name)
 	 */
 	for (i = ZS_SIZE_CLASSES - 1; i >= 0; i--) {
 		int size;
-		int pages_per_zspage;
+		int zpdescs_per_zspage;
 		int objs_per_zspage;
 		struct size_class *class;
 		int fullness;
+		int idx = ZSMALLOC_TYPE_BASEPAGE;
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+		int order = 0;
+#endif
 
-		size = ZS_MIN_ALLOC_SIZE + i * ZS_SIZE_CLASS_DELTA;
+		if (i < ZS_PAGE_SIZE_CLASSES)
+			size = ZS_MIN_ALLOC_SIZE + i * ZS_SIZE_CLASS_DELTA;
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+		if (i >= ZS_PAGE_SIZE_CLASSES)
+			size = PAGE_SIZE + (i - ZS_PAGE_SIZE_CLASSES) *
+			       ZS_MULTI_PAGES_SIZE_CLASS_DELTA;
+#endif
 		if (size > ZS_MAX_ALLOC_SIZE)
 			size = ZS_MAX_ALLOC_SIZE;
-		pages_per_zspage = calculate_zspage_chain_size(size);
-		objs_per_zspage = pages_per_zspage * PAGE_SIZE / size;
+#ifdef CONFIG_ZSMALLOC_MULTI_PAGES
+		order = class_size_to_zs_order(size);
+		if (order == ZSMALLOC_MULTI_PAGES_ORDER)
+			idx = ZSMALLOC_TYPE_MULTI_PAGES;
+#endif
+		zpdescs_per_zspage = calculate_zspage_chain_size(size);
+		objs_per_zspage = zpdescs_per_zspage *
+				  class_size_to_zs_size(size) / size;
 
 		/*
 		 * We iterate from biggest down to smallest classes,
@@ -2232,9 +2325,9 @@ struct zs_pool *zs_create_pool(const char *name)
 		 * class. Any object bigger than or equal to that will
 		 * endup in the huge class.
 		 */
-		if (pages_per_zspage != 1 && objs_per_zspage != 1 &&
-				!huge_class_size) {
-			huge_class_size = size;
+		if (zpdescs_per_zspage != 1 && objs_per_zspage != 1 &&
+				!huge_class_size[idx]) {
+			huge_class_size[idx] = size;
 			/*
 			 * The object uses ZS_HANDLE_SIZE bytes to store the
 			 * handle. We need to subtract it, because zs_store()
@@ -2244,7 +2337,7 @@ struct zs_pool *zs_create_pool(const char *name)
 			 * class because it grows by ZS_HANDLE_SIZE extra bytes
 			 * right before class lookup.
 			 */
-			huge_class_size -= (ZS_HANDLE_SIZE - 1);
+			huge_class_size[idx] -= (ZS_HANDLE_SIZE - 1);
 		}
 
 		/*
@@ -2257,7 +2350,7 @@ struct zs_pool *zs_create_pool(const char *name)
 		 * previous size_class if possible.
 		 */
 		if (prev_class) {
-			if (can_merge(prev_class, pages_per_zspage, objs_per_zspage)) {
+			if (can_merge(prev_class, zpdescs_per_zspage, objs_per_zspage)) {
 				pool->size_class[i] = prev_class;
 				continue;
 			}
@@ -2269,7 +2362,7 @@ struct zs_pool *zs_create_pool(const char *name)
 
 		class->size = size;
 		class->index = i;
-		class->pages_per_zspage = pages_per_zspage;
+		class->zpdescs_per_zspage = zpdescs_per_zspage;
 		class->objs_per_zspage = objs_per_zspage;
 		spin_lock_init(&class->lock);
 		pool->size_class[i] = class;
-- 
2.43.0

From: Tangquan Zheng <zhengtangquan@oppo.com>

Currently, when a large folio with nr_pages is submitted to zram, it is
divided into nr_pages parts for compression and storage individually.
By transitioning to a higher granularity, we can notably enhance
compression rates while simultaneously reducing CPU consumption.

This patch introduces the capability for large folios to be divided
based on the granularity specified by ZSMALLOC_MULTI_PAGES_ORDER, which
defaults to 2. For instance, for folios sized at 128KiB, compression
will occur in eight 16KiB multi-pages.

This modification will notably reduce CPU consumption and enhance
compression ratios. The following data illustrates the time and
compressed data for typical anonymous pages gathered from Android
phones.

Signed-off-by: Tangquan Zheng <zhengtangquan@oppo.com>
Co-developed-by: Barry Song <baohua@kernel.org>
Signed-off-by: Barry Song <baohua@kernel.org>
Co-developed-by: Nanzhe Zhao <zhaonanzhe@xiaomi.com>
Signed-off-by: Nanzhe Zhao <zhaonanzhe@xiaomi.com>
---
 drivers/block/zram/Kconfig    |   9 +
 drivers/block/zram/zcomp.c    |  31 ++-
 drivers/block/zram/zcomp.h    |  19 +-
 drivers/block/zram/zram_drv.c | 500 ++++++++++++++++++++++++++++++----
 drivers/block/zram/zram_drv.h |  43 +++
 mm/zsmalloc.c                 |   3 +-
 6 files changed, 549 insertions(+), 56 deletions(-)

diff --git a/drivers/block/zram/Kconfig b/drivers/block/zram/Kconfig
index 402b7b175863..716e92c5fdfe 100644
--- a/drivers/block/zram/Kconfig
+++ b/drivers/block/zram/Kconfig
@@ -145,3 +145,12 @@ config ZRAM_MULTI_COMP
 	  re-compress pages using a potentially slower but more effective
 	  compression algorithm. Note, that IDLE page recompression
 	  requires ZRAM_TRACK_ENTRY_ACTIME.
+
+config ZRAM_MULTI_PAGES
+	bool "Enable multiple pages compression and decompression"
+	depends on ZRAM && ZSMALLOC_MULTI_PAGES
+	help
+	  Initially, zram divided large folios into blocks of nr_pages, each sized
+	  equal to PAGE_SIZE, for compression. This option fine-tunes zram to
+	  improve compression granularity by dividing large folios into larger
+	  parts defined by the configuration option ZSMALLOC_MULTI_PAGES_ORDER.
diff --git a/drivers/block/zram/zcomp.c b/drivers/block/zram/zcomp.c
index ad8c7dc73922..8fa2a44adaa3 100644
--- a/drivers/block/zram/zcomp.c
+++ b/drivers/block/zram/zcomp.c
@@ -71,7 +71,7 @@ static int zcomp_cstrm_init(struct zcomp *comp, struct zcomp_cstrm *zstrm)
 	 * allocate 2 pages. 1 for compressed data, plus 1 extra for the
 	 * case when compressed size is larger than the original one
 	 */
-	zstrm->buffer = vzalloc(2 * PAGE_SIZE);
+	zstrm->buffer = vzalloc(2 * ZCOMP_MULTI_PAGES_SIZE);
 	if (!zstrm->buffer) {
 		zcomp_cstrm_free(comp, zstrm);
 		return -ENOMEM;
@@ -83,6 +83,10 @@ static void zcomp_dstrm_free(struct zcomp *comp, struct zcomp_dstrm *zstrm)
 {
 	if (comp->ops->destroy_dctx)
 		comp->ops->destroy_dctx(&zstrm->ctx);
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	vfree(zstrm->decomp_buf);
+	zstrm->decomp_buf = NULL;
+#endif
 	vfree(zstrm->local_copy);
 	zstrm->local_copy = NULL;
 	zstrm->ctx.context = NULL;
@@ -101,11 +105,23 @@ static int zcomp_dstrm_init(struct zcomp *comp, struct zcomp_dstrm *zstrm)
 		}
 	}
 
-	zstrm->local_copy = vzalloc(PAGE_SIZE);
+	zstrm->local_copy = vzalloc(ZCOMP_MULTI_PAGES_SIZE);
 	if (!zstrm->local_copy) {
 		zcomp_dstrm_free(comp, zstrm);
 		return -ENOMEM;
 	}
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	/*
+	 * Preallocate a decompression buffer for partial multi-page reads to
+	 * avoid memory allocation in the swapin path.
+	 */
+	zstrm->decomp_buf = vzalloc(ZCOMP_MULTI_PAGES_SIZE);
+	if (!zstrm->decomp_buf) {
+		zcomp_dstrm_free(comp, zstrm);
+		return -ENOMEM;
+	}
+	mutex_init(&zstrm->buf_lock);
+#endif
 	return 0;
 }
 
@@ -255,13 +271,13 @@ void zcomp_dstrm_put(struct zcomp *comp, struct zcomp_dstrm *zstrm,
 }
 
 int zcomp_compress(struct zcomp *comp, struct zcomp_cstrm *zstrm,
-		   const void *src, unsigned int *dst_len)
+		   const void *src, unsigned int src_len, unsigned int *dst_len)
 {
 	struct zcomp_req req = {
 		.src = src,
 		.dst = zstrm->buffer,
-		.src_len = PAGE_SIZE,
-		.dst_len = 2 * PAGE_SIZE,
+		.src_len = src_len,
+		.dst_len = 2 * src_len,
 	};
 	int ret;
 
@@ -282,12 +298,13 @@ static int zcomp_decompress_buf(struct zcomp *comp, struct zcomp_dstrm *zstrm,
 }
 
 int zcomp_decompress(struct zcomp *comp, struct zcomp_dstrm *zstrm,
-		     struct scatterlist *sg, unsigned int src_len, void *dst)
+		     struct scatterlist *sg, unsigned int src_len, void *dst,
+		     unsigned int dst_len)
 {
 	struct zcomp_req req = {
 		.dst = dst,
 		.src_len = src_len,
-		.dst_len = PAGE_SIZE,
+		.dst_len = dst_len,
 	};
 	void *src;
 	int ret;
diff --git a/drivers/block/zram/zcomp.h b/drivers/block/zram/zcomp.h
index 39efefc43a13..fee105f876b4 100644
--- a/drivers/block/zram/zcomp.h
+++ b/drivers/block/zram/zcomp.h
@@ -47,6 +47,15 @@ struct zcomp_ctx {
 	void *context;
 };
 
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+#define ZCOMP_MULTI_PAGES_ORDER	(_AC(CONFIG_ZSMALLOC_MULTI_PAGES_ORDER, UL))
+#define ZCOMP_MULTI_PAGES_NR	(1 << ZCOMP_MULTI_PAGES_ORDER)
+#define ZCOMP_MULTI_PAGES_SIZE	(PAGE_SIZE * ZCOMP_MULTI_PAGES_NR)
+#define MULTI_PAGE_SHIFT (ZCOMP_MULTI_PAGES_ORDER + PAGE_SHIFT)
+#else
+#define ZCOMP_MULTI_PAGES_SIZE	PAGE_SIZE
+#endif
+
 struct zcomp_cstrm {
 	struct mutex lock;
 	/* compression buffer */
@@ -65,6 +74,11 @@ struct zcomp_dstrm {
 	struct mutex excl_lock;
 	/* local copy of handle memory */
 	void *local_copy;
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	/* Temporary buffer and lock for partial multi-page decompression */
+	void *decomp_buf;
+	struct mutex buf_lock;
+#endif
 	struct zcomp_ctx ctx;
 };
 
@@ -124,8 +138,9 @@ void zcomp_dstrm_put(struct zcomp *comp, struct zcomp_dstrm *zstrm,
 		     enum zcomp_dstrm_mode mode);
 
 int zcomp_compress(struct zcomp *comp, struct zcomp_cstrm *zstrm,
-		   const void *src, unsigned int *dst_len);
+		   const void *src, unsigned int src_len, unsigned int *dst_len);
 int zcomp_decompress(struct zcomp *comp, struct zcomp_dstrm *zstrm,
-		     struct scatterlist *sg, unsigned int src_len, void *dst);
+		     struct scatterlist *sg, unsigned int src_len, void *dst,
+		     unsigned int dst_len);
 
 #endif /* _ZCOMP_H_ */
diff --git a/drivers/block/zram/zram_drv.c b/drivers/block/zram/zram_drv.c
index 24eefcf143de..0d33220f4b22 100644
--- a/drivers/block/zram/zram_drv.c
+++ b/drivers/block/zram/zram_drv.c
@@ -53,11 +53,12 @@ static unsigned int num_devices = 1;
  * Pages that compress to sizes equals or greater than this are stored
  * uncompressed in memory.
  */
-static size_t huge_class_size;
+static size_t huge_class_size[ZSMALLOC_TYPE_MAX];
 
 static const struct block_device_operations zram_devops;
 
 static void slot_free(struct zram *zram, unsigned long index);
+static void slot_free_npages(struct zram *zram, unsigned long index, int nr_pages);
 
 /*
  * entry locking rules:
@@ -321,11 +322,12 @@ static inline void zram_fill_page(void *ptr, unsigned long len,
 	memset_l(ptr, value, len / sizeof(unsigned long));
 }
 
-static bool page_same_filled(void *ptr, unsigned long *element)
+static bool page_same_filled(void *ptr, unsigned long *element,
+			     unsigned int page_size)
 {
 	unsigned long *page;
 	unsigned long val;
-	unsigned int pos, last_pos = PAGE_SIZE / sizeof(*page) - 1;
+	unsigned int pos, last_pos = page_size / sizeof(*page) - 1;
 
 	page = (unsigned long *)ptr;
 	val = page[0];
@@ -484,7 +486,7 @@ static ssize_t idle_store(struct device *dev, struct device_attribute *attr,
 static int read_from_zspool_raw(struct zram *zram, struct page *page,
 				unsigned long index);
 static int read_from_zspool(struct zram *zram, struct page *page,
-			    unsigned long index);
+			    unsigned long index, unsigned long page_size);
 
 struct zram_wb_ctl {
 	/* idle list is accessed only by the writeback task, no concurency */
@@ -1092,7 +1094,7 @@ static int zram_writeback_slots(struct zram *zram,
 		if (zram->compressed_wb)
 			err = read_from_zspool_raw(zram, req->page, index);
 		else
-			err = read_from_zspool(zram, req->page, index);
+			err = read_from_zspool(zram, req->page, index, PAGE_SIZE);
 		if (err)
 			goto next;
 		slot_unlock(zram, index);
@@ -1379,7 +1381,7 @@ static int decompress_bdev_page(struct zram *zram, struct page *page,
 
 	zstrm = zcomp_dstrm_get(zram->comps[prio], ZCOMP_DSTRM_EXCLUSIVE);
 	ret = zcomp_decompress(zram->comps[prio], zstrm, sg, size,
-			       zstrm->local_copy);
+			       zstrm->local_copy, PAGE_SIZE);
 	if (!ret) {
 		void *dst = kmap_local_page(page);
 
@@ -1989,6 +1991,31 @@ static ssize_t debug_stat_show(struct device *dev,
 	return ret;
 }
 
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+static ssize_t multi_pages_debug_stat_show(struct device *dev,
+		struct device_attribute *attr, char *buf)
+{
+	struct zram *zram = dev_to_zram(dev);
+	ssize_t ret;
+
+	guard(rwsem_read)(&zram->dev_lock);
+	ret = sysfs_emit(buf,
+			"zram_bio write/read multi_pages count:%8llu %8llu\n"
+			"zram_bio failed write/read multi_pages count%8llu %8llu\n"
+			"zram_bio partial write/read multi_pages count%8llu %8llu\n"
+			"multi_pages_miss_free %8llu\n",
+			(u64)atomic64_read(&zram->stats.zram_bio_write_multi_pages_count),
+			(u64)atomic64_read(&zram->stats.zram_bio_read_multi_pages_count),
+			(u64)atomic64_read(&zram->stats.multi_pages_failed_writes),
+			(u64)atomic64_read(&zram->stats.multi_pages_failed_reads),
+			(u64)atomic64_read(&zram->stats.zram_bio_write_multi_pages_partial_count),
+			(u64)atomic64_read(&zram->stats.zram_bio_read_multi_pages_partial_count),
+			(u64)atomic64_read(&zram->stats.multi_pages_miss_free));
+
+	return ret;
+}
+#endif
+
 static void zram_meta_free(struct zram *zram, u64 disksize)
 {
 	unsigned long num_pages = disksize >> PAGE_SHIFT;
@@ -2010,6 +2037,7 @@ static void zram_meta_free(struct zram *zram, u64 disksize)
 static bool zram_meta_alloc(struct zram *zram, u64 disksize)
 {
 	unsigned long num_pages;
+	int i;
 
 	num_pages = disksize >> PAGE_SHIFT;
 	zram->table = vzalloc(array_size(num_pages, sizeof(*zram->table)));
@@ -2023,9 +2051,10 @@ static bool zram_meta_alloc(struct zram *zram, u64 disksize)
 		return false;
 	}
 
-	if (!huge_class_size)
-		huge_class_size = zs_huge_class_size(zram->mem_pool,
-						ZSMALLOC_TYPE_BASEPAGE);
+	for (i = 0; i < ZSMALLOC_TYPE_MAX; i++) {
+		if (!huge_class_size[i])
+			huge_class_size[i] = zs_huge_class_size(zram->mem_pool, i);
+	}
 
 	lockdep_register_key(&zram->table_lock_key);
 	lockdep_init_map(&zram->table_lock_map, "zram->table[index].lock", &zram->table_lock_key, 0);
@@ -2033,10 +2062,18 @@ static bool zram_meta_alloc(struct zram *zram, u64 disksize)
 	return true;
 }
 
-static void slot_free(struct zram *zram, unsigned long index)
+static void slot_free_npages(struct zram *zram, unsigned long index,
+			     int nr_pages)
 {
 	struct zs_handle *handle;
 
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	if (test_slot_flag(zram, index, ZRAM_COMP_MULTI_PAGES)) {
+		clear_slot_flag(zram, index, ZRAM_COMP_MULTI_PAGES);
+		nr_pages = max(nr_pages, ZCOMP_MULTI_PAGES_NR);
+	}
+#endif
+
 #ifdef CONFIG_ZRAM_TRACK_ENTRY_ACTIME
 	zram->table[index].attr.ac_time = 0;
 #endif
@@ -2068,7 +2105,7 @@ static void slot_free(struct zram *zram, unsigned long index)
 	 */
 	if (test_slot_flag(zram, index, ZRAM_SAME)) {
 		clear_slot_flag(zram, index, ZRAM_SAME);
-		atomic64_dec(&zram->stats.same_pages);
+		atomic64_sub(nr_pages, &zram->stats.same_pages);
 		goto out;
 	}
 
@@ -2081,44 +2118,49 @@ static void slot_free(struct zram *zram, unsigned long index)
 	atomic64_sub(get_slot_size(zram, index),
 		     &zram->stats.compr_data_size);
 out:
-	atomic64_dec(&zram->stats.pages_stored);
+	atomic64_sub(nr_pages, &zram->stats.pages_stored);
 	zram->table[index].element = 0;
 	set_slot_size(zram, index, 0);
 }
 
+static void slot_free(struct zram *zram, unsigned long index)
+{
+	slot_free_npages(zram, index, 1);
+}
+
 static int read_same_filled_page(struct zram *zram, struct page *page,
-				 unsigned long index)
+				 unsigned long index, unsigned long page_size)
 {
 	void *mem;
 
 	mem = kmap_local_page(page);
-	zram_fill_page(mem, PAGE_SIZE, zram->table[index].element);
+	zram_fill_page(mem, page_size, zram->table[index].element);
 	kunmap_local(mem);
 	return 0;
 }
 
 static int read_incompressible_page(struct zram *zram, struct page *page,
-				    unsigned long index)
+				    unsigned long index,
+				    unsigned long page_size)
 {
 	struct scatterlist sg[2];
 	struct zs_handle *handle;
-	void *src, *dst;
+	void *dst;
 
 	handle = &zram->table[index].handle;
-	zs_obj_read_sg_begin(zram->mem_pool, handle, sg, PAGE_SIZE);
+	zs_obj_read_sg_begin(zram->mem_pool, handle, sg, page_size);
 	/* an incompressible object never spans two pages */
-	src = kmap_local_page(sg_page(sg));
 	dst = kmap_local_page(page);
-	copy_page(dst, src);
+	sg_copy_to_buffer(sg, sg_nents(sg), dst, page_size);
 	kunmap_local(dst);
-	kunmap_local(src);
 	zs_obj_read_sg_end(zram->mem_pool, handle);
 
 	return 0;
 }
 
 static int read_compressed_page(struct zram *zram, struct page *page,
-				unsigned long index)
+				unsigned long index,
+				unsigned long page_size)
 {
 	enum zcomp_dstrm_mode mode = ZCOMP_DSTRM_EXCLUSIVE;
 	struct zcomp_dstrm *zstrm;
@@ -2138,7 +2180,8 @@ static int read_compressed_page(struct zram *zram, struct page *page,
 
 	zstrm = zcomp_dstrm_get(zram->comps[prio], mode);
 	dst = kmap_local_page(page);
-	ret = zcomp_decompress(zram->comps[prio], zstrm, sg, size, dst);
+	ret = zcomp_decompress(zram->comps[prio], zstrm, sg, size, dst,
+			       page_size);
 	kunmap_local(dst);
 	zcomp_dstrm_put(zram->comps[prio], zstrm, mode);
 	zs_obj_read_sg_end(zram->mem_pool, handle);
@@ -2179,20 +2222,21 @@ static int read_from_zspool_raw(struct zram *zram, struct page *page,
  * Corresponding ZRAM slot should be locked.
  */
 static int read_from_zspool(struct zram *zram, struct page *page,
-			    unsigned long index)
+			    unsigned long index, unsigned long page_size)
 {
 	if (test_slot_flag(zram, index, ZRAM_SAME) ||
 	    !zs_handle_is_allocated(&zram->table[index].handle))
-		return read_same_filled_page(zram, page, index);
+		return read_same_filled_page(zram, page, index, page_size);
 
 	if (!test_slot_flag(zram, index, ZRAM_HUGE))
-		return read_compressed_page(zram, page, index);
+		return read_compressed_page(zram, page, index, page_size);
 	else
-		return read_incompressible_page(zram, page, index);
+		return read_incompressible_page(zram, page, index, page_size);
 }
 
-static int zram_read_page(struct zram *zram, struct page *page,
-			  unsigned long index, struct bio *parent)
+static int zram_read_page_size(struct zram *zram, struct page *page,
+			       unsigned long index, struct bio *parent,
+			       unsigned long page_size)
 {
 	int ret;
 
@@ -2200,7 +2244,7 @@ static int zram_read_page(struct zram *zram, struct page *page,
 	mark_slot_accessed(zram, index);
 	if (!test_slot_flag(zram, index, ZRAM_WB)) {
 		/* Slot should be locked through out the function call */
-		ret = read_from_zspool(zram, page, index);
+		ret = read_from_zspool(zram, page, index, page_size);
 		slot_unlock(zram, index);
 	} else {
 		unsigned long blk_idx = zram->table[index].blk_idx;
@@ -2216,6 +2260,12 @@ static int zram_read_page(struct zram *zram, struct page *page,
 	return ret;
 }
 
+static int zram_read_page(struct zram *zram, struct page *page,
+			  unsigned long index, struct bio *parent)
+{
+	return zram_read_page_size(zram, page, index, parent, PAGE_SIZE);
+}
+
 /*
  * Use a temporary buffer to decompress the page, as the decompressor
  * always expects a full page for the output.
@@ -2244,23 +2294,32 @@ static int zram_bvec_read(struct zram *zram, struct bio_vec *bvec,
 }
 
 static int write_same_filled_page(struct zram *zram, unsigned long fill,
-				  unsigned long index)
+				  unsigned long index, int nr_pages)
 {
 	slot_lock(zram, index);
 	slot_free(zram, index);
 	set_slot_flag(zram, index, ZRAM_SAME);
 	zram->table[index].element = fill;
 	mark_slot_accessed(zram, index);
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	if (nr_pages > 1) {
+		int i;
+
+		for (i = 0; i < nr_pages; i++)
+			set_slot_flag(zram, index + i, ZRAM_COMP_MULTI_PAGES);
+	}
+#endif
 	slot_unlock(zram, index);
 
-	atomic64_inc(&zram->stats.same_pages);
-	atomic64_inc(&zram->stats.pages_stored);
+	atomic64_add(nr_pages, &zram->stats.same_pages);
+	atomic64_add(nr_pages, &zram->stats.pages_stored);
 
 	return 0;
 }
 
 static int write_incompressible_page(struct zram *zram, struct page *page,
-				     unsigned long index)
+				     unsigned long index, unsigned int page_size,
+				     int nr_pages)
 {
 	void *src;
 	int ret;
@@ -2269,7 +2328,7 @@ static int write_incompressible_page(struct zram *zram, struct page *page,
 	slot_free(zram, index);
 	src = kmap_local_page(page);
 	ret = zs_store(zram->mem_pool, &zram->table[index].handle, src,
-		       PAGE_SIZE, GFP_NOIO | __GFP_NOWARN |
+		       page_size, GFP_NOIO | __GFP_NOWARN |
 		       __GFP_HIGHMEM | __GFP_MOVABLE, page_to_nid(page));
 	kunmap_local(src);
 	if (ret) {
@@ -2280,14 +2339,22 @@ static int write_incompressible_page(struct zram *zram, struct page *page,
 	update_used_max(zram);
 
 	set_slot_flag(zram, index, ZRAM_HUGE);
-	set_slot_size(zram, index, PAGE_SIZE);
+	set_slot_size(zram, index, page_size);
 	mark_slot_accessed(zram, index);
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	if (nr_pages > 1) {
+		int i;
+
+		for (i = 0; i < nr_pages; i++)
+			set_slot_flag(zram, index + i, ZRAM_COMP_MULTI_PAGES);
+	}
+#endif
 	slot_unlock(zram, index);
 
-	atomic64_add(PAGE_SIZE, &zram->stats.compr_data_size);
+	atomic64_add(page_size, &zram->stats.compr_data_size);
 	atomic64_inc(&zram->stats.huge_pages);
 	atomic64_inc(&zram->stats.huge_pages_since);
-	atomic64_inc(&zram->stats.pages_stored);
+	atomic64_add(nr_pages, &zram->stats.pages_stored);
 
 	return 0;
 }
@@ -2301,12 +2368,27 @@ static int zram_write_page(struct zram *zram, struct page *page,
 	struct zcomp_cstrm *zstrm;
 	unsigned long element;
 	bool same_filled;
+	unsigned long page_size = PAGE_SIZE;
+	int huge_class_idx = ZSMALLOC_TYPE_BASEPAGE;
+	int nr_pages = 1;
+
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	/*
+	 * Multi-page compression packs ZCOMP_MULTI_PAGES_NR pages into one
+	 * compressed object.
+	 */
+	if (folio_size(page_folio(page)) >= ZCOMP_MULTI_PAGES_SIZE) {
+		page_size = ZCOMP_MULTI_PAGES_SIZE;
+		huge_class_idx = ZSMALLOC_TYPE_MULTI_PAGES;
+		nr_pages = ZCOMP_MULTI_PAGES_NR;
+	}
+#endif
 
 	mem = kmap_local_page(page);
-	same_filled = page_same_filled(mem, &element);
+	same_filled = page_same_filled(mem, &element, page_size);
 	kunmap_local(mem);
 	if (same_filled)
-		return write_same_filled_page(zram, element, index);
+		return write_same_filled_page(zram, element, index, nr_pages);
 
 	if (!zram_can_store_page(zram))
 		return -ENOMEM;
@@ -2314,7 +2396,7 @@ static int zram_write_page(struct zram *zram, struct page *page,
 	zstrm = zcomp_cstrm_get(zram->comps[ZRAM_PRIMARY_COMP]);
 	mem = kmap_local_page(page);
 	ret = zcomp_compress(zram->comps[ZRAM_PRIMARY_COMP], zstrm,
-			     mem, &comp_len);
+			     mem, page_size, &comp_len);
 	kunmap_local(mem);
 
 	if (unlikely(ret)) {
@@ -2323,9 +2405,10 @@ static int zram_write_page(struct zram *zram, struct page *page,
 		return ret;
 	}
 
-	if (comp_len >= huge_class_size) {
+	if (comp_len >= huge_class_size[huge_class_idx]) {
 		zcomp_cstrm_put(zstrm);
-		return write_incompressible_page(zram, page, index);
+		return write_incompressible_page(zram, page, index, page_size,
+						 nr_pages);
 	}
 
 	slot_lock(zram, index);
@@ -2344,10 +2427,18 @@ static int zram_write_page(struct zram *zram, struct page *page,
 
 	set_slot_size(zram, index, comp_len);
 	mark_slot_accessed(zram, index);
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	if (nr_pages > 1) {
+		int i;
+
+		for (i = 0; i < nr_pages; i++)
+			set_slot_flag(zram, index + i, ZRAM_COMP_MULTI_PAGES);
+	}
+#endif
 	slot_unlock(zram, index);
 
 	/* Update stats */
-	atomic64_inc(&zram->stats.pages_stored);
+	atomic64_add(nr_pages, &zram->stats.pages_stored);
 	atomic64_add(comp_len, &zram->stats.compr_data_size);
 
 	return ret;
@@ -2467,7 +2558,7 @@ static int recompress_slot(struct zram *zram, unsigned long index,
 	if (comp_len_old < threshold)
 		return 0;
 
-	ret = read_from_zspool(zram, page, index);
+	ret = read_from_zspool(zram, page, index, PAGE_SIZE);
 	if (ret)
 		return ret;
 
@@ -2480,7 +2571,8 @@ static int recompress_slot(struct zram *zram, unsigned long index,
 
 	zstrm = zcomp_cstrm_get(zram->comps[prio]);
 	src = kmap_local_page(page);
-	ret = zcomp_compress(zram->comps[prio], zstrm, src, &comp_len_new);
+	ret = zcomp_compress(zram->comps[prio], zstrm, src, PAGE_SIZE,
+			      &comp_len_new);
 	kunmap_local(src);
 
 	/*
@@ -2615,7 +2707,7 @@ static ssize_t recompress_store(struct device *dev,
 		}
 	}
 
-	if (threshold >= huge_class_size)
+	if (threshold >= huge_class_size[ZSMALLOC_TYPE_BASEPAGE])
 		return -EINVAL;
 
 	guard(rwsem_write)(&zram->dev_lock);
@@ -2727,8 +2819,283 @@ static void zram_bio_discard(struct zram *zram, struct bio *bio)
 	bio_endio(bio);
 }
 
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+
+static inline unsigned long zram_multi_pages_head_index(unsigned long index)
+{
+	return index & ~((unsigned long)ZCOMP_MULTI_PAGES_NR - 1);
+}
+
+static inline int __test_multi_pages_comp(struct zram *zram,
+					 unsigned long index)
+{
+	unsigned long head = zram_multi_pages_head_index(index);
+	int i, count = 0;
+
+	for (i = 0; i < ZCOMP_MULTI_PAGES_NR; i++) {
+		if (test_slot_flag(zram, head + i, ZRAM_COMP_MULTI_PAGES))
+			count++;
+	}
+
+	return count;
+}
+
+static inline bool want_multi_pages_comp(struct zram *zram, struct bio *bio)
+{
+	unsigned long index = bio->bi_iter.bi_sector >> SECTORS_PER_PAGE_SHIFT;
+
+	if (bio->bi_io_vec->bv_len >= ZCOMP_MULTI_PAGES_SIZE)
+		return true;
+
+	slot_lock(zram, index);
+	if (__test_multi_pages_comp(zram, index)) {
+		slot_unlock(zram, index);
+		return true;
+	}
+	slot_unlock(zram, index);
+
+	return false;
+}
+
+static inline bool test_multi_pages_comp(struct zram *zram, struct bio *bio)
+{
+	unsigned long index = bio->bi_iter.bi_sector >> SECTORS_PER_PAGE_SHIFT;
+
+	return !!__test_multi_pages_comp(zram, index);
+}
+
+static inline bool is_multi_pages_partial_io(struct bio_vec *bvec)
+{
+	return bvec->bv_len != ZCOMP_MULTI_PAGES_SIZE;
+}
+
+static int zram_read_partial_from_zspool(struct zram *zram, struct page *page,
+					unsigned long index, int offset)
+{
+	struct zs_handle *handle = &zram->table[index].handle;
+	struct zcomp_dstrm *zstrm;
+	struct scatterlist sg[2];
+	enum zcomp_dstrm_mode mode;
+	unsigned int size;
+	void *src, *dst;
+	u32 prio;
+	int ret;
+	unsigned long page_size = ZCOMP_MULTI_PAGES_SIZE;
+
+	if (test_slot_flag(zram, index, ZRAM_SAME) ||
+	    !zs_handle_is_allocated(handle)) {
+		void *mem;
+
+		mem = kmap_local_page(page);
+		zram_fill_page(mem, PAGE_SIZE, zram->table[index].element);
+		kunmap_local(mem);
+		return 0;
+	}
+
+	size = get_slot_size(zram, index);
+	zs_obj_read_sg_begin(zram->mem_pool, handle, sg, size);
+
+	if (size == page_size) {
+		src = kmap_local_page(sg_page(sg));
+		dst = kmap_local_page(page);
+		memcpy(dst, src + offset, PAGE_SIZE);
+		kunmap_local(dst);
+		kunmap_local(src);
+		ret = 0;
+	} else {
+		prio = get_slot_comp_priority(zram, index);
+		mode = ZCOMP_DSTRM_EXCLUSIVE;
+		if (zcomp_shared_decomp(zram->comps[prio]))
+			mode = ZCOMP_DSTRM_SHARED;
+
+		zstrm = zcomp_dstrm_get(zram->comps[prio], mode);
+		mutex_lock(&zstrm->buf_lock);
+		ret = zcomp_decompress(zram->comps[prio], zstrm, sg, size,
+				       zstrm->decomp_buf, page_size);
+		if (!ret) {
+			dst = kmap_local_page(page);
+			memcpy(dst, zstrm->decomp_buf + offset, PAGE_SIZE);
+			kunmap_local(dst);
+		}
+		mutex_unlock(&zstrm->buf_lock);
+		zcomp_dstrm_put(zram->comps[prio], zstrm, mode);
+	}
+	zs_obj_read_sg_end(zram->mem_pool, handle);
+
+	return ret;
+}
+
+static int zram_bvec_read_multi_pages_partial(struct zram *zram,
+					      struct bio_vec *bvec,
+					      unsigned long index, int offset,
+					      struct bio *bio)
+{
+	int ret;
+
+	slot_lock(zram, index);
+	if (!test_slot_flag(zram, index, ZRAM_WB)) {
+		ret = zram_read_partial_from_zspool(zram, bvec->bv_page,
+						    index, offset);
+		slot_unlock(zram, index);
+	} else {
+		unsigned long blk_idx = zram->table[index].blk_idx;
+
+		slot_unlock(zram, index);
+		ret = read_from_bdev(zram, bvec->bv_page, index, blk_idx, bio);
+	}
+
+	if (WARN_ON(ret < 0))
+		pr_err("Decompression failed! err=%d, page=%lu offset=%d\n",
+		       ret, index, offset);
+	if (!ret)
+		atomic64_inc(&zram->stats.zram_bio_read_multi_pages_partial_count);
+
+	return ret;
+}
+
+static int zram_bvec_read_multi_pages(struct zram *zram, struct bio_vec *bvec,
+				      unsigned long index, int offset,
+				      struct bio *bio)
+{
+	if (is_multi_pages_partial_io(bvec))
+		return zram_bvec_read_multi_pages_partial(zram, bvec, index,
+							  offset, bio);
+
+	return zram_read_page_size(zram, bvec->bv_page, index, bio,
+				   ZCOMP_MULTI_PAGES_SIZE);
+}
+
+static int zram_bvec_write_multi_pages_partial(struct zram *zram,
+					       struct bio_vec *bvec,
+					       unsigned long index, int offset,
+					       struct bio *bio)
+{
+	struct page *page;
+	void *src, *dst;
+	int ret;
+
+	page = alloc_pages(GFP_NOIO | __GFP_COMP, ZCOMP_MULTI_PAGES_ORDER);
+	if (!page)
+		return -ENOMEM;
+
+	ret = zram_read_page_size(zram, page, index, bio,
+				  ZCOMP_MULTI_PAGES_SIZE);
+	if (!ret) {
+		src = kmap_local_page(bvec->bv_page);
+		dst = kmap_local_page(page);
+		memcpy(dst + offset, src + bvec->bv_offset, bvec->bv_len);
+		kunmap_local(dst);
+		kunmap_local(src);
+
+		atomic64_inc(&zram->stats.zram_bio_write_multi_pages_partial_count);
+		ret = zram_write_page(zram, page, index);
+	}
+
+	__free_pages(page, ZCOMP_MULTI_PAGES_ORDER);
+	return ret;
+}
+
+static int zram_bvec_write_multi_pages(struct zram *zram, struct bio_vec *bvec,
+				       unsigned long index, int offset,
+				       struct bio *bio)
+{
+	if (is_multi_pages_partial_io(bvec))
+		return zram_bvec_write_multi_pages_partial(zram, bvec, index,
+							   offset, bio);
+
+	return zram_write_page(zram, bvec->bv_page, index);
+}
+
+static void zram_bio_read_multi_pages(struct zram *zram, struct bio *bio)
+{
+	unsigned long start_time = bio_start_io_acct(bio);
+	struct bvec_iter iter = bio->bi_iter;
+
+	do {
+		unsigned long index = zram_multi_pages_head_index(
+			iter.bi_sector >> SECTORS_PER_PAGE_SHIFT);
+		u32 offset = (iter.bi_sector & (SECTORS_PER_MULTI_PAGE - 1)) <<
+			     SECTOR_SHIFT;
+		struct bio_vec bv = multi_pages_bio_iter_iovec(bio, iter);
+
+		atomic64_inc(&zram->stats.zram_bio_read_multi_pages_count);
+		bv.bv_len = min_t(u32, bv.bv_len,
+				  ZCOMP_MULTI_PAGES_SIZE - offset);
+
+		if (zram_bvec_read_multi_pages(zram, &bv, index, offset, bio) < 0) {
+			atomic64_inc(&zram->stats.multi_pages_failed_reads);
+			bio->bi_status = BLK_STS_IOERR;
+			break;
+		}
+		flush_dcache_page(bv.bv_page);
+
+		slot_lock(zram, index);
+		mark_slot_accessed(zram, index);
+		slot_unlock(zram, index);
+
+		bio_advance_iter_single(bio, &iter, bv.bv_len);
+	} while (iter.bi_size);
+
+	bio_end_io_acct(bio, start_time);
+	bio_endio(bio);
+}
+
+static void zram_bio_write_multi_pages(struct zram *zram, struct bio *bio)
+{
+	unsigned long start_time = bio_start_io_acct(bio);
+	struct bvec_iter iter = bio->bi_iter;
+
+	do {
+		unsigned long index = zram_multi_pages_head_index(
+			iter.bi_sector >> SECTORS_PER_PAGE_SHIFT);
+		u32 offset = (iter.bi_sector & (SECTORS_PER_MULTI_PAGE - 1)) <<
+			     SECTOR_SHIFT;
+		struct bio_vec bv = multi_pages_bio_iter_iovec(bio, iter);
+
+		bv.bv_len = min_t(u32, bv.bv_len,
+				  ZCOMP_MULTI_PAGES_SIZE - offset);
+
+		atomic64_inc(&zram->stats.zram_bio_write_multi_pages_count);
+		if (zram_bvec_write_multi_pages(zram, &bv, index, offset, bio) < 0) {
+			atomic64_inc(&zram->stats.multi_pages_failed_writes);
+			bio->bi_status = BLK_STS_IOERR;
+			break;
+		}
+
+		slot_lock(zram, index);
+		mark_slot_accessed(zram, index);
+		slot_unlock(zram, index);
+
+		bio_advance_iter_single(bio, &iter, bv.bv_len);
+	} while (iter.bi_size);
+
+	bio_end_io_acct(bio, start_time);
+	bio_endio(bio);
+}
+#else /* CONFIG_ZRAM_MULTI_PAGES */
+static inline bool test_multi_pages_comp(struct zram *zram, struct bio *bio)
+{
+	return false;
+}
+
+static inline bool want_multi_pages_comp(struct zram *zram, struct bio *bio)
+{
+	return false;
+}
+
+static inline void zram_bio_read_multi_pages(struct zram *zram,
+					     struct bio *bio) { }
+static inline void zram_bio_write_multi_pages(struct zram *zram,
+					      struct bio *bio) { }
+#endif
+
 static void zram_bio_read(struct zram *zram, struct bio *bio)
 {
+	if (test_multi_pages_comp(zram, bio)) {
+		zram_bio_read_multi_pages(zram, bio);
+		return;
+	}
+
 	unsigned long start_time = bio_start_io_acct(bio);
 	struct bvec_iter iter = bio->bi_iter;
 
@@ -2756,6 +3123,11 @@ static void zram_bio_read(struct zram *zram, struct bio *bio)
 
 static void zram_bio_write(struct zram *zram, struct bio *bio)
 {
+	if (want_multi_pages_comp(zram, bio)) {
+		zram_bio_write_multi_pages(zram, bio);
+		return;
+	}
+
 	unsigned long start_time = bio_start_io_acct(bio);
 	struct bvec_iter iter = bio->bi_iter;
 
@@ -2808,6 +3180,9 @@ static void zram_slot_free_notify(struct block_device *bdev,
 				unsigned long index)
 {
 	struct zram *zram;
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	int nr_pages_to_free = 1;
+#endif
 
 	zram = bdev->bd_disk->private_data;
 
@@ -2817,7 +3192,34 @@ static void zram_slot_free_notify(struct block_device *bdev,
 		return;
 	}
 
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	{
+		int comp_count = __test_multi_pages_comp(zram, index);
+
+		if (comp_count > 1) {
+			clear_slot_flag(zram, index, ZRAM_COMP_MULTI_PAGES);
+			slot_unlock(zram, index);
+			return;
+		}
+
+		if (comp_count == 1) {
+			clear_slot_flag(zram, index, ZRAM_COMP_MULTI_PAGES);
+			slot_unlock(zram, index);
+			index = zram_multi_pages_head_index(index);
+			nr_pages_to_free = ZCOMP_MULTI_PAGES_NR;
+			if (!slot_trylock(zram, index)) {
+				atomic64_inc(&zram->stats.multi_pages_miss_free);
+				return;
+			}
+		}
+	}
+#endif
+
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	slot_free_npages(zram, index, nr_pages_to_free);
+#else
 	slot_free(zram, index);
+#endif
 	slot_unlock(zram, index);
 }
 
@@ -2991,6 +3393,9 @@ static const struct block_device_operations zram_devops = {
 static DEVICE_ATTR_RO(io_stat);
 static DEVICE_ATTR_RO(mm_stat);
 static DEVICE_ATTR_RO(debug_stat);
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+static DEVICE_ATTR_RO(multi_pages_debug_stat);
+#endif
 static DEVICE_ATTR_WO(compact);
 static DEVICE_ATTR_RW(disksize);
 static DEVICE_ATTR_RO(initstate);
@@ -3035,6 +3440,9 @@ static struct attribute *zram_disk_attrs[] = {
 	&dev_attr_io_stat.attr,
 	&dev_attr_mm_stat.attr,
 	&dev_attr_debug_stat.attr,
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	&dev_attr_multi_pages_debug_stat.attr,
+#endif
 #ifdef CONFIG_ZRAM_MULTI_COMP
 	&dev_attr_recomp_algorithm.attr,
 	&dev_attr_recompress.attr,
diff --git a/drivers/block/zram/zram_drv.h b/drivers/block/zram/zram_drv.h
index d64c99b74214..1e3ebfec972e 100644
--- a/drivers/block/zram/zram_drv.h
+++ b/drivers/block/zram/zram_drv.h
@@ -28,6 +28,11 @@
 #define ZRAM_SECTOR_PER_LOGICAL_BLOCK	\
 	(1 << (ZRAM_LOGICAL_BLOCK_SHIFT - SECTOR_SHIFT))
 
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+#define SECTORS_PER_MULTI_PAGE_SHIFT	(MULTI_PAGE_SHIFT - SECTOR_SHIFT)
+#define SECTORS_PER_MULTI_PAGE		(1 << SECTORS_PER_MULTI_PAGE_SHIFT)
+#endif
+
 /*
  * ZRAM is mainly used for memory efficiency so we want to keep memory
  * footprint small and thus squeeze size and zram pageflags into a flags
@@ -37,7 +42,12 @@
  *
  * We use BUILD_BUG_ON() to make sure that zram pageflags don't overflow.
  */
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+#define ZRAM_FLAG_SHIFT (PAGE_SHIFT + CONFIG_ZSMALLOC_MULTI_PAGES_ORDER + 1)
+#else
 #define ZRAM_FLAG_SHIFT (PAGE_SHIFT + 1)
+#endif
+
 
 /* Only 2 bits are allowed for comp priority index */
 #define ZRAM_COMP_PRIORITY_MASK	0x3
@@ -55,6 +65,10 @@ enum zram_pageflags {
 	ZRAM_COMP_PRIORITY_BIT1, /* First bit of comp priority index */
 	ZRAM_COMP_PRIORITY_BIT2, /* Second bit of comp priority index */
 
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	ZRAM_COMP_MULTI_PAGES,	/* Compressed by multi-pages */
+#endif
+
 	__NR_ZRAM_PAGEFLAGS,
 };
 
@@ -110,6 +124,15 @@ struct zram_stats {
 	atomic64_t bd_reads;		/* no. of reads from backing device */
 	atomic64_t bd_writes;		/* no. of writes from backing device */
 #endif
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+	atomic64_t zram_bio_write_multi_pages_count;
+	atomic64_t zram_bio_read_multi_pages_count;
+	atomic64_t multi_pages_failed_writes;
+	atomic64_t multi_pages_failed_reads;
+	atomic64_t zram_bio_write_multi_pages_partial_count;
+	atomic64_t zram_bio_read_multi_pages_partial_count;
+	atomic64_t multi_pages_miss_free;
+#endif
 };
 
 #ifdef CONFIG_ZRAM_MULTI_COMP
@@ -167,4 +190,24 @@ struct zram {
 	struct dentry *debugfs_dir;
 #endif
 };
+
+#ifdef CONFIG_ZRAM_MULTI_PAGES
+#define multi_pages_bvec_iter_offset(bvec, iter)				\
+	(mp_bvec_iter_offset((bvec), (iter)) % ZCOMP_MULTI_PAGES_SIZE)
+
+#define multi_pages_bvec_iter_len(bvec, iter)				\
+	min_t(unsigned int, mp_bvec_iter_len((bvec), (iter)),		\
+	      ZCOMP_MULTI_PAGES_SIZE -					\
+	      multi_pages_bvec_iter_offset((bvec), (iter)))
+
+#define multi_pages_bvec_iter_bvec(bvec, iter)				\
+((struct bio_vec) {							\
+	.bv_page	= mp_bvec_iter_page((bvec), (iter)),		\
+	.bv_len		= multi_pages_bvec_iter_len((bvec), (iter)),	\
+	.bv_offset	= multi_pages_bvec_iter_offset((bvec), (iter)),	\
+})
+
+#define multi_pages_bio_iter_iovec(bio, iter)				\
+	multi_pages_bvec_iter_bvec((bio)->bi_io_vec, (iter))
+#endif
 #endif
diff --git a/mm/zsmalloc.c b/mm/zsmalloc.c
index 2d83fcc8610d..81d908d6d5a8 100644
--- a/mm/zsmalloc.c
+++ b/mm/zsmalloc.c
@@ -1697,7 +1697,8 @@ static struct zs_handle *find_alloced_obj(struct size_class *class,
 	offset += class->size * index;
 
 	while (offset < page_size) {
-		if (obj_allocated(zpdesc, addr + offset, &handle))
+		handle = obj_allocated(zpdesc, addr + offset);
+		if (handle)
 			break;
 
 		offset += class->size;
-- 
2.43.0

From: Barry Song <baohua@kernel.org>

zstd_setup_params() passes PAGE_SIZE to zstd_get_params() and
zstd_get_cparams() as the source size hint. When multi-page compression
is enabled, compressing 16 KiB blocks with 4 KiB-tuned parameters limits
zstd's match window to a single page and prevents it from finding
cross-page redundancy.

Pass ZCOMP_MULTI_PAGES_SIZE instead of PAGE_SIZE so that zstd configures
its parameters for the multi-page block size.

Signed-off-by: Barry Song <baohua@kernel.org>
Co-developed-by: Nanzhe Zhao <zhaonanzhe@xiaomi.com>
Signed-off-by: Nanzhe Zhao <zhaonanzhe@xiaomi.com>
---
 drivers/block/zram/backend_zstd.c | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/drivers/block/zram/backend_zstd.c b/drivers/block/zram/backend_zstd.c
index d55615f34300..9438691791ef 100644
--- a/drivers/block/zram/backend_zstd.c
+++ b/drivers/block/zram/backend_zstd.c
@@ -73,12 +73,12 @@ static int zstd_setup_params(struct zcomp_params *params)
 		goto error;
 	}
 
-	zp->cprm = zstd_get_params(params->level, PAGE_SIZE);
+	zp->cprm = zstd_get_params(params->level, ZCOMP_MULTI_PAGES_SIZE);
 
 	zp->custom_mem.customAlloc = zstd_custom_alloc;
 	zp->custom_mem.customFree = zstd_custom_free;
 
-	prm = zstd_get_cparams(params->level, PAGE_SIZE,
+	prm = zstd_get_cparams(params->level, ZCOMP_MULTI_PAGES_SIZE,
 			       params->dict_sz);
 
 	if (!params->dict || !params->dict_sz)
-- 
2.43.0

Synchronous swap devices such as zram complete swap I/O in the faulting
context. With multi-page compression enabled, a 16 KiB block is
stored and decompressed as a single object.

The swapin allocation inherits the VMA's THP gfp mask, which drops
__GFP_DIRECT_RECLAIM for VMAs without MADV_HUGEPAGE. Under memory
pressure, the order-2 allocation fails and falls back to order-0,
requiring the entire 16 KiB block to be decompressed just to swap in
an order-0 page.

Allow direct reclaim for order-2 swapin allocation to improve the
allocation success rate and avoid partial reads.

Signed-off-by: Nanzhe Zhao <zhaonanzhe@xiaomi.com>
---
 mm/swap_state.c | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mm/swap_state.c b/mm/swap_state.c
index da65d75ceead..912b0e21a781 100644
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -462,6 +462,13 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 	if (mpol || !vmf) {
 		folio = folio_alloc_mpol(gfp, order, mpol, ilx, numa_node_id());
 	} else {
+		/*
+		 * zsmalloc uses order-2 pages to store compressed data.
+		 * Allow direct reclaim to improve allocation success and
+		 * avoid partial reads.
+		 */
+		if (order == 2)
+			gfp |= __GFP_RECLAIM;
 		address = round_down(vmf->address, PAGE_SIZE << order);
 		folio = vma_alloc_folio(gfp, order, vmf->vma, address);
 	}
-- 
2.43.0