{
        unsigned long next, *table, *new;
        struct page *page;
+       spinlock_t *ptl;
        pmd_t *pmd;
 
        pmd = pmd_offset(pud, addr);
                if (!new)
                        return -ENOMEM;
 
-               spin_lock(&mm->page_table_lock);
+               ptl = pmd_lock(mm, pmd);
                if (likely((unsigned long *) pmd_deref(*pmd) == table)) {
                        /* Nuke pmd entry pointing to the "short" page table */
                        pmdp_flush_lazy(mm, addr, pmd);
                        page_table_free_rcu(tlb, table);
                        new = NULL;
                }
-               spin_unlock(&mm->page_table_lock);
+               spin_unlock(ptl);
                if (new) {
                        page_table_free_pgste(new);
                        goto again;