Merge tag 'sched_ext-for-7.3' of git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext

Pull sched_ext updates from Tejun Heo:
 "Most of this cycle completes the enqueue-path support for hierarchical
  sub-scheduling, which makes sub-scheduler support feature complete: a
  root BPF scheduler can now hand a cgroup subtree over to a nested
  sub-scheduler together with revocable CPU grants, and the
  sub-scheduler owns all scheduling decisions for its tasks on those
  CPUs.

  Development volume was high and a number of changes plugging holes in
  the new support landed late in the cycle. Also included are core
  scheduling fixes that were completed too late for the v7.2 release and
  are routed through this pull request.

  Sub-scheduler CPU delegation:

   - Parent schedulers now grant and revoke per-CPU capabilities
     (enqueueing, preemption, CPU frequency control) on their children,
     enforced on every path a scheduler can reach a CPU through.
     Previously only dispatching could be delegated; this lets
     sub-schedulers fully schedule their CPUs.

   - Rescue execution: a task whose scheduler doesn't have access to the
     CPUs the task needs to run on starved until the watchdog ejected
     the whole scheduler. The kernel now runs such tasks directly on a
     small bandwidth budget, turning a scheduler-killing failure into
     bounded degradation.

   - Cgroup integration: tasks migrating across a sub-scheduler boundary
     weren't re-homed to the new owner, causing wrong-scheduler
     scheduling and a use-after-free. Sub-schedulers now take over their
     cgroup subtree and receive its cgroup callbacks.

   - Arena objects now cross the kernel/BPF boundary as typed pointer
     arguments, translated transparently by the BPF tree's new arena
     argument support, replacing untyped arguments with manual
     translation.

   - scx_qmap now demonstrates full hierarchical sub-scheduling.

  Other fixes and updates:

   - Robustness improvements: the abort path is now NMI-safe, fixing
     deadlocks when errors are raised from NMI context and making
     hardlockup recovery direct. Reenqueue loops that could monopolize a
     CPU ahead of the watchdog now eject the offending scheduler, and
     stalls are blamed on the scheduler actually responsible.

   - Hardening: BPF-writable arena memory is validated before kernel
     use, and task slice and vtime writes got explicit synchronization
     rules, closing corruption vectors open to buggy or malicious
     schedulers.

   - Core scheduling: sched_ext dispatching can drop the rq lock inside
     the core-wide pick, which let interleaving selections corrupt each
     other's state and hard-hang the machine. The selection now restarts
     when the lock was released. The task ordering callback was also
     invoked with its arguments swapped, and the default ordering is
     updated to work across sub-scheduler boundaries. The fixes are
     marked for stable.

   - Other fixes headed for stable: a task init leak on fork failure
     during enable, tooling compat macros that silently failed to detect
     newer kernels, and a crash on reenqueueing against a destroyed
     dispatch queue.

   - Tooling: scx_pair moves off deprecated callbacks, and the
     deprecated scx_bpf_cpu_rq() kfunc is removed"

* tag 'sched_ext-for-7.3' of git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext: (144 commits)
  sched_ext: Drop the dead SCX_DEQ_CORE_SCHED_EXEC test in dequeue_task_scx()
  sched_ext: Make core-sched task ordering hierarchy-aware
  sched_ext: Use runnable_at for the default core-sched task ordering
  sched_ext: Fix inverted ops.core_sched_before() invocation
  sched_ext: Move the config-off sub-cap kfunc stubs into sub.c
  sched_ext: Rename balance-era identifiers to dispatch terms
  sched_ext: Drop the stale keep_prev fixup in dispatch_pick()
  sched_ext: Keep kick_sync waiting on the rq's own CPU
  sched_ext: Make SCHED_CLASS_EXT select GENERIC_ALLOCATOR
  sched_ext/scx_flatcg: Fix cvtime true-up on slice expiry
  sched_ext: Don't BUG_ON a destroyed DSQ in process_deferred_reenq_users
  sched_ext: Fix scx_bpf_dsq_move_to_local___v2 compat detection
  sched_ext: Make scx_bpf_events() read the calling scheduler's counters
  sched_ext: Drop unlocked scx_rq_clock_invalidate() from scx_root_disable()
  selftests/sched_ext: Fix flaky ddsp failure tests on busy systems
  selftests/sched_ext: Make numa idle validation race-free
  sched_ext: Fix scx_bpf_dsq_reenq___compat kfunc extern prototype
  sched_ext/scx_flatcg: expire cached hweights on weight changes
  sched_ext: Fix exit_task leak on fork failure during enable
  sched_ext: fix stale references in doc comments
  ...

19 files changed, 6886 insertions, 2228 deletions

diff --git a/kernel/Kconfig.preempt b/kernel/Kconfig.preempt
index 35f546a042b1..f294dad43bd7 100644
--- a/kernel/Kconfig.preempt
+++ b/kernel/Kconfig.preempt
@@ -173,6 +173,7 @@ config SCHED_CORE
 config SCHED_CLASS_EXT
 	bool "Extensible Scheduling Class"
 	depends on BPF_SYSCALL && BPF_JIT && DEBUG_INFO_BTF
