[RFC PATCH 4/7] sched/cache: Add prctl to manage per process cache scheduling groups

@@ -884,12 +884,22 @@ static int exec_mmap(struct linux_binprm *bprm)
 	{
 		struct sched_cache_group *old_grp, *new_grp;
 
-		old_grp = rcu_dereference_protected(tsk->sched_cache_grp, true);
-
-		/* Acquire the reference before publishing the pointer. */
+		/*
+		 * Acquire the reference before publishing the pointer: once
+		 * tsk->sched_cache_grp is visible, a concurrent
+		 * prctl(PR_SCHED_CACHE) writer may pick the group up as its
+		 * old_grp and drop a reference we have not taken yet.
+		 *
+		 * pi_lock serializes the exchange against such a writer. IRQs
+		 * are already disabled here, so a plain raw_spin_lock()
+		 * suffices.
+		 */
 		new_grp = sched_cache_group_get(mm->sched_cache_grp);
 
-		rcu_assign_pointer(tsk->sched_cache_grp, new_grp);
+		raw_spin_lock(&tsk->pi_lock);
+		old_grp = sched_cache_grp_replace(tsk, new_grp);
+		raw_spin_unlock(&tsk->pi_lock);
+
 		if (old_grp)
 			sched_cache_group_put(old_grp);
 	}

@@ -2407,7 +2407,11 @@ void sched_cache_group_put(struct sched_cache_group *grp);
 struct sched_cache_group *sched_cache_group_get(struct sched_cache_group *grp);
 struct sched_cache_group *task_cache_group_get(struct task_struct *p);
 struct sched_cache_group *
+sched_cache_grp_replace(struct task_struct *p, struct sched_cache_group *grp);
+struct sched_cache_group *
 sched_cache_alloc_group(struct sched_cache_time __percpu *pcpu_sched);
+int sched_cache_prctl(int option, unsigned long arg2, unsigned long arg3,
+		      unsigned long arg4, unsigned long arg5);
 
 #else
 

