[RFC PATCH 1/3] riscv: vector: Add a guest owner invalidation hook

[slack up]

From: Xiaochun Yu <yxcnew@xxxxxxxxxxx>

KVM can avoid restoring a guest Vector register file when it remains
resident on the same hart, but only if every generic Vector user
invalidates that ownership before touching the registers.

Add a single RCU-protected invalidation callback and guard its call sites
with a static key. Hook task save, restore and discard, kernel Vector
entry, and the resume path for preemptible kernel Vector contexts. With
no registered consumer, the static key leaves the generic paths disabled.

This patch only establishes the invalidation boundary; it does not change
KVM behavior.

Development note: OpenAI Codex assisted with source analysis,
implementation drafting, test orchestration and changelog preparation.

Signed-off-by: Xiaochun Yu <yxcnew@xxxxxxxxxxx>
---
 arch/riscv/include/asm/vector.h        | 31 +++++++++++++++++++++
 arch/riscv/kernel/kernel_mode_vector.c | 38 ++++++++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/arch/riscv/include/asm/vector.h b/arch/riscv/include/asm/vector.h
index 1f8dcae3af51..967dd70f4ea4 100644
--- a/arch/riscv/include/asm/vector.h
+++ b/arch/riscv/include/asm/vector.h
@@ -9,8 +9,17 @@
 #include <linux/types.h>
 #include <uapi/asm-generic/errno.h>
 
+enum riscv_v_invalidate_reason {
+	RISCV_V_INVALIDATE_TASK_SAVE,
+	RISCV_V_INVALIDATE_TASK_RESTORE,
+	RISCV_V_INVALIDATE_TASK_DISCARD,
+	RISCV_V_INVALIDATE_KERNEL_VECTOR,
+	RISCV_V_INVALIDATE_CPU_STATE,
+};
+
 #ifdef CONFIG_RISCV_ISA_V
 
+#include <linux/jump_label.h>
 #include <linux/stringify.h>
 #include <linux/sched.h>
 #include <linux/sched/task_stack.h>
@@ -52,6 +61,17 @@ void riscv_v_thread_free(struct task_struct *tsk);
 void __init riscv_v_setup_ctx_cache(void);
 void riscv_v_thread_alloc(struct task_struct *tsk);
 void __init update_regset_vector_info(unsigned long size);
