memcg: nmi-safe kmem charging

To enable memcg charged kernel memory allocations from nmi context, consume_obj_stock() and refill_obj_stock() needs to be nmi safe. With the simple in_nmi() check, take the slow path of the objcg charging which handles the charging and memcg stats updates correctly for the nmi context.

Signed-off-by Shakeel Butt <[email protected]>

---
 mm/memcontrol.c | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

mm/memcontrol.c +3 -0

index bba549c1f18c..6cfa3550f300 100644
@@ -2965,6 +2965,9 @@ static bool consume_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
        unsigned long flags;
        bool ret = false;
 
+       if (unlikely(in_nmi()))
+               return ret;
+
        local_lock_irqsave(&obj_stock.lock, flags);
 
        stock = this_cpu_ptr(&obj_stock);

@@ -3068,6 +3071,15 @@ static void refill_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
        unsigned long flags;

Quoted 28 lines · Shakeel Butt

To enable memcg charged kernel memory allocations from nmi context, consume_obj_stock() and refill_obj_stock() needs to be nmi safe. With the simple in_nmi() check, take the slow path of the objcg charging which handles the charging and memcg stats updates correctly for the nmi context.

Signed-off-by Shakeel Butt <[email protected]>

---
 mm/memcontrol.c | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

mm/memcontrol.c +3 -0

index bba549c1f18c..6cfa3550f300 100644
@@ -2965,6 +2965,9 @@ static bool consume_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
        unsigned long flags;
        bool ret = false;
 
+       if (unlikely(in_nmi()))
+               return ret;
+
        local_lock_irqsave(&obj_stock.lock, flags);
 
        stock = this_cpu_ptr(&obj_stock);

@@ -3068,6 +3071,15 @@ static void refill_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
        unsigned long flags;

mm/memcontrol.c static void refill_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes, · line 3081

Quoted 10 lines

mm/memcontrol.c +8 -0

        unsigned int nr_pages = 0;
 
+       if (unlikely(in_nmi())) {
+               if (pgdat)
+                       __mod_objcg_mlstate(objcg, pgdat, idx, nr_bytes);
+               nr_pages = nr_bytes >> PAGE_SHIFT;
+               nr_bytes = nr_bytes & (PAGE_SIZE - 1);
+               atomic_add(nr_bytes, &objcg->nr_charged_bytes);
+               goto out;
+       }

Now I see what I did incorrectly in my series and how this patch 4 combined with patch 3 is doing accounting properly.

The only issue here and in other patches is that in_nmi() is an incomplete condition to check for.
The reentrance is possible through kprobe or tracepoint. In PREEMP_RT we will be fully preemptible, but obj_stock.lock will be already taken by the current task. To fix it you need to use local_lock_is_locked(&obj_stock.lock) instead of in_nmi() or use local_trylock_irqsave(&obj_stock.lock).

local_trylock_irqsave() is cleaner and works today,
while local_lock_is_locked() hasn't landed yet, but if we go is_locked route we can decouple reentrant obj_stock operation vs normal. Like the if (!local_lock_is_locked(&obj_stock.lock))
can be done much higher up the stack from __memcg_slab_post_alloc_hook() the way I did in my series, and if locked it can do atomic_add()-style charging.
So refill_obj_stock() and friends won't need to change.

BPF programs can trigger memcg charged kernel allocations in nmi context. However memcg charging infra for kernel memory is not equipped to handle nmi context. This series adds support for kernel memory charging for nmi context.

Quoted 4 lines · On Fri, 9 May 2025 16:28:55 -0700 Shakeel Butt <[email protected]> wrote:

BPF programs can trigger memcg charged kernel allocations in nmi context. However memcg charging infra for kernel memory is not equipped to handle nmi context. This series adds support for kernel memory charging for nmi context.

The patchset adds quite a bit of material to core MM on behalf of a single caller. So can we please take a close look at why BPF is doing this?

What would be involved in changing BPF to avoid doing this, or of changing BPF to handle things locally? What would be the end-user impact of such an alteration? IOW, what is the value to our users of the present BPF behavior?

Hi Andrew,

Quoted 7 lines · Andrew Morton

Quoted 4 lines · On Fri, 9 May 2025 16:28:55 -0700 Shakeel Butt <[email protected]> wrote:

