4 files changed, 51 insertions, 6 deletions

diff --git a/arch/s390/include/asm/page.h b/arch/s390/include/asm/page.h
index 27ab3c7c1e8b..6d5367060a56 100644
--- a/arch/s390/include/asm/page.h
+++ b/arch/s390/include/asm/page.h
@@ -30,12 +30,20 @@
 #include <asm/setup.h>
 #ifndef __ASSEMBLY__
 
+static unsigned long pfmf(unsigned long function, unsigned long address)
+{
+        asm volatile(
+                "       .insn   rre,0xb9af0000,%[function],%[address]"
+                : [address] "+a" (address)
+                : [function] "d" (function)
+                : "memory");
+        return address;
+}
+
 static inline void clear_page(void *page)
 {
         if (MACHINE_HAS_PFMF) {
-                asm volatile(
-                        "       .insn   rre,0xb9af0000,%0,%1"
-                        : : "d" (0x10000), "a" (page) : "memory", "cc");
+                pfmf(0x10000, (unsigned long)page);
         } else {
                 register unsigned long reg1 asm ("1") = 0;
                 register void *reg2 asm ("2") = page;
diff --git a/arch/s390/include/asm/setup.h b/arch/s390/include/asm/setup.h
index a798985515b7..bc1b87b7453a 100644
--- a/arch/s390/include/asm/setup.h
+++ b/arch/s390/include/asm/setup.h
@@ -76,6 +76,7 @@ extern unsigned int s390_user_mode;
 #define MACHINE_FLAG_MVCOS      (1UL << 8)
 #define MACHINE_FLAG_KVM        (1UL << 9)
 #define MACHINE_FLAG_EDAT1      (1UL << 10)
+#define MACHINE_FLAG_EDAT2      (1UL << 11)
 #define MACHINE_FLAG_LPAR       (1UL << 12)
 #define MACHINE_FLAG_SPP        (1UL << 13)
 #define MACHINE_FLAG_TOPOLOGY   (1UL << 14)
@@ -98,6 +99,7 @@ extern unsigned int s390_user_mode;
 #define MACHINE_HAS_MVPG        (S390_lowcore.machine_flags & MACHINE_FLAG_MVPG)
 #define MACHINE_HAS_MVCOS       (0)
 #define MACHINE_HAS_EDAT1       (0)
+#define MACHINE_HAS_EDAT2       (0)
 #define MACHINE_HAS_SPP         (0)
 #define MACHINE_HAS_TOPOLOGY    (0)
 #define MACHINE_HAS_TE          (0)
@@ -110,6 +112,7 @@ extern unsigned int s390_user_mode;
 #define MACHINE_HAS_MVPG        (1)
 #define MACHINE_HAS_MVCOS       (S390_lowcore.machine_flags & MACHINE_FLAG_MVCOS)
 #define MACHINE_HAS_EDAT1       (S390_lowcore.machine_flags & MACHINE_FLAG_EDAT1)
+#define MACHINE_HAS_EDAT2       (S390_lowcore.machine_flags & MACHINE_FLAG_EDAT2)
 #define MACHINE_HAS_SPP         (S390_lowcore.machine_flags & MACHINE_FLAG_SPP)
 #define MACHINE_HAS_TOPOLOGY    (S390_lowcore.machine_flags & MACHINE_FLAG_TOPOLOGY)
 #define MACHINE_HAS_TE          (S390_lowcore.machine_flags & MACHINE_FLAG_TE)
diff --git a/arch/s390/kernel/early.c b/arch/s390/kernel/early.c
index 4c91d078d091..1f0eee9e7daa 100644
--- a/arch/s390/kernel/early.c
+++ b/arch/s390/kernel/early.c
@@ -374,6 +374,8 @@ static __init void detect_machine_facilities(void)
                 S390_lowcore.machine_flags |= MACHINE_FLAG_EDAT1;
                 __ctl_set_bit(0, 23);
         }
+        if (test_facility(78))
+                S390_lowcore.machine_flags |= MACHINE_FLAG_EDAT2;
         if (test_facility(3))
                 S390_lowcore.machine_flags |= MACHINE_FLAG_IDTE;
         if (test_facility(27))
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);