Re: [PATCH v3 2/5] arm64: fpsimd: Preserve/restore kernel mode NEON at context switch

[Mark Rutland <mark.rutland@arm.com>]

Hi Ard,

On Mon, Nov 27, 2023 at 01:23:02PM +0100, Ard Biesheuvel wrote:
> From: Ard Biesheuvel <ardb@kernel.org>
> 
> Currently, the FPSIMD register file is not preserved and restored along
> with the general registers on exception entry/exit or context switch.
> For this reason, we disable preemption when enabling FPSIMD for kernel
> mode use in task context, and suspend the processing of softirqs so that
> there are no concurrent uses in the kernel. (Kernel mode FPSIMD may not
> be used at all in other contexts).
> 
> Disabling preemption while doing CPU intensive work on inputs of
> potentially unbounded size is bad for real-time performance, which is
> why we try and ensure that SIMD crypto code does not operate on more
> than ~4k at a time, which is an arbitrary limit and requires assembler
> code to implement efficiently.
> 
> We can avoid the need for disabling preemption if we can ensure that any
> in-kernel users of the NEON will not lose the FPSIMD register state
> across a context switch. And given that disabling softirqs implicitly
> disables preemption as well, we will also have to ensure that a softirq
> that runs code using FPSIMD can safely interrupt an in-kernel user.
> 
> So introduce a thread_info flag TIF_USING_KMODE_FPSIMD, and modify the
> context switch hook for FPSIMD to preserve and restore the kernel mode
> FPSIMD to/from struct thread_struct when it is set. This avoids any
> scheduling blackouts due to prolonged use of FPSIMD in kernel mode,
> without the need for manual yielding.
> 
> In order to support softirq processing while FPSIMD is being used in
> kernel task context, use the same flag to decide whether the kernel mode
> FPSIMD state needs to be preserved and restored before allowing FPSIMD
> to be used in softirq context.
> 
> Signed-off-by: Ard Biesheuvel <ardb@kernel.org>
> Reviewed-by: Mark Brown <broonie@kernel.org>

I have a couple of naming/structural comments below, but this looks
functionally good to me. I appreciate those are arguable bikeshedding, so
either way:

Reviewed-by: Mark Rutland <mark.rutland@arm.com>

