if (pvmw->pte)
                goto next_pte;
 restart:
-       pgd = pgd_offset(mm, pvmw->address);
-       if (!pgd_present(*pgd))
-               return false;
-       p4d = p4d_offset(pgd, pvmw->address);
-       if (!p4d_present(*p4d))
-               return false;
-       pud = pud_offset(p4d, pvmw->address);
-       if (!pud_present(*pud))
-               return false;
-       pvmw->pmd = pmd_offset(pud, pvmw->address);
-       /*
-        * Make sure the pmd value isn't cached in a register by the
-        * compiler and used as a stale value after we've observed a
-        * subsequent update.
-        */
-       pmde = READ_ONCE(*pvmw->pmd);
-       if (pmd_trans_huge(pmde) || is_pmd_migration_entry(pmde)) {
-               pvmw->ptl = pmd_lock(mm, pvmw->pmd);
-               pmde = *pvmw->pmd;
-               if (likely(pmd_trans_huge(pmde))) {
-                       if (pvmw->flags & PVMW_MIGRATION)
-                               return not_found(pvmw);
-                       if (pmd_page(pmde) != page)
-                               return not_found(pvmw);
-                       return true;
-               }
-               if (!pmd_present(pmde)) {
-                       swp_entry_t entry;
+       {
+               pgd = pgd_offset(mm, pvmw->address);
+               if (!pgd_present(*pgd))
+                       return false;
+               p4d = p4d_offset(pgd, pvmw->address);
+               if (!p4d_present(*p4d))
+                       return false;
+               pud = pud_offset(p4d, pvmw->address);
+               if (!pud_present(*pud))
+                       return false;
 
-                       if (!thp_migration_supported() ||
-                           !(pvmw->flags & PVMW_MIGRATION))
-                               return not_found(pvmw);
-                       entry = pmd_to_swp_entry(pmde);
-                       if (!is_migration_entry(entry) ||
-                           migration_entry_to_page(entry) != page)
-                               return not_found(pvmw);
-                       return true;
-               }
-               /* THP pmd was split under us: handle on pte level */
-               spin_unlock(pvmw->ptl);
-               pvmw->ptl = NULL;
-       } else if (!pmd_present(pmde)) {
+               pvmw->pmd = pmd_offset(pud, pvmw->address);
                /*
-                * If PVMW_SYNC, take and drop THP pmd lock so that we
-                * cannot return prematurely, while zap_huge_pmd() has
-                * cleared *pmd but not decremented compound_mapcount().
+                * Make sure the pmd value isn't cached in a register by the
+                * compiler and used as a stale value after we've observed a
+                * subsequent update.
                 */
-               if ((pvmw->flags & PVMW_SYNC) && PageTransCompound(page)) {
-                       spinlock_t *ptl = pmd_lock(mm, pvmw->pmd);
+               pmde = READ_ONCE(*pvmw->pmd);
 
-                       spin_unlock(ptl);
+               if (pmd_trans_huge(pmde) || is_pmd_migration_entry(pmde)) {
+                       pvmw->ptl = pmd_lock(mm, pvmw->pmd);
+                       pmde = *pvmw->pmd;
+                       if (likely(pmd_trans_huge(pmde))) {
+                               if (pvmw->flags & PVMW_MIGRATION)
+                                       return not_found(pvmw);
+                               if (pmd_page(pmde) != page)
+                                       return not_found(pvmw);
+                               return true;
+                       }
+                       if (!pmd_present(pmde)) {
+                               swp_entry_t entry;
+
+                               if (!thp_migration_supported() ||
+                                   !(pvmw->flags & PVMW_MIGRATION))
+                                       return not_found(pvmw);
+                               entry = pmd_to_swp_entry(pmde);
+                               if (!is_migration_entry(entry) ||
+                                   migration_entry_to_page(entry) != page)
+                                       return not_found(pvmw);
+                               return true;
+                       }
+                       /* THP pmd was split under us: handle on pte level */
+                       spin_unlock(pvmw->ptl);
+                       pvmw->ptl = NULL;
+               } else if (!pmd_present(pmde)) {
+                       /*
+                        * If PVMW_SYNC, take and drop THP pmd lock so that we
+                        * cannot return prematurely, while zap_huge_pmd() has
+                        * cleared *pmd but not decremented compound_mapcount().
+                        */
+                       if ((pvmw->flags & PVMW_SYNC) &&
+                           PageTransCompound(page)) {
+                               spinlock_t *ptl = pmd_lock(mm, pvmw->pmd);
+
+                               spin_unlock(ptl);
+                       }
+                       return false;
                }
-               return false;
+               if (!map_pte(pvmw))
+                       goto next_pte;
        }
-       if (!map_pte(pvmw))
-               goto next_pte;
        while (1) {
                unsigned long end;