#define KVM_REQ_SLEEP \
        KVM_ARCH_REQ_FLAGS(0, KVM_REQUEST_WAIT | KVM_REQUEST_NO_WAKEUP)
 #define KVM_REQ_IRQ_PENDING    KVM_ARCH_REQ(1)
+#define KVM_REQ_VCPU_RESET     KVM_ARCH_REQ(2)
 
 DECLARE_STATIC_KEY_FALSE(userspace_irqchip_in_use);
 
 
 typedef struct kvm_cpu_context kvm_cpu_context_t;
 
+struct vcpu_reset_state {
+       unsigned long   pc;
+       unsigned long   r0;
+       bool            be;
+       bool            reset;
+};
+
 struct kvm_vcpu_arch {
        struct kvm_cpu_context ctxt;
 
        /* Cache some mmu pages needed inside spinlock regions */
        struct kvm_mmu_memory_cache mmu_page_cache;
 
+       struct vcpu_reset_state reset_state;
+
        /* Detect first run of a vcpu */
        bool has_run_once;
 };
 
 #include <asm/cputype.h>
 #include <asm/kvm_arm.h>
 #include <asm/kvm_coproc.h>
+#include <asm/kvm_emulate.h>
 
 #include <kvm/arm_arch_timer.h>
 
        /* Reset CP15 registers */
        kvm_reset_coprocs(vcpu);
 
+       /*
+        * Additional reset state handling that PSCI may have imposed on us.
+        * Must be done after all the sys_reg reset.
+        */
+       if (READ_ONCE(vcpu->arch.reset_state.reset)) {
+               unsigned long target_pc = vcpu->arch.reset_state.pc;
+
+               /* Gracefully handle Thumb2 entry point */
+               if (target_pc & 1) {
+                       target_pc &= ~1UL;
+                       vcpu_set_thumb(vcpu);
+               }
+
+               /* Propagate caller endianness */
+               if (vcpu->arch.reset_state.be)
+                       kvm_vcpu_set_be(vcpu);
+
+               *vcpu_pc(vcpu) = target_pc;
+               vcpu_set_reg(vcpu, 0, vcpu->arch.reset_state.r0);
+
+               vcpu->arch.reset_state.reset = false;
+       }
+
        /* Reset arch_timer context */
        return kvm_timer_vcpu_reset(vcpu);
 }
 
 #define KVM_REQ_SLEEP \
        KVM_ARCH_REQ_FLAGS(0, KVM_REQUEST_WAIT | KVM_REQUEST_NO_WAKEUP)
 #define KVM_REQ_IRQ_PENDING    KVM_ARCH_REQ(1)
+#define KVM_REQ_VCPU_RESET     KVM_ARCH_REQ(2)
 
 DECLARE_STATIC_KEY_FALSE(userspace_irqchip_in_use);
 
 
 typedef struct kvm_cpu_context kvm_cpu_context_t;
 
+struct vcpu_reset_state {
+       unsigned long   pc;
+       unsigned long   r0;
+       bool            be;
+       bool            reset;
+};
+
 struct kvm_vcpu_arch {
        struct kvm_cpu_context ctxt;
 
        /* Virtual SError ESR to restore when HCR_EL2.VSE is set */
        u64 vsesr_el2;
 
+       /* Additional reset state */
+       struct vcpu_reset_state reset_state;
+
        /* True when deferrable sysregs are loaded on the physical CPU,
         * see kvm_vcpu_load_sysregs and kvm_vcpu_put_sysregs. */
        bool sysregs_loaded_on_cpu;
 
 #include <asm/kvm_arm.h>
 #include <asm/kvm_asm.h>
 #include <asm/kvm_coproc.h>
+#include <asm/kvm_emulate.h>
 #include <asm/kvm_mmu.h>
 
 /* Maximum phys_shift supported for any VM on this host */
        /* Reset system registers */
        kvm_reset_sys_regs(vcpu);
 
+       /*
+        * Additional reset state handling that PSCI may have imposed on us.
+        * Must be done after all the sys_reg reset.
+        */
+       if (vcpu->arch.reset_state.reset) {
+               unsigned long target_pc = vcpu->arch.reset_state.pc;
+
+               /* Gracefully handle Thumb2 entry point */
+               if (vcpu_mode_is_32bit(vcpu) && (target_pc & 1)) {
+                       target_pc &= ~1UL;
+                       vcpu_set_thumb(vcpu);
+               }
+
+               /* Propagate caller endianness */
+               if (vcpu->arch.reset_state.be)
+                       kvm_vcpu_set_be(vcpu);
+
+               *vcpu_pc(vcpu) = target_pc;
+               vcpu_set_reg(vcpu, 0, vcpu->arch.reset_state.r0);
+
+               vcpu->arch.reset_state.reset = false;
+       }
+
        /* Reset PMU */
        kvm_pmu_vcpu_reset(vcpu);
 
 
                /* Awaken to handle a signal, request we sleep again later. */
                kvm_make_request(KVM_REQ_SLEEP, vcpu);
        }
