#include <linux/rcupdate.h>
 #include <linux/ratelimit.h>
 #include <linux/err.h>
+#include <linux/irqflags.h>
 #include <asm/signal.h>
 
 #include <linux/kvm.h>
 }
 #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;
+}
+
+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
 
 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);
 }
 
 /*
 
 #ifdef CONFIG_VIRT_CPU_ACCOUNTING
 extern void vtime_task_switch(struct task_struct *prev);
 extern void vtime_account_system(struct task_struct *tsk);
-extern void vtime_account_system_irqsafe(struct task_struct *tsk);
 extern void vtime_account_idle(struct task_struct *tsk);
 extern void vtime_account_user(struct task_struct *tsk);
 extern void vtime_account(struct task_struct *tsk);
 #else /* !CONFIG_VIRT_CPU_ACCOUNTING */
 static inline void vtime_task_switch(struct task_struct *prev) { }
 static inline void vtime_account_system(struct task_struct *tsk) { }
-static inline void vtime_account_system_irqsafe(struct task_struct *tsk) { }
 static inline void vtime_account_user(struct task_struct *tsk) { }
 static inline void vtime_account(struct task_struct *tsk) { }
 static inline bool vtime_accounting_enabled(void) { return false; }
 
        *st = cputime.stime;
 }
 
-void vtime_account_system_irqsafe(struct task_struct *tsk)
-{
-       unsigned long flags;
-
-       local_irq_save(flags);
-       vtime_account_system(tsk);
-       local_irq_restore(flags);
-}
-EXPORT_SYMBOL_GPL(vtime_account_system_irqsafe);
-
 #ifndef __ARCH_HAS_VTIME_TASK_SWITCH
 void vtime_task_switch(struct task_struct *prev)
 {