Hey Andy,

On Fri, Jul 21, 2023 at 11:28:55AM +0000, Andy Chiu wrote:
> Add kernel_vstate to keep track of kernel-mode Vector registers when
> trap introduced context switch happens. Also, provide trap_pt_regs to
> let context save/restore routine reference status.VS at which the trap
> takes place. The thread flag TIF_RISCV_V_KERNEL_MODE indicates whether
> a task is running in kernel-mode Vector with preemption 'ON'. So context
> switch routines know and would save V-regs to kernel_vstate and restore
> V-regs immediately from kernel_vstate if the bit is set.
> 
> Apart from a task's preemption status, the capability of
> running preemptive kernel-mode Vector is jointly controlled by the
> RISCV_V_VSTATE_CTRL_PREEMPTIBLE mask in the task's
> thread.vstate_ctrl. This bit is masked whenever a trap takes place in
> kernel mode while executing preemptive Vector code.
> 
> Also, provide a config CONFIG_RISCV_ISA_V_PREEMPTIVE to give users an
> option to disable preemptible kernel-mode Vector at build time. Users
> with constraint memory may want to disable this config as preemptible
> kernel-mode Vector needs extra space for tracking per thread's
> kernel-mode V context. Or, users might as well want to disable it if all
> kernel-mode Vector code is time sensitive and cannot tolerate context
> swicth overhead.
> 
> Signed-off-by: Andy Chiu <andy.chiu@sifive.com>
> ---
> Changelog v2:
>  - fix build fail when compiling without RISCV_ISA_V (Conor)
>  - 's/TIF_RISCV_V_KMV/TIF_RISCV_V_KERNEL_MODE' and add comment (Conor)
>  - merge Kconfig patch into this oine (Conor).
>  - 's/CONFIG_RISCV_ISA_V_PREEMPTIVE_KMV/CONFIG_RISCV_ISA_V_PREEMPTIVE/'
>    (Conor)
>  - fix some typos (Conor)
>  - enclose assembly with RISCV_ISA_V_PREEMPTIVE.
>  - change riscv_v_vstate_ctrl_config_kmv() to
>    kernel_vector_allow_preemption() for better understanding. (Conor)
>  - 's/riscv_v_kmv_preempitble/kernel_vector_preemptible/'
> ---
>  arch/riscv/Kconfig                     | 10 +++++
>  arch/riscv/include/asm/processor.h     |  2 +
>  arch/riscv/include/asm/simd.h          |  4 +-
>  arch/riscv/include/asm/thread_info.h   |  4 ++
>  arch/riscv/include/asm/vector.h        | 27 +++++++++++--
>  arch/riscv/kernel/asm-offsets.c        |  2 +
>  arch/riscv/kernel/entry.S              | 45 ++++++++++++++++++++++
>  arch/riscv/kernel/kernel_mode_vector.c | 53 ++++++++++++++++++++++++--
>  arch/riscv/kernel/process.c            |  8 +++-
>  arch/riscv/kernel/vector.c             |  3 +-
>  10 files changed, 148 insertions(+), 10 deletions(-)
> 
> diff --git a/arch/riscv/Kconfig b/arch/riscv/Kconfig
> index 4c07b9189c86..0622951b15dd 100644
> --- a/arch/riscv/Kconfig
> +++ b/arch/riscv/Kconfig
> @@ -507,6 +507,16 @@ config RISCV_ISA_V_DEFAULT_ENABLE
>  
>  	  If you don't know what to do here, say Y.
>  
> +config RISCV_ISA_V_PREEMPTIVE
> +	bool "Run kernel-mode Vector with kernel preemption"
> +	depends on PREEMPTION
> +	depends on RISCV_ISA_V
> +	default y
> +	help
> +	  Ordinarily the kernel disables preemption before running in-kernel
> +	  Vector code. This config frees the kernel from disabling preemption
> +	  by adding memory on demand for tracking kernel's V-context.
> +
>  config TOOLCHAIN_HAS_ZBB
>  	bool
>  	default y
> diff --git a/arch/riscv/include/asm/processor.h b/arch/riscv/include/asm/processor.h
> index c950a8d9edef..497c0dd30b2a 100644
> --- a/arch/riscv/include/asm/processor.h
> +++ b/arch/riscv/include/asm/processor.h
> @@ -42,6 +42,8 @@ struct thread_struct {
>  	unsigned long bad_cause;
>  	unsigned long vstate_ctrl;
>  	struct __riscv_v_ext_state vstate;
> +	struct pt_regs *trap_pt_regs;
> +	struct __riscv_v_ext_state kernel_vstate;
>  };
>  
>  /* Whitelist the fstate from the task_struct for hardened usercopy */
> diff --git a/arch/riscv/include/asm/simd.h b/arch/riscv/include/asm/simd.h
> index ef70af78005d..a54a0ce58f4d 100644
> --- a/arch/riscv/include/asm/simd.h
> +++ b/arch/riscv/include/asm/simd.h
> @@ -12,6 +12,7 @@
>  #include <linux/percpu.h>
>  #include <linux/preempt.h>
>  #include <linux/types.h>
> +#include <linux/thread_info.h>
>  
>  #ifdef CONFIG_RISCV_ISA_V
>  
> @@ -35,7 +36,8 @@ static __must_check inline bool may_use_simd(void)
>  	 * where it is set.
>  	 */
>  	return !in_irq() && !irqs_disabled() && !in_nmi() &&
> -	       !this_cpu_read(vector_context_busy);
> +	       !this_cpu_read(vector_context_busy) &&
> +	       !test_thread_flag(TIF_RISCV_V_KERNEL_MODE);
>  }
>  
>  #else /* ! CONFIG_RISCV_ISA_V */
> diff --git a/arch/riscv/include/asm/thread_info.h b/arch/riscv/include/asm/thread_info.h
> index b182f2d03e25..8797d520e8ef 100644
> --- a/arch/riscv/include/asm/thread_info.h
> +++ b/arch/riscv/include/asm/thread_info.h
> @@ -94,6 +94,7 @@ int arch_dup_task_struct(struct task_struct *dst, struct task_struct *src);
>  #define TIF_UPROBE		10	/* uprobe breakpoint or singlestep */
>  #define TIF_32BIT		11	/* compat-mode 32bit process */
>  #define TIF_RISCV_V_DEFER_RESTORE	12 /* restore Vector before returing to user */
> +#define TIF_RISCV_V_KERNEL_MODE			13 /* kernel-mode Vector run with preemption-on */
>  
>  #define _TIF_NOTIFY_RESUME	(1 << TIF_NOTIFY_RESUME)
>  #define _TIF_SIGPENDING		(1 << TIF_SIGPENDING)
> @@ -101,9 +102,12 @@ int arch_dup_task_struct(struct task_struct *dst, struct task_struct *src);
>  #define _TIF_NOTIFY_SIGNAL	(1 << TIF_NOTIFY_SIGNAL)
>  #define _TIF_UPROBE		(1 << TIF_UPROBE)
>  #define _TIF_RISCV_V_DEFER_RESTORE	(1 << TIF_RISCV_V_DEFER_RESTORE)
> +#define _TIF_RISCV_V_KERNEL_MODE	(1 << TIF_RISCV_V_KERNEL_MODE)
>  
>  #define _TIF_WORK_MASK \
>  	(_TIF_NOTIFY_RESUME | _TIF_SIGPENDING | _TIF_NEED_RESCHED | \
>  	 _TIF_NOTIFY_SIGNAL | _TIF_UPROBE)
>  
> +#define RISCV_V_VSTATE_CTRL_PREEMPTIBLE	0x20
> +
>  #endif /* _ASM_RISCV_THREAD_INFO_H */
> diff --git a/arch/riscv/include/asm/vector.h b/arch/riscv/include/asm/vector.h
> index 3b783b317112..c2776851d50d 100644
> --- a/arch/riscv/include/asm/vector.h
> +++ b/arch/riscv/include/asm/vector.h
> @@ -195,9 +195,24 @@ static inline void __switch_to_vector(struct task_struct *prev,
>  {
>  	struct pt_regs *regs;
>  
> -	regs = task_pt_regs(prev);
> -	riscv_v_vstate_save(&prev->thread.vstate, regs);
> -	riscv_v_vstate_set_restore(next, task_pt_regs(next));
> +	if (IS_ENABLED(CONFIG_RISCV_ISA_V_PREEMPTIVE) &&
> +	    test_tsk_thread_flag(prev, TIF_RISCV_V_KERNEL_MODE)) {
> +		regs = prev->thread.trap_pt_regs;
> +		WARN_ON(!regs);

In what cases could these WARN_ON()s be triggered?

> +		riscv_v_vstate_save(&prev->thread.kernel_vstate, regs);
> +	} else {
> +		regs = task_pt_regs(prev);
> +		riscv_v_vstate_save(&prev->thread.vstate, regs);
> +	}
> +
> +	if (IS_ENABLED(CONFIG_RISCV_ISA_V_PREEMPTIVE) &&
> +	    test_tsk_thread_flag(next, TIF_RISCV_V_KERNEL_MODE)) {
> +		regs = next->thread.trap_pt_regs;
> +		WARN_ON(!regs);
> +		riscv_v_vstate_restore(&next->thread.kernel_vstate, regs);
> +	} else {
> +		riscv_v_vstate_set_restore(next, task_pt_regs(next));
> +	}
>  }


>  /*
>   * kernel_vector_begin(): obtain the CPU vector registers for use by the calling
>   * context
> @@ -70,11 +109,14 @@ void kernel_vector_begin(void)
>  
>  	riscv_v_vstate_save(&current->thread.vstate, task_pt_regs(current));
>  
> -	get_cpu_vector_context();
> +	if (!preemptible() || !kernel_vector_preemptible()) {
> +		get_cpu_vector_context();
> +	} else {
> +		if (riscv_v_start_kernel_context())
> +			get_cpu_vector_context();

What happens here if riscv_v_start_kernel_context() fails w/ -ENOMEM?

> +	}
>  
>  	riscv_v_enable();
> -
> -	return 0;
>  }
>  EXPORT_SYMBOL_GPL(kernel_vector_begin);
>  
> @@ -96,6 +138,9 @@  void kernel_vector_end(void)
>  
>  	riscv_v_disable();
>  
> -	put_cpu_vector_context();
> +	if (!test_thread_flag(TIF_RISCV_V_KERNEL_MODE))
> +		put_cpu_vector_context();
> +	else
> +		riscv_v_stop_kernel_context();
>  }

Probably just missing something here, but how come we don't need to call
put_cpu_vector_context() here. I'm just a little confused, since, in
kernel_vector_begin, get_cpu_vector_context() is called.