+DECLARE_STATIC_KEY_FALSE(riscv_v_owner_invalidate_key);
+void __riscv_v_owner_invalidate(enum riscv_v_invalidate_reason reason);
+void kvm_riscv_register_vctx_invalidate(void (*func)(enum riscv_v_invalidate_reason reason));
+void kvm_riscv_unregister_vctx_invalidate(void);
+
+static __always_inline void
+riscv_v_owner_invalidate(enum riscv_v_invalidate_reason reason)
+{
+	if (static_branch_unlikely(&riscv_v_owner_invalidate_key))
+		__riscv_v_owner_invalidate(reason);
+}
 
 static inline u32 riscv_v_flags(void)
 {
@@ -303,6 +323,7 @@ static inline void __riscv_v_vstate_discard(void)
 static inline void riscv_v_vstate_discard(struct pt_regs *regs)
 {
 	if (riscv_v_vstate_query(regs)) {
+		riscv_v_owner_invalidate(RISCV_V_INVALIDATE_TASK_DISCARD);
 		__riscv_v_vstate_discard();
 		__riscv_v_vstate_dirty(regs);
 	}
@@ -312,6 +333,7 @@ static inline void riscv_v_vstate_save(struct __riscv_v_ext_state *vstate,
 				       struct pt_regs *regs)
 {
 	if (__riscv_v_vstate_check(regs->status, DIRTY)) {
+		riscv_v_owner_invalidate(RISCV_V_INVALIDATE_TASK_SAVE);
 		__riscv_v_vstate_save(vstate, vstate->datap);
 		__riscv_v_vstate_clean(regs);
 	}
@@ -321,6 +343,7 @@ static inline void riscv_v_vstate_restore(struct __riscv_v_ext_state *vstate,
 					  struct pt_regs *regs)
 {
 	if (riscv_v_vstate_query(regs)) {
+		riscv_v_owner_invalidate(RISCV_V_INVALIDATE_TASK_RESTORE);
 		__riscv_v_vstate_restore(vstate, vstate->datap);
 		__riscv_v_vstate_clean(regs);
 	}
@@ -393,6 +416,13 @@ static inline void __switch_to_vector(struct task_struct *prev,
 	}
 
 	if (riscv_preempt_v_started(next)) {
+		/*
+		 * A preemptible kernel Vector user can be scheduled back after a
+		 * vCPU published a guest owner on this hart.  Drop that owner before
+		 * the kernel context resumes or restores and overwrites the register
+		 * file.
+		 */
+		riscv_v_owner_invalidate(RISCV_V_INVALIDATE_KERNEL_VECTOR);
 		if (next->thread.riscv_v_flags & RISCV_PREEMPT_V_IN_SCHEDULE) {
 			next->thread.riscv_v_flags &= ~RISCV_PREEMPT_V_IN_SCHEDULE;
 			riscv_v_enable();
@@ -423,6 +453,7 @@ static inline bool riscv_v_vstate_ctrl_user_allowed(void) { return false; }
 #define riscv_v_vstate_discard(regs)		do {} while (0)
 #define riscv_v_vstate_save(vstate, regs)	do {} while (0)
 #define riscv_v_vstate_restore(vstate, regs)	do {} while (0)
+#define riscv_v_owner_invalidate(reason)	((void)(reason))
 #define __switch_to_vector(__prev, __next)	do {} while (0)
 #define riscv_v_vstate_off(regs)		do {} while (0)
 #define riscv_v_vstate_on(regs)			do {} while (0)
diff --git a/arch/riscv/kernel/kernel_mode_vector.c b/arch/riscv/kernel/kernel_mode_vector.c
index 77e98b504485..9e87bdf02028 100644
--- a/arch/riscv/kernel/kernel_mode_vector.c
+++ b/arch/riscv/kernel/kernel_mode_vector.c
@@ -22,6 +22,13 @@
 
 static void (* __rcu kvm_flush_vector_ctx_callback)(void);
 
+/*
+ * KVM may leave a guest Vector register file resident on a hart.  Keep the
+ * generic Vector paths unchanged until KVM registers the invalidation hook.
+ */
+static void (* __rcu kvm_invalidate_vector_ctx_callback)(enum riscv_v_invalidate_reason reason);
+DEFINE_STATIC_KEY_FALSE(riscv_v_owner_invalidate_key);
+
 void kvm_riscv_register_vctx_callback(void (*func)(void))
 {
 	if (WARN_ON_ONCE(rcu_access_pointer(kvm_flush_vector_ctx_callback)))
@@ -38,6 +45,35 @@ void kvm_riscv_unregister_vctx_callback(void)
 }
 EXPORT_SYMBOL_GPL(kvm_riscv_unregister_vctx_callback);
 
+void kvm_riscv_register_vctx_invalidate(void (*func)(enum riscv_v_invalidate_reason reason))
+{
+	if (WARN_ON_ONCE(rcu_access_pointer(kvm_invalidate_vector_ctx_callback)))
+		return;
+
+	rcu_assign_pointer(kvm_invalidate_vector_ctx_callback, func);
+	static_branch_enable(&riscv_v_owner_invalidate_key);
+}
+EXPORT_SYMBOL_GPL(kvm_riscv_register_vctx_invalidate);
+
+void kvm_riscv_unregister_vctx_invalidate(void)
+{
+	static_branch_disable(&riscv_v_owner_invalidate_key);
+	rcu_assign_pointer(kvm_invalidate_vector_ctx_callback, NULL);
+	synchronize_rcu();
+}
+EXPORT_SYMBOL_GPL(kvm_riscv_unregister_vctx_invalidate);
+
+void __riscv_v_owner_invalidate(enum riscv_v_invalidate_reason reason)
+{
+	void (*invalidate_callback)(enum riscv_v_invalidate_reason reason);
+
+	rcu_read_lock();
+	invalidate_callback = rcu_dereference(kvm_invalidate_vector_ctx_callback);
+	if (invalidate_callback)
+		invalidate_callback(reason);
+	rcu_read_unlock();
+}
+
 
 static inline void riscv_v_start(u32 flags)
 {
@@ -106,6 +142,8 @@ static void __riscv_flush_vector_context(void)
 {
 	void (*vcpu_flush_v_callback)(void);
 
+	riscv_v_owner_invalidate(RISCV_V_INVALIDATE_KERNEL_VECTOR);
+
 	if (riscv_v_flags() & RISCV_V_VCPU_CTX) {
 		rcu_read_lock();
 		vcpu_flush_v_callback = rcu_dereference(kvm_flush_vector_ctx_callback);
-- 
2.43.0