1 files changed, 13 insertions, 7 deletions

diff --git a/arch/x86/kernel/vmi_32.c b/arch/x86/kernel/vmi_32.c
index 95deb9f2211e..b263423fbe2a 100644
--- a/arch/x86/kernel/vmi_32.c
+++ b/arch/x86/kernel/vmi_32.c
@@ -462,22 +462,28 @@ vmi_startup_ipi_hook(int phys_apicid, unsigned long start_eip,
 }
 #endif
 
-static void vmi_enter_lazy_cpu(void)
+static void vmi_start_context_switch(struct task_struct *prev)
 {
-        paravirt_enter_lazy_cpu();
+        paravirt_start_context_switch(prev);
         vmi_ops.set_lazy_mode(2);
 }
 
+static void vmi_end_context_switch(struct task_struct *next)
+{
+        vmi_ops.set_lazy_mode(0);
+        paravirt_end_context_switch(next);
+}
+
 static void vmi_enter_lazy_mmu(void)
 {
         paravirt_enter_lazy_mmu();
         vmi_ops.set_lazy_mode(1);
 }
 
-static void vmi_leave_lazy(void)
+static void vmi_leave_lazy_mmu(void)
 {
-        paravirt_leave_lazy(paravirt_get_lazy_mode());
         vmi_ops.set_lazy_mode(0);
+        paravirt_leave_lazy_mmu();
 }
 
 static inline int __init check_vmi_rom(struct vrom_header *rom)
@@ -711,14 +717,14 @@ static inline int __init activate_vmi(void)
         para_fill(pv_cpu_ops.set_iopl_mask, SetIOPLMask);
         para_fill(pv_cpu_ops.io_delay, IODelay);
 
-        para_wrap(pv_cpu_ops.lazy_mode.enter, vmi_enter_lazy_cpu,
+        para_wrap(pv_cpu_ops.start_context_switch, vmi_start_context_switch,
                   set_lazy_mode, SetLazyMode);
-        para_wrap(pv_cpu_ops.lazy_mode.leave, vmi_leave_lazy,
+        para_wrap(pv_cpu_ops.end_context_switch, vmi_end_context_switch,
                   set_lazy_mode, SetLazyMode);
 
         para_wrap(pv_mmu_ops.lazy_mode.enter, vmi_enter_lazy_mmu,
                   set_lazy_mode, SetLazyMode);
-        para_wrap(pv_mmu_ops.lazy_mode.leave, vmi_leave_lazy,
+        para_wrap(pv_mmu_ops.lazy_mode.leave, vmi_leave_lazy_mmu,
                   set_lazy_mode, SetLazyMode);
 
         /* user and kernel flush are just handled with different flags to FlushTLB */