BPF programs can trigger memcg charged kernel allocations in nmi context. However memcg charging infra for kernel memory is not equipped to handle nmi context. This series adds support for kernel memory charging for nmi context.

Quoted 9 lines

The patchset adds quite a bit of material to core MM on behalf of a single caller. So can we please take a close look at why BPF is doing this?

What would be involved in changing BPF to avoid doing this, or of changing BPF to handle things locally? What would be the end-user impact of such an alteration? IOW, what is the value to our users of the present BPF behavior?

Before answering the questions, let me clarify that this series is continuation of the work which added similar support for page allocator and related memcg charging and now the work is happening for kmalloc/slab allocations. Alexei has a proposal on reentrant kmalloc and here I am providing how memcg charging for that (reentrant kmalloc) should work.

Next let me take a stab in answering the questions and BPF folks can correct me if I am wrong. From what I understand, users can attach BPF programs at almost any place in kernel and those BPF programs can do memory allocations. This line of work is to make those allocations work if the any such BPF attach point is triggered in mni context.

Before this line of work (reentrant page and slab allocators), I think BPF had its internal cache but it was very limited and can easily fail allocation requests (please BPF folks correct me if I am wrong). This was discussed in LSFMM this year as well.

Now regarding the impact to the users. First there will not be any negative impact on the non-users of this feature. For the value this feature will provide to users, I think this line of work will make BPF programs of the users more reliable with better allocation behavior. BPF folks, please add more comments for the value of these features.

Before answering the questions, let me clarify that this series is continuation of the work which added similar support for page allocator and related memcg charging and now the work is happening for kmalloc/slab allocations. Alexei has a proposal on reentrant kmalloc and here I am providing how memcg charging for that (reentrant kmalloc) should work.

Next let me take a stab in answering the questions and BPF folks can correct me if I am wrong. From what I understand, users can attach BPF programs at almost any place in kernel and those BPF programs can do memory allocations. This line of work is to make those allocations work if the any such BPF attach point is triggered in mni context.

Before this line of work (reentrant page and slab allocators), I think BPF had its internal cache but it was very limited and can easily fail allocation requests (please BPF folks correct me if I am wrong). This was discussed in LSFMM this year as well.

Now regarding the impact to the users. First there will not be any negative impact on the non-users of this feature. For the value this feature will provide to users, I think this line of work will make BPF programs of the users more reliable with better allocation behavior. BPF folks, please add more comments for the value of these features.

There will be a followup series which will make kernel memory charging reentrant for irq and will be able to do without disabling irqs.

BPF programs can trigger memcg charged kernel allocations in nmi context. However memcg charging infra for kernel memory is not equipped to handle nmi context. This series adds support for kernel memory charging for nmi context.

The initial prototype tried to make memcg charging infra for kernel memory re-entrant against irq and nmi. However upon realizing that this_cpu_* operations are not safe on all architectures (Tejun), this

series took a different approach targeting only nmi context. Since the number of stats that are updated in kernel memory charging path are 3, this series added special handling of those stats in nmi context rather than making all >100 memcg stats nmi safe.

There will be a followup series which will make kernel memory charging reentrant for irq and will be able to do without disabling irqs.

We ran network intensive workload on this series and have not seen any significant performance differences with and without the series.

Shakeel Butt (4):
memcg: add infra for nmi safe memcg stats
memcg: add nmi-safe update for MEMCG_KMEM
memcg: nmi-safe slab stats updates
memcg: make objcg charging nmi safe

 include/linux/memcontrol.h |  6 +++
 mm/memcontrol.c            | 87 +++++++++++++++++++++++++++++++++++---
 2 files changed, 88 insertions(+), 5 deletions(-)

Quoted 5 lines · Shakeel Butt

BPF programs can trigger memcg charged kernel allocations in nmi context. However memcg charging infra for kernel memory is not equipped to handle nmi context. This series adds support for kernel memory charging for nmi context.

Quoted 3 lines

The initial prototype tried to make memcg charging infra for kernel memory re-entrant against irq and nmi. However upon realizing that this_cpu_* operations are not safe on all architectures (Tejun), this

I assume it was an off-list discussion?
Could we avoid this for the architectures where these are safe, which should be the major ones I hope?

