#endif
 }
 
+/*
+ * Remove all empty zero pages from the mapping for lazy refaulting
+ * - This must be called after mm->context.has_pgste is set, to avoid
+ *   future creation of zero pages
+ * - This must be called after THP was enabled
+ */
+static int __zap_zero_pages(pmd_t *pmd, unsigned long start,
+                          unsigned long end, struct mm_walk *walk)
+{
+       unsigned long addr;
+
+       for (addr = start; addr != end; addr += PAGE_SIZE) {
+               pte_t *ptep;
+               spinlock_t *ptl;
+
+               ptep = pte_offset_map_lock(walk->mm, pmd, addr, &ptl);
+               if (is_zero_pfn(pte_pfn(*ptep)))
+                       ptep_xchg_direct(walk->mm, addr, ptep, __pte(_PAGE_INVALID));
+               pte_unmap_unlock(ptep, ptl);
+       }
+       return 0;
+}
+
+static inline void zap_zero_pages(struct mm_struct *mm)
+{
+       struct mm_walk walk = { .pmd_entry = __zap_zero_pages };
+
+       walk.mm = mm;
+       walk_page_range(0, TASK_SIZE, &walk);
+}
+
 /*
  * switch on pgstes for its userspace process (for kvm)
  */
        mm->context.has_pgste = 1;
        /* split thp mappings and disable thp for future mappings */
        thp_split_mm(mm);
+       zap_zero_pages(mm);
        up_write(&mm->mmap_sem);
        return 0;
 }
 static int __s390_enable_skey(pte_t *pte, unsigned long addr,
                              unsigned long next, struct mm_walk *walk)
 {
-       /*
-        * Remove all zero page mappings,
-        * after establishing a policy to forbid zero page mappings
-        * following faults for that page will get fresh anonymous pages
-        */
-       if (is_zero_pfn(pte_pfn(*pte)))
-               ptep_xchg_direct(walk->mm, addr, pte, __pte(_PAGE_INVALID));
        /* Clear storage key */
        ptep_zap_key(walk->mm, addr, pte);
        return 0;