3 files changed, 31 insertions, 9 deletions

diff --git a/arch/powerpc/include/asm/pgtable-ppc64.h b/arch/powerpc/include/asm/pgtable-ppc64.h
index 4a191c472867..d27960c89a71 100644
--- a/arch/powerpc/include/asm/pgtable-ppc64.h
+++ b/arch/powerpc/include/asm/pgtable-ppc64.h
@@ -558,5 +558,19 @@ extern pgtable_t pgtable_trans_huge_withdraw(struct mm_struct *mm, pmd_t *pmdp);
 #define __HAVE_ARCH_PMDP_INVALIDATE
 extern void pmdp_invalidate(struct vm_area_struct *vma, unsigned long address,
                             pmd_t *pmdp);
+
+#define pmd_move_must_withdraw pmd_move_must_withdraw
+typedef struct spinlock spinlock_t;
+static inline int pmd_move_must_withdraw(spinlock_t *new_pmd_ptl,
+                                         spinlock_t *old_pmd_ptl)
+{
+        /*
+         * Archs like ppc64 use pgtable to store per pmd
+         * specific information. So when we switch the pmd,
+         * we should also withdraw and deposit the pgtable
+         */
+        return true;
+}
+
 #endif /* __ASSEMBLY__ */
 #endif /* _ASM_POWERPC_PGTABLE_PPC64_H_ */
diff --git a/include/asm-generic/pgtable.h b/include/asm-generic/pgtable.h
index db0923458940..8e4f41d9af4d 100644
--- a/include/asm-generic/pgtable.h
+++ b/include/asm-generic/pgtable.h
@@ -558,6 +558,18 @@ static inline pmd_t pmd_read_atomic(pmd_t *pmdp)
 }
 #endif
 
+#ifndef pmd_move_must_withdraw
+static inline int pmd_move_must_withdraw(spinlock_t *new_pmd_ptl,
+                                         spinlock_t *old_pmd_ptl)
+{
+        /*
+         * With split pmd lock we also need to move preallocated
+         * PTE page table if new_pmd is on different PMD page table.
+         */
+        return new_pmd_ptl != old_pmd_ptl;
+}
+#endif
+
 /*
  * This function is meant to be used by sites walking pagetables with
  * the mmap_sem hold in read mode to protect against MADV_DONTNEED and
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 65c98eb5483c..82166bf974e1 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -1508,19 +1508,15 @@ int move_huge_pmd(struct vm_area_struct *vma, struct vm_area_struct *new_vma,
                         spin_lock_nested(new_ptl, SINGLE_DEPTH_NESTING);
                 pmd = pmdp_get_and_clear(mm, old_addr, old_pmd);
                 VM_BUG_ON(!pmd_none(*new_pmd));
-                set_pmd_at(mm, new_addr, new_pmd, pmd_mksoft_dirty(pmd));
-                if (new_ptl != old_ptl) {
-                        pgtable_t pgtable;
 
-                        /*
-                         * Move preallocated PTE page table if new_pmd is on
-                         * different PMD page table.
-                         */
+                if (pmd_move_must_withdraw(new_ptl, old_ptl)) {
+                        pgtable_t pgtable;
                         pgtable = pgtable_trans_huge_withdraw(mm, old_pmd);
                         pgtable_trans_huge_deposit(mm, new_pmd, pgtable);
-
-                        spin_unlock(new_ptl);
                 }
+                set_pmd_at(mm, new_addr, new_pmd, pmd_mksoft_dirty(pmd));
+                if (new_ptl != old_ptl)
+                        spin_unlock(new_ptl);
                 spin_unlock(old_ptl);
         }
 out: