3 files changed, 20 insertions, 5 deletions

diff --git a/arch/s390/Kconfig b/arch/s390/Kconfig
index 65a07750f4f9..86db5a8e6586 100644
--- a/arch/s390/Kconfig
+++ b/arch/s390/Kconfig
@@ -415,6 +415,10 @@ config ARCH_ENABLE_MEMORY_HOTPLUG
 config ARCH_ENABLE_MEMORY_HOTREMOVE
         def_bool y
 
+config ARCH_ENABLE_SPLIT_PMD_PTLOCK
+        def_bool y
+        depends on 64BIT
+
 config FORCE_MAX_ZONEORDER
         int
         default "9"
diff --git a/arch/s390/include/asm/pgalloc.h b/arch/s390/include/asm/pgalloc.h
index 14d43c77d6cf..884017cbfa9f 100644
--- a/arch/s390/include/asm/pgalloc.h
+++ b/arch/s390/include/asm/pgalloc.h
@@ -92,11 +92,22 @@ static inline pud_t *pud_alloc_one(struct mm_struct *mm, unsigned long address)
 static inline pmd_t *pmd_alloc_one(struct mm_struct *mm, unsigned long vmaddr)
 {
         unsigned long *table = crst_table_alloc(mm);
-        if (table)
-                crst_table_init(table, _SEGMENT_ENTRY_EMPTY);
+
+        if (!table)
+                return NULL;
+        crst_table_init(table, _SEGMENT_ENTRY_EMPTY);
+        if (!pgtable_pmd_page_ctor(virt_to_page(table))) {
+                crst_table_free(mm, table);
+                return NULL;
+        }
         return (pmd_t *) table;
 }
-#define pmd_free(mm, pmd) crst_table_free(mm, (unsigned long *) pmd)
+
+static inline void pmd_free(struct mm_struct *mm, pmd_t *pmd)
+{
+        pgtable_pmd_page_dtor(virt_to_page(pmd));
+        crst_table_free(mm, (unsigned long *) pmd);
+}
 
 static inline void pgd_populate(struct mm_struct *mm, pgd_t *pgd, pud_t *pud)
 {
diff --git a/arch/s390/mm/pgtable.c b/arch/s390/mm/pgtable.c
index 9c26b7aa96d9..f8b58a7a3048 100644
--- a/arch/s390/mm/pgtable.c
+++ b/arch/s390/mm/pgtable.c
@@ -1397,7 +1397,7 @@ void pgtable_trans_huge_deposit(struct mm_struct *mm, pmd_t *pmdp,
 {
         struct list_head *lh = (struct list_head *) pgtable;
 
-        assert_spin_locked(&mm->page_table_lock);
+        assert_spin_locked(pmd_lockptr(mm, pmdp));
 
         /* FIFO */
         if (!pmd_huge_pte(mm, pmdp))
@@ -1413,7 +1413,7 @@ pgtable_t pgtable_trans_huge_withdraw(struct mm_struct *mm, pmd_t *pmdp)
         pgtable_t pgtable;
         pte_t *ptep;
 
-        assert_spin_locked(&mm->page_table_lock);
+        assert_spin_locked(pmd_lockptr(mm, pmdp));
 
         /* FIFO */
         pgtable = pmd_huge_pte(mm, pmdp);