[PATCH v2 16/16] fs/resctrl: Add per-task kmode enable support via rdtgroup

[Babu Moger]

Introduce support for enabling kmode on a per-task basis through the
resctrl control-group interface.

Add an architecture helper to set the kmode state in the task structure and
extend the rdtgroup task handling path to apply kmode (e.g. PLZA) when
associating a task with a CTRL_MON or MON group.

Proper memory ordering is enforced to ensure that task closid and rmid
updates are visible before determining whether the task is currently
running. If the task is active on a CPU, the relevant MSRs are updated
immediately; otherwise, PLZA state is programmed on the next context
switch.

Signed-off-by: Babu Moger <babu.moger@xxxxxxx>
---
v2: Few name changes to refer PLZA as kmode.
---
 arch/x86/include/asm/resctrl.h | 13 +++++
 fs/resctrl/rdtgroup.c          | 98 +++++++++++++++++++++++++++++++++-
 2 files changed, 110 insertions(+), 1 deletion(-)

diff --git a/arch/x86/include/asm/resctrl.h b/arch/x86/include/asm/resctrl.h
index ccfd95b98bac..f48d1279e33d 100644
--- a/arch/x86/include/asm/resctrl.h
+++ b/arch/x86/include/asm/resctrl.h
@@ -238,6 +238,19 @@ static inline void resctrl_arch_set_cpu_kmode(int cpu, u32 closid, u32 rmid, u32
 	WRITE_ONCE(per_cpu(pqr_state.kmode_rmid, cpu), rmid);
 }
 