@@ -416,4 +416,11 @@ struct prctl_mm_map {
 # define PR_CFI_DISABLE		_BITUL(1)
 # define PR_CFI_LOCK		_BITUL(2)
 
+/* Cache-aware scheduling */
+#define PR_SCHED_CACHE			82
+# define PR_SCHED_CACHE_GET		0
+# define PR_SCHED_CACHE_CREATE		1
+# define PR_SCHED_CACHE_SHARE_FROM	2
+# define PR_SCHED_CACHE_MAX		3
+
 #endif /* _LINUX_PRCTL_H */

@@ -592,10 +592,17 @@ static void exit_mm(void)
 
 #ifdef CONFIG_SCHED_CACHE
 	{
-		struct sched_cache_group *grp =
-			rcu_dereference_protected(current->sched_cache_grp, true);
+		struct sched_cache_group *grp;
+		unsigned long flags;
 
-		rcu_assign_pointer(current->sched_cache_grp, NULL);
+		/*
+		 * pi_lock serializes the clear against a concurrent
+		 * prctl(PR_SCHED_CACHE) writer targeting this task, so the
+		 * reference is dropped exactly once.
+		 */
+		raw_spin_lock_irqsave(&current->pi_lock, flags);
+		grp = sched_cache_grp_replace(current, NULL);
+		raw_spin_unlock_irqrestore(&current->pi_lock, flags);
 
 		if (grp)
 			sched_cache_group_put(grp);

@@ -1601,13 +1601,28 @@ static int copy_mm(u64 clone_flags, struct task_struct *tsk)
 	tsk->active_mm = mm;
 #ifdef CONFIG_SCHED_CACHE
 	{
+		struct sched_cache_group *grp;
+
 		/*
-		 * A task holds its own reference on the group, separate from
-		 * the reference held by its mm_struct. Acquire it before
-		 * publishing the pointer.
+		 * For CLONE_VM (threads): inherit the parent's
+		 * sched_cache_grp, which may be a cookie-based group different
+		 * from mm->sched_cache_grp. For kernel thread, don't assign
+		 * it a sched cache group.
+		 *
+		 * For new processes (own mm): use the new mm's group. The mm
+		 * has just been created and is not visible to anyone else, so
+		 * there is no race against a concurrent writer and the mm's
+		 * own reference keeps the group alive. Take the refcount
+		 * directly, without RCU.
 		 */
-		struct sched_cache_group *grp =
-			sched_cache_group_get(mm->sched_cache_grp);
+		if (clone_flags & CLONE_VM) {
+			if (tsk->flags & PF_KTHREAD)
+				grp = NULL;
+			else
+				grp = task_cache_group_get(current);
+		} else {
+			grp = sched_cache_group_get(mm->sched_cache_grp);
+		}
 
 		rcu_assign_pointer(tsk->sched_cache_grp, grp);
 	}

@@ -83,3 +83,277 @@ sched_cache_alloc_group(struct sched_cache_time __percpu *_pcpu_sched)
 	sched_cache_group_init(grp, _pcpu_sched);
 	return grp;
 }
+
+static struct sched_cache_group *sched_cache_alloc_all(void)
+{
+	struct sched_cache_time __percpu *pcpu_sched;
+
+	pcpu_sched = alloc_percpu(struct sched_cache_time);
+	if (!pcpu_sched)
+		return NULL;
+
+	return sched_cache_alloc_group(pcpu_sched);
+}
+
+#ifdef CONFIG_NUMA_BALANCING
+/*
+ * When prctl() moves a task between groups, move its footprint estimate too:
+ * otherwise the group it leaves over-estimates its footprint forever, and the
+ * group it joins under-estimates it once the task exits.
+ */
+static void sched_cache_xfer_footprint(struct task_struct *p,
+				       struct sched_cache_group *old,
+				       struct sched_cache_group *new)
+{
+	unsigned long fp, sub;
+
+	if (!old || !new || old == new)
+		return;
+
+	sub = READ_ONCE(p->total_numa_faults);
+	if (!sub)
+		return;
+
+	fp = READ_ONCE(old->footprint);
+	sub = min(fp, sub);
+	WRITE_ONCE(old->footprint, fp - sub);
+
+	fp = READ_ONCE(new->footprint);
+	WRITE_ONCE(new->footprint, fp + sub);
+}
+#else
+static inline void sched_cache_xfer_footprint(struct task_struct *p,
+					      struct sched_cache_group *old,
+					      struct sched_cache_group *new)
+{
+}
+#endif /* CONFIG_NUMA_BALANCING */
+
+/* Swap a task's cache group pointer and return the previous one. */
+struct sched_cache_group *
+sched_cache_grp_replace(struct task_struct *p, struct sched_cache_group *grp)
+{
+	struct sched_cache_group *old;
+
+	lockdep_assert_held(&p->pi_lock);
+	old = rcu_dereference_protected(p->sched_cache_grp,
+					lockdep_is_held(&p->pi_lock));
+	rcu_assign_pointer(p->sched_cache_grp, grp);
+
+	return old;
+}
+
+static void __sched_cache_set(struct task_struct *p,
+			      struct sched_cache_group *grp)
+{
+	struct sched_cache_group *old_grp;
+	unsigned long flags;
+
+	grp = sched_cache_group_get(grp);
+
+	/*
+	 * p->pi_lock serializes the exchange against concurrent writers
+	 * (other prctl callers as well as exec_mmap()/exit_mm()). Without
+	 * it two writers could fetch the same old pointer and each drop a
+	 * reference, over-decrementing the refcount.
+	 */
+	raw_spin_lock_irqsave(&p->pi_lock, flags);
+
+	/*
+	 * Avoid increasing the refcount for an exiting task,
+	 * otherwise the grp can not be put after the task exits,
+	 * and cause memory leak:
+	 * CPU0 (prctl)                    CPU1 (task exiting)
+	 * ----                            ----
+	 *                                 exit_mm()
+	 *                                   put(p->sched_cache_grp)
+	 *                                    free(old_grp)
+	 *
+	 * __sched_cache_set(p, new_grp)
+	 *
+	 *                                   free_task()
+	 *                                     free(p)
+	 * the "freed" p is unable to put new_grp
+	 *
+	 * PF_EXITING is set in exit_signals() without p->pi_lock held, so this
+	 * test on its own would only narrow the window. What closes it is that
+	 * exit_mm() takes p->pi_lock to clear the pointer, and exit_signals()
+	 * runs before exit_mm(). So of the two possible orders:
+	 *
+	 *  - we get pi_lock before exit_mm() does: we install new_grp and its
+	 *    reference, and exit_mm() subsequently drops it, because it puts
+	 *    whatever sched_cache_grp_replace() hands back.
+	 *
+	 *  - we get pi_lock after exit_mm() released it: the lock's release/
+	 *    acquire ordering makes the PF_EXITING store visible to us, so the
+	 *    test below fires and we install nothing.
+	 *
+	 * Either way the reference is dropped exactly once. The test is thus
+	 * about not leaving a reference behind on a task that is already past
+	 * exit_mm(), not about racing with the PF_EXITING store itself.
+	 */
+	if (p->flags & PF_EXITING) {
+		raw_spin_unlock_irqrestore(&p->pi_lock, flags);
+		if (grp)
+			sched_cache_group_put(grp);
+		return;
+	}
+
+	old_grp = sched_cache_grp_replace(p, grp);
+	raw_spin_unlock_irqrestore(&p->pi_lock, flags);
+
+	/* Carry this task's footprint estimate to the group it just joined. */
+	sched_cache_xfer_footprint(p, old_grp, grp);
+
+	if (old_grp)
+		sched_cache_group_put(old_grp);
+}
+
+static struct task_struct *sched_cache_find_get_task(unsigned long vpid)
+{
+	struct task_struct *p;
+
+	guard(rcu)();
+	p = vpid ? find_task_by_vpid(vpid) : current;
+	if (p)
+		get_task_struct(p);
+
+	return p;
+}
+
+/*
+ * arg2: subcommand,
+ * arg3: destination pid, SHARE_FROM copies the group of arg4 to arg3
+ * arg4: source pid, or cookie out ptr for GET
+ * arg5: pid type, the scope of the destination
+ *
+ */
+int sched_cache_prctl(int option, unsigned long arg2, unsigned long arg3,
+		      unsigned long arg4, unsigned long arg5)
+{
+	struct task_struct *dst = NULL, *src = NULL, *p;
+	struct sched_cache_group *grp = NULL;
+	enum pid_type type = arg5;
+	struct pid *pid_grp;
+	int err = 0;
+
+	if (arg2 >= PR_SCHED_CACHE_MAX || arg3 > INT_MAX ||
+	    arg5 > PIDTYPE_PGID)
+		return -EINVAL;
+
+	/* only GET and SHARE_FROM take a 4th argument */
+	if (arg4 && arg2 != PR_SCHED_CACHE_GET &&
+	    arg2 != PR_SCHED_CACHE_SHARE_FROM)
+		return -EINVAL;
+
+	dst = sched_cache_find_get_task(arg3);
+	if (!dst)
+		return -ESRCH;
+
+	if (dst->flags & PF_KTHREAD) {
+		err = -EINVAL;
+		goto out_task;
+	}
+
+	if (!ptrace_may_access(dst, PTRACE_MODE_READ_REALCREDS)) {
+		err = -EPERM;
+		goto out_task;
+	}
+
+	switch (arg2) {
+	case PR_SCHED_CACHE_GET: {
+		unsigned long id = 0;
+
+		if (type != PIDTYPE_PID || arg4 & 7) {
+			err = -EINVAL;
+			goto out_task;
+		}
+
+		grp = task_cache_group_get(dst);
+		if (grp)
+			ptr_to_hashval((void *)grp, &id);
+
+		if (arg4)
+			err = put_user((u64)id, (u64 __user *)arg4);
+
+		goto out_group;
+	}
+
+	case PR_SCHED_CACHE_CREATE:
+		/*
+		 * Allocator owns ref 1, __sched_cache_set() acquires ref 2.
+		 * The sched_cache_group_put() at out_group: drops ref 1, leaving
+		 * ref 1 held by the task.
+		 */
+		grp = sched_cache_alloc_all();
+		if (!grp) {
+			err = -ENOMEM;
+			goto out_task;
+		}
+		break;
+
+	case PR_SCHED_CACHE_SHARE_FROM:
+		if (arg4 > INT_MAX) {
+			err = -EINVAL;
+			goto out_task;
+		}
+
+		src = sched_cache_find_get_task(arg4);
+		if (!src) {
+			err = -ESRCH;
+			goto out_task;
+		}
+
+		if (src->flags & PF_KTHREAD) {
+			err = -EINVAL;
+			goto out_task;
+		}
+
+		if (!ptrace_may_access(src, PTRACE_MODE_READ_REALCREDS)) {
+			err = -EPERM;
+			goto out_task;
+		}
+
+		/* copy the group of src to dst */
+		grp = task_cache_group_get(src);
+		if (!grp) {
+			err = -ENOENT;
+			goto out_task;
+		}
+		break;
+
+	default:
+		err = -EINVAL;
+		goto out_task;
+	}
+
+	if (type == PIDTYPE_PID) {
+		__sched_cache_set(dst, grp);
+		goto out_group;
+	}
+
+	read_lock(&tasklist_lock);
+	pid_grp = task_pid_type(dst, type);
+
+	do_each_pid_thread(pid_grp, type, p) {
+		if (!ptrace_may_access(p, PTRACE_MODE_READ_REALCREDS)) {
+			err = -EPERM;
+			goto out_tasklist;
+		}
+	} while_each_pid_thread(pid_grp, type, p);
+
+	do_each_pid_thread(pid_grp, type, p) {
+		__sched_cache_set(p, grp);
+	} while_each_pid_thread(pid_grp, type, p);
+out_tasklist:
+	read_unlock(&tasklist_lock);
+
+out_group:
+	sched_cache_group_put(grp);
+out_task:
+	if (src)
+		put_task_struct(src);
+	if (dst)
+		put_task_struct(dst);
+	return err;
+}

@@ -1484,7 +1484,13 @@ static bool invalid_llc_nr(struct sched_cache_group *grp, struct task_struct *p,
 {
 	int scale;
 
-	if (get_nr_threads(p) <= 1)
+	/*
+	 * Single threaded process that is not grouped with other threads
+	 * do not need cache aware scheduling.
+	 * A single-threaded process has a refcount of 2: from the mm and task
+	 * respectively.
+	 */
+	if (refcount_read(&grp->refcnt) <= 2 && get_nr_threads(p) <= 1)
 		return true;
 
 	/*

@@ -2907,6 +2907,11 @@ SYSCALL_DEFINE5(prctl, int, option, unsigned long, arg2, unsigned long, arg3,
 		if (arg3 & PR_CFI_LOCK && !(arg3 & PR_CFI_DISABLE))
 			error = arch_prctl_lock_branch_landing_pad_state(me);
 		break;
+#ifdef CONFIG_SCHED_CACHE
+	case PR_SCHED_CACHE:
+		error = sched_cache_prctl(option, arg2, arg3, arg4, arg5);
+		break;
+#endif
 	default:
 		trace_task_prctl_unknown(option, arg2, arg3, arg4, arg5);
 		error = -EINVAL;