set_pte_at() should not be used to set a pte entry at locations that
already holds a valid pte entry. Architectures like ppc64 don't do TLB
invalidate in set_pte_at() and hence expect it to be used to set locations
that are not a valid PTE.

Signed-off-by: Aneesh Kumar K.V <redacted>
---
  mm/debug_vm_pgtable.c | 35 +++++++++++++++--------------------
  1 file changed, 15 insertions(+), 20 deletions(-)

diff --git a/mm/debug_vm_pgtable.c b/mm/debug_vm_pgtable.c
index 76f4c713e5a3..9c7e2c9cfc76 100644
--- a/mm/debug_vm_pgtable.c
+++ b/mm/debug_vm_pgtable.c

@@ -74,15 +74,18 @@ static void __init pte_advanced_tests(struct mm_struct *mm,
  {
  	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));

@@ -96,13 +99,11 @@ static void __init pte_advanced_tests(struct mm_struct *mm,
  	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);

@@ -164,9 +165,6 @@ static void __init pmd_advanced_tests(struct mm_struct *mm,
  	pmdp_set_wrprotect(mm, vaddr, pmdp);
  	pmd = READ_ONCE(*pmdp);
  	WARN_ON(pmd_write(pmd));
-
-	pmd = pmd_mkhuge(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));

@@ -180,13 +178,11 @@ static void __init pmd_advanced_tests(struct mm_struct *mm,
  	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);

@@ -283,18 +279,10 @@ static void __init pud_advanced_tests(struct mm_struct *mm,
  	WARN_ON(pud_write(pud));
  
  #ifndef __PAGETABLE_PMD_FOLDED

-
-	pud = pud_mkhuge(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 = pud_mkhuge(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 = pud_mkhuge(pfn_pud(pfn, prot));

@@ -307,6 +295,13 @@ static void __init pud_advanced_tests(struct mm_struct *mm,
  	pud = READ_ONCE(*pudp);
  	WARN_ON(!(pud_write(pud) && pud_dirty(pud)));
  
+#ifndef __PAGETABLE_PMD_FOLDED
+	pudp_huge_get_and_clear_full(vma, vaddr, pudp, 1);
+	pud = READ_ONCE(*pudp);
+	WARN_ON(!pud_none(pud));
+#endif /* __PAGETABLE_PMD_FOLDED */