#define KVM_ARM64_DEBUG_DIRTY_SHIFT    0
 #define KVM_ARM64_DEBUG_DIRTY          (1 << KVM_ARM64_DEBUG_DIRTY_SHIFT)
 
+/* Translate a kernel address of @sym into its equivalent linear mapping */
 #define kvm_ksym_ref(sym)                                              \
        ({                                                              \
                void *val = &sym;                                       \
 
 extern void __qcom_hyp_sanitize_btac_predictors(void);
 
+#else /* __ASSEMBLY__ */
+
+.macro get_host_ctxt reg, tmp
+       adr_l   \reg, kvm_host_cpu_state
+       mrs     \tmp, tpidr_el2
+       add     \reg, \reg, \tmp
+.endm
+
+.macro get_vcpu_ptr vcpu, ctxt
+       get_host_ctxt \ctxt, \vcpu
+       ldr     \vcpu, [\ctxt, #HOST_CONTEXT_VCPU]
+       kern_hyp_va     \vcpu
+.endm
+
 #endif
 
 #endif /* __ARM_KVM_ASM_H__ */
 
 
 struct kvm_vcpu *kvm_mpidr_to_vcpu(struct kvm *kvm, unsigned long mpidr);
 
+void __kvm_set_tpidr_el2(u64 tpidr_el2);
+DECLARE_PER_CPU(kvm_cpu_context_t, kvm_host_cpu_state);
+
 static inline void __cpu_init_hyp_mode(phys_addr_t pgd_ptr,
                                       unsigned long hyp_stack_ptr,
                                       unsigned long vector_ptr)
 {
+       u64 tpidr_el2;
+
        /*
         * Call initialization code, and switch to the full blown HYP code.
         * If the cpucaps haven't been finalized yet, something has gone very
         */
        BUG_ON(!static_branch_likely(&arm64_const_caps_ready));
        __kvm_call_hyp((void *)pgd_ptr, hyp_stack_ptr, vector_ptr);
+
+       /*
+        * Calculate the raw per-cpu offset without a translation from the
+        * kernel's mapping to the linear mapping, and store it in tpidr_el2
+        * so that we can use adr_l to access per-cpu variables in EL2.
+        */
+       tpidr_el2 = (u64)this_cpu_ptr(&kvm_host_cpu_state)
+               - (u64)kvm_ksym_ref(kvm_host_cpu_state);
+
+       kvm_call_hyp(__kvm_set_tpidr_el2, tpidr_el2);
 }
 
 static inline void kvm_arch_hardware_unsetup(void) {}
 
   DEFINE(CPU_FP_REGS,          offsetof(struct kvm_regs, fp_regs));
   DEFINE(VCPU_FPEXC32_EL2,     offsetof(struct kvm_vcpu, arch.ctxt.sys_regs[FPEXC32_EL2]));
   DEFINE(VCPU_HOST_CONTEXT,    offsetof(struct kvm_vcpu, arch.host_cpu_context));
+  DEFINE(HOST_CONTEXT_VCPU,    offsetof(struct kvm_cpu_context, __hyp_running_vcpu));
 #endif
 #ifdef CONFIG_CPU_PM
   DEFINE(CPU_SUSPEND_SZ,       sizeof(struct cpu_suspend_ctx));
 
        // Store the host regs
        save_callee_saved_regs x1
 
-       // Store host_ctxt and vcpu for use at exit time
-       stp     x1, x0, [sp, #-16]!
-
        add     x18, x0, #VCPU_CONTEXT
 
        // Restore guest regs x0-x17
        // Store the guest regs x19-x29, lr
        save_callee_saved_regs x1
 
-       // Restore the host_ctxt from the stack
-       ldr     x2, [sp], #16
+       get_host_ctxt   x2, x3
 
        // Now restore the host regs
        restore_callee_saved_regs x2
 
 el1_sync:                              // Guest trapped into EL2
        stp     x0, x1, [sp, #-16]!
 
-alternative_if_not ARM64_HAS_VIRT_HOST_EXTN
-       mrs     x1, esr_el2
-alternative_else
-       mrs     x1, esr_el1
-alternative_endif
-       lsr     x0, x1, #ESR_ELx_EC_SHIFT
-
+       mrs     x0, esr_el2
+       lsr     x0, x0, #ESR_ELx_EC_SHIFT
        cmp     x0, #ESR_ELx_EC_HVC64
        ccmp    x0, #ESR_ELx_EC_HVC32, #4, ne
        b.ne    el1_trap
        eret
 
 el1_trap:
+       get_vcpu_ptr    x1, x0
+
+       mrs             x0, esr_el2
+       lsr             x0, x0, #ESR_ELx_EC_SHIFT
        /*
         * x0: ESR_EC
+        * x1: vcpu pointer
         */
-       ldr     x1, [sp, #16 + 8]       // vcpu stored by __guest_enter
 
        /*
         * We trap the first access to the FP/SIMD to save the host context
 
 el1_irq:
        stp     x0, x1, [sp, #-16]!
-       ldr     x1, [sp, #16 + 8]
+       get_vcpu_ptr    x1, x0
        mov     x0, #ARM_EXCEPTION_IRQ
        b       __guest_exit
 
 el1_error:
        stp     x0, x1, [sp, #-16]!
-       ldr     x1, [sp, #16 + 8]
+       get_vcpu_ptr    x1, x0
        mov     x0, #ARM_EXCEPTION_EL1_SERROR
        b       __guest_exit
 
 ENDPROC(__hyp_do_panic)
 
 ENTRY(__hyp_panic)
-       /*
-        * '=kvm_host_cpu_state' is a host VA from the constant pool, it may
-        * not be accessible by this address from EL2, hyp_panic() converts
-        * it with kern_hyp_va() before use.
-        */
-       ldr     x0, =kvm_host_cpu_state
-       mrs     x1, tpidr_el2
-       add     x0, x0, x1
+       get_host_ctxt x0, x1
        b       hyp_panic
 ENDPROC(__hyp_panic)
 
 
                            __hyp_call_panic_nvhe, __hyp_call_panic_vhe,
                            ARM64_HAS_VIRT_HOST_EXTN);
 
-void __hyp_text __noreturn hyp_panic(struct kvm_cpu_context *__host_ctxt)
+void __hyp_text __noreturn hyp_panic(struct kvm_cpu_context *host_ctxt)
 {
        struct kvm_vcpu *vcpu = NULL;
 
        u64 par = read_sysreg(par_el1);
 
        if (read_sysreg(vttbr_el2)) {
-               struct kvm_cpu_context *host_ctxt;
-
-               host_ctxt = kern_hyp_va(__host_ctxt);
                vcpu = host_ctxt->__hyp_running_vcpu;
                __timer_disable_traps(vcpu);
                __deactivate_traps(vcpu);
 
        if (vcpu->arch.debug_flags & KVM_ARM64_DEBUG_DIRTY)
                write_sysreg(sysreg[DBGVCR32_EL2], dbgvcr32_el2);
 }
+
+void __hyp_text __kvm_set_tpidr_el2(u64 tpidr_el2)
+{
+       asm("msr tpidr_el2, %0": : "r" (tpidr_el2));
+}