diff --git a/arch/arm64/kvm/Kconfig b/arch/arm64/kvm/Kconfig
index 2257dfcc44cc..5b2c8d8c9722 100644
--- a/arch/arm64/kvm/Kconfig
+++ b/arch/arm64/kvm/Kconfig
@@ -39,6 +39,7 @@ config KVM
 	select HAVE_KVM_IRQ_ROUTING
 	select IRQ_BYPASS_MANAGER
 	select HAVE_KVM_IRQ_BYPASS
+	select HAVE_KVM_VCPU_RUN_PID_CHANGE
 	---help---
 	  Support hosting virtualized guest machines.
 	  We don't support KVM with 16K page tables yet, due to the multiple
diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index ac0062b74aed..10a37b122f6f 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -1272,4 +1272,13 @@ static inline long kvm_arch_vcpu_async_ioctl(struct file *filp,
 }
 #endif /* CONFIG_HAVE_KVM_VCPU_ASYNC_IOCTL */
 
+#ifdef CONFIG_HAVE_KVM_VCPU_RUN_PID_CHANGE
+int kvm_arch_vcpu_run_pid_change(struct kvm_vcpu *vcpu);
+#else
+static inline int kvm_arch_vcpu_run_pid_change(struct kvm_vcpu *vcpu)
+{
+	return 0;
+}
+#endif /* CONFIG_HAVE_KVM_VCPU_RUN_PID_CHANGE */
+
 #endif
diff --git a/virt/kvm/Kconfig b/virt/kvm/Kconfig
index cca7e065a075..72143cfaf6ec 100644
--- a/virt/kvm/Kconfig
+++ b/virt/kvm/Kconfig
@@ -54,3 +54,6 @@ config HAVE_KVM_IRQ_BYPASS
 
 config HAVE_KVM_VCPU_ASYNC_IOCTL
        bool
+
+config HAVE_KVM_VCPU_RUN_PID_CHANGE
+       bool
diff --git a/virt/kvm/arm/arm.c b/virt/kvm/arm/arm.c
index 0330e1f8fb09..99eb52559f24 100644
--- a/virt/kvm/arm/arm.c
+++ b/virt/kvm/arm/arm.c
@@ -867,6 +867,23 @@ int kvm_arch_vcpu_ioctl_run(struct kvm_vcpu *vcpu, struct kvm_run *run)
 	return ret;
 }
 
