diff --git a/arch/s390/include/asm/page.h b/arch/s390/include/asm/page.h
index 6d5367060a5656a92d5d1b95f28dc3097fa47e78..24afd079e5cee6cc331e6bf577e8f87e5b9e9662 100644
--- a/arch/s390/include/asm/page.h
+++ b/arch/s390/include/asm/page.h
@@ -30,6 +30,8 @@
 #include <asm/setup.h>
 #ifndef __ASSEMBLY__
 
+void storage_key_init_range(unsigned long start, unsigned long end);
+
 static unsigned long pfmf(unsigned long function, unsigned long address)
 {
 	asm volatile(
diff --git a/arch/s390/kernel/setup.c b/arch/s390/kernel/setup.c
index b1f2be9aaaad7459b339af793e6e95d754543723..cfdad57957febb6f0d931ace4c103fb4bb352b3c 100644
--- a/arch/s390/kernel/setup.c
+++ b/arch/s390/kernel/setup.c
@@ -777,40 +777,6 @@ static void __init reserve_crashkernel(void)
 #endif
 }
 
-static void __init init_storage_keys(unsigned long start, unsigned long end)
-{
-	unsigned long boundary, function, size;
-
-	while (start < end) {
-		if (MACHINE_HAS_EDAT2) {
-			/* set storage keys for a 2GB frame */
-			function = 0x22000 | PAGE_DEFAULT_KEY;
-			size = 1UL << 31;
-			boundary = (start + size) & ~(size - 1);
-			if (boundary <= end) {
-				do {
-					start = pfmf(function, start);
-				} while (start < boundary);
-				continue;
-			}
-		}
-		if (MACHINE_HAS_EDAT1) {
-			/* set storage keys for a 1MB frame */
-			function = 0x21000 | PAGE_DEFAULT_KEY;
-			size = 1UL << 20;
-			boundary = (start + size) & ~(size - 1);
-			if (boundary <= end) {
-				do {
-					start = pfmf(function, start);
-				} while (start < boundary);
-				continue;
-			}
-		}
-		page_set_storage_key(start, PAGE_DEFAULT_KEY, 0);
-		start += PAGE_SIZE;
-	}
-}
-
 static void __init setup_memory(void)
 {
         unsigned long bootmap_size;
@@ -889,7 +855,7 @@ static void __init setup_memory(void)
 		memblock_add_node(PFN_PHYS(start_chunk),
 				  PFN_PHYS(end_chunk - start_chunk), 0);
 		pfn = max(start_chunk, start_pfn);
-		init_storage_keys(PFN_PHYS(pfn), PFN_PHYS(end_chunk));
+		storage_key_init_range(PFN_PHYS(pfn), PFN_PHYS(end_chunk));
 	}
 
 	psw_set_key(PAGE_DEFAULT_KEY);
diff --git a/arch/s390/mm/Makefile b/arch/s390/mm/Makefile
index 1bea6d1f55ab334b86db2042d84c84e5f01231df..640bea12303c0a659bcbb235bf191aaee8da4c62 100644
--- a/arch/s390/mm/Makefile
+++ b/arch/s390/mm/Makefile
@@ -2,9 +2,9 @@
 # Makefile for the linux s390-specific parts of the memory manager.
 #
 
-obj-y	 := init.o fault.o extmem.o mmap.o vmem.o pgtable.o maccess.o \
-	    page-states.o gup.o extable.o
-obj-$(CONFIG_CMM) += cmm.o
-obj-$(CONFIG_HUGETLB_PAGE) += hugetlbpage.o
-obj-$(CONFIG_DEBUG_SET_MODULE_RONX) += pageattr.o
-obj-$(CONFIG_S390_PTDUMP) += dump_pagetables.o
+obj-y		:= init.o fault.o extmem.o mmap.o vmem.o pgtable.o maccess.o
+obj-y		+= page-states.o gup.o extable.o pageattr.o
+
+obj-$(CONFIG_CMM)		+= cmm.o
+obj-$(CONFIG_HUGETLB_PAGE)	+= hugetlbpage.o
+obj-$(CONFIG_S390_PTDUMP)	+= dump_pagetables.o
diff --git a/arch/s390/mm/pageattr.c b/arch/s390/mm/pageattr.c
index c7ec7c2e46b0864d7010bce697c2df0f0691cf26..73c3da393e88dfd332fbb8314397ae69d6c623cc 100644
--- a/arch/s390/mm/pageattr.c
+++ b/arch/s390/mm/pageattr.c
@@ -2,11 +2,46 @@
  * Copyright IBM Corp. 2011
  * Author(s): Jan Glauber <jang@linux.vnet.ibm.com>
  */
+#include <linux/hugetlb.h>
 #include <linux/module.h>
 #include <linux/mm.h>
-#include <linux/hugetlb.h>
 #include <asm/cacheflush.h>
 #include <asm/pgtable.h>
+#include <asm/page.h>
+
+void storage_key_init_range(unsigned long start, unsigned long end)
+{
+	unsigned long boundary, function, size;
+
+	while (start < end) {
+		if (MACHINE_HAS_EDAT2) {
+			/* set storage keys for a 2GB frame */
+			function = 0x22000 | PAGE_DEFAULT_KEY;
+			size = 1UL << 31;
+			boundary = (start + size) & ~(size - 1);
+			if (boundary <= end) {
+				do {
+					start = pfmf(function, start);
+				} while (start < boundary);
+				continue;
+			}
+		}
+		if (MACHINE_HAS_EDAT1) {
+			/* set storage keys for a 1MB frame */
+			function = 0x21000 | PAGE_DEFAULT_KEY;
+			size = 1UL << 20;
+			boundary = (start + size) & ~(size - 1);
+			if (boundary <= end) {
+				do {
+					start = pfmf(function, start);
+				} while (start < boundary);
+				continue;
+			}
+		}
+		page_set_storage_key(start, PAGE_DEFAULT_KEY, 0);
+		start += PAGE_SIZE;
+	}
+}
 
 static pte_t *walk_page_table(unsigned long addr)
 {
diff --git a/drivers/s390/char/sclp_cmd.c b/drivers/s390/char/sclp_cmd.c
index 71ea923c322d4330f2272dac818e0db327457199..0dfa88a30118667bb2b7063ccf175e8ed4d1352b 100644
--- a/drivers/s390/char/sclp_cmd.c
+++ b/drivers/s390/char/sclp_cmd.c
@@ -19,10 +19,11 @@
 #include <linux/memory.h>
 #include <linux/module.h>
 #include <linux/platform_device.h>
+#include <asm/ctl_reg.h>
 #include <asm/chpid.h>
-#include <asm/sclp.h>
 #include <asm/setup.h>
-#include <asm/ctl_reg.h>
+#include <asm/page.h>
+#include <asm/sclp.h>
 
 #include "sclp.h"
 
@@ -400,17 +401,15 @@ static int do_assign_storage(sclp_cmdw_t cmd, u16 rn)
 
 static int sclp_assign_storage(u16 rn)
 {
-	unsigned long long start, address;
+	unsigned long long start;
 	int rc;
 
 	rc = do_assign_storage(0x000d0001, rn);
 	if (rc)
-		goto out;
-	start = address = rn2addr(rn);
-	for (; address < start + rzm; address += PAGE_SIZE)
-		page_set_storage_key(address, PAGE_DEFAULT_KEY, 0);
-out:
-	return rc;
+		return rc;
+	start = rn2addr(rn);
+	storage_key_init_range(start, start + rzm);
+	return 0;
 }
 
 static int sclp_unassign_storage(u16 rn)