1 files changed, 46 insertions, 9 deletions

diff --git a/include/linux/kvm_host.h b/include/linux/kvm_host.h
index 2c497ab0d03d..b7996a768eb2 100644
--- a/include/linux/kvm_host.h
+++ b/include/linux/kvm_host.h
@@ -22,6 +22,7 @@
 #include <linux/rcupdate.h>
 #include <linux/ratelimit.h>
 #include <linux/err.h>
+#include <linux/irqflags.h>
 #include <asm/signal.h>
 
 #include <linux/kvm.h>
@@ -740,15 +741,52 @@ static inline int kvm_deassign_device(struct kvm *kvm,
 }
 #endif /* CONFIG_IOMMU_API */
 
-static inline void kvm_guest_enter(void)
+static inline void __guest_enter(void)
 {
-        BUG_ON(preemptible());
         /*
          * This is running in ioctl context so we can avoid
          * the call to vtime_account() with its unnecessary idle check.
          */
-        vtime_account_system_irqsafe(current);
+        vtime_account_system(current);
         current->flags |= PF_VCPU;
+}
+
+static inline void __guest_exit(void)
+{
+        /*
+         * This is running in ioctl context so we can avoid
+         * the call to vtime_account() with its unnecessary idle check.
+         */
+        vtime_account_system(current);
+        current->flags &= ~PF_VCPU;
+}
+
+#ifdef CONFIG_CONTEXT_TRACKING
+extern void guest_enter(void);
+extern void guest_exit(void);
+
+#else /* !CONFIG_CONTEXT_TRACKING */
+static inline void guest_enter(void)
+{
+        __guest_enter();
+}
+
+static inline void guest_exit(void)
+{
+        __guest_exit();
+}
+#endif /* !CONFIG_CONTEXT_TRACKING */
+
+static inline void kvm_guest_enter(void)
+{
+        unsigned long flags;
+
+        BUG_ON(preemptible());
+
+        local_irq_save(flags);
+        guest_enter();
+        local_irq_restore(flags);
+
         /* KVM does not hold any references to rcu protected data when it
          * switches CPU into a guest mode. In fact switching to a guest mode
          * is very similar to exiting to userspase from rcu point of view. In
@@ -761,12 +799,11 @@ static inline void kvm_guest_enter(void)
 
 static inline void kvm_guest_exit(void)
 {
-        /*
-         * This is running in ioctl context so we can avoid
-         * the call to vtime_account() with its unnecessary idle check.
-         */
-        vtime_account_system_irqsafe(current);
-        current->flags &= ~PF_VCPU;
+        unsigned long flags;
+
+        local_irq_save(flags);
+        guest_exit();
+        local_irq_restore(flags);
 }
 
 /*