return ret;
 }
 
-void __avic_vcpu_load(struct kvm_vcpu *vcpu, int cpu)
+void avic_vcpu_load(struct kvm_vcpu *vcpu, int cpu)
 {
        u64 entry;
        int h_physical_id = kvm_cpu_get_apicid(cpu);
        avic_update_iommu_vcpu_affinity(vcpu, h_physical_id, true);
 }
 
-void __avic_vcpu_put(struct kvm_vcpu *vcpu)
+void avic_vcpu_put(struct kvm_vcpu *vcpu)
 {
        u64 entry;
        struct vcpu_svm *svm = to_svm(vcpu);
        WRITE_ONCE(*(svm->avic_physical_id_cache), entry);
 }
 
-static void avic_vcpu_load(struct kvm_vcpu *vcpu)
-{
-       int cpu = get_cpu();
-
-       WARN_ON(cpu != vcpu->cpu);
-
-       __avic_vcpu_load(vcpu, cpu);
-
-       put_cpu();
-}
-
-static void avic_vcpu_put(struct kvm_vcpu *vcpu)
-{
-       preempt_disable();
-
-       __avic_vcpu_put(vcpu);
-
-       preempt_enable();
-}
 
 void avic_refresh_apicv_exec_ctrl(struct kvm_vcpu *vcpu)
 {
        vmcb_mark_dirty(vmcb, VMCB_AVIC);
 
        if (activated)
-               avic_vcpu_load(vcpu);
+               avic_vcpu_load(vcpu, vcpu->cpu);
        else
                avic_vcpu_put(vcpu);
 
        if (!kvm_vcpu_apicv_active(vcpu))
                return;
 
-       avic_vcpu_load(vcpu);
+       avic_vcpu_load(vcpu, vcpu->cpu);
 }
 
                indirect_branch_prediction_barrier();
        }
        if (kvm_vcpu_apicv_active(vcpu))
-               __avic_vcpu_load(vcpu, cpu);
+               avic_vcpu_load(vcpu, cpu);
 }
 
 static void svm_vcpu_put(struct kvm_vcpu *vcpu)
 {
        if (kvm_vcpu_apicv_active(vcpu))
-               __avic_vcpu_put(vcpu);
+               avic_vcpu_put(vcpu);
 
        svm_prepare_host_switch(vcpu);
 
 
 int avic_incomplete_ipi_interception(struct kvm_vcpu *vcpu);
 int avic_unaccelerated_access_interception(struct kvm_vcpu *vcpu);
 int avic_init_vcpu(struct vcpu_svm *svm);
-void __avic_vcpu_load(struct kvm_vcpu *vcpu, int cpu);
-void __avic_vcpu_put(struct kvm_vcpu *vcpu);
+void avic_vcpu_load(struct kvm_vcpu *vcpu, int cpu);
+void avic_vcpu_put(struct kvm_vcpu *vcpu);
 void avic_apicv_post_state_restore(struct kvm_vcpu *vcpu);
 void avic_set_virtual_apic_mode(struct kvm_vcpu *vcpu);
 void avic_refresh_apicv_exec_ctrl(struct kvm_vcpu *vcpu);