#include <linux/start_kernel.h>
 #include <linux/sched/mm.h>
 #include <linux/io.h>
+
+#include <asm/cacheflush.h>
 #include <asm/pgalloc.h>
 #include <asm/tlbflush.h>
 
 
 static void __init pte_advanced_tests(struct pgtable_debug_args *args)
 {
+       struct page *page;
        pte_t pte;
 
        /*
         * Architectures optimize set_pte_at by avoiding TLB flush.
         * This requires set_pte_at to be not used to update an
         * existing pte entry. Clear pte before we do set_pte_at
+        *
+        * flush_dcache_page() is called after set_pte_at() to clear
+        * PG_arch_1 for the page on ARM64. The page flag isn't cleared
+        * when it's released and page allocation check will fail when
+        * the page is allocated again. For architectures other than ARM64,
+        * the unexpected overhead of cache flushing is acceptable.
         */
-       if (args->pte_pfn == ULONG_MAX)
+       page = (args->pte_pfn != ULONG_MAX) ? pfn_to_page(args->pte_pfn) : NULL;
+       if (!page)
                return;
 
        pr_debug("Validating PTE advanced\n");
        pte = pfn_pte(args->pte_pfn, args->page_prot);
        set_pte_at(args->mm, args->vaddr, args->ptep, pte);
+       flush_dcache_page(page);
        ptep_set_wrprotect(args->mm, args->vaddr, args->ptep);
        pte = ptep_get(args->ptep);
        WARN_ON(pte_write(pte));
        pte = pte_wrprotect(pte);
        pte = pte_mkclean(pte);
        set_pte_at(args->mm, args->vaddr, args->ptep, pte);
+       flush_dcache_page(page);
        pte = pte_mkwrite(pte);
        pte = pte_mkdirty(pte);
        ptep_set_access_flags(args->vma, args->vaddr, args->ptep, pte, 1);
        pte = pfn_pte(args->pte_pfn, args->page_prot);
        pte = pte_mkyoung(pte);
        set_pte_at(args->mm, args->vaddr, args->ptep, pte);
+       flush_dcache_page(page);
        ptep_test_and_clear_young(args->vma, args->vaddr, args->ptep);
        pte = ptep_get(args->ptep);
        WARN_ON(pte_young(pte));
 
 static void __init pmd_advanced_tests(struct pgtable_debug_args *args)
 {
+       struct page *page;
        pmd_t pmd;
        unsigned long vaddr = args->vaddr;
 
        if (!has_transparent_hugepage())
                return;
 
-       if (args->pmd_pfn == ULONG_MAX)
+       page = (args->pmd_pfn != ULONG_MAX) ? pfn_to_page(args->pmd_pfn) : NULL;
+       if (!page)
                return;
 
+       /*
+        * flush_dcache_page() is called after set_pmd_at() to clear
+        * PG_arch_1 for the page on ARM64. The page flag isn't cleared
+        * when it's released and page allocation check will fail when
+        * the page is allocated again. For architectures other than ARM64,
+        * the unexpected overhead of cache flushing is acceptable.
+        */
        pr_debug("Validating PMD advanced\n");
        /* Align the address wrt HPAGE_PMD_SIZE */
        vaddr &= HPAGE_PMD_MASK;
 
        pmd = pfn_pmd(args->pmd_pfn, args->page_prot);
        set_pmd_at(args->mm, vaddr, args->pmdp, pmd);
+       flush_dcache_page(page);
        pmdp_set_wrprotect(args->mm, vaddr, args->pmdp);
        pmd = READ_ONCE(*args->pmdp);
        WARN_ON(pmd_write(pmd));
        pmd = pmd_wrprotect(pmd);
        pmd = pmd_mkclean(pmd);
        set_pmd_at(args->mm, vaddr, args->pmdp, pmd);
+       flush_dcache_page(page);
        pmd = pmd_mkwrite(pmd);
        pmd = pmd_mkdirty(pmd);
        pmdp_set_access_flags(args->vma, vaddr, args->pmdp, pmd, 1);
        pmd = pmd_mkhuge(pfn_pmd(args->pmd_pfn, args->page_prot));
        pmd = pmd_mkyoung(pmd);
        set_pmd_at(args->mm, vaddr, args->pmdp, pmd);
+       flush_dcache_page(page);
        pmdp_test_and_clear_young(args->vma, vaddr, args->pmdp);
        pmd = READ_ONCE(*args->pmdp);
        WARN_ON(pmd_young(pmd));
 
 static void __init pud_advanced_tests(struct pgtable_debug_args *args)
 {
+       struct page *page;
        unsigned long vaddr = args->vaddr;
        pud_t pud;
 
        if (!has_transparent_hugepage())
                return;
 
-       if (args->pud_pfn == ULONG_MAX)
+       page = (args->pud_pfn != ULONG_MAX) ? pfn_to_page(args->pud_pfn) : NULL;
+       if (!page)
                return;
 
+       /*
+        * flush_dcache_page() is called after set_pud_at() to clear
+        * PG_arch_1 for the page on ARM64. The page flag isn't cleared
+        * when it's released and page allocation check will fail when
+        * the page is allocated again. For architectures other than ARM64,
+        * the unexpected overhead of cache flushing is acceptable.
+        */
        pr_debug("Validating PUD advanced\n");
        /* Align the address wrt HPAGE_PUD_SIZE */
        vaddr &= HPAGE_PUD_MASK;
 
        pud = pfn_pud(args->pud_pfn, args->page_prot);
        set_pud_at(args->mm, vaddr, args->pudp, pud);
+       flush_dcache_page(page);
        pudp_set_wrprotect(args->mm, vaddr, args->pudp);
        pud = READ_ONCE(*args->pudp);
        WARN_ON(pud_write(pud));
        pud = pud_wrprotect(pud);
        pud = pud_mkclean(pud);
        set_pud_at(args->mm, vaddr, args->pudp, pud);
+       flush_dcache_page(page);
        pud = pud_mkwrite(pud);
        pud = pud_mkdirty(pud);
        pudp_set_access_flags(args->vma, vaddr, args->pudp, pud, 1);
        pud = pfn_pud(args->pud_pfn, args->page_prot);
        pud = pud_mkyoung(pud);
        set_pud_at(args->mm, vaddr, args->pudp, pud);
+       flush_dcache_page(page);
        pudp_test_and_clear_young(args->vma, vaddr, args->pudp);
        pud = READ_ONCE(*args->pudp);
        WARN_ON(pud_young(pud));
 
 static void __init pte_clear_tests(struct pgtable_debug_args *args)
 {
+       struct page *page;
        pte_t pte = pfn_pte(args->pte_pfn, args->page_prot);
 
-       if (args->pte_pfn == ULONG_MAX)
+       page = (args->pte_pfn != ULONG_MAX) ? pfn_to_page(args->pte_pfn) : NULL;
+       if (!page)
                return;
 
+       /*
+        * flush_dcache_page() is called after set_pte_at() to clear
+        * PG_arch_1 for the page on ARM64. The page flag isn't cleared
+        * when it's released and page allocation check will fail when
+        * the page is allocated again. For architectures other than ARM64,
+        * the unexpected overhead of cache flushing is acceptable.
+        */
        pr_debug("Validating PTE clear\n");
 #ifndef CONFIG_RISCV
        pte = __pte(pte_val(pte) | RANDOM_ORVALUE);
 #endif
        set_pte_at(args->mm, args->vaddr, args->ptep, pte);
+       flush_dcache_page(page);
        barrier();
        pte_clear(args->mm, args->vaddr, args->ptep);
        pte = ptep_get(args->ptep);