return pte;
 }
 
+#define ptep_set_numa ptep_set_numa
+static inline void ptep_set_numa(struct mm_struct *mm, unsigned long addr,
+                                pte_t *ptep)
+{
+       if ((pte_val(*ptep) & _PAGE_PRESENT) == 0)
+               VM_BUG_ON(1);
+
+       pte_update(mm, addr, ptep, _PAGE_PRESENT, _PAGE_NUMA, 0);
+       return;
+}
+
 #define pmd_numa pmd_numa
 static inline int pmd_numa(pmd_t pmd)
 {
        return pte_numa(pmd_pte(pmd));
 }
 
+#define pmdp_set_numa pmdp_set_numa
+static inline void pmdp_set_numa(struct mm_struct *mm, unsigned long addr,
+                                pmd_t *pmdp)
+{
+       if ((pmd_val(*pmdp) & _PAGE_PRESENT) == 0)
+               VM_BUG_ON(1);
+
+       pmd_hugepage_update(mm, addr, pmdp, _PAGE_PRESENT, _PAGE_NUMA);
+       return;
+}
+
 #define pmd_mknonnuma pmd_mknonnuma
 static inline pmd_t pmd_mknonnuma(pmd_t pmd)
 {
 
 }
 #endif
 
+#ifndef ptep_set_numa
+static inline void ptep_set_numa(struct mm_struct *mm, unsigned long addr,
+                                pte_t *ptep)
+{
+       pte_t ptent = *ptep;
+
+       ptent = pte_mknuma(ptent);
+       set_pte_at(mm, addr, ptep, ptent);
+       return;
+}
+#endif
+
 #ifndef pmd_mknuma
 static inline pmd_t pmd_mknuma(pmd_t pmd)
 {
        return pmd_clear_flags(pmd, _PAGE_PRESENT);
 }
 #endif
+
+#ifndef pmdp_set_numa
+static inline void pmdp_set_numa(struct mm_struct *mm, unsigned long addr,
+                                pmd_t *pmdp)
+{
+       pmd_t pmd = *pmdp;
+
+       pmd = pmd_mknuma(pmd);
+       set_pmd_at(mm, addr, pmdp, pmd);
+       return;
+}
+#endif
 #else
 extern int pte_numa(pte_t pte);
 extern int pmd_numa(pmd_t pmd);
 extern pmd_t pmd_mknonnuma(pmd_t pmd);
 extern pte_t pte_mknuma(pte_t pte);
 extern pmd_t pmd_mknuma(pmd_t pmd);
+extern void ptep_set_numa(struct mm_struct *mm, unsigned long addr, pte_t *ptep);
+extern void pmdp_set_numa(struct mm_struct *mm, unsigned long addr, pmd_t *pmdp);
 #endif /* CONFIG_ARCH_USES_NUMA_PROT_NONE */
 #else
 static inline int pmd_numa(pmd_t pmd)
        return pte;
 }
 
+static inline void ptep_set_numa(struct mm_struct *mm, unsigned long addr,
+                                pte_t *ptep)
+{
+       return;
+}
+
+
 static inline pmd_t pmd_mknuma(pmd_t pmd)
 {
        return pmd;
 }
+
+static inline void pmdp_set_numa(struct mm_struct *mm, unsigned long addr,
+                                pmd_t *pmdp)
+{
+       return ;
+}
 #endif /* CONFIG_NUMA_BALANCING */
 
 #endif /* CONFIG_MMU */
 
                                entry = pmd_mknonnuma(entry);
                        entry = pmd_modify(entry, newprot);
                        ret = HPAGE_PMD_NR;
+                       set_pmd_at(mm, addr, pmd, entry);
                        BUG_ON(pmd_write(entry));
                } else {
                        struct page *page = pmd_page(*pmd);
                         */
                        if (!is_huge_zero_page(page) &&
                            !pmd_numa(*pmd)) {
-                               entry = *pmd;
-                               entry = pmd_mknuma(entry);
+                               pmdp_set_numa(mm, addr, pmd);
                                ret = HPAGE_PMD_NR;
                        }
                }
-
-               /* Set PMD if cleared earlier */
-               if (ret == HPAGE_PMD_NR)
-                       set_pmd_at(mm, addr, pmd, entry);
-
                spin_unlock(ptl);
        }
 
 
                        } else {
                                struct page *page;
 
-                               ptent = *pte;
                                page = vm_normal_page(vma, addr, oldpte);
                                if (page && !PageKsm(page)) {
                                        if (!pte_numa(oldpte)) {
-                                               ptent = pte_mknuma(ptent);
-                                               set_pte_at(mm, addr, pte, ptent);
+                                               ptep_set_numa(mm, addr, pte);
                                                updated = true;
                                        }
                                }