sched_ext: Add per-shard scx_sched storage scaffolding

Add struct scx_pshard and sch->pshard[] indexed by shard_idx, each entry
allocated on its shard's NUMA node from scx_shard_node[si]. The struct
starts empty (one dummy field). Follow-up patches will grow it as
shard-local state lands. Only cid-type schedulers with an arena pool get
pshards.

Allocation happens after ops.init_cids() returns so any
scx_bpf_cid_override() it issues has finalized scx_nr_cid_shards and
scx_shard_node[]. sch->nr_pshards records the array size for the async RCU
free path, which may run after a later scheduler's scx_cid_init() has
rewritten the global.

v3: Build and publish pshard[] fully-formed here rather than a later patch.
v2: Free the partially-allocated pshard array on alloc failure. (sashiko AI)

Signed-off-by: Tejun Heo <tj@kernel.org>
Reviewed-by: Andrea Righi <arighi@nvidia.com>

4 files changed, 84 insertions, 0 deletions

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index b70170c3de18..7a4e007e4ed7 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -4681,6 +4681,8 @@ static void scx_sched_free_rcu_work(struct work_struct *work)
 		free_pnode(sch->pnode[node]);
 	kfree(sch->pnode);
 
+	scx_free_pshards(sch);
+
 	rhashtable_walk_enter(&sch->dsq_hash, &rht_iter);
 	do {
 		rhashtable_walk_start(&rht_iter);
@@ -6763,6 +6765,12 @@ static void scx_root_enable_workfn(struct kthread_work *work)
 		goto err_disable;
 	}
 
+	ret = scx_alloc_pshards(sch);
+	if (ret) {
+		cpus_read_unlock();
+		goto err_disable;
+	}
+
 	if (sch->ops.init) {
 		ret = SCX_CALL_OP_RET(sch, init, NULL);
 		if (ret) {
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 160ff79faedc..5b18c4192c62 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1183,6 +1183,12 @@ struct scx_sched_pnode {
 	struct scx_dispatch_q	global_dsq;
 };
 
+#ifdef CONFIG_EXT_SUB_SCHED
+struct scx_pshard {
+	int			_dummy;		/* until the first real field lands */
+};
+#endif
+
 struct scx_sched {
 	/*
 	 * cpu-form and cid-form ops share field offsets up to .priv (verified
@@ -1230,6 +1236,9 @@ struct scx_sched {
 	 */
 	struct rhashtable	dsq_hash;
 	struct scx_sched_pnode	**pnode;
+#ifdef CONFIG_EXT_SUB_SCHED
+	struct scx_pshard	**pshard;	/* indexed by shard_idx */
+#endif
 	struct scx_sched_pcpu __percpu *pcpu;
 
 	u64			slice_dfl;
@@ -1245,6 +1254,15 @@ struct scx_sched {
 	u32			dsp_max_batch;
 	s32			level;
 
+#ifdef CONFIG_EXT_SUB_SCHED
+	/*
+	 * pshard[] size captured at enable for the async RCU free path -
+	 * scx_nr_cid_shards may be rewritten by a later scx_cid_init() before
+	 * free runs. While sch is active, use the global.
+	 */
+	u32			nr_pshards;
+#endif
+
 	/*
 	 * Updates to the following warned bitfields can race causing RMW issues
 	 * but it doesn't really matter.
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 9855a9a4e709..3adec9343e46 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -82,6 +82,60 @@ void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch)
 		rcu_assign_pointer(pos->scx_sched, sch);
 }
 
+static void free_pshard(struct scx_pshard *pshard)
+{
+	kfree(pshard);
+}
+
+void scx_free_pshards(struct scx_sched *sch)
+{
+	s32 si;
+
+	if (!sch->pshard)
+		return;
+	for (si = 0; si < sch->nr_pshards; si++)
+		free_pshard(sch->pshard[si]);
+	kfree(sch->pshard);
+}
+
+static struct scx_pshard *alloc_pshard(struct scx_sched *sch, s32 shard_idx, s32 node)
+{
+	return kzalloc_node(sizeof(struct scx_pshard), GFP_KERNEL, node);
+}
+
+s32 scx_alloc_pshards(struct scx_sched *sch)
+{
+	struct scx_pshard **pshard;
+	s32 si;
+
+	if (!sch->is_cid_type || !sch->arena_pool)
+		return 0;
+
+	pshard = kzalloc_objs(pshard[0], scx_nr_cid_shards, GFP_KERNEL);
+	if (!pshard)
+		return -ENOMEM;
+
+	for (si = 0; si < scx_nr_cid_shards; si++) {
+		pshard[si] = alloc_pshard(sch, si, scx_shard_node[si]);
+		if (!pshard[si]) {
+			while (--si >= 0)
+				free_pshard(pshard[si]);
+			kfree(pshard);
+			return -ENOMEM;
+		}
+	}
+
+	sch->nr_pshards = scx_nr_cid_shards;
+	/*
+	 * Publish only after every entry is built so a reader observing
+	 * @sch->pshard never sees a partially-filled array. Pair the store
+	 * with a barrier and READ_ONCE() on the read side.
+	 */
+	smp_wmb();
+	WRITE_ONCE(sch->pshard, pshard);
+	return 0;
+}
+
 static DECLARE_WAIT_QUEUE_HEAD(scx_unlink_waitq);
 
 void drain_descendants(struct scx_sched *sch)
diff --git a/kernel/sched/ext/sub.h b/kernel/sched/ext/sub.h
index 9b5ac07e5e76..c06d0bfafbd8 100644
--- a/kernel/sched/ext/sub.h
+++ b/kernel/sched/ext/sub.h
@@ -24,6 +24,8 @@ void drain_descendants(struct scx_sched *sch);
 void scx_sub_disable(struct scx_sched *sch);
 void scx_sub_enable_workfn(struct kthread_work *work);
 bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux);
+void scx_free_pshards(struct scx_sched *sch);
+s32 scx_alloc_pshards(struct scx_sched *sch);
 
 static inline const char *sch_cgrp_path(struct scx_sched *sch)
 {
@@ -39,6 +41,8 @@ static inline const char *sch_cgrp_path(struct scx_sched *sch) { return "/"; }
 static inline void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch) {}
 static inline void drain_descendants(struct scx_sched *sch) { }
 static inline void scx_sub_disable(struct scx_sched *sch) { }
+static inline void scx_free_pshards(struct scx_sched *sch) {}
+static inline s32 scx_alloc_pshards(struct scx_sched *sch) { return 0; }
 
 #endif	/* CONFIG_EXT_SUB_SCHED */