5 files changed, 42 insertions, 40 deletions

diff --git a/arch/x86/kernel/kvm.c b/arch/x86/kernel/kvm.c
index 33019ddb56b4..6551dedee20c 100644
--- a/arch/x86/kernel/kvm.c
+++ b/arch/x86/kernel/kvm.c
@@ -195,7 +195,7 @@ static void kvm_leave_lazy_mmu(void)
         struct kvm_para_state *state = kvm_para_state();
 
         mmu_queue_flush(state);
-        paravirt_leave_lazy(paravirt_get_lazy_mode());
+        paravirt_leave_lazy_mmu();
         state->mode = paravirt_get_lazy_mode();
 }
 
diff --git a/arch/x86/kernel/paravirt.c b/arch/x86/kernel/paravirt.c
index 8e45f4464880..aa3442340705 100644
--- a/arch/x86/kernel/paravirt.c
+++ b/arch/x86/kernel/paravirt.c
@@ -246,18 +246,16 @@ static DEFINE_PER_CPU(enum paravirt_lazy_mode, paravirt_lazy_mode) = PARAVIRT_LA
 
 static inline void enter_lazy(enum paravirt_lazy_mode mode)
 {
-        BUG_ON(__get_cpu_var(paravirt_lazy_mode) != PARAVIRT_LAZY_NONE);
-        BUG_ON(preemptible());
+        BUG_ON(percpu_read(paravirt_lazy_mode) != PARAVIRT_LAZY_NONE);
 
-        __get_cpu_var(paravirt_lazy_mode) = mode;
+        percpu_write(paravirt_lazy_mode, mode);
 }
 
-void paravirt_leave_lazy(enum paravirt_lazy_mode mode)
+static void leave_lazy(enum paravirt_lazy_mode mode)
 {
-        BUG_ON(__get_cpu_var(paravirt_lazy_mode) != mode);
-        BUG_ON(preemptible());
+        BUG_ON(percpu_read(paravirt_lazy_mode) != mode);
 
-        __get_cpu_var(paravirt_lazy_mode) = PARAVIRT_LAZY_NONE;
+        percpu_write(paravirt_lazy_mode, PARAVIRT_LAZY_NONE);
 }
 
 void paravirt_enter_lazy_mmu(void)
@@ -267,22 +265,36 @@ void paravirt_enter_lazy_mmu(void)
 
 void paravirt_leave_lazy_mmu(void)
 {
-        paravirt_leave_lazy(PARAVIRT_LAZY_MMU);
+        leave_lazy(PARAVIRT_LAZY_MMU);
 }
 
-void paravirt_enter_lazy_cpu(void)
+void paravirt_start_context_switch(struct task_struct *prev)
 {
+        BUG_ON(preemptible());
+
+        if (percpu_read(paravirt_lazy_mode) == PARAVIRT_LAZY_MMU) {
+                arch_leave_lazy_mmu_mode();
+                set_ti_thread_flag(task_thread_info(prev), TIF_LAZY_MMU_UPDATES);
+        }
         enter_lazy(PARAVIRT_LAZY_CPU);
 }
 
-void paravirt_leave_lazy_cpu(void)
+void paravirt_end_context_switch(struct task_struct *next)
 {
-        paravirt_leave_lazy(PARAVIRT_LAZY_CPU);
+        BUG_ON(preemptible());
+
+        leave_lazy(PARAVIRT_LAZY_CPU);
+
+        if (test_and_clear_ti_thread_flag(task_thread_info(next), TIF_LAZY_MMU_UPDATES))
+                arch_enter_lazy_mmu_mode();
 }
 
 enum paravirt_lazy_mode paravirt_get_lazy_mode(void)
 {
-        return __get_cpu_var(paravirt_lazy_mode);
+        if (in_interrupt())
+                return PARAVIRT_LAZY_NONE;
+
+        return percpu_read(paravirt_lazy_mode);
 }
 
 void arch_flush_lazy_mmu_mode(void)
@@ -290,7 +302,6 @@ void arch_flush_lazy_mmu_mode(void)
         preempt_disable();
 
         if (paravirt_get_lazy_mode() == PARAVIRT_LAZY_MMU) {
-                WARN_ON(preempt_count() == 1);
                 arch_leave_lazy_mmu_mode();
                 arch_enter_lazy_mmu_mode();
         }
@@ -298,19 +309,6 @@ void arch_flush_lazy_mmu_mode(void)
         preempt_enable();
 }
 
-void arch_flush_lazy_cpu_mode(void)
-{
-        preempt_disable();
-
-        if (paravirt_get_lazy_mode() == PARAVIRT_LAZY_CPU) {
-                WARN_ON(preempt_count() == 1);
-                arch_leave_lazy_cpu_mode();
-                arch_enter_lazy_cpu_mode();
-        }
-
-        preempt_enable();
-}
-
 struct pv_info pv_info = {
         .name = "bare hardware",
         .paravirt_enabled = 0,
@@ -402,10 +400,8 @@ struct pv_cpu_ops pv_cpu_ops = {
         .set_iopl_mask = native_set_iopl_mask,
         .io_delay = native_io_delay,
 
-        .lazy_mode = {
-                .enter = paravirt_nop,
-                .leave = paravirt_nop,
-        },
+        .start_context_switch = paravirt_nop,
+        .end_context_switch = paravirt_nop,
 };
 
 struct pv_apic_ops pv_apic_ops = {
diff --git a/arch/x86/kernel/process_32.c b/arch/x86/kernel/process_32.c
index 76f8f84043a2..5de30f0960fb 100644
--- a/arch/x86/kernel/process_32.c
+++ b/arch/x86/kernel/process_32.c
@@ -407,7 +407,7 @@ __switch_to(struct task_struct *prev_p, struct task_struct *next_p)
          * done before math_state_restore, so the TS bit is up
          * to date.
          */
-        arch_leave_lazy_cpu_mode();
+        arch_end_context_switch(next_p);
 
         /* If the task has used fpu the last 5 timeslices, just do a full
          * restore of the math state immediately to avoid the trap; the
diff --git a/arch/x86/kernel/process_64.c b/arch/x86/kernel/process_64.c
index b751a41392b1..66ad06791d6f 100644
--- a/arch/x86/kernel/process_64.c
+++ b/arch/x86/kernel/process_64.c
@@ -428,7 +428,7 @@ __switch_to(struct task_struct *prev_p, struct task_struct *next_p)
          * done before math_state_restore, so the TS bit is up
          * to date.
          */
-        arch_leave_lazy_cpu_mode();
+        arch_end_context_switch(next_p);
 
         /*
          * Switch FS and GS.
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 */