put_cpu_fpsimd_context();
 }
 
-/*
- * Like fpsimd_preserve_current_state(), but ensure that
- * current->thread.uw.fpsimd_state is updated so that it can be copied to
- * the signal frame.
- */
-void fpsimd_signal_preserve_current_state(void)
-{
-       fpsimd_preserve_current_state();
-       if (current->thread.fp_type == FP_STATE_SVE)
-               sve_to_fpsimd(current);
-}
-
 /*
  * Associate current's FPSIMD context with this cpu
  * The caller must have ownership of the cpu FPSIMD context before calling
        put_cpu_fpsimd_context();
 }
 
-/*
- * Load an updated userland FPSIMD state for 'current' from memory and set the
- * flag that indicates that the FPSIMD register contents are the most recent
- * FPSIMD state of 'current'. This is used by the signal code to restore the
- * register state when returning from a signal handler in FPSIMD only cases,
- * any SVE context will be discarded.
- */
 void fpsimd_update_current_state(struct user_fpsimd_state const *state)
 {
        if (WARN_ON(!system_supports_fpsimd()))
                return;
 
-       get_cpu_fpsimd_context();
-
        current->thread.uw.fpsimd_state = *state;
        if (current->thread.fp_type == FP_STATE_SVE)
                fpsimd_to_sve(current);
-
-       task_fpsimd_load();
-       fpsimd_bind_task_to_cpu();
-
-       clear_thread_flag(TIF_FOREIGN_FPSTATE);
-
-       put_cpu_fpsimd_context();
 }
 
 /*
 
                ¤t->thread.uw.fpsimd_state;
        int err;
 
+       sve_sync_to_fpsimd(current);
+
        /* copy the FP and status/control registers */
        err = __copy_to_user(ctx->vregs, fpsimd->vregs, sizeof(fpsimd->vregs));
        __put_user_error(fpsimd->fpsr, &ctx->fpsr, err);
 {
        int err = 0;
 
-       current->thread.uw.fpmr = read_sysreg_s(SYS_FPMR);
-
        __put_user_error(FPMR_MAGIC, &ctx->head.magic, err);
        __put_user_error(sizeof(*ctx), &ctx->head.size, err);
        __put_user_error(current->thread.uw.fpmr, &ctx->fpmr, err);
 
        __get_user_error(fpmr, &user->fpmr->fpmr, err);
        if (!err)
-               write_sysreg_s(fpmr, SYS_FPMR);
+               current->thread.uw.fpmr = fpmr;
 
        return err;
 }
        err |= __copy_to_user(&ctx->__reserved, reserved, sizeof(reserved));
 
        if (vq) {
-               /*
-                * This assumes that the SVE state has already been saved to
-                * the task struct by calling the function
-                * fpsimd_signal_preserve_current_state().
-                */
                err |= __copy_to_user((char __user *)ctx + SVE_SIG_REGS_OFFSET,
                                      current->thread.sve_state,
                                      SVE_SIG_REGS_SIZE(vq));
        if (user->sve_size < SVE_SIG_CONTEXT_SIZE(vq))
                return -EINVAL;
 
-       /*
-        * Careful: we are about __copy_from_user() directly into
-        * thread.sve_state with preemption enabled, so protection is
-        * needed to prevent a racing context switch from writing stale
-        * registers back over the new data.
-        */
-
-       fpsimd_flush_task_state(current);
-       /* From now, fpsimd_thread_switch() won't touch thread.sve_state */
-
        sve_alloc(current, true);
        if (!current->thread.sve_state) {
                clear_thread_flag(TIF_SVE);
        err |= __copy_to_user(&ctx->__reserved, reserved, sizeof(reserved));
 
        if (vq) {
-               /*
-                * This assumes that the ZA state has already been saved to
-                * the task struct by calling the function
-                * fpsimd_signal_preserve_current_state().
-                */
                err |= __copy_to_user((char __user *)ctx + ZA_SIG_REGS_OFFSET,
                                      current->thread.sme_state,
                                      ZA_SIG_REGS_SIZE(vq));
        if (user->za_size < ZA_SIG_CONTEXT_SIZE(vq))
                return -EINVAL;
 
-       /*
-        * Careful: we are about __copy_from_user() directly into
-        * thread.sme_state with preemption enabled, so protection is
-        * needed to prevent a racing context switch from writing stale
-        * registers back over the new data.
-        */
-
-       fpsimd_flush_task_state(current);
-       /* From now, fpsimd_thread_switch() won't touch thread.sve_state */
-
        sme_alloc(current, true);
        if (!current->thread.sme_state) {
                current->thread.svcr &= ~SVCR_ZA_MASK;
        BUILD_BUG_ON(sizeof(ctx->__reserved) != sizeof(reserved));
        err |= __copy_to_user(&ctx->__reserved, reserved, sizeof(reserved));
 
-       /*
-        * This assumes that the ZT state has already been saved to
-        * the task struct by calling the function
-        * fpsimd_signal_preserve_current_state().
-        */
        err |= __copy_to_user((char __user *)ctx + ZT_SIG_REGS_OFFSET,
                              thread_zt_state(¤t->thread),
                              ZT_SIG_REGS_SIZE(1));
        if (nregs != 1)
                return -EINVAL;
 
-       /*
-        * Careful: we are about __copy_from_user() directly into
-        * thread.zt_state with preemption enabled, so protection is
-        * needed to prevent a racing context switch from writing stale
-        * registers back over the new data.
-        */
-
-       fpsimd_flush_task_state(current);
-       /* From now, fpsimd_thread_switch() won't touch ZT in thread state */
-
        err = __copy_from_user(thread_zt_state(¤t->thread),
                               (char __user const *)user->zt +
                                        ZT_SIG_REGS_OFFSET,
         */
        forget_syscall(regs);
 
+       fpsimd_save_and_flush_current_state();
+
        err |= !valid_user_regs(®s->user_regs, current);
        if (err == 0)
                err = parse_user_sigframe(&user, sf);
                /*
                 * If we were in streaming mode the saved register
                 * state was SVE but we will exit SM and use the
-                * FPSIMD register state - flush the saved FPSIMD
-                * register state in case it gets loaded.
+                * FPSIMD register state.
+                *
+                * TODO: decide if this should behave as SMSTOP (e.g. reset
+                * FPSR + FPMR), or whether this should only clear the scalable
+                * registers + ZA state.
                 */
                if (current->thread.svcr & SVCR_SM_MASK) {
                        memset(¤t->thread.uw.fpsimd_state, 0,
                        current->thread.fp_type = FP_STATE_FPSIMD;
                }
 
-               current->thread.svcr &= ~(SVCR_ZA_MASK |
-                                         SVCR_SM_MASK);
-               sme_smstop();
+               current->thread.svcr &= ~(SVCR_ZA_MASK | SVCR_SM_MASK);
        }
 
        return 0;
        struct user_access_state ua_state;
        int err = 0;
 
-       fpsimd_signal_preserve_current_state();
+       fpsimd_save_and_flush_current_state();
 
        if (get_sigframe(&user, ksig, regs))
                return 1;