Re: [PATCH v4] sched/fair: Prefer fully idle cores for NOHZ balancing

Hi Andrea, Peter, Prateek,

>> BTW, Andrea, are you planning to do that refactoring of select_rq_mask
>> or you want me to pick it up?
> 
> I was waiting for this patch to land before looking at the select_rq_mask
> refactoring. If you have time and would like to pick it up, please go ahead,
> I'll be happy to review it.
> 
> -Andrea

"Copying peter's reply from v1"
>> Peter, do you have a preference on whether the mask renaming/accessor work
>> should be addressed separately, or included as a preparatory patch for this?

>Separate would be fine. Perhaps another way to do it is using the fancy
>new clang context analysis.

I am not sure if below is what peter meant.
Below diff uses scope based stuff we have in kernel.

This is only a draft, this may need to cleaned up and split into 5 patches.
If the patch below makes sense, please let me know I can work on sending the
series.

Further Notes:
- I checked all the current usecases in sched, all have irq disabled. So
   have lockdep_assert_irqs_disabled in the constructor.
- Its usage is limited to scheduler as for now, there may be few such callsites
   outside. But i don't think its worth making it as generic api usecase.
- Variables names are not well thought off. I just scribbled something. I am sure
   it can be better. Please suggest.
- May need PROVE_LOCKING/debug build check for in_use and those WARN_ON's so that
   it remains minimal. or let it be as is?


I have run hackbench on small 40 CPU system, didn;t see any difference in performance
numbers, not any warnings. Please try and see if it works.


Subject: [PATCH] sched: introduce generic scratch cpumask

---
  kernel/sched/core.c     | 18 +++++++++--
  kernel/sched/deadline.c | 18 +++--------
  kernel/sched/ext/idle.c | 28 +++++------------
  kernel/sched/fair.c     | 42 +++++++++-----------------
  kernel/sched/rt.c       | 15 ++-------
  kernel/sched/sched.h    | 67 +++++++++++++++++++++++++++++++++++++++--
  6 files changed, 109 insertions(+), 79 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 2e7cde033a31..eb44b562abfb 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -130,6 +130,7 @@ EXPORT_TRACEPOINT_SYMBOL_GPL(sched_dl_server_stop_tp);
  
  DEFINE_PER_CPU_SHARED_ALIGNED(struct rq, runqueues);
  DEFINE_PER_CPU(struct rnd_state, sched_rnd_state);
+DEFINE_PER_CPU(struct sched_scratchmask_pool, sched_scratchmask_irq);
  
  #ifdef CONFIG_SCHED_PROXY_EXEC
  DEFINE_STATIC_KEY_TRUE(__sched_proxy_exec);
@@ -8875,9 +8876,6 @@ void __init sched_init_smp(void)
  	current->flags &= ~PF_NO_SETAFFINITY;
  	sched_init_granularity();
  
-	init_sched_rt_class();
-	init_sched_dl_class();
-
  	sched_init_dl_servers();
  
  	sched_smp_initialized = true;
@@ -8897,6 +8895,18 @@ int in_sched_functions(unsigned long addr)
  		&& addr < (unsigned long)__sched_text_end);
  }
  
+void __init init_scratchmasks(void)
+{
+	struct sched_scratchmask_pool *pool;
+	int cpu, slot;
+
+	for_each_possible_cpu(cpu) {
+		pool = per_cpu_ptr(&sched_scratchmask_irq, cpu);
+		for (slot = 0; slot < SCHED_SCRATCHMASK_IRQ_SLOTS; slot++)
+			zalloc_cpumask_var_node(&pool->mask[slot], GFP_KERNEL, cpu_to_node(cpu));
+	}
+}
+
  #ifdef CONFIG_CGROUP_SCHED
  /*
   * Default task group.
@@ -8946,6 +8956,8 @@ void __init sched_init(void)
  
  #endif /* CONFIG_RT_GROUP_SCHED */
  
+	init_scratchmasks();
+
  	init_defrootdomain();
  
  #ifdef CONFIG_RT_GROUP_SCHED
