diff --git a/arch/s390/mm/pageattr.c b/arch/s390/mm/pageattr.c
index 0f33bab3e98483d6ac5df2f430cfa868d82517b7..b56953dd95f82fdf1ff17e4c28b1c08f8d041e84 100644
--- a/arch/s390/mm/pageattr.c
+++ b/arch/s390/mm/pageattr.c
@@ -8,25 +8,38 @@
 #include <asm/cacheflush.h>
 #include <asm/pgtable.h>
 
+static pte_t *walk_page_table(unsigned long addr)
+{
+	pgd_t *pgdp;
+	pud_t *pudp;
+	pmd_t *pmdp;
+	pte_t *ptep;
+
+	pgdp = pgd_offset_k(addr);
+	if (pgd_none(*pgdp))
+		return NULL;
+	pudp = pud_offset(pgdp, addr);
+	if (pud_none(*pudp))
+		return NULL;
+	pmdp = pmd_offset(pudp, addr);
+	if (pmd_none(*pmdp) || pmd_large(*pmdp))
+		return NULL;
+	ptep = pte_offset_kernel(pmdp, addr);
+	if (pte_none(*ptep))
+		return NULL;
+	return ptep;
+}
+
 static void change_page_attr(unsigned long addr, int numpages,
 			     pte_t (*set) (pte_t))
 {
 	pte_t *ptep, pte;
-	pmd_t *pmdp;
-	pud_t *pudp;
-	pgd_t *pgdp;
 	int i;
 
 	for (i = 0; i < numpages; i++) {
-		pgdp = pgd_offset(&init_mm, addr);
-		pudp = pud_offset(pgdp, addr);
-		pmdp = pmd_offset(pudp, addr);
-		if (pmd_large(*pmdp)) {
-			WARN_ON_ONCE(1);
-			continue;
-		}
-		ptep = pte_offset_kernel(pmdp, addr);
-
+		ptep = walk_page_table(addr);
+		if (WARN_ON_ONCE(!ptep))
+			break;
 		pte = *ptep;
 		pte = set(pte);
 		__ptep_ipte(addr, ptep);