Quoted 5 lines

series took a different approach targeting only nmi context. Since the number of stats that are updated in kernel memory charging path are 3, this series added special handling of those stats in nmi context rather than making all >100 memcg stats nmi safe.

Hmm so from patches 2 and 3 I see this relies on atomic64_add(). But AFAIU lib/atomic64.c has the generic fallback implementation for architectures that don't know better, and that would be using the "void generic_atomic64_##op" macro, which AFAICS is doing:

        local_irq_save(flags);                                          \
        arch_spin_lock(lock);                                           \
        v->counter c_op a;                                              \
        arch_spin_unlock(lock);                                         \
        local_irq_restore(flags);                                       \

so in case of a nmi hitting after the spin_lock this can still deadlock?

Hm or is there some assumption that we only use these paths when already in_nmi() and then another nmi can't come in that context?

But even then, flush_nmi_stats() in patch 1 isn't done in_nmi() and uses atomic64_xchg() which in generic_atomic64_xchg() implementation also has the irq_save+spin_lock. So can't we deadlock there?

I forgot to CC Tejun, so doing it now.

Quoted 6 lines · Vlastimil Babka

Quoted 5 lines · Shakeel Butt

BPF programs can trigger memcg charged kernel allocations in nmi context. However memcg charging infra for kernel memory is not equipped to handle nmi context. This series adds support for kernel memory charging for nmi context.

Quoted 7 lines

Quoted 3 lines

The initial prototype tried to make memcg charging infra for kernel memory re-entrant against irq and nmi. However upon realizing that this_cpu_* operations are not safe on all architectures (Tejun), this

I assume it was an off-list discussion?
Could we avoid this for the architectures where these are safe, which should be the major ones I hope?

Yes it was an off-list discussion. The discussion was more about the this_cpu_* ops vs atomic_* ops as on x86 this_cpu_* does not have lock prefix and how I should prefer this_cpu_* over atomic_* for my series on objcg charging without disabling irqs. Tejun pointed out this_cpu_* are not nmi safe for some archs and it would be better to handle nmi context separately. So, I am not that worried about optimizing for NMI context

but your next comment on generic_atomic64_* ops is giving me headache.

Quoted 17 lines

Quoted 5 lines

series took a different approach targeting only nmi context. Since the number of stats that are updated in kernel memory charging path are 3, this series added special handling of those stats in nmi context rather than making all >100 memcg stats nmi safe.

Hmm so from patches 2 and 3 I see this relies on atomic64_add(). But AFAIU lib/atomic64.c has the generic fallback implementation for architectures that don't know better, and that would be using the "void generic_atomic64_##op" macro, which AFAICS is doing:

        local_irq_save(flags);                                          \
        arch_spin_lock(lock);                                           \
        v->counter c_op a;                                              \
        arch_spin_unlock(lock);                                         \
        local_irq_restore(flags);                                       \

Quoted 8 lines

so in case of a nmi hitting after the spin_lock this can still deadlock?

Hm or is there some assumption that we only use these paths when already in_nmi() and then another nmi can't come in that context?

But even then, flush_nmi_stats() in patch 1 isn't done in_nmi() and uses atomic64_xchg() which in generic_atomic64_xchg() implementation also has the irq_save+spin_lock. So can't we deadlock there?

I was actually assuming that atomic_* ops are safe against nmis for all archs. I looked at atomic_* ops in include/asm-generic/atomic.h and it is using arch_cmpxchg() for CONFIG_SMP and it seems like for archs with cmpxchg should be fine against nmi. I am not sure why atomic64_* are not using arch_cmpxchg() instead. I will dig more.

I also have the followup series on objcg charging without irq almost ready. I will send it out as rfc soon.

Thanks a lot for awesome and insightful comments.
Shakeel

Quoted 7 lines

Quoted 7 lines

Quoted 3 lines

The initial prototype tried to make memcg charging infra for kernel memory re-entrant against irq and nmi. However upon realizing that this_cpu_* operations are not safe on all architectures (Tejun), this

I assume it was an off-list discussion?
Could we avoid this for the architectures where these are safe, which should be the major ones I hope?

Quoted 6 lines