> ---
>  arch/arm64/include/asm/processor.h   |  2 +
>  arch/arm64/include/asm/thread_info.h |  1 +
>  arch/arm64/kernel/fpsimd.c           | 92 ++++++++++++++++----
>  3 files changed, 77 insertions(+), 18 deletions(-)
> 
> diff --git a/arch/arm64/include/asm/processor.h b/arch/arm64/include/asm/processor.h
> index e5bc54522e71..dcb51c0571af 100644
> --- a/arch/arm64/include/asm/processor.h
> +++ b/arch/arm64/include/asm/processor.h
> @@ -167,6 +167,8 @@ struct thread_struct {
>  	unsigned long		fault_address;	/* fault info */
>  	unsigned long		fault_code;	/* ESR_EL1 value */
>  	struct debug_info	debug;		/* debugging */
> +
> +	struct user_fpsimd_state	kmode_fpsimd_state;
>  #ifdef CONFIG_ARM64_PTR_AUTH
>  	struct ptrauth_keys_user	keys_user;
>  #ifdef CONFIG_ARM64_PTR_AUTH_KERNEL
> diff --git a/arch/arm64/include/asm/thread_info.h b/arch/arm64/include/asm/thread_info.h
> index 553d1bc559c6..6b254cf90e8b 100644
> --- a/arch/arm64/include/asm/thread_info.h
> +++ b/arch/arm64/include/asm/thread_info.h
> @@ -80,6 +80,7 @@ void arch_setup_new_exec(void);
>  #define TIF_TAGGED_ADDR		26	/* Allow tagged user addresses */
>  #define TIF_SME			27	/* SME in use */
>  #define TIF_SME_VL_INHERIT	28	/* Inherit SME vl_onexec across exec */
> +#define TIF_USING_KMODE_FPSIMD	29	/* Task is in a kernel mode FPSIMD section */

Sorry for the bikeshedding, but "KMODE" isn't a term we use elsewhere, and I
think it'd be nicer/clearer if this had "KERNEL" spelled out in full,
especially as it's only 1 additional character.

Could this be TIF_FPSIMD_KERNEL, or maybe TIF_KERNEL_FPSTATE to align with
TIF_FORIEGN_FPSTATE?

>  
>  #define _TIF_SIGPENDING		(1 << TIF_SIGPENDING)
>  #define _TIF_NEED_RESCHED	(1 << TIF_NEED_RESCHED)
> diff --git a/arch/arm64/kernel/fpsimd.c b/arch/arm64/kernel/fpsimd.c
> index ccc4a78a70e4..198918805bf6 100644
> --- a/arch/arm64/kernel/fpsimd.c
> +++ b/arch/arm64/kernel/fpsimd.c
> @@ -357,6 +357,7 @@ static void task_fpsimd_load(void)
>  
>  	WARN_ON(!system_supports_fpsimd());
>  	WARN_ON(preemptible());
> +	WARN_ON(test_thread_flag(TIF_USING_KMODE_FPSIMD));
>  
>  	if (system_supports_sve() || system_supports_sme()) {
>  		switch (current->thread.fp_type) {
> @@ -379,7 +380,7 @@ static void task_fpsimd_load(void)
>  		default:
>  			/*
>  			 * This indicates either a bug in
> -			 * fpsimd_save() or memory corruption, we
> +			 * fpsimd_save_user_state() or memory corruption, we
>  			 * should always record an explicit format
>  			 * when we save. We always at least have the
>  			 * memory allocated for FPSMID registers so
> @@ -430,7 +431,7 @@ static void task_fpsimd_load(void)
>   * than via current, if we are saving KVM state then it will have
>   * ensured that the type of registers to save is set in last->to_save.
>   */
> -static void fpsimd_save(void)
> +static void fpsimd_save_user_state(void)
>  {
>  	struct cpu_fp_state const *last =
>  		this_cpu_ptr(&fpsimd_last_state);
> @@ -861,7 +862,7 @@ int vec_set_vector_length(struct task_struct *task, enum vec_type type,
>  	if (task == current) {
>  		get_cpu_fpsimd_context();
>  
> -		fpsimd_save();
> +		fpsimd_save_user_state();
>  	}
>  
>  	fpsimd_flush_task_state(task);
> @@ -1473,6 +1474,16 @@ void do_fpsimd_exc(unsigned long esr, struct pt_regs *regs)
>  		       current);
>  }
>  
> +static void fpsimd_load_kernel_state(struct task_struct *task)
> +{
> +	fpsimd_load_state(&task->thread.kmode_fpsimd_state);
> +}
> +
> +static void fpsimd_save_kernel_state(struct task_struct *task)
> +{
> +	fpsimd_save_state(&task->thread.kmode_fpsimd_state);
> +}
> +
>  void fpsimd_thread_switch(struct task_struct *next)
>  {
>  	bool wrong_task, wrong_cpu;
> @@ -1483,19 +1494,28 @@ void fpsimd_thread_switch(struct task_struct *next)
>  	WARN_ON_ONCE(!irqs_disabled());
>  
>  	/* Save unsaved fpsimd state, if any: */
> -	fpsimd_save();
> +	if (!test_thread_flag(TIF_USING_KMODE_FPSIMD))
> +		fpsimd_save_user_state();
> +	else
> +		fpsimd_save_kernel_state(current);

Minor nit: I find this condition inversion slightly hard to read since in prose
it'd be "if there's not kernel state, save the user state; else save the kernel
state", whereas:

	if (test_thread_flag(TIF_USING_KMODE_FPSIMD))
		fpsimd_save_kernel_state(current);
	else
		fpsimd_save_user_state();

... is more clearly "if there's kernel state, save it; else save the user
state", and I think that'd be preferable.

>  
> -	/*
> -	 * Fix up TIF_FOREIGN_FPSTATE to correctly describe next's
> -	 * state.  For kernel threads, FPSIMD registers are never loaded
> -	 * and wrong_task and wrong_cpu will always be true.
> -	 */
> -	wrong_task = __this_cpu_read(fpsimd_last_state.st) !=
> -					&next->thread.uw.fpsimd_state;
> -	wrong_cpu = next->thread.fpsimd_cpu != smp_processor_id();
> +	if (test_tsk_thread_flag(next, TIF_USING_KMODE_FPSIMD)) {
> +		fpsimd_load_kernel_state(next);
> +		set_tsk_thread_flag(next, TIF_FOREIGN_FPSTATE);
> +	} else {
> +		/*
> +		 * Fix up TIF_FOREIGN_FPSTATE to correctly describe next's
> +		 * state.  For kernel threads, FPSIMD registers are never
> +		 * loaded with user mode FPSIMD state and so wrong_task and
> +		 * wrong_cpu will always be true.
> +		 */
> +		wrong_task = __this_cpu_read(fpsimd_last_state.st) !=
> +			&next->thread.uw.fpsimd_state;
> +		wrong_cpu = next->thread.fpsimd_cpu != smp_processor_id();
>  
> -	update_tsk_thread_flag(next, TIF_FOREIGN_FPSTATE,
> -			       wrong_task || wrong_cpu);
> +		update_tsk_thread_flag(next, TIF_FOREIGN_FPSTATE,
> +				       wrong_task || wrong_cpu);
> +	}
>  }
>  
>  static void fpsimd_flush_thread_vl(enum vec_type type)
> @@ -1585,7 +1605,7 @@ void fpsimd_preserve_current_state(void)
>  		return;
>  
>  	get_cpu_fpsimd_context();
> -	fpsimd_save();
> +	fpsimd_save_user_state();
>  	put_cpu_fpsimd_context();
>  }
>  
> @@ -1801,7 +1821,7 @@ void fpsimd_save_and_flush_cpu_state(void)
>  		return;
>  	WARN_ON(preemptible());
>  	get_cpu_fpsimd_context();
> -	fpsimd_save();
> +	fpsimd_save_user_state();
>  	fpsimd_flush_cpu_state();
>  	put_cpu_fpsimd_context();
>  }
> @@ -1835,10 +1855,37 @@ void kernel_neon_begin(void)
>  	get_cpu_fpsimd_context();
>  
>  	/* Save unsaved fpsimd state, if any: */
> -	fpsimd_save();
> +	if (!test_thread_flag(TIF_USING_KMODE_FPSIMD)) {
> +		fpsimd_save_user_state();
> +
> +		/*
> +		 * Set the thread flag so that the kernel mode FPSIMD state
> +		 * will be context switched along with the rest of the task
> +		 * state.
> +		 *
> +		 * On non-PREEMPT_RT, softirqs may interrupt task level kernel
> +		 * mode FPSIMD, but the task will not be preemptible so setting
> +		 * TIF_USING_KMODE_FPSIMD for those would be both wrong (as it
> +		 * would mark the task context FPSIMD state as requiring a
> +		 * context switch) and unnecessary.
> +		 *
> +		 * On PREEMPT_RT, softirqs are serviced from a separate thread,
> +		 * which is scheduled as usual, and this guarantees that these
> +		 * softirqs are not interrupting use of the FPSIMD in kernel
> +		 * mode in task context. So in this case, setting the flag here
> +		 * is always appropriate.
> +		 */
> +		if (IS_ENABLED(CONFIG_PREEMPT_RT) || !in_serving_softirq())
> +			set_thread_flag(TIF_USING_KMODE_FPSIMD);
> +	} else {
> +		BUG_ON(IS_ENABLED(CONFIG_PREEMPT_RT) || !in_serving_softirq());
> +		fpsimd_save_kernel_state(current);
> +	}

Same comment as above for condition inversion here.

Mark.

>  
>  	/* Invalidate any task state remaining in the fpsimd regs: */
>  	fpsimd_flush_cpu_state();
> +
> +	put_cpu_fpsimd_context();
>  }
>  EXPORT_SYMBOL_GPL(kernel_neon_begin);
>  
> @@ -1856,7 +1903,16 @@ void kernel_neon_end(void)
>  	if (!system_supports_fpsimd())
>  		return;
>  
> -	put_cpu_fpsimd_context();
> +	/*
> +	 * If we are returning from a nested use of kernel mode FPSIMD, restore
> +	 * the task context kernel mode FPSIMD state. This can only happen when
> +	 * running in softirq context on non-PREEMPT_RT.
> +	 */
> +	if (!IS_ENABLED(CONFIG_PREEMPT_RT) && in_serving_softirq() &&
> +	    test_thread_flag(TIF_USING_KMODE_FPSIMD))
> +		fpsimd_load_kernel_state(current);
> +	else
> +		clear_thread_flag(TIF_USING_KMODE_FPSIMD);
>  }
>  EXPORT_SYMBOL_GPL(kernel_neon_end);
>  
> -- 
> 2.43.0.rc1.413.gea7ed67945-goog
> 

_______________________________________________
linux-arm-kernel mailing list
linux-arm-kernel@lists.infradead.org
http://lists.infradead.org/mailman/listinfo/linux-arm-kernel