diff --git a/kernel/sched/deadline.c b/kernel/sched/deadline.c
index 857dbe3519a8..5d64536acfef 100644
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -2927,13 +2927,11 @@ static struct task_struct *pick_earliest_pushable_dl_task(struct rq *rq, int cpu
  	return NULL;
  }
  
-/* Access rule: must be called on local CPU with preemption disabled */
-static DEFINE_PER_CPU(cpumask_var_t, local_cpu_mask_dl);
-
  static int find_later_rq(struct task_struct *task)
  {
  	struct sched_domain *sd;
-	struct cpumask *later_mask = this_cpu_cpumask_var_ptr(local_cpu_mask_dl);
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	struct cpumask *later_mask = scratch.mask;
  	int this_cpu = smp_processor_id();
  	int cpu = task_cpu(task);
  
@@ -3382,15 +3380,6 @@ static void rq_offline_dl(struct rq *rq)
  	cpudl_clear(&rq->rd->cpudl, rq->cpu, false);
  }
  
-void __init init_sched_dl_class(void)
-{
-	unsigned int i;
-
-	for_each_possible_cpu(i)
-		zalloc_cpumask_var_node(&per_cpu(local_cpu_mask_dl, i),
-					GFP_KERNEL, cpu_to_node(i));
-}
-
  /*
   * This function always returns a non-empty bitmap in @cpus. This is because
   * if a root domain has reserved bandwidth for DL tasks, the DL bandwidth
@@ -3435,7 +3424,8 @@ void dl_add_task_root_domain(struct task_struct *p)
  		return;
  	}
  
-	msk = this_cpu_cpumask_var_ptr(local_cpu_mask_dl);
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	msk = scratch.mask;
  	dl_get_task_effective_cpus(p, msk);
  	cpu = cpumask_first_and(cpu_active_mask, msk);
  	BUG_ON(cpu >= nr_cpu_ids);
diff --git a/kernel/sched/ext/idle.c b/kernel/sched/ext/idle.c
index 6f93cc32b650..056cb9016129 100644
--- a/kernel/sched/ext/idle.c
+++ b/kernel/sched/ext/idle.c
@@ -47,13 +47,6 @@ static struct scx_idle_cpus scx_idle_global_masks;
   */
  static struct scx_idle_cpus **scx_idle_node_masks;
  
-/*
- * Local per-CPU cpumasks (used to generate temporary idle cpumasks).
- */
-static DEFINE_PER_CPU(cpumask_var_t, local_idle_cpumask);
-static DEFINE_PER_CPU(cpumask_var_t, local_llc_idle_cpumask);
-static DEFINE_PER_CPU(cpumask_var_t, local_numa_idle_cpumask);
-
  /*
   * Return the idle masks associated to a target @node.
   *
@@ -468,8 +461,13 @@ s32 scx_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags,
  	/*
  	 * Determine the subset of CPUs usable by @p within @cpus_allowed.
  	 */
