Re: [RFC][PATCH] sched: Cache aware load-balancing

Tim Chen
2025-03-31T21:06:19Z

On Tue, 2025-03-25 at 13:09 +0100, Peter Zijlstra wrote:
> Hi all,
>
> One of the many things on the eternal todo list has been finishing the
> below hackery.
>
> It is an attempt at modelling cache affinity -- and while the patch
> really only targets LLC, it could very well be extended to also apply to
> clusters (L2). Specifically any case of multiple cache domains inside a
> node.
>
> Anyway, I wrote this about a year ago, and I mentioned this at the
> recent OSPM conf where Gautham and Prateek expressed interest in playing
> with this code.
>
> So here goes, very rough and largely unproven code ahead :-)
>
> It applies to current tip/master, but I know it will fail the __percpu
> validation that sits in -next, although that shouldn't be terribly hard
> to fix up.
>
> As is, it only computes a CPU inside the LLC that has the highest recent
> runtime, this CPU is then used in the wake-up path to steer towards this
> LLC and in task_hot() to limit migrations away from it.
>
> More elaborate things could be done, notably there is an XXX in there
> somewhere about finding the best LLC inside a NODE (interaction with
> NUMA_BALANCING).

Thanks for sharing the patch.

> +
> +static void task_cache_work(struct callback_head *work)
> +{
> + struct task_struct *p = current;
> + struct mm_struct *mm = p->mm;
> + unsigned long m_a_occ = 0;
> + int cpu, m_a_cpu = -1;
> + cpumask_var_t cpus;
> +
> + WARN_ON_ONCE(work != &p->cache_work);
> +
> + work->next = work;
> +
> + if (p->flags & PF_EXITING)
> + return;
> +
> + if (!alloc_cpumask_var(&cpus, GFP_KERNEL))
> + return;
> +
> + scoped_guard (cpus_read_lock) {
> + cpumask_copy(cpus, cpu_online_mask);
> +

We can constrain the preferred LLC in the same preferred node
as that from numa balancing. Then numa balancing and preferred LLC
won't fight each other if preferred LLC falls on a different node.

Perhaps something like this here

+#ifdef CONFIG_NUMA_BALANCING
+ if (static_branch_likely(&sched_numa_balancing) &&
+ p->numa_preferred_nid != NUMA_NO_NODE)
+ cpumask_and(cpus, cpus, cpumask_of_node(p->numa_preferred_nid));
+#endif

> + for_each_cpu(cpu, cpus) {
> + /* XXX sched_cluster_active */
> + struct sched_domain *sd = per_cpu(sd_llc, cpu);
> + unsigned long occ, m_occ = 0, a_occ = 0;
> + int m_cpu = -1, nr = 0, i;
> +
> + for_each_cpu(i, sched_domain_span(sd)) {
> + occ = fraction_mm_sched(cpu_rq(i),
> + per_cpu_ptr(mm->pcpu_sched, i));
> + a_occ += occ;
> + if (occ > m_occ) {
> + m_occ = occ;
> + m_cpu = i;
> + }
> + nr++;
> + trace_printk("(%d) occ: %ld m_occ: %ld m_cpu: %d nr: %d\n",
> + per_cpu(sd_llc_id, i), occ, m_occ, m_cpu, nr);
> + }
> +
> + a_occ /= nr;

Is it necessary to divide by nr (#CPUs in LLC)? Suppose we have uneven
number of CPUs between LLC, but some total occupancy, we will skew towards
the smaller LLC with this division, which may not be desirable. Could happen
if some CPUs are offlined. I think the preferred LLC can be the one with
most accumulated occupancy.

> + if (a_occ > m_a_occ) {
> + m_a_occ = a_occ;
> + m_a_cpu = m_cpu;
> + }
> +
> + trace_printk("(%d) a_occ: %ld m_a_occ: %ld\n",
> + per_cpu(sd_llc_id, cpu), a_occ, m_a_occ);
> +
> + for_each_cpu(i, sched_domain_span(sd)) {
> + /* XXX threshold ? */
> + per_cpu_ptr(mm->pcpu_sched, i)->occ = a_occ;
> + }
> +
> + cpumask_andnot(cpus, cpus, sched_domain_span(sd));
> + }
> + }
> +
> + /*
> + * If the max average cache occupancy is 'small' we don't care.
> + */
> + if (m_a_occ < (NICE_0_LOAD >> EPOCH_OLD))
> + m_a_cpu = -1;
> +
> + mm->mm_sched_cpu = m_a_cpu;
> +
> + free_cpumask_var(cpus);
> +}
> +
> +void init_sched_mm(struct task_struct *p)
> +{
> + struct callback_head *work = &p->cache_work;
> + init_task_work(work, task_cache_work);
> + work->next = work;
> +}
> +
> +#else
> +
> +static inline void account_mm_sched(struct rq *rq, struct task_struct *p,
> + s64 delta_exec) { }
> +
> +
> +void init_sched_mm(struct task_struct *p) { }
> +
> +static void task_tick_cache(struct rq *rq, struct task_struct *p) { }
> +
> +#endif
> +
> +static inline
> +void update_curr_task(struct rq *rq, struct task_struct *p, s64 delta_exec)
> {
> trace_sched_stat_runtime(p, delta_exec);
> account_group_exec_runtime(p, delta_exec);
> + account_mm_sched(rq, p, delta_exec);
> cgroup_account_cputime(p, delta_exec);
> }
>
> @@ -1215,7 +1434,7 @@ s64 update_curr_common(struct rq *rq)
>
> delta_exec = update_curr_se(rq, &donor->se);
> if (likely(delta_exec > 0))
> - update_curr_task(donor, delta_exec);
> + update_curr_task(rq, donor, delta_exec);
>
> return delta_exec;
> }
> @@ -1244,7 +1463,7 @@ static void update_curr(struct cfs_rq *cfs_rq)
> if (entity_is_task(curr)) {
> struct task_struct *p = task_of(curr);
>
> - update_curr_task(p, delta_exec);
> + update_curr_task(rq, p, delta_exec);
>
> /*
> * If the fair_server is active, we need to account for the
> @@ -7850,7 +8069,7 @@ static int select_idle_sibling(struct task_struct *p, int prev, int target)
> * per-cpu select_rq_mask usage
> */
> lockdep_assert_irqs_disabled();
> -
> +again:
> if ((available_idle_cpu(target) || sched_idle_cpu(target)) &&
> asym_fits_cpu(task_util, util_min, util_max, target))
> return target;
> @@ -7888,7 +8107,8 @@ static int select_idle_sibling(struct task_struct *p, int prev, int target)
> /* Check a recently used CPU as a potential idle candidate: */
> recent_used_cpu = p->recent_used_cpu;
> p->recent_used_cpu = prev;
> - if (recent_used_cpu != prev &&
> + if (prev == p->wake_cpu &&
> + recent_used_cpu != prev &&
> recent_used_cpu != target &&
> cpus_share_cache(recent_used_cpu, target) &&
> (available_idle_cpu(recent_used_cpu) || sched_idle_cpu(recent_used_cpu)) &&
> @@ -7941,6 +8161,18 @@ static int select_idle_sibling(struct task_struct *p, int prev, int target)
> if ((unsigned)i < nr_cpumask_bits)
> return i;
>
> + if (prev != p->wake_cpu && !cpus_share_cache(prev, p->wake_cpu)) {
> + /*
> + * Most likely select_cache_cpu() will have re-directed
> + * the wakeup, but getting here means the preferred cache is
> + * too busy, so re-try with the actual previous.
> + *
> + * XXX wake_affine is lost for this pass.
> + */
> + prev = target = p->wake_cpu;
> + goto again;
> + }
> +
> /*
> * For cluster machines which have lower sharing cache like L2 or
> * LLC Tag, we tend to find an idle CPU in the target's cluster
> @@ -8563,6 +8795,40 @@ static int find_energy_efficient_cpu(struct task_struct *p, int prev_cpu)
> return target;
> }
>
> +#ifdef CONFIG_SCHED_CACHE
> +static long __migrate_degrades_locality(struct task_struct *p, int src_cpu, int dst_cpu, bool idle);
> +
> +static int select_cache_cpu(struct task_struct *p, int prev_cpu)
> +{
> + struct mm_struct *mm = p->mm;
> + int cpu;
> +
> + if (!mm || p->nr_cpus_allowed == 1)
> + return prev_cpu;
> +
> + cpu = mm->mm_sched_cpu;

Some regressions seen when the preferred LLC has significant load, causing frequent task
migrations between preferred and other LLCs. For those cases, tasks should
stay with prev_cpu to avoid migration overhead. Perhaps we could
consider the load in preferred LLC to see if we should move a task there.
Say if preferred LLC is heavily loaded (say nr_running in LLC >= cpus in LLC),
we can stick with prev_cpu and not attempt to switch LLC.

> + if (cpu < 0)
> + return prev_cpu;
> +
> +
> + if (static_branch_likely(&sched_numa_balancing) &&
> + __migrate_degrades_locality(p, prev_cpu, cpu, false) > 0) {
> + /*
> + * XXX look for max occupancy inside prev_cpu's node
> + */
> + return prev_cpu;
> + }
> +
> + return cpu;
> +}
> +#else
> +static int select_cache_cpu(struct task_struct *p, int prev_cpu)
> +{
> + return prev_cpu;
> +}
> +#endif
> +
> +
> /*
> * select_task_rq_fair: Select target runqueue for the waking task in domains
> * that have the relevant SD flag set. In practice, this is SD_BALANCE_WAKE,
> @@ -8588,6 +8854,8 @@ select_task_rq_fair(struct task_struct *p, int prev_cpu, int wake_flags)
> * required for stable ->cpus_allowed
> */
> lockdep_assert_held(&p->pi_lock);
> + guard(rcu)();

Tim