+#ifdef CONFIG_ARM64
+int kvm_arch_vcpu_run_pid_change(struct kvm_vcpu *vcpu)
+{
+	struct task_struct *tsk = current;
+	int ret;
+
+	/*
+	 * Make sure struct thread_info (and TIF flags) and the fpsimd state
+	 * are visible to hyp.
+	 */
+	ret = create_hyp_mappings(tsk, tsk + 1, PAGE_HYP);
+	if (!ret)
+		vcpu->arch.hyp_current = kern_hyp_va(current);
+	return ret;
+}
+#endif
+
 static int vcpu_interrupt_line(struct kvm_vcpu *vcpu, int number, bool level)
 {
 	int bit_index;
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 4501e658e8d6..dbd35abe7d9c 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -2551,8 +2551,13 @@ static long kvm_vcpu_ioctl(struct file *filp,
 		oldpid = rcu_access_pointer(vcpu->pid);
 		if (unlikely(oldpid != current->pids[PIDTYPE_PID].pid)) {
 			/* The thread running this VCPU changed. */
-			struct pid *newpid = get_task_pid(current, PIDTYPE_PID);
+			struct pid *newpid;
 
+			r = kvm_arch_vcpu_run_pid_change(vcpu);
+			if (r)
+				break;
+
+			newpid = get_task_pid(current, PIDTYPE_PID);
 			rcu_assign_pointer(vcpu->pid, newpid);
 			if (oldpid)
 				synchronize_rcu();

commit 6bb55488489d69885b51819add3690da523be12a (HEAD -> kvm-vfp-integration-rfc)
Author: Christoffer Dall <christoffer.dall@linaro.org>
Date:   Fri Feb 23 17:58:17 2018 +0100

    KVM: arm64: Be more lazy with switching KVM guest FPSIMD state
    
    We currently save the FPSIMD state back from the CPU on every exit, when
    the guest has touched the FPSIMD state.
    
    We can try to avoid this by changing the state that is tracked by the
    kernel FPSIMD mechanism to the KVM guest state, and keep track of this
    using additional thread flag.  Whenever we go back to userspace from the
    KVM_RUN ioctl, we check if we switched to the KVM state, and make sure
    the state is copied back.
    
    Signed-off-by: Christoffer Dall <christoffer.dall@linaro.org>

diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index 95ffb54daec2..df819376ae9a 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -241,8 +241,14 @@ struct kvm_vcpu_arch {
 
 	/* Pointer to host CPU context */
 	kvm_cpu_context_t *host_cpu_context;
-	struct user_fpsimd_state *host_fpsimd_state; /* hyp va */
-	bool guest_fpsimd_loaded;
+	struct task_struct *hyp_current;
+
+	/*
+	 * If FPSIMD registers are valid when entering the guest, this is
+	 * where we store the host userspace register state.
+	 */
+	struct user_fpsimd_state host_fpsimd_state;
+
 	struct {
 		/* {Break,watch}point registers */
 		struct kvm_guest_debug_arch regs;
diff --git a/arch/arm64/include/asm/thread_info.h b/arch/arm64/include/asm/thread_info.h
index 9f1fa1a49bb4..6ec3c8b51898 100644
--- a/arch/arm64/include/asm/thread_info.h
+++ b/arch/arm64/include/asm/thread_info.h
@@ -94,7 +94,7 @@ void arch_release_task_struct(struct task_struct *tsk);
 #define TIF_32BIT		22	/* 32bit process */
 #define TIF_SVE			23	/* Scalable Vector Extension in use */
 #define TIF_SVE_VL_INHERIT	24	/* Inherit sve_vl_onexec across exec */
-#define TIF_MAPPED_TO_HYP	25	/* task_struct mapped to Hyp (KVM) */
+#define TIF_KVM_GUEST_FPSTATE	25	/* current's FP state belongs to KVM */
 
 #define _TIF_SIGPENDING		(1 << TIF_SIGPENDING)
 #define _TIF_NEED_RESCHED	(1 << TIF_NEED_RESCHED)
diff --git a/arch/arm64/kvm/hyp/switch.c b/arch/arm64/kvm/hyp/switch.c
index b88e83fc76c8..a1034e880d6e 100644
--- a/arch/arm64/kvm/hyp/switch.c
+++ b/arch/arm64/kvm/hyp/switch.c
@@ -17,6 +17,7 @@
 
 #include <linux/types.h>
 #include <linux/jump_label.h>
+#include <linux/thread_info.h>
 #include <uapi/linux/psci.h>
 
 #include <kvm/arm_psci.h>
@@ -28,6 +29,15 @@
 #include <asm/fpsimd.h>
 #include <asm/debug-monitors.h>
 
+#define hyp_current(vcpu) ((vcpu)->arch.hyp_current)
+
+#define hyp_set_thread_flag(vcpu, flag) \
+	set_ti_thread_flag(&hyp_current(vcpu)->thread_info, flag)
+#define hyp_clear_thread_flag(vcpu, flag) \
+	clear_ti_thread_flag(&hyp_current(vcpu)->thread_info, flag)
+#define hyp_test_thread_flag(vcpu, flag) \
+	test_ti_thread_flag(&hyp_current(vcpu)->thread_info, flag)
+
 static bool __hyp_text __fpsimd_enabled_nvhe(void)
 {
 	return !(read_sysreg(cptr_el2) & CPTR_EL2_TFP);
@@ -85,7 +95,7 @@ static void __hyp_text __deactivate_traps_common(void)
 	write_sysreg(0, pmuserenr_el0);
 }
 
-static inline void activate_traps_vhe(struct kvm_vcpu *vcpu)
+static void __hyp_text activate_traps_vhe(struct kvm_vcpu *vcpu)
 {
 	u64 val;
 
@@ -93,7 +103,8 @@ static inline void activate_traps_vhe(struct kvm_vcpu *vcpu)
 	val |= CPACR_EL1_TTA;
 
 	val &= ~CPACR_EL1_ZEN;
-	if (!vcpu->arch.guest_fpsimd_loaded)
+	if (!hyp_test_thread_flag(vcpu, TIF_KVM_GUEST_FPSTATE) ||
+	    hyp_test_thread_flag(vcpu, TIF_FOREIGN_FPSTATE))
 		val &= ~CPACR_EL1_FPEN;
 
 	write_sysreg(val, cpacr_el1);
@@ -109,7 +120,8 @@ static inline void __hyp_text __activate_traps_nvhe(struct kvm_vcpu *vcpu)
 
 	val = CPTR_EL2_DEFAULT;
 	val |= CPTR_EL2_TTA | CPTR_EL2_TZ;
-	if (!vcpu->arch.guest_fpsimd_loaded)
+	if (!hyp_test_thread_flag(vcpu, TIF_KVM_GUEST_FPSTATE) ||
+	    hyp_test_thread_flag(vcpu, TIF_FOREIGN_FPSTATE))
 		val |= CPTR_EL2_TFP;
 
 	write_sysreg(val, cptr_el2);
@@ -512,6 +524,13 @@ int __hyp_text __kvm_vcpu_run_nvhe(struct kvm_vcpu *vcpu)
 void __hyp_text __hyp_switch_fpsimd(u64 esr __always_unused,
 				    struct kvm_vcpu *vcpu)
 {
+	struct user_fpsimd_state *current_fpsimd =
+		&hyp_current(vcpu)->thread.fpsimd_state.user_fpsimd;
+	struct user_fpsimd_state *guest_fpsimd =
+		&vcpu->arch.ctxt.gp_regs.fp_regs;
+	struct user_fpsimd_state *host_fpsimd =
+		&vcpu->arch.host_fpsimd_state;
+
 	if (has_vhe())
 		write_sysreg(read_sysreg(cpacr_el1) | CPACR_EL1_FPEN,
 			     cpacr_el1);
@@ -521,11 +540,27 @@ void __hyp_text __hyp_switch_fpsimd(u64 esr __always_unused,
 
 	isb();
 
-	if (vcpu->arch.host_fpsimd_state)
-		__fpsimd_save_state(vcpu->arch.host_fpsimd_state);
+	/*
+	 * We trapped on guest FPSIMD access.  There are two situations:
+	 *   (1) This is the first use of FPSIMD by the guest for this ioctl
+	 *       invocation.  We make sure the host userspace state is backed
+	 *       up (either from the CPU or from memory).
+	 *   (2) We were preempted or a softirq called kernel_neon_being.  We
+	 *       rely on the kernel fpsimd machinery to have saved our state
+	 *       and we simply restore it.
+	 */
+	if (!hyp_test_thread_flag(vcpu, TIF_KVM_GUEST_FPSTATE)) {
+		if (!hyp_test_thread_flag(vcpu, TIF_FOREIGN_FPSTATE))
+			__fpsimd_save_state(host_fpsimd);
+		else
+			memcpy(host_fpsimd, current_fpsimd, sizeof(*host_fpsimd));
+		__fpsimd_restore_state(guest_fpsimd);
+	} else {
+		__fpsimd_restore_state(current_fpsimd);
+	}
 
-	__fpsimd_restore_state(&vcpu->arch.ctxt.gp_regs.fp_regs);
-	vcpu->arch.guest_fpsimd_loaded = true;
+	hyp_clear_thread_flag(vcpu, TIF_FOREIGN_FPSTATE);
+	hyp_set_thread_flag(vcpu, TIF_KVM_GUEST_FPSTATE);
 
 	/* Skip restoring fpexc32 for AArch64 guests */
 	if (!(read_sysreg(hcr_el2) & HCR_RW))
diff --git a/virt/kvm/arm/arm.c b/virt/kvm/arm/arm.c
index 99eb52559f24..2fe59aff2099 100644
--- a/virt/kvm/arm/arm.c
+++ b/virt/kvm/arm/arm.c
@@ -329,10 +329,6 @@ void kvm_arch_vcpu_unblocking(struct kvm_vcpu *vcpu)
 
 int kvm_arch_vcpu_init(struct kvm_vcpu *vcpu)
 {
-	/* Mark this vcpu's FPSIMD state as non-live initially: */
-	fpsimd_flush_state(&vcpu->arch.ctxt.fpsimd_state);
-	vcpu->arch.guest_fpsimd_loaded = false;
-
 	/* Force users to call KVM_ARM_VCPU_INIT */
 	vcpu->arch.target = -1;
 	bitmap_zero(vcpu->arch.features, KVM_VCPU_MAX_FEATURES);
@@ -635,9 +631,6 @@ static void check_vcpu_requests(struct kvm_vcpu *vcpu)
 int kvm_arch_vcpu_ioctl_run(struct kvm_vcpu *vcpu, struct kvm_run *run)
 {
 	int ret;
-	struct fpsimd_state *guest_fpsimd = &vcpu->arch.ctxt.fpsimd_state;
-	struct user_fpsimd_state *host_fpsimd =
-		&current->thread.fpsimd_state.user_fpsimd;
 
 	if (unlikely(!kvm_vcpu_initialized(vcpu)))
 		return -ENOEXEC;
@@ -659,15 +652,6 @@ int kvm_arch_vcpu_ioctl_run(struct kvm_vcpu *vcpu, struct kvm_run *run)
 
 	WARN_ON(!current->mm);
 
-	if (!test_thread_flag(TIF_MAPPED_TO_HYP)) {
-		ret = create_hyp_mappings(host_fpsimd, host_fpsimd + 1,
-					  PAGE_HYP);
-		if (ret)
-			return ret;
-
-		set_thread_flag(TIF_MAPPED_TO_HYP);
-	}
-
 	vcpu_load(vcpu);
 
 	kvm_sigset_activate(vcpu);
@@ -698,23 +682,8 @@ int kvm_arch_vcpu_ioctl_run(struct kvm_vcpu *vcpu, struct kvm_run *run)
 
 		local_irq_disable();
 
-		/*
-		 * host_fpsimd_state indicates to hyp that there is host state
-		 * to save, and where to save it:
-		 */
-		if (test_thread_flag(TIF_FOREIGN_FPSTATE))
-			vcpu->arch.host_fpsimd_state = NULL;
-		else
-			vcpu->arch.host_fpsimd_state = kern_hyp_va(host_fpsimd);
-
-		vcpu->arch.guest_fpsimd_loaded =
-			!fpsimd_foreign_fpstate(guest_fpsimd);
-
 		BUG_ON(system_supports_sve());
 
-		BUG_ON(vcpu->arch.guest_fpsimd_loaded &&
-		       vcpu->arch.host_fpsimd_state);
-
 		kvm_vgic_flush_hwstate(vcpu);
 
 		/*
@@ -809,9 +778,6 @@ int kvm_arch_vcpu_ioctl_run(struct kvm_vcpu *vcpu, struct kvm_run *run)
 		if (static_branch_unlikely(&userspace_irqchip_in_use))
 			kvm_timer_sync_hwstate(vcpu);
 
-		/* defend against kernel-mode NEON in softirq */
-		local_bh_disable();
-
 		/*
 		 * We may have taken a host interrupt in HYP mode (ie
 		 * while executing the guest). This interrupt is still
@@ -824,18 +790,6 @@ int kvm_arch_vcpu_ioctl_run(struct kvm_vcpu *vcpu, struct kvm_run *run)
 		 */
 		local_irq_enable();
 
-		if (vcpu->arch.guest_fpsimd_loaded) {
-			set_thread_flag(TIF_FOREIGN_FPSTATE);
-			fpsimd_bind_state_to_cpu(&vcpu->arch.ctxt.fpsimd_state);
-
-			/*
-			 * Protect ourselves against a softirq splatting the
-			 * FPSIMD state once irqs are enabled:
-			 */
-			fpsimd_save_state(guest_fpsimd);
-		}
-		local_bh_enable();
-
 		/*
 		 * We do local_irq_enable() before calling guest_exit() so
 		 * that if a timer interrupt hits while running the guest we
@@ -863,6 +817,25 @@ int kvm_arch_vcpu_ioctl_run(struct kvm_vcpu *vcpu, struct kvm_run *run)
 
 	kvm_sigset_deactivate(vcpu);
 
+	if (test_thread_flag(TIF_KVM_GUEST_FPSTATE)) {
+		struct user_fpsimd_state *current_fpsimd =
+			&current->thread.fpsimd_state.user_fpsimd;
+		struct user_fpsimd_state *guest_fpsimd =
+			&vcpu->arch.ctxt.gp_regs.fp_regs;
+		struct user_fpsimd_state *host_fpsimd =
+			&vcpu->arch.host_fpsimd_state;
+
+		local_bh_disable();
+		if (!test_thread_flag(TIF_FOREIGN_FPSTATE))
+			__fpsimd_save_state(guest_fpsimd);
+		else
+			memcpy(guest_fpsimd, current_fpsimd, sizeof(*guest_fpsimd));
+
+		memcpy(current_fpsimd, host_fpsimd, sizeof(*current_fpsimd));
+		set_thread_flag(TIF_FOREIGN_FPSTATE);
+		local_bh_enable();
+	}
+
 	vcpu_put(vcpu);
 	return ret;
 }
