__r;                                                    \
 })
 
+static inline int pte_write(pte_t pte)
+{
+       return !!(pte_raw(pte) & cpu_to_be64(_PAGE_WRITE));
+}
+
+#ifdef CONFIG_NUMA_BALANCING
+#define pte_savedwrite pte_savedwrite
+static inline bool pte_savedwrite(pte_t pte)
+{
+       /*
+        * Saved write ptes are prot none ptes that doesn't have
+        * privileged bit sit. We mark prot none as one which has
+        * present and pviliged bit set and RWX cleared. To mark
+        * protnone which used to have _PAGE_WRITE set we clear
+        * the privileged bit.
+        */
+       return !(pte_raw(pte) & cpu_to_be64(_PAGE_RWX | _PAGE_PRIVILEGED));
+}
+#else
+#define pte_savedwrite pte_savedwrite
+static inline bool pte_savedwrite(pte_t pte)
+{
+       return false;
+}
+#endif
+
 #define __HAVE_ARCH_PTEP_SET_WRPROTECT
 static inline void ptep_set_wrprotect(struct mm_struct *mm, unsigned long addr,
                                      pte_t *ptep)
 {
-       if ((pte_raw(*ptep) & cpu_to_be64(_PAGE_WRITE)) == 0)
-               return;
-
-       pte_update(mm, addr, ptep, _PAGE_WRITE, 0, 0);
+       if (pte_write(*ptep))
+               pte_update(mm, addr, ptep, _PAGE_WRITE, 0, 0);
+       else if (unlikely(pte_savedwrite(*ptep)))
+               pte_update(mm, addr, ptep, 0, _PAGE_PRIVILEGED, 0);
 }
 
 static inline void huge_ptep_set_wrprotect(struct mm_struct *mm,
                                           unsigned long addr, pte_t *ptep)
 {
-       if ((pte_raw(*ptep) & cpu_to_be64(_PAGE_WRITE)) == 0)
-               return;
-
-       pte_update(mm, addr, ptep, _PAGE_WRITE, 0, 1);
+       /*
+        * We should not find protnone for hugetlb, but this complete the
+        * interface.
+        */
+       if (pte_write(*ptep))
+               pte_update(mm, addr, ptep, _PAGE_WRITE, 0, 1);
+       else if (unlikely(pte_savedwrite(*ptep)))
+               pte_update(mm, addr, ptep, 0, _PAGE_PRIVILEGED, 1);
 }
 
 #define __HAVE_ARCH_PTEP_GET_AND_CLEAR
        pte_update(mm, addr, ptep, ~0UL, 0, 0);
 }
 
-static inline int pte_write(pte_t pte)
-{
-       return !!(pte_raw(pte) & cpu_to_be64(_PAGE_WRITE));
-}
-
 static inline int pte_dirty(pte_t pte)
 {
        return !!(pte_raw(pte) & cpu_to_be64(_PAGE_DIRTY));
        return __pte(pte_val(pte) | _PAGE_PRIVILEGED);
 }
 
-#define pte_savedwrite pte_savedwrite
-static inline bool pte_savedwrite(pte_t pte)
-{
-       /*
-        * Saved write ptes are prot none ptes that doesn't have
-        * privileged bit sit. We mark prot none as one which has
-        * present and pviliged bit set and RWX cleared. To mark
-        * protnone which used to have _PAGE_WRITE set we clear
-        * the privileged bit.
-        */
-       VM_BUG_ON(!pte_protnone(pte));
-       return !(pte_raw(pte) & cpu_to_be64(_PAGE_RWX | _PAGE_PRIVILEGED));
-}
 #endif /* CONFIG_NUMA_BALANCING */
 
 static inline int pte_present(pte_t pte)
 static inline void pmdp_set_wrprotect(struct mm_struct *mm, unsigned long addr,
                                      pmd_t *pmdp)
 {
-
-       if ((pmd_raw(*pmdp) & cpu_to_be64(_PAGE_WRITE)) == 0)
-               return;
-
-       pmd_hugepage_update(mm, addr, pmdp, _PAGE_WRITE, 0);
+       if (pmd_write((*pmdp)))
+               pmd_hugepage_update(mm, addr, pmdp, _PAGE_WRITE, 0);
+       else if (unlikely(pmd_savedwrite(*pmdp)))
+               pmd_hugepage_update(mm, addr, pmdp, 0, _PAGE_PRIVILEGED);
 }
 
 static inline int pmd_trans_huge(pmd_t pmd)