diff --git a/arch/arm64/kvm/arm.c b/arch/arm64/kvm/arm.c
index 2adca153ad68..348bdef382c1 100644
--- a/arch/arm64/kvm/arm.c
+++ b/arch/arm64/kvm/arm.c
@@ -576,7 +576,9 @@ nommu:
 	kvm_arch_vcpu_load_debug_state_flags(vcpu);
 
 	if (is_protected_kvm_enabled()) {
-		kvm_call_hyp_nvhe(__pkvm_vcpu_load, vcpu);
+		kvm_call_hyp_nvhe(__pkvm_vcpu_load,
+				  vcpu->kvm->arch.pkvm.shadow_handle,
+				  vcpu->vcpu_idx, vcpu->arch.hcr_el2);
 		kvm_call_hyp(__vgic_v3_restore_vmcr_aprs,
 			     &vcpu->arch.vgic_cpu.vgic_v3);
 	}
diff --git a/arch/arm64/kvm/hyp/nvhe/hyp-main.c b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
index 1086ff2e4245..5946fb5462c0 100644
--- a/arch/arm64/kvm/hyp/nvhe/hyp-main.c
+++ b/arch/arm64/kvm/hyp/nvhe/hyp-main.c
@@ -616,9 +616,10 @@ static void fpsimd_host_restore(void)
 
 static void handle___pkvm_vcpu_load(struct kvm_cpu_context *host_ctxt)
 {
-	DECLARE_REG(struct kvm_vcpu *, vcpu, host_ctxt, 1);
+	DECLARE_REG(int, shadow_handle, host_ctxt, 1);
+	DECLARE_REG(int, vcpu_idx, host_ctxt, 2);
+	DECLARE_REG(u64, hcr_el2, host_ctxt, 3);
 	struct pkvm_loaded_state *state;
-	int handle;
 
 	/* Why did you bother? */
 	if (!is_protected_kvm_enabled())
@@ -630,10 +631,7 @@ static void handle___pkvm_vcpu_load(struct kvm_cpu_context *host_ctxt)
 	if (state->vcpu)
 		return;
 
-	vcpu = kern_hyp_va(vcpu);
-
-	handle = READ_ONCE(vcpu->arch.pkvm.shadow_handle);
-	state->vcpu = get_shadow_vcpu(handle, vcpu->vcpu_idx);
+	state->vcpu = get_shadow_vcpu(shadow_handle, vcpu_idx);
 
 	if (!state->vcpu)
 		return;
@@ -647,8 +645,7 @@ static void handle___pkvm_vcpu_load(struct kvm_cpu_context *host_ctxt)
 		/* Propagate WFx trapping flags, trap ptrauth */
 		state->vcpu->arch.hcr_el2 &= ~(HCR_TWE | HCR_TWI |
 					       HCR_API | HCR_APK);
-		state->vcpu->arch.hcr_el2 |= vcpu->arch.hcr_el2 & (HCR_TWE |
-								   HCR_TWI);
+		state->vcpu->arch.hcr_el2 |= hcr_el2 & (HCR_TWE | HCR_TWI);
 	}
 }