+};
+
  #endif /* _LINUX_MM_TYPES_TASK_H */

diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index 194b2c3e7576..269225a733de 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h

@@ -228,28 +228,89 @@ static inline int pmd_dirty(pmd_t pmd)
   * of the lazy mode. So the implementation must assume preemption may be enabled
   * and cpu migration is possible; it must take steps to be robust against this.
   * (In practice, for user PTE updates, the appropriate page table lock(s) are
- * held, but for kernel PTE updates, no lock is held). Nesting is not permitted
- * and the mode cannot be used in interrupt context.
+ * held, but for kernel PTE updates, no lock is held). The mode cannot be used
+ * in interrupt context.
+ *
+ * The lazy MMU mode is enabled for a given block of code using:
+ *
+ *   lazy_mmu_mode_enable();
+ *   <code>
+ *   lazy_mmu_mode_disable();
+ *
+ * Nesting is permitted: <code> may itself use an enable()/disable() pair.
+ * A nested call to enable() has no functional effect; however disable() causes
+ * any batched architectural state to be flushed regardless of nesting. After a
+ * call to disable(), the caller can therefore rely on all previous page table
+ * modifications to have taken effect, but the lazy MMU mode may still be
+ * enabled.
+ *
+ * In certain cases, it may be desirable to temporarily pause the lazy MMU mode.
+ * This can be done using:
+ *
+ *   lazy_mmu_mode_pause();
+ *   <code>
+ *   lazy_mmu_mode_resume();
+ *
+ * This sequence must only be used if the lazy MMU mode is already enabled.
+ * pause() ensures that the mode is exited regardless of the nesting level;
+ * resume() re-enters the mode at the same nesting level. <code> must not modify
+ * the lazy MMU state (i.e. it must not call any of the lazy_mmu_mode_*
+ * helpers).
+ *
+ * in_lazy_mmu_mode() can be used to check whether the lazy MMU mode is
+ * currently enabled.
   */
  #ifdef CONFIG_ARCH_LAZY_MMU
  static inline void lazy_mmu_mode_enable(void)
  {
-	arch_enter_lazy_mmu_mode();
+	struct lazy_mmu_state *state = &current->lazy_mmu_state;
+
+	VM_BUG_ON(state->count == U8_MAX);

+		state->enabled = false;
+		arch_leave_lazy_mmu_mode();
+	} else {
+		/* Exiting a nested section */
+		arch_flush_lazy_mmu_mode();
+	}
  }
  
  static inline void lazy_mmu_mode_pause(void)
  {
+	struct lazy_mmu_state *state = &current->lazy_mmu_state;
+
+	VM_WARN_ON(state->count == 0 || !state->enabled);
+
+	state->enabled = false;
  	arch_leave_lazy_mmu_mode();
  }
  
  static inline void lazy_mmu_mode_resume(void)
  {
+	struct lazy_mmu_state *state = &current->lazy_mmu_state;
+
+	VM_WARN_ON(state->count == 0 || state->enabled);
+
  	arch_enter_lazy_mmu_mode();
+	state->enabled = true;
  }
  #else
  static inline void lazy_mmu_mode_enable(void) {}

diff --git a/include/linux/sched.h b/include/linux/sched.h
index cbb7340c5866..2862d8bf2160 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h

@@ -1441,6 +1441,10 @@ struct task_struct {
  
  	struct page_frag		task_frag;
  
+#ifdef CONFIG_ARCH_LAZY_MMU
+	struct lazy_mmu_state		lazy_mmu_state;
+#endif
+
  #ifdef CONFIG_TASK_DELAY_ACCT
  	struct task_delay_info		*delays;
  #endif

@@ -1724,6 +1728,18 @@ static inline char task_state_to_char(struct task_struct *tsk)
  	return task_index_to_char(task_state_index(tsk));
  }
  
+#ifdef CONFIG_ARCH_LAZY_MMU
+static inline bool in_lazy_mmu_mode(void)