1 files changed, 18 insertions, 1 deletions

diff --git a/arch/s390/mm/pgtable.c b/arch/s390/mm/pgtable.c
index 8345ae1f117d..05ae254f84cf 100644
--- a/arch/s390/mm/pgtable.c
+++ b/arch/s390/mm/pgtable.c
@@ -1237,11 +1237,28 @@ EXPORT_SYMBOL_GPL(s390_reset_cmma);
  */
 bool gmap_test_and_clear_dirty(unsigned long address, struct gmap *gmap)
 {
+        pgd_t *pgd;
+        pud_t *pud;
+        pmd_t *pmd;
         pte_t *pte;
         spinlock_t *ptl;
         bool dirty = false;
 
-        pte = get_locked_pte(gmap->mm, address, &ptl);
+        pgd = pgd_offset(gmap->mm, address);
+        pud = pud_alloc(gmap->mm, pgd, address);
+        if (!pud)
+                return false;
+        pmd = pmd_alloc(gmap->mm, pud, address);
+        if (!pmd)
+                return false;
+        /* We can't run guests backed by huge pages, but userspace can
+         * still set them up and then try to migrate them without any
+         * migration support.
+         */
+        if (pmd_large(*pmd))
+                return true;
+
+        pte = pte_alloc_map_lock(gmap->mm, pmd, address, &ptl);
         if (unlikely(!pte))
                 return false;