The resctrl subsystem writes the task's RMID/CLOSID to IA32_PQR_ASSOC in
__resctrl_sched_in(). With PLZA support being introduced and guarded by
rdt_plza_enable_key, the kernel needs a way to track and program the PLZA
association independently of the regular RMID/CLOSID path.

Extend the per-CPU resctrl_pqr_state to track PLZA-related state, including
the current and default PLZA values along with the associated RMID and
CLOSID.

Update the resctrl scheduling-in path to program the PLZA MSR when PLZA
support is enabled. During the context switch, the task-specific PLZA
setting is applied if present; otherwise, the per-CPU default PLZA value is
used. The MSR is only written when the PLZA state changes, avoiding
unnecessary writes.

PLZA programming is guarded by a static key to ensure there is no overhead
when the feature is disabled.

Signed-off-by: Babu Moger <babu.moger@amd.com>
---
 arch/x86/include/asm/resctrl.h | 19 +++++++++++++++++++
 include/linux/sched.h          |  1 +
 2 files changed, 20 insertions(+)

diff --git a/arch/x86/include/asm/resctrl.h b/arch/x86/include/asm/resctrl.h
index fc0a7f64649e..76de7d6051b7 100644
--- a/arch/x86/include/asm/resctrl.h
+++ b/arch/x86/include/asm/resctrl.h
@@ -38,6 +38,10 @@ struct resctrl_pqr_state {
 	u32			cur_closid;
 	u32			default_rmid;
 	u32			default_closid;
+	u32			cur_plza;
+	u32			default_plza;
+	u32			plza_rmid;
+	u32			plza_closid;
 };
 
 DECLARE_PER_CPU(struct resctrl_pqr_state, pqr_state);
@@ -115,6 +119,7 @@ static inline void __resctrl_sched_in(struct task_struct *tsk)
 	struct resctrl_pqr_state *state = this_cpu_ptr(&pqr_state);
 	u32 closid = READ_ONCE(state->default_closid);
 	u32 rmid = READ_ONCE(state->default_rmid);
+	u32 plza = READ_ONCE(state->default_plza);
 	u32 tmp;
 
 	/*
@@ -138,6 +143,20 @@ static inline void __resctrl_sched_in(struct task_struct *tsk)
 		state->cur_rmid = rmid;
 		wrmsr(MSR_IA32_PQR_ASSOC, rmid, closid);
 	}
+
+	if (static_branch_likely(&rdt_plza_enable_key)) {
+		tmp = READ_ONCE(tsk->plza);
+		if (tmp)
+			plza = tmp;
+
+		if (plza != state->cur_plza) {
+			state->cur_plza = plza;
+			wrmsr(MSR_IA32_PQR_PLZA_ASSOC,
+			      RMID_EN | state->plza_rmid,
+			      (plza ? PLZA_EN : 0) | CLOSID_EN | state->plza_closid);
+		}
+	}
+
 }
 
 static inline unsigned int resctrl_arch_round_mon_val(unsigned int val)
diff --git a/include/linux/sched.h b/include/linux/sched.h
index 8f3a60f13393..d573163865ae 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1326,6 +1326,7 @@ struct task_struct {
 #ifdef CONFIG_X86_CPU_RESCTRL
 	u32				closid;
 	u32				rmid;
+	u32				plza;
 #endif
 #ifdef CONFIG_FUTEX
 	struct robust_list_head __user	*robust_list;
-- 
2.34.1