mm/slab: introduce kmalloc_flags()

With alloc_flags usage in slab, we can replace __GFP_NO_OBJ_EXT with an
alloc flag that prevents kmalloc recursion. For that we need a version
of kmalloc() that takes alloc_flags and use it in places that perform
these potentially recursive kmalloc allocations (of sheaves or obj_ext
arrays).

Add this function, named kmalloc_flags(). Right now it's only useful for
these nested allocations, so it doesn't need to optimize build-time
constant sizes like kmalloc() or kmalloc_buckets.

Since we need it to support both normal and non-spinning
kmalloc_nolock() context through the SLAB_ALLOC_NOLOCK flag, split out
most of the special _kmalloc_nolock_noprof() implementation to
__kmalloc_nolock_noprof() that takes a slab_alloc_context, and make
_kmalloc_nolock_noprof() a simple tail calling wrapper with the proper
context.

kmalloc_flags() can thus determine whether to call
__kmalloc_nolock_noprof() or __do_kmalloc_node(), based on the
given alloc_flags.

Link: https://patch.msgid.link/20260610-slab_alloc_flags-v2-14-7190909db118@kernel.org
Reviewed-by: Hao Li <hao.li@linux.dev>
Reviewed-by: Suren Baghdasaryan <surenb@google.com>
Reviewed-by: Harry Yoo (Oracle) <harry@kernel.org>
Signed-off-by: Vlastimil Babka (SUSE) <vbabka@kernel.org>

diff --git a/mm/slab.h b/mm/slab.h
index d86203131f588ef64c9bcdbebf966bf53a3da16e..482b8e0fe797b033f6adb1d22b64035fed14c88d 100644 (file)
--- a/mm/slab.h
+++ b/mm/slab.h
@@ -11,6 +11,7 @@
 #include <linux/memcontrol.h>
 #include <linux/kfence.h>
 #include <linux/kasan.h>
+#include <linux/slab.h>
 
 /*
  * Internal slab definitions
@@ -26,6 +27,18 @@ static inline bool alloc_flags_allow_spinning(const unsigned int alloc_flags)
        return !(alloc_flags & SLAB_ALLOC_NOLOCK);
 }
 
+void *__kmalloc_flags_noprof(DECL_TOKEN_PARAMS(size, token), gfp_t flags,
+                                 unsigned int alloc_flags, int node)
+                                 __assume_kmalloc_alignment __alloc_size(1);
+
+static __always_inline __alloc_size(1) void *_kmalloc_flags_noprof(size_t size,
+               gfp_t flags, unsigned int alloc_flags, int node, kmalloc_token_t token)
+{
+       return __kmalloc_flags_noprof(PASS_TOKEN_PARAMS(size, token), flags, alloc_flags, node);
+}
+#define kmalloc_flags_noprof(...)      _kmalloc_flags_noprof(__VA_ARGS__, __kmalloc_token(__VA_ARGS__))
+#define kmalloc_flags(...)             alloc_hooks(kmalloc_flags_noprof(__VA_ARGS__))
+
 #ifdef CONFIG_64BIT
 # ifdef system_has_cmpxchg128
 # define system_has_freelist_aba()     system_has_cmpxchg128()

diff --git a/mm/slub.c b/mm/slub.c
index 8769083bec8180a19494385d2b6be852ee667b54..383d39a22561c3e8322167b50fcd91e2d3cf4b74 100644 (file)
--- a/mm/slub.c
+++ b/mm/slub.c
@@ -5385,19 +5385,14 @@ void *__kmalloc_noprof(DECL_TOKEN_PARAMS(size, token), gfp_t flags)
 }
 EXPORT_SYMBOL(__kmalloc_noprof);
 
-void *_kmalloc_nolock_noprof(DECL_TOKEN_PARAMS(size, token), gfp_t gfp_flags, int node)
+static void *__kmalloc_nolock_noprof(DECL_TOKEN_PARAMS(size, token), gfp_t gfp_flags,
+                                    int node, const struct slab_alloc_context *ac)
 {
-       size_t orig_size = size;
-       unsigned int alloc_flags = SLAB_ALLOC_NOLOCK;
        struct kmem_cache *s;
        bool can_retry = true;
        void *ret;
-       const struct slab_alloc_context ac = {
-               .caller_addr = _RET_IP_,
-               .orig_size = orig_size,
-               .alloc_flags = alloc_flags,
-       };
 
+       VM_WARN_ON_ONCE(alloc_flags_allow_spinning(ac->alloc_flags));
        VM_WARN_ON_ONCE(gfp_flags & ~(__GFP_ACCOUNT | __GFP_ZERO |
                        __GFP_NO_OBJ_EXT | __GFP_NOWARN | __GFP_NOMEMALLOC));
 
@@ -5434,7 +5429,7 @@ retry:
                 */
                return NULL;
 