+	select GENERIC_ALLOCATOR
 	select STACKTRACE if STACKTRACE_SUPPORT
 	help
 	  This option enables a new scheduler class sched_ext (SCX), which
diff --git a/kernel/cgroup/cgroup.c b/kernel/cgroup/cgroup.c
index 98c536f8b666..c3a12fee7528 100644
--- a/kernel/cgroup/cgroup.c
+++ b/kernel/cgroup/cgroup.c
@@ -88,6 +88,8 @@ EXPORT_SYMBOL_GPL(css_set_lock);
 
 struct blocking_notifier_head cgroup_lifetime_notifier =
 	BLOCKING_NOTIFIER_INIT(cgroup_lifetime_notifier);
+struct blocking_notifier_head cgroup_task_notifier =
+	BLOCKING_NOTIFIER_INIT(cgroup_task_notifier);
 
 DEFINE_SPINLOCK(trace_cgroup_path_lock);
 char trace_cgroup_path[TRACE_CGROUP_PATH_LEN];
@@ -2676,14 +2678,27 @@ struct task_struct *cgroup_taskset_next(struct cgroup_taskset *tset,
 	return NULL;
 }
 
+static void cgroup_migrate_notify_canceled(struct css_set *src_cset,
+					   struct task_struct *task)
+{
+	struct cgroup_task_migrate_ctx ctx = {
+		.task = task,
+		.src_dcgrp = src_cset->dfl_cgrp,
+		.dst_dcgrp = src_cset->mg_dst_cset->dfl_cgrp,
+	};
+
+	blocking_notifier_call_chain(&cgroup_task_notifier,
+				     CGROUP_TASK_MIGRATE_CANCELED, &ctx);
+}
+
 /**
  * cgroup_migrate_execute - migrate a taskset
  * @mgctx: migration context
  *
- * Migrate tasks in @mgctx as setup by migration preparation functions.
- * This function fails iff one of the ->can_attach callbacks fails and
- * guarantees that either all or none of the tasks in @mgctx are migrated.
- * @mgctx is consumed regardless of success.
+ * Migrate tasks in @mgctx as setup by migration preparation functions. This
+ * function fails iff one of the ->can_attach callbacks or CGROUP_TASK_MIGRATING
+ * notifications fails and guarantees that either all or none of the tasks in
+ * @mgctx are migrated. @mgctx is consumed regardless of success.
  */
 static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 {
@@ -2691,6 +2706,7 @@ static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 	struct cgroup_subsys *ss;
 	struct task_struct *task, *tmp_task;
 	struct css_set *cset, *tmp_cset;
+	bool dfl_migration = false;
 	int ssid, failed_ssid, ret;
 
 	/* check that we can legitimately attach to the cgroup */
@@ -2708,6 +2724,33 @@ static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 	}
 
 	/*
+	 * Notify each task about the impending migration. An error return fails
+	 * the migration. Only migrations on the default hierarchy are reported:
+	 * a migration modifies either every moved task's dfl cgroup or, on
+	 * cgroup1 or for subtree_control writes, none.
+	 */
+	list_for_each_entry(cset, &tset->src_csets, mg_node) {
+		if (cset->dfl_cgrp == cset->mg_dst_cset->dfl_cgrp)
+			continue;
+		dfl_migration = true;
+		list_for_each_entry(task, &cset->mg_tasks, cg_list) {
+			struct cgroup_task_migrate_ctx ctx = {
+				.task = task,
+				.src_dcgrp = cset->dfl_cgrp,
+				.dst_dcgrp = cset->mg_dst_cset->dfl_cgrp,
+			};
+
+			ret = blocking_notifier_call_chain_robust(&cgroup_task_notifier,
+								  CGROUP_TASK_MIGRATING,
+								  CGROUP_TASK_MIGRATE_CANCELED,
+								  &ctx);
+			ret = notifier_to_errno(ret);
+			if (ret)
+				goto out_cancel_migrating;
+		}
+	}
+
+	/*
 	 * Now that we're guaranteed success, proceed to move all tasks to
 	 * the new cgroup.  There are no failure cases after here, so this
 	 * is the commit point.
@@ -2750,9 +2793,41 @@ static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 		} while_each_subsys_mask();
 	}
 
+	/*
+	 * Notify each task after successful migration. The operation can no
+	 * longer fail and the return value is ignored. The MIGRATING loop
+	 * above explains why only dfl migrations are reported. Per-task
+	 * sources are not tracked past the commit point, so src_dcgrp is
+	 * NULL.
+	 */
+	if (dfl_migration) {
+		list_for_each_entry(cset, &tset->dst_csets, mg_node) {
+			list_for_each_entry(task, &cset->mg_tasks, cg_list) {
+				struct cgroup_task_migrate_ctx ctx = {
+					.task = task,
+					.dst_dcgrp = cset->dfl_cgrp,
+				};
+
+				blocking_notifier_call_chain(
+					&cgroup_task_notifier,
+					CGROUP_TASK_MIGRATED, &ctx);
+			}
+		}
+	}
+
 	ret = 0;
 	goto out_release_tset;
 
