mm/slub: use empty sheaf helpers for oversized sheaves

Oversized prefilled sheaves are allocated separately because their
capacity can be larger than the cache's regular sheaf capacity. After
they are flushed, however, they are empty sheaves as well, and should be
released through the same empty-sheaf helper.

Allocate oversized prefilled sheaves with __alloc_empty_sheaf() and free
them with free_empty_sheaf() after a failed prefill or after they are
returned and flushed. This keeps the oversized and pfmemalloc return paths
consistent, including the SLAB_KMALLOC-specific __GFP_NO_OBJ_EXT and
mark_obj_codetag_empty() handling.

Keep the caller-GFP filtering in alloc_empty_sheaf() instead of
__alloc_empty_sheaf(). In particular, do not clear OBJCGS_CLEAR_MASK in
the raw helper, so the oversized prefill path does not unexpectedly drop
caller-provided flags such as __GFP_NOFAIL. The SLAB_KMALLOC-specific
addition of __GFP_NO_OBJ_EXT remains in __alloc_empty_sheaf(), matching
the free_empty_sheaf() assumption.

Since oversized sheaves are now allocated and freed through the empty
sheaf helpers, SHEAF_ALLOC and SHEAF_FREE also account for oversized
sheaves. Update the stat comments accordingly.

Keep the capacity initialization in the oversized prefill path, since
capacity is currently only used for prefilled sheaves

Signed-off-by: Shengming Hu <hu.shengming@zte.com.cn>
Link: https://patch.msgid.link/20260528193537623nAo-xYBNYBysGKSBjREuO@zte.com.cn
Reviewed-by: Harry Yoo (Oracle) <harry@kernel.org>
Reviewed-by: Hao Li <hao.li@linux.dev>
Signed-off-by: Vlastimil Babka (SUSE) <vbabka@kernel.org>

Author: HSM6236

mm/slub.c

@@ -362,8 +362,8 @@ enum stat_item {
 	CMPXCHG_DOUBLE_FAIL,	/* Failures of slab freelist update */
 	SHEAF_FLUSH,		/* Objects flushed from a sheaf */
 	SHEAF_REFILL,		/* Objects refilled to a sheaf */
-	SHEAF_ALLOC,		/* Allocation of an empty sheaf */
-	SHEAF_FREE,		/* Freeing of an empty sheaf */
+	SHEAF_ALLOC,		/* Allocation of an empty sheaf including oversized ones */
+	SHEAF_FREE,		/* Freeing of an empty sheaf including oversized ones */
 	BARN_GET,		/* Got full sheaf from barn */
 	BARN_GET_FAIL,		/* Failed to get full sheaf from barn */
 	BARN_PUT,		/* Put full sheaf to barn */
@@ -2762,11 +2762,6 @@ static struct slab_sheaf *__alloc_empty_sheaf(struct kmem_cache *s, gfp_t gfp,
 	struct slab_sheaf *sheaf;
 	size_t sheaf_size;
 
-	if (gfp & __GFP_NO_OBJ_EXT)
-		return NULL;
-
-	gfp &= ~OBJCGS_CLEAR_MASK;
-
 	/*
 	 * Prevent recursion to the same cache, or a deep stack of kmallocs of
 	 * varying sizes (sheaf capacity might differ for each kmalloc size
@@ -2791,6 +2786,11 @@ static struct slab_sheaf *__alloc_empty_sheaf(struct kmem_cache *s, gfp_t gfp,
 static inline struct slab_sheaf *alloc_empty_sheaf(struct kmem_cache *s,
 						   gfp_t gfp)
 {
+	if (gfp & __GFP_NO_OBJ_EXT)
+		return NULL;
+
+	gfp &= ~OBJCGS_CLEAR_MASK;
+
 	return __alloc_empty_sheaf(s, gfp, s->sheaf_capacity);
 }
 
@@ -5014,12 +5014,11 @@ kmem_cache_prefill_sheaf(struct kmem_cache *s, gfp_t gfp, unsigned int size)
 
 	if (unlikely(size > s->sheaf_capacity)) {
 
-		sheaf = kzalloc_flex(*sheaf, objects, size, gfp);
+		sheaf = __alloc_empty_sheaf(s, gfp, size);
 		if (!sheaf)
 			return NULL;
 
 		stat(s, SHEAF_PREFILL_OVERSIZE);
-		sheaf->cache = s;
 		sheaf->capacity = size;
 
 		/*
@@ -5028,7 +5027,7 @@ kmem_cache_prefill_sheaf(struct kmem_cache *s, gfp_t gfp, unsigned int size)
 		 */
 		if (!__kmem_cache_alloc_bulk(s, gfp, size,
 					     &sheaf->objects[0])) {
-			kfree(sheaf);
+			free_empty_sheaf(s, sheaf);
 			return NULL;
 		}
 
@@ -5096,7 +5095,7 @@ void kmem_cache_return_sheaf(struct kmem_cache *s, gfp_t gfp,
 	if (unlikely((sheaf->capacity != s->sheaf_capacity)
 		     || sheaf->pfmemalloc)) {
 		sheaf_flush_unused(s, sheaf);
-		kfree(sheaf);
+		free_empty_sheaf(s, sheaf);
 		return;
 	}