-       ret = alloc_from_pcs(s, gfp_flags, alloc_flags, node);
+       ret = alloc_from_pcs(s, gfp_flags, ac->alloc_flags, node);
        if (ret)
                goto success;
 
@@ -5444,7 +5439,7 @@ retry:
         * kfence_alloc. Hence call __slab_alloc_node() (at most twice)
         * and slab_post_alloc_hook() directly.
         */
-       ret = __slab_alloc_node(s, gfp_flags, node, &ac);
+       ret = __slab_alloc_node(s, gfp_flags, node, ac);
 
        /*
         * It's possible we failed due to trylock as we preempted someone with
@@ -5467,11 +5462,23 @@ retry:
 
 success:
        maybe_wipe_obj_freeptr(s, ret);
-       slab_post_alloc_hook(s, gfp_flags, 1, &ret, &ac);
+       slab_post_alloc_hook(s, gfp_flags, 1, &ret, ac);
 
-       ret = kasan_kmalloc(s, ret, orig_size, gfp_flags);
+       ret = kasan_kmalloc(s, ret, ac->orig_size, gfp_flags);
        return ret;
 }
+
+void *_kmalloc_nolock_noprof(DECL_TOKEN_PARAMS(size, token), gfp_t gfp_flags, int node)
+{
+       const struct slab_alloc_context ac = {
+               .caller_addr = _RET_IP_,
+               .orig_size = size,
+               .alloc_flags = SLAB_ALLOC_NOLOCK,
+       };
+
+       return __kmalloc_nolock_noprof(PASS_TOKEN_PARAMS(size, token),
+                                      gfp_flags, node, &ac);
+}
 EXPORT_SYMBOL_GPL(_kmalloc_nolock_noprof);
 
 void *__kmalloc_node_track_caller_noprof(DECL_KMALLOC_PARAMS(size, b, token), gfp_t flags,
@@ -5525,6 +5532,30 @@ void *__kmalloc_cache_node_noprof(struct kmem_cache *s, gfp_t gfpflags,
 }
 EXPORT_SYMBOL(__kmalloc_cache_node_noprof);
 
+/*
+ * The only version of kmalloc_node() that takes alloc_flags and thus can
+ * determine on its own whether to handle the allocation via kmalloc_nolock() or
+ * normally
+ */
+void *__kmalloc_flags_noprof(DECL_TOKEN_PARAMS(size, token), gfp_t flags,
+                            unsigned int alloc_flags, int node)
+{
+       const struct slab_alloc_context ac = {
+               .caller_addr = _RET_IP_,
+               .orig_size = size,
+               .alloc_flags = alloc_flags,
+       };
+
+       if (alloc_flags_allow_spinning(alloc_flags)) {
+               return __do_kmalloc_node(NULL, flags, node,
+                               PASS_TOKEN_PARAM(token), &ac);
+       } else {
+               return __kmalloc_nolock_noprof(PASS_TOKEN_PARAMS(size, token),
+                                              flags, node, &ac);
+       }
+}
+
+
 static noinline void free_to_partial_list(
        struct kmem_cache *s, struct slab *slab,
        void *head, void *tail, int bulk_cnt,