+
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	CLASS(sched_scratchmask_irq, scratch1)(1);
+	CLASS(sched_scratchmask_irq, scratch2)(2);
+
  	if (allowed != p->cpus_ptr) {
-		struct cpumask *local_cpus = this_cpu_cpumask_var_ptr(local_idle_cpumask);
+		struct cpumask *local_cpus = scratch.mask;
  
  		if (task_affinity_all(p)) {
  			allowed = cpus_allowed;
@@ -500,7 +498,7 @@ s32 scx_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags,
  	 * directly.
  	 */
  	if (static_branch_maybe(CONFIG_NUMA, &scx_selcpu_topo_numa)) {
-		struct cpumask *local_cpus = this_cpu_cpumask_var_ptr(local_numa_idle_cpumask);
+		struct cpumask *local_cpus = scratch1.mask;
  		const struct cpumask *cpus = numa_span(prev_cpu);
  
  		if (allowed == p->cpus_ptr && task_affinity_all(p))
@@ -510,7 +508,7 @@ s32 scx_select_cpu_dfl(struct task_struct *p, s32 prev_cpu, u64 wake_flags,
  	}
  
  	if (static_branch_maybe(CONFIG_SCHED_MC, &scx_selcpu_topo_llc)) {
-		struct cpumask *local_cpus = this_cpu_cpumask_var_ptr(local_llc_idle_cpumask);
+		struct cpumask *local_cpus = scratch2.mask;
  		const struct cpumask *cpus = llc_span(prev_cpu);
  
  		if (allowed == p->cpus_ptr && task_affinity_all(p))
@@ -695,16 +693,6 @@ void scx_idle_init_masks(void)
  		BUG_ON(!alloc_cpumask_var_node(&scx_idle_node_masks[i]->cpu, GFP_KERNEL, i));
  		BUG_ON(!alloc_cpumask_var_node(&scx_idle_node_masks[i]->smt, GFP_KERNEL, i));
  	}
-
-	/* Allocate local per-cpu idle cpumasks */
-	for_each_possible_cpu(i) {
-		BUG_ON(!alloc_cpumask_var_node(&per_cpu(local_idle_cpumask, i),
-					       GFP_KERNEL, cpu_to_node(i)));
-		BUG_ON(!alloc_cpumask_var_node(&per_cpu(local_llc_idle_cpumask, i),
-					       GFP_KERNEL, cpu_to_node(i)));
-		BUG_ON(!alloc_cpumask_var_node(&per_cpu(local_numa_idle_cpumask, i),
-					       GFP_KERNEL, cpu_to_node(i)));
-	}
  }
  
  static void update_builtin_idle(int cpu, bool idle)
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 9a975a684b48..770363a59b0c 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -8190,11 +8190,6 @@ static inline unsigned int cfs_h_nr_delayed(struct rq *rq)
  	return (rq->cfs.h_nr_queued - rq->cfs.h_nr_runnable);
  }
  
-/* Working cpumask for: sched_balance_rq(), sched_balance_newidle(). */
-static DEFINE_PER_CPU(cpumask_var_t, load_balance_mask);
-static DEFINE_PER_CPU(cpumask_var_t, select_rq_mask);
-static DEFINE_PER_CPU(cpumask_var_t, should_we_balance_tmpmask);
-
  #ifdef CONFIG_NO_HZ_COMMON
  
  static struct {
@@ -8665,7 +8660,8 @@ static int select_idle_smt(struct task_struct *p, struct sched_domain *sd, int t
   */
  static int select_idle_cpu(struct task_struct *p, struct sched_domain *sd, bool has_idle_core, int target)
  {
-	struct cpumask *cpus = this_cpu_cpumask_var_ptr(select_rq_mask);
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	struct cpumask *cpus = scratch.mask;
  	int i, cpu, idle_cpu = -1, nr = INT_MAX;
  
  	if (sched_feat(SIS_UTIL) && sd->shared) {
@@ -8800,7 +8796,8 @@ select_idle_capacity(struct task_struct *p, struct sched_domain *sd, int target)
  	struct cpumask *cpus;
  	int nr = INT_MAX;
  
-	cpus = this_cpu_cpumask_var_ptr(select_rq_mask);
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	cpus = scratch.mask;
  	cpumask_and(cpus, sched_domain_span(sd), p->cpus_ptr);
  
  	task_util = task_util_est(p);
@@ -8948,9 +8945,6 @@ static int select_idle_sibling(struct task_struct *p, int prev, int target)
  		util_max = uclamp_eff_value(p, UCLAMP_MAX);
  	}
  
-	/*
-	 * per-cpu select_rq_mask usage
-	 */
  	lockdep_assert_irqs_disabled();
  
  	if (choose_idle_cpu(target, p) &&
@@ -9487,7 +9481,8 @@ compute_energy(struct energy_env *eenv, struct perf_domain *pd,
   */
  static int find_energy_efficient_cpu(struct task_struct *p, int prev_cpu)
  {
-	struct cpumask *cpus = this_cpu_cpumask_var_ptr(select_rq_mask);
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	struct cpumask *cpus = scratch.mask;
  	unsigned long prev_delta = ULONG_MAX, best_delta = ULONG_MAX;
  	unsigned long p_util_min = uclamp_is_used() ? uclamp_eff_value(p, UCLAMP_MIN) : 0;
  	unsigned long p_util_max = uclamp_is_used() ? uclamp_eff_value(p, UCLAMP_MAX) : 1024;
@@ -13271,7 +13266,9 @@ static int active_load_balance_cpu_stop(void *data);
  
  static int should_we_balance(struct lb_env *env)
  {
-	struct cpumask *swb_cpus = this_cpu_cpumask_var_ptr(should_we_balance_tmpmask);
+	/* 0 is used already in sched_balance_rq */
+	CLASS(sched_scratchmask_irq, scratch1)(1);
+	struct cpumask *swb_cpus = scratch1.mask;
  	struct sched_group *sg = env->sd->groups;
  	int cpu, idle_smt = -1;
  
@@ -13387,7 +13384,8 @@ static int sched_balance_rq(int this_cpu, struct rq *this_rq,
  	struct sched_group *group;
  	struct rq *busiest;
  	struct rq_flags rf;
-	struct cpumask *cpus = this_cpu_cpumask_var_ptr(load_balance_mask);
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	struct cpumask *cpus = scratch.mask;
  	struct lb_env env = {
  		.sd		= sd,
  		.dst_cpu	= this_cpu,
@@ -13967,13 +13965,9 @@ static inline int find_new_ilb(void)
  	struct cpumask *ilb_cpus;
  	int ilb_cpu, fallback = -1;
  
-	lockdep_assert_irqs_disabled();
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	ilb_cpus = scratch.mask;
  
-	/*
-	 * Reuse the per-CPU select_rq_mask, which is protected from concurrent
-	 * use on this CPU by having interrupts disabled.
-	 */
-	ilb_cpus = this_cpu_cpumask_var_ptr(select_rq_mask);
  	cpumask_and(ilb_cpus, nohz.idle_cpus_mask,
  		    housekeeping_cpumask(HK_TYPE_KERNEL_NOISE));
  
@@ -15574,18 +15568,12 @@ __init void init_sched_fair_class(void)
  {
  	int i;
  
-	for_each_possible_cpu(i) {
-		zalloc_cpumask_var_node(&per_cpu(load_balance_mask, i), GFP_KERNEL, cpu_to_node(i));
-		zalloc_cpumask_var_node(&per_cpu(select_rq_mask,    i), GFP_KERNEL, cpu_to_node(i));
-		zalloc_cpumask_var_node(&per_cpu(should_we_balance_tmpmask, i),
-					GFP_KERNEL, cpu_to_node(i));
-
  #ifdef CONFIG_CFS_BANDWIDTH
+	for_each_possible_cpu(i) {
  		INIT_CSD(&cpu_rq(i)->cfsb_csd, __cfsb_csd_unthrottle, cpu_rq(i));
  		INIT_LIST_HEAD(&cpu_rq(i)->cfsb_csd_list);
-#endif
  	}
-
+#endif
  	open_softirq(SCHED_SOFTIRQ, sched_balance_softirq);
  
  #ifdef CONFIG_NO_HZ_COMMON
diff --git a/kernel/sched/rt.c b/kernel/sched/rt.c
index e6e5f8a2caaf..95b2b20a661d 100644
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -1770,12 +1770,11 @@ static struct task_struct *pick_highest_pushable_task(struct rq *rq, int cpu)
  	return NULL;
  }
  
-static DEFINE_PER_CPU(cpumask_var_t, local_cpu_mask);
-
  static int find_lowest_rq(struct task_struct *task)
  {
  	struct sched_domain *sd;
-	struct cpumask *lowest_mask = this_cpu_cpumask_var_ptr(local_cpu_mask);
+	CLASS(sched_scratchmask_irq, scratch)(0);
+	struct cpumask *lowest_mask = scratch.mask;
  	int this_cpu = smp_processor_id();
  	int cpu      = task_cpu(task);
  	int ret;
@@ -2425,16 +2424,6 @@ static void switched_from_rt(struct rq *rq, struct task_struct *p)
  	rt_queue_pull_task(rq);
  }
  
-void __init init_sched_rt_class(void)
-{
-	unsigned int i;
-
-	for_each_possible_cpu(i) {
-		zalloc_cpumask_var_node(&per_cpu(local_cpu_mask, i),
-					GFP_KERNEL, cpu_to_node(i));
-	}
-}
-
  /*
   * When switching a task to RT, we may overload the runqueue
   * with RT tasks. In this case we try to push them off to
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 26ae13c86b69..afa44303b0da 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1064,6 +1064,7 @@ struct root_domain {
  };
  
  extern void init_defrootdomain(void);
+extern void init_scratchmasks(void);
  extern int sched_init_domains(const struct cpumask *cpu_map);
  extern void rq_attach_root(struct rq *rq, struct root_domain *rd);
  extern void sched_get_rd(struct root_domain *rd);
@@ -2968,8 +2969,6 @@ extern void sysrq_sched_debug_show(void);
  extern void sched_init_granularity(void);
  extern void update_max_interval(void);
  
-extern void init_sched_dl_class(void);
-extern void init_sched_rt_class(void);
  extern void init_sched_fair_class(void);
  
  extern void resched_curr(struct rq *rq);
@@ -4230,4 +4229,68 @@ DEFINE_CLASS_IS_UNCONDITIONAL(sched_change)
  
  #include "ext/ext.h"
  
+#define SCHED_SCRATCHMASK_IRQ_SLOTS	3
+
+struct sched_scratchmask_pool {
+	cpumask_var_t mask[SCHED_SCRATCHMASK_IRQ_SLOTS];
+	unsigned long in_use;	/* bitwise slot usage */
+};
+
+DECLARE_PER_CPU(struct sched_scratchmask_pool, sched_scratchmask_irq);
+
+static __always_inline struct cpumask *
+sched_get_scratchmask_irq(unsigned int slot)
+{
+	struct sched_scratchmask_pool *pool = this_cpu_ptr(&sched_scratchmask_irq);
+
+	lockdep_assert_irqs_disabled();
+
+	if (WARN_ON_ONCE(slot >= SCHED_SCRATCHMASK_IRQ_SLOTS))
+		return NULL;
+
+	if (WARN_ON_ONCE(test_and_set_bit(slot, &pool->in_use)))
+		return NULL;
+
+	return pool->mask[slot];
+}
+
+static __always_inline void
+sched_put_scratchmask_irq(unsigned int slot)
+{
+	struct sched_scratchmask_pool *pool = this_cpu_ptr(&sched_scratchmask_irq);
+
+	lockdep_assert_irqs_disabled();
+
+	if (WARN_ON_ONCE(slot >= SCHED_SCRATCHMASK_IRQ_SLOTS))
+		return;
+
+	WARN_ON_ONCE(!test_and_clear_bit(slot, &pool->in_use));
+}
+
+struct sched_scratchmask {
+	struct cpumask	*mask;
+	unsigned int	slot;
+};
+
+static __always_inline struct sched_scratchmask
+sched_scratchmask_acquire_irq(unsigned int slot)
+{
+	return (struct sched_scratchmask) {
+		.mask = sched_get_scratchmask_irq(slot),
+		.slot = slot,
+	};
+}
+
+static __always_inline void
+sched_scratchmask_release_irq(struct sched_scratchmask *scratch)
+{
+	if (scratch->mask)
+		sched_put_scratchmask_irq(scratch->slot);
+}
+
+DEFINE_CLASS(sched_scratchmask_irq, struct sched_scratchmask,
+	     sched_scratchmask_release_irq(&_T),
+	     sched_scratchmask_acquire_irq(_slot),
+	     unsigned int _slot);
+
  #endif /* _KERNEL_SCHED_SCHED_H */
-- 
2.47.3