Re: [PATCH RFC] alloc_tag: add option to pick the first codetag along callchain

[Suren Baghdasaryan <surenb@google.com>]

On Mon, Dec 15, 2025 at 10:44 PM David Wang <00107082@163.com> wrote:
>
> When tracking memory allocation for some specific function,
> picking the first codetag is more desired, because there
> is no need to track down all allocation sites in the call
> graph and change them to _noprof version, which is quite
> inflexible when the call graph is complex.
>
> For example, consider a simple graph:
>
> A ---> B ---> C ===> D
>        E ---> C
>
> ===> means a call with codetag
> ---> means a call without codetag
>
> To profiling memory allocation for A, the call graph needs
> to be changed to
> A ===> B ---> C ---> D
>        E ===> C
> Three call sites needs to be changed.
>
> But if pick the first codetag, only one change is needed.
> A ===> B ---> C ===> D
>        E ---> C
>
> The drawback is some accounting for C is splited to A,
> making the number not accurate for C. (But the overall
> accounting is still the same.)
>
> This is useful when debug memory problems, not meant for
> production usage though.

Hi David,
Sorry for the delay. Do you have specific examples when allocation
needs to be accounted at the highest level?
Our usual approach is that we account allocation at the lowest
"allocator" level and if that allocator uses lowel level allocators it
should use _noprof versions so that allocation is still done at the
right level. I would like to keep that simple approach but if there
are cases when that's not enough, I would like to know more about them
before trying to address them.
Thanks,
Suren.

>
> Signed-off-by: David Wang <00107082@163.com>
> ---
>  include/linux/sched.h |  6 ++++++
>  lib/Kconfig.debug     | 12 ++++++++++++
>  2 files changed, 18 insertions(+)
>
> diff --git a/include/linux/sched.h b/include/linux/sched.h
> index d395f2810fac..4a4f7000737e 100644
> --- a/include/linux/sched.h
> +++ b/include/linux/sched.h
> @@ -2288,14 +2288,20 @@ extern void sched_set_stop_task(int cpu, struct task_struct *stop);
>  #ifdef CONFIG_MEM_ALLOC_PROFILING
>  static __always_inline struct alloc_tag *alloc_tag_save(struct alloc_tag *tag)
>  {
> +#ifdef CONFIG_MEM_ALLOC_PROFILING_PICK_FIRST_CODETAG
> +       if (current->alloc_tag)
> +               return current->alloc_tag;
> +#endif
>         swap(current->alloc_tag, tag);
>         return tag;
>  }
>
>  static __always_inline void alloc_tag_restore(struct alloc_tag *tag, struct alloc_tag *old)
>  {
> +#ifndef CONFIG_MEM_ALLOC_PROFILING_PICK_FIRST_CODETAG
>  #ifdef CONFIG_MEM_ALLOC_PROFILING_DEBUG
>         WARN(current->alloc_tag != tag, "current->alloc_tag was changed:\n");
> +#endif
>  #endif
>         current->alloc_tag = old;
>  }
> diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
> index ba36939fda79..6e6f3a12033a 100644
> --- a/lib/Kconfig.debug
> +++ b/lib/Kconfig.debug
> @@ -1038,6 +1038,18 @@ config MEM_ALLOC_PROFILING_DEBUG
>           Adds warnings with helpful error messages for memory allocation
>           profiling.
>
> +config MEM_ALLOC_PROFILING_PICK_FIRST_CODETAG
> +       bool "Use the first tag along the call chain"
> +       default n
> +       depends on MEM_ALLOC_PROFILING
> +       select MEM_ALLOC_PROFILING_ENABLED_BY_DEFAULT
> +       help
> +         Make memory allocation profiling store counters to the first
> +         codetag along the call chain. This help profiling memory allocation
> +         for specific function by simply adding codetag to the function,
> +         without clearup all the codetag down the callchain.
> +         It is used for debug purpose.
> +
>  source "lib/Kconfig.kasan"
>  source "lib/Kconfig.kfence"
>  source "lib/Kconfig.kmsan"
> --
> 2.47.3
>