sched: add sched_numa_find_nth_cpu()

The function finds Nth set CPU in a given cpumask starting from a given
node.

Leveraging the fact that each hop in sched_domains_numa_masks includes the
same or greater number of CPUs than the previous one, we can use binary
search on hops instead of linear walk, which makes the overall complexity
of O(log n) in terms of number of cpumask_weight() calls.

Signed-off-by: Yury Norov <yury.norov@gmail.com>
Acked-by: Tariq Toukan <tariqt@nvidia.com>
Reviewed-by: Jacob Keller <jacob.e.keller@intel.com>
Reviewed-by: Peter Lafreniere <peter@n8pjl.ca>
Signed-off-by: Jakub Kicinski <kuba@kernel.org>

diff --git a/include/linux/topology.h b/include/linux/topology.h
index 4564faafd0e127eeee0439b093c5f4207daeac32..72f264575698286c1b194e62d7c84fecf43c4cc0 100644 (file)
--- a/include/linux/topology.h
+++ b/include/linux/topology.h
@@ -245,5 +245,13 @@ static inline const struct cpumask *cpu_cpu_mask(int cpu)
        return cpumask_of_node(cpu_to_node(cpu));
 }
 
+#ifdef CONFIG_NUMA
+int sched_numa_find_nth_cpu(const struct cpumask *cpus, int cpu, int node);
+#else
+static __always_inline int sched_numa_find_nth_cpu(const struct cpumask *cpus, int cpu, int node)
+{
+       return cpumask_nth(cpu, cpus);
+}
+#endif /* CONFIG_NUMA */
 
 #endif /* _LINUX_TOPOLOGY_H */

diff --git a/kernel/sched/topology.c b/kernel/sched/topology.c
index 8739c2a5a54eabbff18bb5c7666454e75a34c663..2bf89186a10fa0082b1a1e5bae4e072dd3dde86c 100644 (file)
--- a/kernel/sched/topology.c
+++ b/kernel/sched/topology.c
@@ -3,6 +3,8 @@
  * Scheduler topology setup/handling methods
  */
 
+#include <linux/bsearch.h>
+
 DEFINE_MUTEX(sched_domains_mutex);
 
 /* Protected by sched_domains_mutex: */
@@ -2067,6 +2069,61 @@ unlock:
        return found;
 }
 
+struct __cmp_key {
+       const struct cpumask *cpus;
+       struct cpumask ***masks;
+       int node;
+       int cpu;
+       int w;
+};
+
+static int hop_cmp(const void *a, const void *b)
+{
+       struct cpumask **prev_hop = *((struct cpumask ***)b - 1);
+       struct cpumask **cur_hop = *(struct cpumask ***)b;
+       struct __cmp_key *k = (struct __cmp_key *)a;
+
+       if (cpumask_weight_and(k->cpus, cur_hop[k->node]) <= k->cpu)
+               return 1;
+
+       k->w = (b == k->masks) ? 0 : cpumask_weight_and(k->cpus, prev_hop[k->node]);
+       if (k->w <= k->cpu)
+               return 0;
+
+       return -1;
+}
+
+/*
+ * sched_numa_find_nth_cpu() - given the NUMA topology, find the Nth next cpu
+ *                             closest to @cpu from @cpumask.
+ * cpumask: cpumask to find a cpu from
+ * cpu: Nth cpu to find
+ *
+ * returns: cpu, or nr_cpu_ids when nothing found.
+ */
+int sched_numa_find_nth_cpu(const struct cpumask *cpus, int cpu, int node)
+{
+       struct __cmp_key k = { .cpus = cpus, .node = node, .cpu = cpu };
+       struct cpumask ***hop_masks;
+       int hop, ret = nr_cpu_ids;
+
+       rcu_read_lock();
+
+       k.masks = rcu_dereference(sched_domains_numa_masks);
+       if (!k.masks)
+               goto unlock;
+
+       hop_masks = bsearch(&k, k.masks, sched_domains_numa_levels, sizeof(k.masks[0]), hop_cmp);
+       hop = hop_masks - k.masks;
+
+       ret = hop ?
+               cpumask_nth_and_andnot(cpu - k.w, cpus, k.masks[hop][node], k.masks[hop-1][node]) :
+               cpumask_nth_and(cpu, cpus, k.masks[0][node]);
+unlock:
+       rcu_read_unlock();
+       return ret;
+}
+EXPORT_SYMBOL_GPL(sched_numa_find_nth_cpu);
 #endif /* CONFIG_NUMA */
 
 static int __sdt_alloc(const struct cpumask *cpu_map)