+/**
+ * resctrl_arch_set_task_kmode() - Set per-task kernel mode (e.g. PLZA) flag
+ * @tsk:	Task to update.
+ * @enable:	1 to enable kmode for this task; 0 to disable.
+ *
+ * When enabled, the task will use the group's CLOSID/RMID for kernel mode
+ * on context switch (see __resctrl_sched_in()).
+ */
+static inline void resctrl_arch_set_task_kmode(struct task_struct *tsk, u32 enable)
+{
+	WRITE_ONCE(tsk->kmode, enable);
+}
+
 static inline void resctrl_arch_sched_in(struct task_struct *tsk)
 {
 	if (static_branch_likely(&rdt_enable_key))
diff --git a/fs/resctrl/rdtgroup.c b/fs/resctrl/rdtgroup.c
index b41e681f6922..74fc942e6a4e 100644
--- a/fs/resctrl/rdtgroup.c
+++ b/fs/resctrl/rdtgroup.c
@@ -827,6 +827,31 @@ static int __rdtgroup_move_task(struct task_struct *tsk,
 	return 0;
 }
 
+/**
+ * __rdtgroup_task_kmode() - Enable kernel mode (e.g. PLZA) for a single task
+ * @tsk:	Task to enable kmode for.
+ * @rdtgrp:	Rdtgroup with kmode enabled (used for context; CLOSID/RMID applied on sched-in).
+ *
+ * Sets t->kmode so that the task uses the group's CLOSID/RMID on context
+ * switch. Memory ordering ensures the store is visible before we check if
+ * the task is current (and thus before any sched-in that may observe it).
+ *
+ * Return: 0.
+ */
+static int __rdtgroup_task_kmode(struct task_struct *tsk, struct rdtgroup *rdtgrp)
+{
+	resctrl_arch_set_task_kmode(tsk, true);
+
+	/*
+	 * Order the task's kmode state stores above before the loads in
+	 * task_curr(). This pairs with the full barrier between the
+	 * rq->curr update and resctrl_arch_sched_in() during context switch.
+	 */
+	smp_mb();
+
+	return 0;
+}
+
 static bool is_closid_match(struct task_struct *t, struct rdtgroup *r)
 {
 	return (resctrl_arch_alloc_capable() && (r->type == RDTCTRL_GROUP) &&
@@ -916,6 +941,48 @@ static int rdtgroup_move_task(pid_t pid, struct rdtgroup *rdtgrp,
 	return ret;
 }
 
+/**
+ * rdtgroup_task_kmode() - Enable kernel mode for a task added to a kmode group
+ * @pid:	PID of the task (0 for current).
+ * @rdtgrp:	Rdtgroup with kmode enabled.
+ * @of:		kernfs file (for permission check).
+ *
+ * Called when a task is written to the "tasks" file of a group that has
+ * kernel mode enabled. Enables kmode for that task so it uses the group's
+ * CLOSID/RMID on context switch. If the task is currently running, MSRs are
+ * updated on next sched-in.
+ *
+ * Return: 0 on success, or -ESRCH/-EPERM on error.
+ */
+static int rdtgroup_task_kmode(pid_t pid, struct rdtgroup *rdtgrp,
+			       struct kernfs_open_file *of)
+{
+	struct task_struct *tsk;
+	int ret;
+
+	rcu_read_lock();
+	if (pid) {
+		tsk = find_task_by_vpid(pid);
+		if (!tsk) {
+			rcu_read_unlock();
+			rdt_last_cmd_printf("No task %d\n", pid);
+			return -ESRCH;
+		}
+	} else {
+		tsk = current;
+	}
+
+	get_task_struct(tsk);
+	rcu_read_unlock();
+
+	ret = rdtgroup_task_write_permission(tsk, of);
+	if (!ret)
+		ret = __rdtgroup_task_kmode(tsk, rdtgrp);
+
+	put_task_struct(tsk);
+	return ret;
+}
+
 static ssize_t rdtgroup_tasks_write(struct kernfs_open_file *of,
 				    char *buf, size_t nbytes, loff_t off)
 {
@@ -953,7 +1020,11 @@ static ssize_t rdtgroup_tasks_write(struct kernfs_open_file *of,
 			break;
 		}
 
-		ret = rdtgroup_move_task(pid, rdtgrp, of);
+		/* Group has kmode: set task kmode; else move task CLOSID/RMID. */
+		if (rdtgrp->kmode)
+			ret = rdtgroup_task_kmode(pid, rdtgrp, of);
+		else
+			ret = rdtgroup_move_task(pid, rdtgrp, of);
 		if (ret) {
 			rdt_last_cmd_printf("Error while processing task %d\n", pid);
 			break;
@@ -1011,6 +1082,28 @@ static void show_rdt_tasks(struct rdtgroup *r, struct seq_file *s)
 	rcu_read_unlock();
 }
 
+/**
+ * rdt_task_set_kmode() - Set or clear kmode for all tasks in the rdtgroup
+ * @r:		Rdtgroup (must have r->kmode set for matching).
+ * @kmode:	True to set t->kmode for each matching task; false to clear.
+ *
+ * Walks all tasks that belong to @r (via rdt_task_match) and updates their
+ * per-task kmode flag. Used when enabling or disabling kernel mode for the
+ * group so existing members get the new state.
+ */
+static void rdt_task_set_kmode(struct rdtgroup *r, bool kmode)
+{
+	struct task_struct *p, *t;
+
+	rcu_read_lock();
+	for_each_process_thread(p, t) {
+		if (!rdt_task_match(t, r, r->kmode))
+			continue;
+		resctrl_arch_set_task_kmode(t, kmode);
+	}
+	rcu_read_unlock();
+}
+
 static int rdtgroup_tasks_show(struct kernfs_open_file *of,
 			       struct seq_file *s, void *v)
 {
@@ -1225,6 +1318,9 @@ static int rdtgroup_config_kmode(struct rdtgroup *rdtgrp, bool enable)
 
 	resctrl_arch_set_kmode(&rdtgrp->cpu_mask, &resctrl_kcfg, closid,
 			       rdtgrp->mon.rmid, enable);
+
+	rdt_task_set_kmode(rdtgrp, enable);
+
 	rdtgrp->kmode = enable;
 	if (enable)
 		resctrl_kcfg.k_rdtgrp = rdtgrp;
-- 
2.43.0