Yes it was an off-list discussion. The discussion was more about the this_cpu_* ops vs atomic_* ops as on x86 this_cpu_* does not have lock prefix and how I should prefer this_cpu_* over atomic_* for my series on objcg charging without disabling irqs. Tejun pointed out this_cpu_* are not nmi safe for some archs and it would be better to handle nmi context separately. So, I am not that worried about optimizing for NMI context

Well, we're introducing in_nmi() check and different execution paths to all charging. This could be e.g. compiled out for architectures where this_cpu* is NMI safe or they don't have NMIs in the first place.

Quoted 24 lines

but your next comment on generic_atomic64_* ops is giving me headache.

Quoted 22 lines

Quoted 5 lines

series took a different approach targeting only nmi context. Since the number of stats that are updated in kernel memory charging path are 3, this series added special handling of those stats in nmi context rather than making all >100 memcg stats nmi safe.

Hmm so from patches 2 and 3 I see this relies on atomic64_add(). But AFAIU lib/atomic64.c has the generic fallback implementation for architectures that don't know better, and that would be using the "void generic_atomic64_##op" macro, which AFAICS is doing:

        local_irq_save(flags);                                          \
        arch_spin_lock(lock);                                           \
        v->counter c_op a;                                              \
        arch_spin_unlock(lock);                                         \
        local_irq_restore(flags);                                       \

so in case of a nmi hitting after the spin_lock this can still deadlock?

Hm or is there some assumption that we only use these paths when already in_nmi() and then another nmi can't come in that context?

Quoted 6 lines

Quoted 3 lines

But even then, flush_nmi_stats() in patch 1 isn't done in_nmi() and uses atomic64_xchg() which in generic_atomic64_xchg() implementation also has the irq_save+spin_lock. So can't we deadlock there?

I was actually assuming that atomic_* ops are safe against nmis for all archs.

I looked at atomic_* ops in include/asm-generic/atomic.h and it

Quoted 3 lines

is using arch_cmpxchg() for CONFIG_SMP and it seems like for archs with cmpxchg should be fine against nmi. I am not sure why atomic64_* are not using arch_cmpxchg() instead. I will dig more.

Yeah I've found https://docs.kernel.org/core-api/local_ops.html and since it listed Mathieu we discussed on IRC and he mentioned the same thing that atomic_ ops are fine, but the later added 64bit variant isn't, which PeterZ (who added it) acknowledged.

But there could be way out if we could somehow compile-time assert that either is true:
- CONFIG_HAVE_NMI=n - we can compile out all the nmi code

- this_cpu is safe on that arch - we can also compile out the nmi code

- (if the above leaves any 64bit arch) its 64bit atomics implementation is safe

- (if there are any 32bit applicable arch left) 32bit atomics should be enough for the nmi counters even with >4GB memory as we flush them? and we know the 32bit ops are safe

Quoted 4 lines · Vlastimil Babka

Quoted 4 lines

Quoted 4 lines

Quoted 3 lines

The initial prototype tried to make memcg charging infra for kernel memory re-entrant against irq and nmi. However upon realizing that this_cpu_* operations are not safe on all architectures (Tejun), this

Quoted 3 lines

Quoted 3 lines

Quoted 3 lines

I assume it was an off-list discussion?
Could we avoid this for the architectures where these are safe, which should be the major ones I hope?

IIRC Power64 has issues here, 'funnily' their local_t is NMI safe. Perhaps we could do the same for their this_cpu_*(), but ideally someone with actual power hardware should do this ;-)

Quoted 10 lines

Quoted 6 lines

Yes it was an off-list discussion. The discussion was more about the this_cpu_* ops vs atomic_* ops as on x86 this_cpu_* does not have lock prefix and how I should prefer this_cpu_* over atomic_* for my series on objcg charging without disabling irqs. Tejun pointed out this_cpu_* are not nmi safe for some archs and it would be better to handle nmi context separately. So, I am not that worried about optimizing for NMI context

Well, we're introducing in_nmi() check and different execution paths to all charging. This could be e.g. compiled out for architectures where this_cpu* is NMI safe or they don't have NMIs in the first place.

Very few architectures one would care about do not have NMIs. Risc-V seems to be the exception here ?!?

Quoted 24 lines

Quoted 24 lines

