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.

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