cpumask: add cpumask_nth_{,and,andnot}

Add cpumask_nth_{,and,andnot} as wrappers around corresponding
find functions, and use it in cpumask_local_spread().

Signed-off-by: Yury Norov <yury.norov@gmail.com>

diff --git a/include/linux/cpumask.h b/include/linux/cpumask.h
index 9a71af8097cadeab2801747faacf166241cc086a..e4f9136a4a630634400b0403eec69b1ecbfb69ae 100644 (file)
--- a/include/linux/cpumask.h
+++ b/include/linux/cpumask.h
@@ -337,6 +337,50 @@ unsigned int cpumask_any_but(const struct cpumask *mask, unsigned int cpu)
        return i;
 }
 
+/**
+ * cpumask_nth - get the first cpu in a cpumask
+ * @srcp: the cpumask pointer
+ * @cpu: the N'th cpu to find, starting from 0
+ *
+ * Returns >= nr_cpu_ids if such cpu doesn't exist.
+ */
+static inline unsigned int cpumask_nth(unsigned int cpu, const struct cpumask *srcp)
+{
+       return find_nth_bit(cpumask_bits(srcp), nr_cpumask_bits, cpumask_check(cpu));
+}
+
+/**
+ * cpumask_nth_and - get the first cpu in 2 cpumasks
+ * @srcp1: the cpumask pointer
+ * @srcp2: the cpumask pointer
+ * @cpu: the N'th cpu to find, starting from 0
+ *
+ * Returns >= nr_cpu_ids if such cpu doesn't exist.
+ */
+static inline
+unsigned int cpumask_nth_and(unsigned int cpu, const struct cpumask *srcp1,
+                                                       const struct cpumask *srcp2)
+{
+       return find_nth_and_bit(cpumask_bits(srcp1), cpumask_bits(srcp2),
+                               nr_cpumask_bits, cpumask_check(cpu));
+}
+
+/**
+ * cpumask_nth_andnot - get the first cpu set in 1st cpumask, and clear in 2nd.
+ * @srcp1: the cpumask pointer
+ * @srcp2: the cpumask pointer
+ * @cpu: the N'th cpu to find, starting from 0
+ *
+ * Returns >= nr_cpu_ids if such cpu doesn't exist.
+ */
+static inline
+unsigned int cpumask_nth_andnot(unsigned int cpu, const struct cpumask *srcp1,
+                                                       const struct cpumask *srcp2)
+{
+       return find_nth_andnot_bit(cpumask_bits(srcp1), cpumask_bits(srcp2),
+                               nr_cpumask_bits, cpumask_check(cpu));
+}
+
 #define CPU_BITS_NONE                                          \
 {                                                              \
        [0 ... BITS_TO_LONGS(NR_CPUS)-1] = 0UL                  \

diff --git a/lib/cpumask.c b/lib/cpumask.c
index f0ae119be8c41e3e1f56233ed122f8e2590916b3..2c4a63b6f03f5e496b7cd984072de065f48bbd96 100644 (file)
--- a/lib/cpumask.c
+++ b/lib/cpumask.c
@@ -128,23 +128,21 @@ unsigned int cpumask_local_spread(unsigned int i, int node)
        i %= num_online_cpus();
 
        if (node == NUMA_NO_NODE) {
-               for_each_cpu(cpu, cpu_online_mask)
-                       if (i-- == 0)
-                               return cpu;
+               cpu = cpumask_nth(i, cpu_online_mask);
+               if (cpu < nr_cpu_ids)
+                       return cpu;
        } else {
                /* NUMA first. */
-               for_each_cpu_and(cpu, cpumask_of_node(node), cpu_online_mask)
-                       if (i-- == 0)
-                               return cpu;
-
-               for_each_cpu(cpu, cpu_online_mask) {
-                       /* Skip NUMA nodes, done above. */
-                       if (cpumask_test_cpu(cpu, cpumask_of_node(node)))
-                               continue;
-
-                       if (i-- == 0)
-                               return cpu;
-               }
+               cpu = cpumask_nth_and(i, cpu_online_mask, cpumask_of_node(node));
+               if (cpu < nr_cpu_ids)
+                       return cpu;
+
+               i -= cpumask_weight_and(cpu_online_mask, cpumask_of_node(node));
+
+               /* Skip NUMA nodes, done above. */
+               cpu = cpumask_nth_andnot(i, cpu_online_mask, cpumask_of_node(node));
+               if (cpu < nr_cpu_ids)
+                       return cpu;
        }
        BUG();
 }