but your next comment on generic_atomic64_* ops is giving me headache.

Quoted 22 lines

Quoted 5 lines

series took a different approach targeting only nmi context. Since the number of stats that are updated in kernel memory charging path are 3, this series added special handling of those stats in nmi context rather than making all >100 memcg stats nmi safe.

Hmm so from patches 2 and 3 I see this relies on atomic64_add(). But AFAIU lib/atomic64.c has the generic fallback implementation for architectures that don't know better, and that would be using the "void generic_atomic64_##op" macro, which AFAICS is doing:

        local_irq_save(flags);                                          \
        arch_spin_lock(lock);                                           \
        v->counter c_op a;                                              \
        arch_spin_unlock(lock);                                         \
        local_irq_restore(flags);                                       \

so in case of a nmi hitting after the spin_lock this can still deadlock?

Hm or is there some assumption that we only use these paths when already in_nmi() and then another nmi can't come in that context?

Quoted 6 lines

Quoted 6 lines

Quoted 3 lines

But even then, flush_nmi_stats() in patch 1 isn't done in_nmi() and uses atomic64_xchg() which in generic_atomic64_xchg() implementation also has the irq_save+spin_lock. So can't we deadlock there?

I was actually assuming that atomic_* ops are safe against nmis for all archs.

We have HAVE_NMI_SAFE_CMPXCHG for this -- there are architectures where this is not the case -- but again, those are typically oddball archs you don't much care about.

But yes, *64 on 32bit archs is generally not NMI safe.

Quoted 4 lines

I looked at atomic_* ops in include/asm-generic/atomic.h and it

Quoted 3 lines

is using arch_cmpxchg() for CONFIG_SMP and it seems like for archs with cmpxchg should be fine against nmi. I am not sure why atomic64_* are not using arch_cmpxchg() instead. I will dig more.

Not many 32bit architectures have 64bit cmpxchg. We're only now dropping support for x86 chips without CMPXCHG8b.

Quoted 8 lines

Yeah I've found https://docs.kernel.org/core-api/local_ops.html and since it listed Mathieu we discussed on IRC and he mentioned the same thing that atomic_ ops are fine, but the later added 64bit variant isn't, which PeterZ (who added it) acknowledged.

But there could be way out if we could somehow compile-time assert that either is true:
- CONFIG_HAVE_NMI=n - we can compile out all the nmi code

Note that in_nmi() is not depending on HAVE_NMI -- nor can it be. Many architectures treat various traps as NMI-like, even though they might not have real NMIs.

Quoted 1 line

- this_cpu is safe on that arch - we can also compile out the nmi code

There is no config symbol for this presently.

Quoted 2 lines

- (if the above leaves any 64bit arch) its 64bit atomics implementation is safe

True, only because HPPA does not in fact have NMIs.

Quoted 4 lines

- (if there are any 32bit applicable arch left) 32bit atomics should be enough for the nmi counters even with >4GB memory as we flush them? and we know the 32bit ops are safe

Older ARM might qualify here.

Quoted 4 lines

IIRC Power64 has issues here, 'funnily' their local_t is NMI safe. Perhaps we could do the same for their this_cpu_*(), but ideally someone with actual power hardware should do this ;-)

Is CONFIG_ARCH_HAS_NMI_SAFE_THIS_CPU_OPS the right config to differentiate between such archs? I see Power64 does not have that enabled.

Quoted 1 line

There is no config symbol for this presently.

Hmm what about CONFIG_ARCH_HAS_NMI_SAFE_THIS_CPU_OPS?

Quoted 4 lines

Quoted 2 lines

- (if the above leaves any 64bit arch) its 64bit atomics implementation is safe

True, only because HPPA does not in fact have NMIs.

What is HPPA?

Thanks a lot Vlastimil & Peter for the suggestions. Let me summarize what I plan to do and please point out if I am doing something wrong:

#if defined(CONFIG_ARCH_HAS_NMI_SAFE_THIS_CPU_OPS) || !defined(CONFIG_HAVE_NMI)

// Do normal this_cpu* ops

#elif defined(ARCH_HAVE_NMI_SAFE_CMPXCHG)

// Do 32 bit atomic ops with in_nmi() checks

#else

// Build error or ignore nmi stats??

#endif