1 files changed, 35 insertions, 3 deletions

diff --git a/arch/s390/kernel/setup.c b/arch/s390/kernel/setup.c
index afa9fdba200e..bfb48f18169c 100644
--- a/arch/s390/kernel/setup.c
+++ b/arch/s390/kernel/setup.c
@@ -768,6 +768,40 @@ 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;
@@ -846,9 +880,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);
-		for (; pfn < end_chunk; pfn++)
-			page_set_storage_key(PFN_PHYS(pfn),
-					     PAGE_DEFAULT_KEY, 0);
+		init_storage_keys(PFN_PHYS(pfn), PFN_PHYS(end_chunk));
 	}
 
 	psw_set_key(PAGE_DEFAULT_KEY);