1 files changed, 18 insertions, 5 deletions

diff --git a/mm/migrate.c b/mm/migrate.c
index b1f57501de9c..6f0c24438bba 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -200,15 +200,14 @@ static void remove_migration_ptes(struct page *old, struct page *new)
  * get to the page and wait until migration is finished.
  * When we return from this function the fault will be retried.
  */
-void migration_entry_wait(struct mm_struct *mm, pmd_t *pmd,
-                                unsigned long address)
+static void __migration_entry_wait(struct mm_struct *mm, pte_t *ptep,
+                                spinlock_t *ptl)
 {
-        pte_t *ptep, pte;
-        spinlock_t *ptl;
+        pte_t pte;
         swp_entry_t entry;
         struct page *page;
 
-        ptep = pte_offset_map_lock(mm, pmd, address, &ptl);
+        spin_lock(ptl);
         pte = *ptep;
         if (!is_swap_pte(pte))
                 goto out;
@@ -236,6 +235,20 @@ out:
         pte_unmap_unlock(ptep, ptl);
 }
 
+void migration_entry_wait(struct mm_struct *mm, pmd_t *pmd,
+                                unsigned long address)
+{
+        spinlock_t *ptl = pte_lockptr(mm, pmd);
+        pte_t *ptep = pte_offset_map(pmd, address);
+        __migration_entry_wait(mm, ptep, ptl);
+}
+
+void migration_entry_wait_huge(struct mm_struct *mm, pte_t *pte)
+{
+        spinlock_t *ptl = &(mm)->page_table_lock;
+        __migration_entry_wait(mm, pte, ptl);
+}
+
 #ifdef CONFIG_BLOCK
 /* Returns true if all buffers are successfully locked */
 static bool buffer_migrate_lock_buffers(struct buffer_head *head,