static inline void enter_lazy(enum paravirt_lazy_mode mode)
 {
        BUG_ON(__get_cpu_var(paravirt_lazy_mode) != PARAVIRT_LAZY_NONE);
-       BUG_ON(preemptible());
 
        __get_cpu_var(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());
 
        __get_cpu_var(paravirt_lazy_mode) = PARAVIRT_LAZY_NONE;
 }
 
 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);
 
 void paravirt_end_context_switch(struct task_struct *next)
 {
+       BUG_ON(preemptible());
+
        leave_lazy(PARAVIRT_LAZY_CPU);
 
        if (test_and_clear_ti_thread_flag(task_thread_info(next), TIF_LAZY_MMU_UPDATES))
        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();
        }
 
 void xen_set_pte_at(struct mm_struct *mm, unsigned long addr,
                    pte_t *ptep, pte_t pteval)
 {
-       /* updates to init_mm may be done without lock */
-       if (mm == &init_mm)
-               preempt_disable();
-
        ADD_STATS(set_pte_at, 1);
 //     ADD_STATS(set_pte_at_pinned, xen_page_pinned(ptep));
        ADD_STATS(set_pte_at_current, mm == current->mm);
        }
        xen_set_pte(ptep, pteval);
 
-out:
-       if (mm == &init_mm)
-               preempt_enable();
+out:   return;
 }
 
 pte_t xen_ptep_modify_prot_start(struct mm_struct *mm,