diff --git a/arch/arm64/include/asm/pgalloc.h b/arch/arm64/include/asm/pgalloc.h
index 172d76fa0245..920da9c5786c 100644
--- a/arch/arm64/include/asm/pgalloc.h
+++ b/arch/arm64/include/asm/pgalloc.h
@@ -30,7 +30,7 @@ static inline pmd_t *pmd_alloc_one(struct mm_struct *mm, unsigned long addr)
 	page = alloc_page(gfp);
 	if (!page)
 		return NULL;
-	if (!pgtable_pmd_page_ctor(page)) {
+	if (!pgtable_pmd_page_ctor(page, mm)) {
 		__free_page(page);
 		return NULL;
 	}
diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index c706bed1e496..b7bdde1990be 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -383,7 +383,7 @@ static phys_addr_t pgd_pgtable_alloc(int shift, struct mm_struct *mm)
 	if (shift == PAGE_SHIFT)
 		BUG_ON(!pgtable_pte_page_ctor(phys_to_page(pa), mm));
 	else if (shift == PMD_SHIFT)
-		BUG_ON(!pgtable_pmd_page_ctor(phys_to_page(pa)));
+		BUG_ON(!pgtable_pmd_page_ctor(phys_to_page(pa), mm));
 
 	return pa;
 }
diff --git a/arch/powerpc/mm/book3s64/pgtable.c b/arch/powerpc/mm/book3s64/pgtable.c
index e0bb69c616e4..9fda5287c197 100644
--- a/arch/powerpc/mm/book3s64/pgtable.c
+++ b/arch/powerpc/mm/book3s64/pgtable.c
@@ -297,7 +297,7 @@ static pmd_t *__alloc_for_pmdcache(struct mm_struct *mm)
 	page = alloc_page(gfp);
 	if (!page)
 		return NULL;
-	if (!pgtable_pmd_page_ctor(page)) {
+	if (!pgtable_pmd_page_ctor(page, mm)) {
 		__free_pages(page, 0);
 		return NULL;
 	}
diff --git a/arch/s390/include/asm/pgalloc.h b/arch/s390/include/asm/pgalloc.h
index 74a352f8c0d1..bebad4e5d42a 100644
--- a/arch/s390/include/asm/pgalloc.h
+++ b/arch/s390/include/asm/pgalloc.h
@@ -86,7 +86,7 @@ static inline pmd_t *pmd_alloc_one(struct mm_struct *mm, unsigned long vmaddr)
 	if (!table)
 		return NULL;
 	crst_table_init(table, _SEGMENT_ENTRY_EMPTY);
-	if (!pgtable_pmd_page_ctor(virt_to_page(table))) {
+	if (!pgtable_pmd_page_ctor(virt_to_page(table), mm)) {
 		crst_table_free(mm, table);
 		return NULL;
 	}
diff --git a/arch/x86/include/asm/pgalloc.h b/arch/x86/include/asm/pgalloc.h
index 29aa7859bdee..33514f0a9e79 100644
--- a/arch/x86/include/asm/pgalloc.h
+++ b/arch/x86/include/asm/pgalloc.h
@@ -96,7 +96,7 @@ static inline pmd_t *pmd_alloc_one(struct mm_struct *mm, unsigned long addr)
 	page = alloc_pages(gfp, 0);
 	if (!page)
 		return NULL;
-	if (!pgtable_pmd_page_ctor(page)) {
+	if (!pgtable_pmd_page_ctor(page, mm)) {
 		__free_pages(page, 0);
 		return NULL;
 	}
diff --git a/arch/x86/mm/pgtable.c b/arch/x86/mm/pgtable.c
index f5f46737aea0..8f4255662c5a 100644
--- a/arch/x86/mm/pgtable.c
+++ b/arch/x86/mm/pgtable.c
@@ -229,7 +229,7 @@ static int preallocate_pmds(struct mm_struct *mm, pmd_t *pmds[], int count)
 		pmd_t *pmd = (pmd_t *)__get_free_page(gfp);
 		if (!pmd)
 			failed = true;
-		if (pmd && !pgtable_pmd_page_ctor(virt_to_page(pmd))) {
+		if (pmd && !pgtable_pmd_page_ctor(virt_to_page(pmd), mm)) {
 			free_page((unsigned long)pmd);
 			pmd = NULL;
 			failed = true;
diff --git a/include/linux/mm.h b/include/linux/mm.h
index 2a98eebeba91..e2924d900fc5 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -2216,11 +2216,14 @@ static inline spinlock_t *pmd_lockptr(struct mm_struct *mm, pmd_t *pmd)
 	return ptlock_ptr(pmd_to_page(pmd));
 }
 
-static inline bool pgtable_pmd_page_ctor(struct page *page)
+static inline
+bool pgtable_pmd_page_ctor(struct page *page, struct mm_struct *mm)
 {
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
 	page->pmd_huge_pte = NULL;
 #endif
+	__SetPageTable(page);
+	page->pt_mm = mm;
 	return ptlock_init(page);
 }
 
@@ -2229,6 +2232,8 @@ static inline void pgtable_pmd_page_dtor(struct page *page)
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
 	VM_BUG_ON_PAGE(page->pmd_huge_pte, page);
 #endif
+	__ClearPageTable(page);
+	page->pt_mm = NULL;
 	ptlock_free(page);
 }
 
@@ -2241,7 +2246,11 @@ static inline spinlock_t *pmd_lockptr(struct mm_struct *mm, pmd_t *pmd)
 	return &mm->page_table_lock;
 }
 
-static inline bool pgtable_pmd_page_ctor(struct page *page) { return true; }
+static inline
+bool pgtable_pmd_page_ctor(struct page *page, struct mm_struct *mm)
+{
+	return true;
+}
 static inline void pgtable_pmd_page_dtor(struct page *page) {}
 
 #define pmd_huge_pte(mm, pmd) ((mm)->pmd_huge_pte)
