bpf: Simplify code by using for_each_cpu_wrap()

In the percpu freelist code, it is a common pattern to iterate over
the possible CPUs mask starting with the current CPU. The pattern is
implemented using a hand rolled while loop with the loop variable
increment being open-coded.

Simplify the code by using for_each_cpu_wrap() helper to iterate over
the possible cpus starting with the current CPU. As a result, some of
the special-casing in the loop also gets simplified.

No functional change intended.

Signed-off-by: Punit Agrawal <punit.agrawal@bytedance.com>
Acked-by: Song Liu <song@kernel.org>
Link: https://lore.kernel.org/r/20220907155746.1750329-1-punit.agrawal@bytedance.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/kernel/bpf/percpu_freelist.c b/kernel/bpf/percpu_freelist.c
index 00b874c8e889f1a019d553fbca2bff34c0527864..b6e7f5c5b9ab5712783b3c7875f6d78994daa038 100644 (file)
--- a/kernel/bpf/percpu_freelist.c
+++ b/kernel/bpf/percpu_freelist.c
@@ -58,23 +58,21 @@ static inline void ___pcpu_freelist_push_nmi(struct pcpu_freelist *s,
 {
        int cpu, orig_cpu;
 
-       orig_cpu = cpu = raw_smp_processor_id();
+       orig_cpu = raw_smp_processor_id();
        while (1) {
-               struct pcpu_freelist_head *head;
+               for_each_cpu_wrap(cpu, cpu_possible_mask, orig_cpu) {
+                       struct pcpu_freelist_head *head;
 
-               head = per_cpu_ptr(s->freelist, cpu);
-               if (raw_spin_trylock(&head->lock)) {
-                       pcpu_freelist_push_node(head, node);
-                       raw_spin_unlock(&head->lock);
-                       return;
+                       head = per_cpu_ptr(s->freelist, cpu);
+                       if (raw_spin_trylock(&head->lock)) {
+                               pcpu_freelist_push_node(head, node);
+                               raw_spin_unlock(&head->lock);
+                               return;
+                       }
                }
-               cpu = cpumask_next(cpu, cpu_possible_mask);
-               if (cpu >= nr_cpu_ids)
-                       cpu = 0;
 
                /* cannot lock any per cpu lock, try extralist */
-               if (cpu == orig_cpu &&
-                   pcpu_freelist_try_push_extra(s, node))
+               if (pcpu_freelist_try_push_extra(s, node))
                        return;
        }
 }
@@ -125,13 +123,12 @@ static struct pcpu_freelist_node *___pcpu_freelist_pop(struct pcpu_freelist *s)
 {
        struct pcpu_freelist_head *head;
        struct pcpu_freelist_node *node;
-       int orig_cpu, cpu;
+       int cpu;
 
-       orig_cpu = cpu = raw_smp_processor_id();
-       while (1) {
+       for_each_cpu_wrap(cpu, cpu_possible_mask, raw_smp_processor_id()) {
                head = per_cpu_ptr(s->freelist, cpu);
                if (!READ_ONCE(head->first))
-                       goto next_cpu;
+                       continue;
                raw_spin_lock(&head->lock);
                node = head->first;
                if (node) {
@@ -140,12 +137,6 @@ static struct pcpu_freelist_node *___pcpu_freelist_pop(struct pcpu_freelist *s)
                        return node;
                }
                raw_spin_unlock(&head->lock);
-next_cpu:
-               cpu = cpumask_next(cpu, cpu_possible_mask);
-               if (cpu >= nr_cpu_ids)
-                       cpu = 0;
-               if (cpu == orig_cpu)
-                       break;
        }
 
        /* per cpu lists are all empty, try extralist */
@@ -164,13 +155,12 @@ ___pcpu_freelist_pop_nmi(struct pcpu_freelist *s)
 {
        struct pcpu_freelist_head *head;
        struct pcpu_freelist_node *node;
-       int orig_cpu, cpu;
+       int cpu;
 
-       orig_cpu = cpu = raw_smp_processor_id();
-       while (1) {
+       for_each_cpu_wrap(cpu, cpu_possible_mask, raw_smp_processor_id()) {
                head = per_cpu_ptr(s->freelist, cpu);
                if (!READ_ONCE(head->first))
-                       goto next_cpu;
+                       continue;
                if (raw_spin_trylock(&head->lock)) {
                        node = head->first;
                        if (node) {
@@ -180,12 +170,6 @@ ___pcpu_freelist_pop_nmi(struct pcpu_freelist *s)
                        }
                        raw_spin_unlock(&head->lock);
                }
-next_cpu:
-               cpu = cpumask_next(cpu, cpu_possible_mask);
-               if (cpu >= nr_cpu_ids)
-                       cpu = 0;
-               if (cpu == orig_cpu)
-                       break;
        }
 
        /* cannot pop from per cpu lists, try extralist */