xref: /linux-6.15/include/linux/alloc_tag.h (revision b951aaff)
1 /* SPDX-License-Identifier: GPL-2.0 */
2 /*
3  * allocation tagging
4  */
5 #ifndef _LINUX_ALLOC_TAG_H
6 #define _LINUX_ALLOC_TAG_H
7 
8 #include <linux/bug.h>
9 #include <linux/codetag.h>
10 #include <linux/container_of.h>
11 #include <linux/preempt.h>
12 #include <asm/percpu.h>
13 #include <linux/cpumask.h>
14 #include <linux/static_key.h>
15 #include <linux/irqflags.h>
16 
17 struct alloc_tag_counters {
18 	u64 bytes;
19 	u64 calls;
20 };
21 
22 /*
23  * An instance of this structure is created in a special ELF section at every
24  * allocation callsite. At runtime, the special section is treated as
25  * an array of these. Embedded codetag utilizes codetag framework.
26  */
27 struct alloc_tag {
28 	struct codetag			ct;
29 	struct alloc_tag_counters __percpu	*counters;
30 } __aligned(8);
31 
32 #ifdef CONFIG_MEM_ALLOC_PROFILING
33 
34 static inline struct alloc_tag *ct_to_alloc_tag(struct codetag *ct)
35 {
36 	return container_of(ct, struct alloc_tag, ct);
37 }
38 
39 #ifdef ARCH_NEEDS_WEAK_PER_CPU
40 /*
41  * When percpu variables are required to be defined as weak, static percpu
42  * variables can't be used inside a function (see comments for DECLARE_PER_CPU_SECTION).
43  * Instead we will accound all module allocations to a single counter.
44  */
45 DECLARE_PER_CPU(struct alloc_tag_counters, _shared_alloc_tag);
46 
47 #define DEFINE_ALLOC_TAG(_alloc_tag)						\
48 	static struct alloc_tag _alloc_tag __used __aligned(8)			\
49 	__section("alloc_tags") = {						\
50 		.ct = CODE_TAG_INIT,						\
51 		.counters = &_shared_alloc_tag };
52 
53 #else /* ARCH_NEEDS_WEAK_PER_CPU */
54 
55 #define DEFINE_ALLOC_TAG(_alloc_tag)						\
56 	static DEFINE_PER_CPU(struct alloc_tag_counters, _alloc_tag_cntr);	\
57 	static struct alloc_tag _alloc_tag __used __aligned(8)			\
58 	__section("alloc_tags") = {						\
59 		.ct = CODE_TAG_INIT,						\
60 		.counters = &_alloc_tag_cntr };
61 
62 #endif /* ARCH_NEEDS_WEAK_PER_CPU */
63 
64 DECLARE_STATIC_KEY_MAYBE(CONFIG_MEM_ALLOC_PROFILING_ENABLED_BY_DEFAULT,
65 			mem_alloc_profiling_key);
66 
67 static inline bool mem_alloc_profiling_enabled(void)
68 {
69 	return static_branch_maybe(CONFIG_MEM_ALLOC_PROFILING_ENABLED_BY_DEFAULT,
70 				   &mem_alloc_profiling_key);
71 }
72 
73 static inline struct alloc_tag_counters alloc_tag_read(struct alloc_tag *tag)
74 {
75 	struct alloc_tag_counters v = { 0, 0 };
76 	struct alloc_tag_counters *counter;
77 	int cpu;
78 
79 	for_each_possible_cpu(cpu) {
80 		counter = per_cpu_ptr(tag->counters, cpu);
81 		v.bytes += counter->bytes;
82 		v.calls += counter->calls;
83 	}
84 
85 	return v;
86 }
87 
88 #ifdef CONFIG_MEM_ALLOC_PROFILING_DEBUG
89 static inline void alloc_tag_add_check(union codetag_ref *ref, struct alloc_tag *tag)
90 {
91 	WARN_ONCE(ref && ref->ct,
92 		  "alloc_tag was not cleared (got tag for %s:%u)\n",
93 		  ref->ct->filename, ref->ct->lineno);
94 
95 	WARN_ONCE(!tag, "current->alloc_tag not set");
96 }
97 
98 static inline void alloc_tag_sub_check(union codetag_ref *ref)
99 {
100 	WARN_ONCE(ref && !ref->ct, "alloc_tag was not set\n");
101 }
102 #else
103 static inline void alloc_tag_add_check(union codetag_ref *ref, struct alloc_tag *tag) {}
104 static inline void alloc_tag_sub_check(union codetag_ref *ref) {}
105 #endif
106 
107 /* Caller should verify both ref and tag to be valid */
108 static inline void __alloc_tag_ref_set(union codetag_ref *ref, struct alloc_tag *tag)
109 {
110 	ref->ct = &tag->ct;
111 	/*
112 	 * We need in increment the call counter every time we have a new
113 	 * allocation or when we split a large allocation into smaller ones.
114 	 * Each new reference for every sub-allocation needs to increment call
115 	 * counter because when we free each part the counter will be decremented.
116 	 */
117 	this_cpu_inc(tag->counters->calls);
118 }
119 
120 static inline void alloc_tag_add(union codetag_ref *ref, struct alloc_tag *tag, size_t bytes)
121 {
122 	alloc_tag_add_check(ref, tag);
123 	if (!ref || !tag)
124 		return;
125 
126 	__alloc_tag_ref_set(ref, tag);
127 	this_cpu_add(tag->counters->bytes, bytes);
128 }
129 
130 static inline void alloc_tag_sub(union codetag_ref *ref, size_t bytes)
131 {
132 	struct alloc_tag *tag;
133 
134 	alloc_tag_sub_check(ref);
135 	if (!ref || !ref->ct)
136 		return;
137 
138 	tag = ct_to_alloc_tag(ref->ct);
139 
140 	this_cpu_sub(tag->counters->bytes, bytes);
141 	this_cpu_dec(tag->counters->calls);
142 
143 	ref->ct = NULL;
144 }
145 
146 #else /* CONFIG_MEM_ALLOC_PROFILING */
147 
148 #define DEFINE_ALLOC_TAG(_alloc_tag)
149 static inline bool mem_alloc_profiling_enabled(void) { return false; }
150 static inline void alloc_tag_add(union codetag_ref *ref, struct alloc_tag *tag,
151 				 size_t bytes) {}
152 static inline void alloc_tag_sub(union codetag_ref *ref, size_t bytes) {}
153 
154 #endif /* CONFIG_MEM_ALLOC_PROFILING */
155 
156 #define alloc_hooks_tag(_tag, _do_alloc)				\
157 ({									\
158 	struct alloc_tag * __maybe_unused _old = alloc_tag_save(_tag);	\
159 	typeof(_do_alloc) _res = _do_alloc;				\
160 	alloc_tag_restore(_tag, _old);					\
161 	_res;								\
162 })
163 
164 #define alloc_hooks(_do_alloc)						\
165 ({									\
166 	DEFINE_ALLOC_TAG(_alloc_tag);					\
167 	alloc_hooks_tag(&_alloc_tag, _do_alloc);			\
168 })
169 
170 #endif /* _LINUX_ALLOC_TAG_H */
171