+
+       /*
+        * Make sure we will observe a potential reset request if we've
+        * observed a change to the power state. Pairs with the smp_wmb() in
+        * kvm_psci_vcpu_on().
+        */
+       smp_rmb();
 }
 
 static int kvm_vcpu_initialized(struct kvm_vcpu *vcpu)
                if (kvm_check_request(KVM_REQ_SLEEP, vcpu))
                        vcpu_req_sleep(vcpu);
 
+               if (kvm_check_request(KVM_REQ_VCPU_RESET, vcpu))
+                       kvm_reset_vcpu(vcpu);
+
                /*
                 * Clear IRQ_PENDING requests that were made to guarantee
                 * that a VCPU sees new virtual interrupts.
 
 
 static unsigned long kvm_psci_vcpu_on(struct kvm_vcpu *source_vcpu)
 {
+       struct vcpu_reset_state *reset_state;
        struct kvm *kvm = source_vcpu->kvm;
        struct kvm_vcpu *vcpu = NULL;
-       struct swait_queue_head *wq;
        unsigned long cpu_id;
-       unsigned long context_id;
-       phys_addr_t target_pc;
 
        cpu_id = smccc_get_arg1(source_vcpu) & MPIDR_HWID_BITMASK;
        if (vcpu_mode_is_32bit(source_vcpu))
                        return PSCI_RET_INVALID_PARAMS;
        }
 
-       target_pc = smccc_get_arg2(source_vcpu);
-       context_id = smccc_get_arg3(source_vcpu);
+       reset_state = &vcpu->arch.reset_state;
 
-       kvm_reset_vcpu(vcpu);
-
-       /* Gracefully handle Thumb2 entry point */
-       if (vcpu_mode_is_32bit(vcpu) && (target_pc & 1)) {
-               target_pc &= ~((phys_addr_t) 1);
-               vcpu_set_thumb(vcpu);
-       }
+       reset_state->pc = smccc_get_arg2(source_vcpu);
 
        /* Propagate caller endianness */
-       if (kvm_vcpu_is_be(source_vcpu))
-               kvm_vcpu_set_be(vcpu);
+       reset_state->be = kvm_vcpu_is_be(source_vcpu);
 
-       *vcpu_pc(vcpu) = target_pc;
        /*
         * NOTE: We always update r0 (or x0) because for PSCI v0.1
         * the general puspose registers are undefined upon CPU_ON.
         */
-       smccc_set_retval(vcpu, context_id, 0, 0, 0);
-       vcpu->arch.power_off = false;
-       smp_mb();               /* Make sure the above is visible */
+       reset_state->r0 = smccc_get_arg3(source_vcpu);
+
+       WRITE_ONCE(reset_state->reset, true);
+       kvm_make_request(KVM_REQ_VCPU_RESET, vcpu);
 
-       wq = kvm_arch_vcpu_wq(vcpu);
-       swake_up_one(wq);
+       /*
+        * Make sure the reset request is observed if the change to
+        * power_state is observed.
+        */
+       smp_wmb();
+
+       vcpu->arch.power_off = false;
+       kvm_vcpu_wake_up(vcpu);
 
        return PSCI_RET_SUCCESS;
 }