mm/slab: add cache_ and slab_needs_objcg() helpers

Slabs of some caches never need the objcg part of struct slabobj_ext.
Introduce helpers to query this for a cache or a slab.

Introduce SLAB_MAY_ACCOUNT flag that is currently only internal and all
caches have it set except:

- KMALLOC_NORMAL caches, as long as KMALLOC_RECLAIM caches are separate
- KMALLOC_NO_OBJ_EXT caches, if they exist

For named caches we currently can't derive SLAB_MAY_ACCOUNT from
SLAB_ACCOUNT because some caches might be created without SLAB_ACCOUNT
and then used both with and without __GFP_ACCOUNT concurrently,
allocating obj_ext arrays on demand. So just add the SLAB_MAY_ACCOUNT
to all kmem caches, unless kmem accounting is disabled.

This can be improved later by finding out all caches used with
__GFP_ACCOUNT, creating them with the SLAB_MAY_ACCOUNT flag explicitly,
and then ignoring __GFP_ACCOUNT for all other caches (possibly with a
warning).

To make the evaluation of slab_needs_objcg() faster in the allocation
and free fast paths, add a obj_exts_needs_objcg flag into slab itself.
This optimization is only available on 64bit architectures where free
bits are available for the flag.

Reviewed-by: Hao Li <hao.li@linux.dev>
Link: https://patch.msgid.link/20260727-b4-objext_split-v3-11-c29ef0f1f257@kernel.org
Reviewed-by: Harry Yoo <harry@kernel.org>
Signed-off-by: Vlastimil Babka (SUSE) <vbabka@kernel.org>

5 files changed, 61 insertions, 8 deletions

diff --git a/include/linux/slab.h b/include/linux/slab.h
index 32c9f8ed7ae2..e1915026e030 100644
--- a/include/linux/slab.h
+++ b/include/linux/slab.h
@@ -45,6 +45,7 @@ enum _slab_flag_bits {
 #endif
 #ifdef CONFIG_MEMCG
 	_SLAB_ACCOUNT,
+	_SLAB_MAY_ACCOUNT,
 #endif
 #ifdef CONFIG_KASAN_GENERIC
 	_SLAB_KASAN,
@@ -204,8 +205,10 @@ enum _slab_flag_bits {
  */
 #ifdef CONFIG_MEMCG
 # define SLAB_ACCOUNT		__SLAB_FLAG_BIT(_SLAB_ACCOUNT)
+# define SLAB_MAY_ACCOUNT	__SLAB_FLAG_BIT(_SLAB_MAY_ACCOUNT)
 #else
 # define SLAB_ACCOUNT		__SLAB_FLAG_UNUSED
+# define SLAB_MAY_ACCOUNT	__SLAB_FLAG_UNUSED
 #endif
 
 #ifdef CONFIG_KASAN_GENERIC
diff --git a/mm/kfence/core.c b/mm/kfence/core.c
index 5608a37f2b4d..7da3be9fbf5d 100644
--- a/mm/kfence/core.c
+++ b/mm/kfence/core.c
@@ -640,6 +640,9 @@ static unsigned long kfence_init_pool(void)
 		struct slab *slab = page_slab(page);
 		slab->obj_exts = (unsigned long)&kfence_metadata_init[i / 2 - 1].obj_exts |
 				 MEMCG_DATA_OBJEXTS;
+#ifdef CONFIG_64BIT
+		slab->obj_exts_needs_objcg = 1;
+#endif
 #endif
 	}
 
diff --git a/mm/slab.h b/mm/slab.h
index 041cda2d7395..69cd3d9631ff 100644
--- a/mm/slab.h
+++ b/mm/slab.h
@@ -81,10 +81,11 @@ struct freelist_counters {
 #ifdef CONFIG_64BIT
 					/*
 					 * Some optimizations use free bits in 'counters' field
-					 * to save memory. If these free bits are not available,
-					 * such optimizations are disabled.
+					 * to save memory or CPU. If these free bits are not
+					 * available, such optimizations are disabled.
 					 */
 					unsigned obj_exts_in_object:1;
+					unsigned obj_exts_needs_objcg:1;
 #endif
 				};
 			};
@@ -584,6 +585,32 @@ static inline bool slab_obj_ext_has_codetag(void)
 }
 #endif
 
