In preparation for wp to support mthp, add the ptep_clear_flush_range()
function to clear and flush TLB for PTE in the specified range.

Signed-off-by: Vernon Yang <vernon2gm@gmail.com>
---
 include/linux/pgtable.h |  3 +++
 mm/memory.c             |  2 +-
 mm/pgtable-generic.c    | 20 ++++++++++++++++++++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/include/linux/pgtable.h b/include/linux/pgtable.h
index 0b6e1f781d86..1ccddcd0098f 100644
--- a/include/linux/pgtable.h
+++ b/include/linux/pgtable.h
@@ -826,6 +826,9 @@ static inline void clear_not_present_full_ptes(struct mm_struct *mm,
 extern pte_t ptep_clear_flush(struct vm_area_struct *vma,
 			      unsigned long address,
 			      pte_t *ptep);
+extern void ptep_clear_flush_range(struct vm_area_struct *vma,
+				   unsigned long address,
+				   pte_t *ptep, unsigned int nr);
 #endif
 
 #ifndef __HAVE_ARCH_PMDP_HUGE_CLEAR_FLUSH
diff --git a/mm/memory.c b/mm/memory.c
index a6bc1db22387..90cbed5ad150 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -3632,7 +3632,7 @@ static vm_fault_t wp_page_copy(struct vm_fault *vmf)
 		 * that left a window where the new PTE could be loaded into
 		 * some TLBs while the old PTE remains in others.
 		 */
-		ptep_clear_flush(vma, vmf->address, vmf->pte);
+		ptep_clear_flush_range(vma, vmf->address, vmf->pte, 1);
 		folio_add_new_anon_rmap(new_folio, vma, vmf->address, RMAP_EXCLUSIVE);
 		folio_add_lru_vma(new_folio, vma);
 		BUG_ON(unshare && pte_write(entry));
diff --git a/mm/pgtable-generic.c b/mm/pgtable-generic.c
index 5a882f2b10f9..cdffec4f54d9 100644
--- a/mm/pgtable-generic.c
+++ b/mm/pgtable-generic.c
@@ -101,6 +101,26 @@ pte_t ptep_clear_flush(struct vm_area_struct *vma, unsigned long address,
 		flush_tlb_page(vma, address);
 	return pte;
 }
+
+void ptep_clear_flush_range(struct vm_area_struct *vma, unsigned long address,
+			    pte_t *ptep, unsigned int nr)
+{
+	struct mm_struct *mm = (vma)->vm_mm;
+	bool accessible = false;
+	pte_t pte;
+	int i;
+
+	for (i = 0; i < nr; i++) {
+		pte = ptep_get_and_clear(mm, address + i * PAGE_SIZE, ptep + i);
+		if (!accessible && pte_accessible(mm, pte))
+			accessible = true;
+	}
+
+	if (accessible) {
+		flush_tlb_mm_range(vma->vm_mm, address, address + nr * PAGE_SIZE,
+				   PAGE_SHIFT, false);
+	}
+}
 #endif
 
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
-- 
2.50.1