return 0;
 }
 
-#define arch_clear_hugepage_flags(page)                do { } while (0)
+static inline void arch_clear_hugepage_flags(struct page *page)
+{
+       clear_bit(PG_arch_1, &page->flags);
+}
 
 static inline void huge_pte_clear(struct mm_struct *mm, unsigned long addr,
                                  pte_t *ptep, unsigned long sz)
 
 {
        pmd_t *pmd = (pmd_t *)pte;
        unsigned long start, end;
+       struct page *page = pmd_page(*pmd);
 
        /*
         * The write check makes sure we do not set a key on shared
        start = pmd_val(*pmd) & HPAGE_MASK;
        end = start + HPAGE_SIZE - 1;
        __storage_key_init_range(start, end);
+       set_bit(PG_arch_1, &page->flags);
        return 0;
 }
 
 
        return pte;
 }
 
+static void clear_huge_pte_skeys(struct mm_struct *mm, unsigned long rste)
+{
+       struct page *page;
+       unsigned long size, paddr;
+
+       if (!mm_uses_skeys(mm) ||
+           rste & _SEGMENT_ENTRY_INVALID)
+               return;
+
+       if ((rste & _REGION_ENTRY_TYPE_MASK) == _REGION_ENTRY_TYPE_R3) {
+               page = pud_page(__pud(rste));
+               size = PUD_SIZE;
+               paddr = rste & PUD_MASK;
+       } else {
+               page = pmd_page(__pmd(rste));
+               size = PMD_SIZE;
+               paddr = rste & PMD_MASK;
+       }
+
+       if (!test_and_set_bit(PG_arch_1, &page->flags))
+               __storage_key_init_range(paddr, paddr + size - 1);
+}
+
 void set_huge_pte_at(struct mm_struct *mm, unsigned long addr,
                     pte_t *ptep, pte_t pte)
 {
                rste |= _REGION_ENTRY_TYPE_R3 | _REGION3_ENTRY_LARGE;
        else
                rste |= _SEGMENT_ENTRY_LARGE;
+       clear_huge_pte_skeys(mm, rste);
        pte_val(*ptep) = rste;
 }