+out_cancel_migrating:
+	list_for_each_entry_continue_reverse(task, &cset->mg_tasks, cg_list)
+		cgroup_migrate_notify_canceled(cset, task);
+	list_for_each_entry_continue_reverse(cset, &tset->src_csets, mg_node) {
+		if (cset->dfl_cgrp == cset->mg_dst_cset->dfl_cgrp)
+			continue;
+		list_for_each_entry_reverse(task, &cset->mg_tasks, cg_list)
+			cgroup_migrate_notify_canceled(cset, task);
+	}
+	failed_ssid = CGROUP_SUBSYS_COUNT;
 out_cancel_attach:
 	if (tset->nr_tasks) {
 		do_each_subsys_mask(ss, ssid, mgctx->ss_mask) {
@@ -2976,11 +3051,11 @@ int cgroup_migrate_prepare_dst(struct cgroup_mgctx *mgctx)
  * cgroup_migrate_prepare_dst() on the targets before invoking this
  * function and following up with cgroup_migrate_finish().
  *
- * As long as a controller's ->can_attach() doesn't fail, this function is
- * guaranteed to succeed.  This means that, excluding ->can_attach()
- * failure, when migrating multiple targets, the success or failure can be
- * decided for all targets by invoking group_migrate_prepare_dst() before
- * actually starting migrating.
+ * As long as a controller's ->can_attach() or a CGROUP_TASK_MIGRATING
+ * notification doesn't fail, this function is guaranteed to succeed.  This
+ * means that, excluding those failures, when migrating multiple targets,
+ * the success or failure can be decided for all targets by invoking
+ * group_migrate_prepare_dst() before actually starting migrating.
  */
 int cgroup_migrate(struct task_struct *leader, bool threadgroup,
 		   struct cgroup_mgctx *mgctx)
diff --git a/kernel/rcu/tree.c b/kernel/rcu/tree.c
index 03a43d3d2616..415583c35f8c 100644
--- a/kernel/rcu/tree.c
+++ b/kernel/rcu/tree.c
@@ -4900,6 +4900,9 @@ static void __init rcu_dump_rcu_node_tree(void)
 
 struct workqueue_struct *rcu_gp_wq;
 
+static struct cpumask rcu_stall_cpumask;
+static struct cpumask rcu_exp_stall_cpumask;
+
 void __init rcu_init(void)
 {
 	int cpu = smp_processor_id();
diff --git a/kernel/rcu/tree_exp.h b/kernel/rcu/tree_exp.h
index 82cada459e5d..46b6907f1b09 100644
--- a/kernel/rcu/tree_exp.h
+++ b/kernel/rcu/tree_exp.h
@@ -578,6 +578,7 @@ static void synchronize_rcu_expedited_stall(unsigned long jiffies_start, unsigne
 			if (!(READ_ONCE(rnp->expmask) & mask))
 				continue;
 			ndetected++;
+			cpumask_set_cpu(cpu, &rcu_exp_stall_cpumask);
 			rdp = per_cpu_ptr(&rcu_data, cpu);
 			pr_cont(" %d-%c%c%c%c", cpu,
 				"O."[!!cpu_online(cpu)],
@@ -665,6 +666,8 @@ static void synchronize_rcu_expedited_wait(void)
 		if (rcu_stall_is_suppressed())
 			continue;
 
+		cpumask_clear(&rcu_exp_stall_cpumask);
+
 		nbcon_cpu_emergency_enter();
 
 		j = jiffies;
@@ -675,7 +678,7 @@ static void synchronize_rcu_expedited_wait(void)
 
 		nbcon_cpu_emergency_exit();
 
-		panic_on_rcu_stall();
+		panic_on_rcu_stall(&rcu_exp_stall_cpumask);
 	}
 }
 
diff --git a/kernel/rcu/tree_stall.h b/kernel/rcu/tree_stall.h
index cf7ae51cba40..ebf381936eb1 100644
--- a/kernel/rcu/tree_stall.h
+++ b/kernel/rcu/tree_stall.h
@@ -159,7 +159,7 @@ static int __init check_cpu_stall_init(void)
 early_initcall(check_cpu_stall_init);
 
 /* If so specified via sysctl, panic, yielding cleaner stall-warning output. */
-static void panic_on_rcu_stall(void)
+static void panic_on_rcu_stall(const struct cpumask *stalled_mask)
 {
 	static int cpu_stall;
 
@@ -167,7 +167,7 @@ static void panic_on_rcu_stall(void)
 	 * Attempt to kick out the BPF scheduler if it's installed and defer
 	 * the panic to give the system a chance to recover.
 	 */
-	if (scx_rcu_cpu_stall())
+	if (scx_rcu_cpu_stall(stalled_mask))
 		return;
 
 	if (++cpu_stall < sysctl_max_rcu_stall_to_panic)
@@ -644,6 +644,8 @@ static void print_other_cpu_stall(unsigned long gp_seq, unsigned long gps)
 	if (rcu_stall_is_suppressed())
 		return;
 
+	cpumask_clear(&rcu_stall_cpumask);
+
 	nbcon_cpu_emergency_enter();
 
 	/*
@@ -659,6 +661,7 @@ static void print_other_cpu_stall(unsigned long gp_seq, unsigned long gps)
 			for_each_leaf_node_possible_cpu(rnp, cpu)
 				if (rnp->qsmask & leaf_node_cpu_bit(rnp, cpu)) {
 					print_cpu_stall_info(cpu);
+					cpumask_set_cpu(cpu, &rcu_stall_cpumask);
 					ndetected++;
 				}
 		}
@@ -700,7 +703,7 @@ static void print_other_cpu_stall(unsigned long gp_seq, unsigned long gps)
 
 	nbcon_cpu_emergency_exit();
 
-	panic_on_rcu_stall();
+	panic_on_rcu_stall(&rcu_stall_cpumask);
 
 	rcu_force_quiescent_state();  /* Kick them all. */
 }
@@ -753,7 +756,9 @@ static void print_cpu_stall(unsigned long gp_seq, unsigned long gps)
 
 	nbcon_cpu_emergency_exit();
 
-	panic_on_rcu_stall();
+	cpumask_clear(&rcu_stall_cpumask);
+	cpumask_set_cpu(smp_processor_id(), &rcu_stall_cpumask);
+	panic_on_rcu_stall(&rcu_stall_cpumask);
 
 	/*
 	 * Attempt to revive the RCU machinery by forcing a context switch.
diff --git a/kernel/sched/build_policy.c b/kernel/sched/build_policy.c
index d74b54f81992..2a828725a7f9 100644
--- a/kernel/sched/build_policy.c
+++ b/kernel/sched/build_policy.c
@@ -66,10 +66,13 @@
 # include "ext/cid.h"
 # include "ext/arena.h"
 # include "ext/idle.h"
+# include "ext/sub.h"
+# include "ext/inlines.h"
 # include "ext/ext.c"
 # include "ext/cid.c"
 # include "ext/arena.c"
 # include "ext/idle.c"
+# include "ext/sub.c"
 #endif
 
 #include "syscalls.c"
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 6544e56925ee..4520d63763d4 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -443,6 +443,17 @@ static void __sched_core_flip(bool enabled)
 
 		sched_core_lock(cpu, &flags);
 
+		/*
+		 * A core-wide selection may have the shared rq lock temporarily
+		 * released by a lock-dropping ->pick_task(). Flipping would
+		 * rebind rq_lockp() under it. Wait it out.
+		 */
+		while (cpu_rq(cpu)->core->core_pick_in_flight) {
+			sched_core_unlock(cpu, &flags);
+			cpu_relax();
+			sched_core_lock(cpu, &flags);
+		}
+
 		for_each_cpu(t, smt_mask)
 			cpu_rq(t)->core_enabled = enabled;
 
@@ -6228,7 +6239,7 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 	unsigned long cookie;
 	int i, cpu, occ = 0;
 	struct rq *rq_i;
-	bool need_sync;
+	bool need_sync = false;
 
 	if (!sched_core_enabled(rq))
 		return __pick_next_task(rq, rf);
@@ -6247,6 +6258,8 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 		return __pick_next_task(rq, rf);
 	}
 
+	rq->core->core_pick_in_flight++;
+
 	/*
 	 * If there were no {en,de}queues since we picked (IOW, the task
 	 * pointers are all still valid), and we haven't scheduled the last
@@ -6271,7 +6284,9 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 	prev_balance(rq, rf);
 
 	smt_mask = cpu_smt_mask(cpu);
-	need_sync = !!rq->core->core_cookie;
+
+restart:
+	need_sync |= !!rq->core->core_cookie;
 
 	/* reset state */
 	rq->core->core_cookie = 0UL;
@@ -6306,10 +6321,15 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 	 * and there are no cookied tasks running on siblings.
 	 */
 	if (!need_sync) {
-restart_single:
 		next = pick_task(rq, rf);
-		if (unlikely(next == RETRY_TASK))
-			goto restart_single;
+		if (unlikely(next == RETRY_TASK)) {
+			/* rq lock may have been dropped, clocks invalidated */
+			core_clock_updated = false;
+			if (!(rq->clock_update_flags & RQCF_UPDATED))
+				update_rq_clock(rq);
+			goto restart;
+		}
+
 		if (!next->core_cookie) {
 			rq->core_pick = NULL;
 			rq->core_dl_server = NULL;
@@ -6329,7 +6349,6 @@ restart_single:
 	 *
 	 * Tie-break prio towards the current CPU
 	 */
-restart_multi:
 	max = NULL;
 	for_each_cpu_wrap(i, smt_mask, cpu) {
 		rq_i = cpu_rq(i);
@@ -6343,8 +6362,13 @@ restart_multi:
 			update_rq_clock(rq_i);
 
 		p = pick_task(rq_i, rf);
-		if (unlikely(p == RETRY_TASK))
-			goto restart_multi;
+		if (unlikely(p == RETRY_TASK)) {
+			/* rq lock may have been dropped, clocks invalidated */
+			core_clock_updated = false;
+			if (!(rq->clock_update_flags & RQCF_UPDATED))
+				update_rq_clock(rq);
+			goto restart;
+		}
 
 		rq_i->core_pick = p;
 		rq_i->core_dl_server = rq_i->dl_server;
@@ -6450,6 +6474,7 @@ restart_multi:
 	}
 
 out_set_next:
+	rq->core->core_pick_in_flight--;
 	put_prev_set_next_task(rq, rq->donor, next);
 	if (rq->core->core_forceidle_count && next == rq->idle)
 		queue_core_balance(rq);
@@ -6645,6 +6670,13 @@ static void sched_core_cpu_deactivate(unsigned int cpu)
 	core_rq->core_forceidle_occupation = rq->core_forceidle_occupation;
 
 	/*
+	 * A stale leftover would bias the count forever if this CPU later
+	 * returns as its own leader. Move, don't copy.
+	 */
+	core_rq->core_pick_in_flight       = rq->core_pick_in_flight;
+	rq->core_pick_in_flight            = 0;
+
+	/*
 	 * Accounting edge for forced idle is handled in pick_next_task().
 	 * Don't need another one here, since the hotplug thread shouldn't
 	 * have a cookie.
@@ -9058,6 +9090,7 @@ void __init sched_init(void)
 		rq->core_forceidle_count = 0;
 		rq->core_forceidle_occupation = 0;
 		rq->core_forceidle_start = 0;
+		rq->core_pick_in_flight = 0;
 
 		rq->core_cookie = 0UL;
 #endif
diff --git a/kernel/sched/ext/arena.c b/kernel/sched/ext/arena.c
index 5783694ec21d..f7a9f633f435 100644
--- a/kernel/sched/ext/arena.c
+++ b/kernel/sched/ext/arena.c
@@ -6,8 +6,7 @@
  *
  * Each chunk added to @sch->arena_pool comes from one
  * bpf_arena_alloc_pages_sleepable() call and is registered at the
- * kernel-side mapping address. Callers translate to the BPF-arena form
- * themselves if needed.
+ * kernel-side mapping address.
  *
  * Allocations grow the pool on demand. Underlying arena pages are released
  * when the arena map itself is torn down.
@@ -70,8 +69,6 @@ void scx_arena_pool_destroy(struct scx_sched *sch)
  */
 static int scx_arena_grow(struct scx_sched *sch, u32 page_cnt)
 {
-	u64 kern_vm_start;
-	u32 uaddr32;
 	void *p;
 	int ret;
 
@@ -83,15 +80,8 @@ static int scx_arena_grow(struct scx_sched *sch, u32 page_cnt)
 	if (!p)
 		return -ENOMEM;
 
-	uaddr32 = (u32)(unsigned long)p;
-	/* arena.o, which defines these, is built only on MMU && 64BIT */
-#if defined(CONFIG_MMU) && defined(CONFIG_64BIT)
-	kern_vm_start = bpf_arena_map_kern_vm_start(sch->arena_map);
-#else
-	kern_vm_start = 0;
-#endif
-
-	ret = gen_pool_add(sch->arena_pool, kern_vm_start + uaddr32,
+	ret = gen_pool_add(sch->arena_pool,
+			   (unsigned long)scx_arena_to_kaddr(sch, p),
 			   page_cnt * PAGE_SIZE, NUMA_NO_NODE);
 	if (ret) {
 		bpf_arena_free_pages_non_sleepable(sch->arena_map, p, page_cnt);
diff --git a/kernel/sched/ext/cid.c b/kernel/sched/ext/cid.c
index af83084ec740..39f88deb94bc 100644
--- a/kernel/sched/ext/cid.c
+++ b/kernel/sched/ext/cid.c
@@ -11,20 +11,26 @@
 #include "cid.h"
 
 /*
- * cid tables.
- *
- * Pointers are published once on first enable and never revoked. The default
- * mapping is populated before ops.init() runs; scx_bpf_cid_override() commits
- * before it returns. As long as the BPF scheduler only uses the tables from
- * those points onward, it sees a consistent view.
+ * cid tables. The cid kfuncs are available whether the root scheduler is
+ * cid-form or cpu-form, the latter to allow gradual migration to cids, so every
+ * root builds a default mapping. Each root enable allocates a fresh set, builds
+ * it privately and publishes the __rcu globals below once the layout is final.
+ * Root disable unpublishes and RCU-frees the set. kfuncs may run before the
+ * tables are published and must check for NULL.
  */
-s16 *scx_cid_to_cpu_tbl;
-s16 *scx_cpu_to_cid_tbl;
-struct scx_cid_topo *scx_cid_topo;
+u32 scx_nr_cid_shards;
+s16 __rcu *scx_cid_to_cpu_tbl;
+s16 __rcu *scx_cpu_to_cid_tbl;
+s32 __rcu *scx_cid_to_shard;
+s32 __rcu *scx_shard_node;
+struct scx_cid_shard __rcu *scx_cid_shard_ranges;
+struct scx_cid_topo __rcu *scx_cid_topo;
+
+static struct scx_cid_tables *scx_cid_tables;	/* used only during alloc/free */
 
 #define SCX_CID_TOPO_NEG	(struct scx_cid_topo) {				\
 	.core_cid = -1, .core_idx = -1, .llc_cid = -1, .llc_idx = -1,		\
-	.node_cid = -1, .node_idx = -1,						\
+	.node_cid = -1, .node_idx = -1, .shard_cid = -1, .shard_idx = -1,	\
 }
 
 /*
@@ -43,37 +49,128 @@ static const struct cpumask *cpu_llc_mask(int cpu, struct cpumask *fallbacks)
 	return &ci->info_list[ci->num_leaves - 1].shared_cpu_map;
 }
 
-/* Allocate the cid tables once on first enable; never freed. */
-static s32 scx_cid_arrays_alloc(void)
+/*
+ * Compute per-LLC shard layout. Each shard holds at most @shard_size cids, and
+ * in any case no more than SCX_CID_SHARD_MAX_CPUS. Cores are spread as evenly
+ * as possible across shards so cpu count is balanced: the first *@nr_large_p
+ * shards get (*@cores_per_shard_p + 1) cores, the rest get *@cores_per_shard_p.
+ */
+static void calc_shard_layout(const struct cpumask *llc_cpus, u32 shard_size,
+			      u32 *cores_per_shard_p, u32 *nr_large_p)
 {
-	u32 npossible = num_possible_cpus();
-	s16 *cid_to_cpu, *cpu_to_cid;
-	struct scx_cid_topo *cid_topo;
+	u32 nr_cores = 0, nr_cpus = 0, nr_shards;
+	int cpu;
 
-	if (scx_cid_to_cpu_tbl)
-		return 0;
+	for_each_cpu(cpu, llc_cpus) {
+		nr_cpus++;
+		if (cpumask_first(topology_sibling_cpumask(cpu)) == cpu)
+			nr_cores++;
+	}
 
-	cid_to_cpu = kzalloc_objs(*scx_cid_to_cpu_tbl, npossible, GFP_KERNEL);
-	cpu_to_cid = kzalloc_objs(*scx_cpu_to_cid_tbl, nr_cpu_ids, GFP_KERNEL);
-	cid_topo = kmalloc_objs(*scx_cid_topo, npossible, GFP_KERNEL);
+	nr_shards = max_t(u32, 1, DIV_ROUND_UP(nr_cpus, shard_size));
+	nr_shards = max_t(u32, nr_shards,
+			  DIV_ROUND_UP(nr_cpus, SCX_CID_SHARD_MAX_CPUS));
 
-	if (!cid_to_cpu || !cpu_to_cid || !cid_topo) {
-		kfree(cid_to_cpu);
-		kfree(cpu_to_cid);
-		kfree(cid_topo);
-		return -ENOMEM;
+	*cores_per_shard_p = nr_cores / nr_shards;
+	*nr_large_p = nr_cores % nr_shards;
+}
+
+static void scx_cid_tables_free(struct scx_cid_tables *tbls)
+{
+	if (!tbls)
+		return;
+	kvfree(tbls->cid_to_cpu);
+	kvfree(tbls->cpu_to_cid);
+	kvfree(tbls->cid_to_shard);
+	kvfree(tbls->shard_node);
+	kvfree(tbls->shard_ranges);
+	kvfree(tbls->topo);
+	kfree(tbls);
+}
+
+static void scx_cid_tables_free_rcufn(struct rcu_head *rcu)
+{
+	scx_cid_tables_free(container_of(rcu, struct scx_cid_tables, rcu));
+}
+
+static struct scx_cid_tables *scx_cid_alloc_tables(void)
+{
+	u32 npossible = num_possible_cpus();
+	struct scx_cid_tables *tbls;
+
+	tbls = kzalloc_obj(*tbls, GFP_KERNEL);
+	if (!tbls)
+		return NULL;
+
+	tbls->cid_to_cpu = kvcalloc(npossible, sizeof(*tbls->cid_to_cpu), GFP_KERNEL);
+	tbls->cpu_to_cid = kvcalloc(nr_cpu_ids, sizeof(*tbls->cpu_to_cid), GFP_KERNEL);
+	tbls->cid_to_shard = kvcalloc(npossible, sizeof(*tbls->cid_to_shard), GFP_KERNEL);
+	tbls->shard_node = kvcalloc(npossible, sizeof(*tbls->shard_node), GFP_KERNEL);
+	tbls->shard_ranges = kvcalloc(npossible, sizeof(*tbls->shard_ranges), GFP_KERNEL);
+	tbls->topo = kvcalloc(npossible, sizeof(*tbls->topo), GFP_KERNEL);
+
+	if (!tbls->cid_to_cpu || !tbls->cpu_to_cid || !tbls->cid_to_shard ||
+	    !tbls->shard_node || !tbls->shard_ranges || !tbls->topo) {
+		scx_cid_tables_free(tbls);
+		return NULL;
 	}
 
-	WRITE_ONCE(scx_cid_to_cpu_tbl, cid_to_cpu);
-	WRITE_ONCE(scx_cpu_to_cid_tbl, cpu_to_cid);
-	WRITE_ONCE(scx_cid_topo, cid_topo);
-	return 0;
+	return tbls;
+}
+
+/**
+ * scx_cid_publish_tables - Publish the tables scx_cid_init() built
+ *
+ * Called after ops.init_cids() where the layout is final.
+ */
+void scx_cid_publish_tables(void)
+{
+	struct scx_cid_tables *tbls = scx_cid_tables;
+
+	lockdep_assert_held(&scx_enable_mutex);
+
+	scx_nr_cid_shards = tbls->nr_shards;
+	rcu_assign_pointer(scx_cid_to_cpu_tbl, tbls->cid_to_cpu);
+	rcu_assign_pointer(scx_cpu_to_cid_tbl, tbls->cpu_to_cid);
+	rcu_assign_pointer(scx_cid_to_shard, tbls->cid_to_shard);
+	rcu_assign_pointer(scx_shard_node, tbls->shard_node);
+	rcu_assign_pointer(scx_cid_shard_ranges, tbls->shard_ranges);
+	rcu_assign_pointer(scx_cid_topo, tbls->topo);
+}
+
+/**
+ * scx_cid_retire_tables - Unpublish and retire the cid tables
+ *
+ * Called by root disable after the readers which dereference without NULL
+ * checks are drained, inside cpus_read_lock() to exclude the hotplug path.
+ */
+void scx_cid_retire_tables(void)
+{
+	struct scx_cid_tables *tbls = scx_cid_tables;
+
+	lockdep_assert_held(&scx_enable_mutex);
+	lockdep_assert_cpus_held();
+
+	if (!tbls)
+		return;
+
+	scx_cid_tables = NULL;
+	RCU_INIT_POINTER(scx_cid_to_cpu_tbl, NULL);
+	RCU_INIT_POINTER(scx_cpu_to_cid_tbl, NULL);
+	RCU_INIT_POINTER(scx_cid_to_shard, NULL);
+	RCU_INIT_POINTER(scx_shard_node, NULL);
+	RCU_INIT_POINTER(scx_cid_shard_ranges, NULL);
+	RCU_INIT_POINTER(scx_cid_topo, NULL);
+	call_rcu(&tbls->rcu, scx_cid_tables_free_rcufn);
 }
 
 /**
  * scx_cid_init - build the cid mapping
  * @sch: the scx_sched being initialized; used as the scx_error() target
  *
+ * Build a fresh table set. It becomes visible through scx_cid_publish_tables()
+ * and is retired by scx_cid_retire_tables() at disable.
+ *
  * See "Topological CPU IDs" in cid.h for the model. Walk online cpus by
  * intersection at each level (parent_scratch & this_level_mask), which keeps
  * containment correct by construction and naturally splits a physical LLC
@@ -88,18 +185,32 @@ s32 scx_cid_init(struct scx_sched *sch)
 	cpumask_var_t core_scratch __free(free_cpumask_var) = CPUMASK_VAR_NULL;
 	cpumask_var_t llc_fallback __free(free_cpumask_var) = CPUMASK_VAR_NULL;
 	cpumask_var_t online_no_topo __free(free_cpumask_var) = CPUMASK_VAR_NULL;
+	struct scx_cid_tables *tbls;
 	u32 next_cid = 0;
 	s32 next_node_idx = 0, next_llc_idx = 0, next_core_idx = 0;
-	s32 cpu, ret;
+	s32 next_shard_idx = 0;
+	u32 shard_size, max_cids;
+	u32 notopo_in_shard;
+	s32 notopo_shard_cid, notopo_shard_idx;
+	s32 cpu, cid, si;
 
 	/* CMASK_MAX_WORDS in cid.bpf.h covers NR_CPUS up to 8192 */
 	BUILD_BUG_ON(NR_CPUS > 8192);
 
 	lockdep_assert_cpus_held();
+	lockdep_assert_held(&scx_enable_mutex);
 
-	ret = scx_cid_arrays_alloc();
-	if (ret)
-		return ret;
+	shard_size = sch->ops.cid_shard_size ?: SCX_CID_SHARD_SIZE_DFL;
+	max_cids = min_t(u32, shard_size, SCX_CID_SHARD_MAX_CPUS);
+
+	tbls = scx_cid_alloc_tables();
+	if (!tbls)
+		return -ENOMEM;
+
+	scx_cid_tables = tbls;
+
+	for (si = 0; si < num_possible_cpus(); si++)
+		tbls->shard_node[si] = NUMA_NO_NODE;
 
 	if (!zalloc_cpumask_var(&to_walk, GFP_KERNEL) ||
 	    !zalloc_cpumask_var(&node_scratch, GFP_KERNEL) ||
@@ -111,7 +222,7 @@ s32 scx_cid_init(struct scx_sched *sch)
 
 	/* -1 sentinels for sparse-possible cpu id holes (0 is a valid cid) */
 	for (cpu = 0; cpu < nr_cpu_ids; cpu++)
-		scx_cpu_to_cid_tbl[cpu] = -1;
+		tbls->cpu_to_cid[cpu] = -1;
 
 	cpumask_copy(to_walk, cpu_online_mask);
 
@@ -142,36 +253,69 @@ s32 scx_cid_init(struct scx_sched *sch)
 			const struct cpumask *llc_mask = cpu_llc_mask(ncpu, llc_fallback);
 			s32 llc_cid = next_cid;
 			s32 llc_idx = next_llc_idx++;
+			u32 cores_per_shard, nr_large;
+			u32 shard_local = 0, cores_in_shard = 0, cids_in_shard = 0;
+			s32 shard_cid, shard_idx;
 
 			/* llc_scratch = node_scratch & this llc */
 			cpumask_and(llc_scratch, node_scratch, llc_mask);
 			if (WARN_ON_ONCE(!cpumask_test_cpu(ncpu, llc_scratch)))
 				return -EINVAL;
 
+			calc_shard_layout(llc_scratch, shard_size, &cores_per_shard, &nr_large);
+			shard_cid = next_cid;
+			shard_idx = next_shard_idx++;
+			tbls->shard_node[shard_idx] = nid;
+
 			while (!cpumask_empty(llc_scratch)) {
 				s32 lcpu = cpumask_first(llc_scratch);
 				const struct cpumask *sib = topology_sibling_cpumask(lcpu);
 				s32 core_cid = next_cid;
 				s32 core_idx = next_core_idx++;
 				s32 ccpu;
+				u32 max_cores, cids_in_core;
 
 				/* core_scratch = llc_scratch & this core */
 				cpumask_and(core_scratch, llc_scratch, sib);
 				if (WARN_ON_ONCE(!cpumask_test_cpu(lcpu, core_scratch)))
 					return -EINVAL;
 
+				/*
+				 * Advance to a new shard when either core or
+				 * cid count reaches max. The latter bounds
+				 * shard sizes under uneven SMT. Never start an
+				 * empty shard.
+				 */
+				cids_in_core = cpumask_weight(core_scratch);
+				max_cores = cores_per_shard + (shard_local < nr_large ? 1 : 0);
+				if (cores_in_shard &&
+				    (cores_in_shard >= max_cores ||
+				     cids_in_shard + cids_in_core > max_cids)) {
+					shard_local++;
+					cores_in_shard = 0;
+					cids_in_shard = 0;
+					shard_cid = next_cid;
+					shard_idx = next_shard_idx++;
+					tbls->shard_node[shard_idx] = nid;
+				}
+				cores_in_shard++;
+				cids_in_shard += cids_in_core;
+
 				for_each_cpu(ccpu, core_scratch) {
 					s32 cid = next_cid++;
 
-					scx_cid_to_cpu_tbl[cid] = ccpu;
-					scx_cpu_to_cid_tbl[ccpu] = cid;
-					scx_cid_topo[cid] = (struct scx_cid_topo){
+					tbls->cid_to_cpu[cid] = ccpu;
+					tbls->cpu_to_cid[ccpu] = cid;
+					tbls->cid_to_shard[cid] = shard_idx;
+					tbls->topo[cid] = (struct scx_cid_topo){
 						.core_cid = core_cid,
 						.core_idx = core_idx,
 						.llc_cid = llc_cid,
 						.llc_idx = llc_idx,
 						.node_cid = node_cid,
 						.node_idx = node_idx,
+						.shard_cid = shard_cid,
+						.shard_idx = shard_idx,
 					};
 
 					cpumask_clear_cpu(ccpu, llc_scratch);
@@ -184,21 +328,37 @@ s32 scx_cid_init(struct scx_sched *sch)
 
 	/*
 	 * No-topo section: any possible cpu without a cid - normally just the
-	 * not-online ones. Collect any currently-online cpus that land here in
-	 * @online_no_topo so we can warn about them at the end.
+	 * not-online ones. Pack into shards of up to min(@shard_size,
+	 * SCX_CID_SHARD_MAX_CPUS) cids so that every cid has a valid shard
+	 * assignment and the hard cap holds even with a large @shard_size.
+	 * Collect any currently-online cpus that land here in @online_no_topo
+	 * so we can warn about them at the end.
 	 */
-	for_each_cpu(cpu, cpu_possible_mask) {
-		s32 cid;
+	notopo_in_shard = min_t(u32, shard_size, SCX_CID_SHARD_MAX_CPUS);
+	notopo_shard_cid = -1;
+	notopo_shard_idx = -1;
 
-		if (__scx_cpu_to_cid(cpu) != -1)
+	for_each_cpu(cpu, cpu_possible_mask) {
+		if (tbls->cpu_to_cid[cpu] != -1)
 			continue;
 		if (cpu_online(cpu))
 			cpumask_set_cpu(cpu, online_no_topo);
 
 		cid = next_cid++;
-		scx_cid_to_cpu_tbl[cid] = cpu;
-		scx_cpu_to_cid_tbl[cpu] = cid;
-		scx_cid_topo[cid] = SCX_CID_TOPO_NEG;
+		tbls->cid_to_cpu[cid] = cpu;