[RFC PATCH 06/12] khugepaged: Generalize __collapse_huge_page_copy_failed()

[Dev Jain <dev.jain@arm.com>]

Upon failure, we repopulate the PMD in case of PMD-THP collapse. Hence, make
this logic specific for PMD case.

Signed-off-by: Dev Jain <dev.jain@arm.com>
---
 mm/khugepaged.c | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index de044b1f83d4..886c76816963 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -766,7 +766,7 @@ static void __collapse_huge_page_copy_failed(pte_t *pte,
 					     pmd_t *pmd,
 					     pmd_t orig_pmd,
 					     struct vm_area_struct *vma,
-					     struct list_head *compound_pagelist)
+					     struct list_head *compound_pagelist, int order)
 {
 	spinlock_t *pmd_ptl;
 
@@ -776,14 +776,16 @@ static void __collapse_huge_page_copy_failed(pte_t *pte,
 	 * pages. Since pages are still isolated and locked here,
 	 * acquiring anon_vma_lock_write is unnecessary.
 	 */
-	pmd_ptl = pmd_lock(vma->vm_mm, pmd);
-	pmd_populate(vma->vm_mm, pmd, pmd_pgtable(orig_pmd));
-	spin_unlock(pmd_ptl);
+	if (order == HPAGE_PMD_ORDER) {
+		pmd_ptl = pmd_lock(vma->vm_mm, pmd);
+		pmd_populate(vma->vm_mm, pmd, pmd_pgtable(orig_pmd));
+		spin_unlock(pmd_ptl);
+	}
 	/*
 	 * Release both raw and compound pages isolated
 	 * in __collapse_huge_page_isolate.
 	 */
-	release_pte_pages(pte, pte + HPAGE_PMD_NR, compound_pagelist);
+	release_pte_pages(pte, pte + (1UL << order), compound_pagelist);
 }
 
 /*
@@ -834,7 +836,7 @@ static int __collapse_huge_page_copy(pte_t *pte, struct folio *folio,
 						    compound_pagelist);
 	else
 		__collapse_huge_page_copy_failed(pte, pmd, orig_pmd, vma,
-						 compound_pagelist);
+						 compound_pagelist, order);
 
 	return result;
 }
-- 
2.30.2