{
        pte_t pte = pfn_pte(pfn, prot);
 
+       /*
+        * 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
+        */
+
        pr_debug("Validating PTE advanced\n");
        pte = pfn_pte(pfn, prot);
        set_pte_at(mm, vaddr, ptep, pte);
        ptep_set_wrprotect(mm, vaddr, ptep);
        pte = ptep_get(ptep);
        WARN_ON(pte_write(pte));
-
-       pte = pfn_pte(pfn, prot);
-       set_pte_at(mm, vaddr, ptep, pte);
        ptep_get_and_clear(mm, vaddr, ptep);
        pte = ptep_get(ptep);
        WARN_ON(!pte_none(pte));
        ptep_set_access_flags(vma, vaddr, ptep, pte, 1);
        pte = ptep_get(ptep);
        WARN_ON(!(pte_write(pte) && pte_dirty(pte)));
-
-       pte = pfn_pte(pfn, prot);
-       set_pte_at(mm, vaddr, ptep, pte);
        ptep_get_and_clear_full(mm, vaddr, ptep, 1);
        pte = ptep_get(ptep);
        WARN_ON(!pte_none(pte));
 
+       pte = pfn_pte(pfn, prot);
        pte = pte_mkyoung(pte);
        set_pte_at(mm, vaddr, ptep, pte);
        ptep_test_and_clear_young(vma, vaddr, ptep);
        pmdp_set_wrprotect(mm, vaddr, pmdp);
        pmd = READ_ONCE(*pmdp);
        WARN_ON(pmd_write(pmd));
-
-       pmd = pfn_pmd(pfn, prot);
-       set_pmd_at(mm, vaddr, pmdp, pmd);
        pmdp_huge_get_and_clear(mm, vaddr, pmdp);
        pmd = READ_ONCE(*pmdp);
        WARN_ON(!pmd_none(pmd));
        pmdp_set_access_flags(vma, vaddr, pmdp, pmd, 1);
        pmd = READ_ONCE(*pmdp);
        WARN_ON(!(pmd_write(pmd) && pmd_dirty(pmd)));
-
-       pmd = pmd_mkhuge(pfn_pmd(pfn, prot));
-       set_pmd_at(mm, vaddr, pmdp, pmd);
        pmdp_huge_get_and_clear_full(vma, vaddr, pmdp, 1);
        pmd = READ_ONCE(*pmdp);
        WARN_ON(!pmd_none(pmd));
 
+       pmd = pmd_mkhuge(pfn_pmd(pfn, prot));
        pmd = pmd_mkyoung(pmd);
        set_pmd_at(mm, vaddr, pmdp, pmd);
        pmdp_test_and_clear_young(vma, vaddr, pmdp);
        WARN_ON(pud_write(pud));
 
 #ifndef __PAGETABLE_PMD_FOLDED
-       pud = pfn_pud(pfn, prot);
-       set_pud_at(mm, vaddr, pudp, pud);
        pudp_huge_get_and_clear(mm, vaddr, pudp);
        pud = READ_ONCE(*pudp);
        WARN_ON(!pud_none(pud));
-
-       pud = pfn_pud(pfn, prot);
-       set_pud_at(mm, vaddr, pudp, pud);
-       pudp_huge_get_and_clear_full(mm, vaddr, pudp, 1);
-       pud = READ_ONCE(*pudp);
-       WARN_ON(!pud_none(pud));
 #endif /* __PAGETABLE_PMD_FOLDED */
        pud = pfn_pud(pfn, prot);
        pud = pud_wrprotect(pud);
        pud = READ_ONCE(*pudp);
        WARN_ON(!(pud_write(pud) && pud_dirty(pud)));
 
+#ifndef __PAGETABLE_PMD_FOLDED
+       pudp_huge_get_and_clear_full(mm, vaddr, pudp, 1);
+       pud = READ_ONCE(*pudp);
+       WARN_ON(!pud_none(pud));
+#endif /* __PAGETABLE_PMD_FOLDED */
+
+       pud = pfn_pud(pfn, prot);
        pud = pud_mkyoung(pud);
        set_pud_at(mm, vaddr, pudp, pud);
        pudp_test_and_clear_young(vma, vaddr, pudp);