#include <linux/highmem.h>
 #include <linux/memblock.h>
 #include <linux/slab.h>
+#include <linux/set_memory.h>
 
 #include <asm/pgalloc.h>
 #include <asm/fixmap.h>
        }
 }
 
-static int __change_page_attr_noflush(struct page *page, pgprot_t prot)
-{
-       pte_t *kpte;
-       unsigned long address;
-
-       BUG_ON(PageHighMem(page));
-       address = (unsigned long)page_address(page);
-
-       if (v_block_mapped(address))
-               return 0;
-       kpte = virt_to_kpte(address);
-       if (!kpte)
-               return -EINVAL;
-       __set_pte_at(&init_mm, address, kpte, mk_pte(page, prot), 0);
-
-       return 0;
-}
-
-/*
- * Change the page attributes of an page in the linear mapping.
- *
- * THIS DOES NOTHING WITH BAT MAPPINGS, DEBUG USE ONLY
- */
-static int change_page_attr(struct page *page, int numpages, pgprot_t prot)
-{
-       int i, err = 0;
-       unsigned long flags;
-       struct page *start = page;
-
-       local_irq_save(flags);
-       for (i = 0; i < numpages; i++, page++) {
-               err = __change_page_attr_noflush(page, prot);
-               if (err)
-                       break;
-       }
-       wmb();
-       local_irq_restore(flags);
-       flush_tlb_kernel_range((unsigned long)page_address(start),
-                              (unsigned long)page_address(page));
-       return err;
-}
-
 void mark_initmem_nx(void)
 {
-       struct page *page = virt_to_page(_sinittext);
        unsigned long numpages = PFN_UP((unsigned long)_einittext) -
                                 PFN_DOWN((unsigned long)_sinittext);
 
        if (v_block_mapped((unsigned long)_sinittext))
                mmu_mark_initmem_nx();
        else
-               change_page_attr(page, numpages, PAGE_KERNEL);
+               set_memory_attr((unsigned long)_sinittext, numpages, PAGE_KERNEL);
 }
 
 #ifdef CONFIG_STRICT_KERNEL_RWX
 void mark_rodata_ro(void)
 {
-       struct page *page;
        unsigned long numpages;
 
        if (v_block_mapped((unsigned long)_stext + 1)) {
                return;
        }
 
-       page = virt_to_page(_stext);
        numpages = PFN_UP((unsigned long)_etext) -
                   PFN_DOWN((unsigned long)_stext);
 
-       change_page_attr(page, numpages, PAGE_KERNEL_ROX);
+       set_memory_attr((unsigned long)_stext, numpages, PAGE_KERNEL_ROX);
        /*
         * mark .rodata as read only. Use __init_begin rather than __end_rodata
         * to cover NOTES and EXCEPTION_TABLE.
         */
-       page = virt_to_page(__start_rodata);
        numpages = PFN_UP((unsigned long)__init_begin) -
                   PFN_DOWN((unsigned long)__start_rodata);
 
-       change_page_attr(page, numpages, PAGE_KERNEL_RO);
+       set_memory_attr((unsigned long)__start_rodata, numpages, PAGE_KERNEL_RO);
 
        // mark_initmem_nx() should have already run by now
        ptdump_check_wx();
 #ifdef CONFIG_DEBUG_PAGEALLOC
 void __kernel_map_pages(struct page *page, int numpages, int enable)
 {
+       unsigned long addr = (unsigned long)page_address(page);
+
        if (PageHighMem(page))
                return;
 
-       change_page_attr(page, numpages, enable ? PAGE_KERNEL : __pgprot(0));
+       if (enable)
+               set_memory_attr(addr, numpages, PAGE_KERNEL);
+       else
+               set_memory_attr(addr, numpages, __pgprot(0));
 }
 #endif /* CONFIG_DEBUG_PAGEALLOC */