+#ifdef CONFIG_MEMCG
+static inline bool cache_needs_objcg(struct kmem_cache *cache)
+{
+	return (cache->flags & SLAB_MAY_ACCOUNT);
+}
+
+static inline bool slab_needs_objcg(struct slab *slab)
+{
+#ifdef CONFIG_64BIT
+	return slab->obj_exts_needs_objcg;
+#else
+	return cache_needs_objcg(slab->slab_cache);
+#endif
+}
+#else
+static inline bool cache_needs_objcg(struct kmem_cache *cache)
+{
+	return false;
+}
+
+static inline bool slab_needs_objcg(struct slab *slab)
+{
+	return false;
+}
+#endif
+
 static inline size_t cache_obj_ext_size(struct kmem_cache *s)
 {
 	size_t sz = 0;
diff --git a/mm/slab_common.c b/mm/slab_common.c
index 03ecac12cd86..1e1d3feec353 100644
--- a/mm/slab_common.c
+++ b/mm/slab_common.c
@@ -52,7 +52,7 @@ struct kmem_cache *kmem_cache;
 		SLAB_OBJ_EXT_IN_OBJ)
 
 #define SLAB_MERGE_SAME (SLAB_RECLAIM_ACCOUNT | SLAB_CACHE_DMA | \
-			 SLAB_CACHE_DMA32 | SLAB_ACCOUNT)
+			 SLAB_CACHE_DMA32 | SLAB_ACCOUNT | SLAB_MAY_ACCOUNT)
 
 /*
  * Merge control. If this is set then no merging of slab caches will occur.
@@ -359,6 +359,13 @@ struct kmem_cache *__kmem_cache_create_args(const char *name,
 		goto out_unlock;
 	}
 
+	/*
+	 * For now we assume any cache can be used with __GFP_ACCOUNT and thus
+	 * may need to store objcg pointers for objects
+	 */
+	if (!mem_cgroup_kmem_disabled())
+		flags |= SLAB_MAY_ACCOUNT;
+
 	/* Fail closed on bad usersize of useroffset values. */
 	if (!IS_ENABLED(CONFIG_HARDENED_USERCOPY) ||
 	    WARN_ON(!args->usersize && args->useroffset) ||
@@ -984,11 +991,20 @@ new_kmalloc_cache(int idx, enum kmalloc_cache_type type)
 #endif
 
 	/*
-	 * If CONFIG_MEMCG is enabled, disable cache merging for
-	 * KMALLOC_NORMAL caches.
+	 * If memcg_kmem is enabled and this is a KMALLOC_NORMAL cache and not
+	 * aliased with any other type, make sure it's never merged with any other
+	 * cache.
+	 *
+	 * In other cases the kmalloc cache may end up being used for a
+	 * __GFP_ACCOUNT allocation so mark it as such. The exception is a
+	 * KMALLOC_NO_OBJ_EXT cache.
 	 */
-	if (IS_ENABLED(CONFIG_MEMCG) && (type == KMALLOC_NORMAL))
-		flags |= SLAB_NO_MERGE;
+	if (!mem_cgroup_kmem_disabled()) {
+		if (type == KMALLOC_NORMAL && KMALLOC_RECLAIM != KMALLOC_NORMAL)
+			flags |= SLAB_NO_MERGE;
+		else if (!(flags & SLAB_NO_OBJ_EXT))
+			flags |= SLAB_MAY_ACCOUNT;
+	}
 
 	if (minalign > ARCH_KMALLOC_MINALIGN) {
 		aligned_size = ALIGN(aligned_size, minalign);
diff --git a/mm/slub.c b/mm/slub.c
index a24428ae32fb..331dd9a9b99d 100644
--- a/mm/slub.c
+++ b/mm/slub.c
@@ -2558,7 +2558,7 @@ bool memcg_slab_post_charge(void *p, gfp_t flags)
 	 * of slab_obj_exts being allocated from the same slab and thus the slab
 	 * becoming effectively unfreeable.
 	 */
-	if (is_kmalloc_normal(s))
+	if (!cache_needs_objcg(s))
 		return true;
 
 	/* Ignore already charged objects. */
@@ -3436,6 +3436,10 @@ static struct slab *allocate_slab(struct kmem_cache *s, gfp_t flags,
 
 	slab->objects = oo_objects(oo);
 
+#ifdef CONFIG_64BIT
+	if (cache_needs_objcg(s))
+		slab->obj_exts_needs_objcg = 1;
+#endif
 	slab->slab_cache = s;
 
 	kasan_poison_slab(slab);