#ifndef __ASSEMBLY__
 #include <linux/sched.h>
 #include <linux/mm_types.h>
+#include <linux/page-flags.h>
 #include <asm/bug.h>
 #include <asm/page.h>
 
 /* Software bits in the page table entry */
 #define _PAGE_SWT      0x001           /* SW pte type bit t */
 #define _PAGE_SWX      0x002           /* SW pte type bit x */
-#define _PAGE_SWC      0x004           /* SW pte changed bit (for KVM) */
-#define _PAGE_SWR      0x008           /* SW pte referenced bit (for KVM) */
-#define _PAGE_SPECIAL  0x010           /* SW associated with special page */
+#define _PAGE_SWC      0x004           /* SW pte changed bit */
+#define _PAGE_SWR      0x008           /* SW pte referenced bit */
+#define _PAGE_SWW      0x010           /* SW pte write bit */
+#define _PAGE_SPECIAL  0x020           /* SW associated with special page */
 #define __HAVE_ARCH_PTE_SPECIAL
 
 /* Set of bits not changed in pte_modify */
-#define _PAGE_CHG_MASK (PAGE_MASK | _PAGE_SPECIAL | _PAGE_SWC | _PAGE_SWR)
+#define _PAGE_CHG_MASK         (PAGE_MASK | _PAGE_SPECIAL | _PAGE_CO | \
+                                _PAGE_SWC | _PAGE_SWR)
 
 /* Six different types of pages. */
 #define _PAGE_TYPE_EMPTY       0x400
 
 /* Bits in the region table entry */
 #define _REGION_ENTRY_ORIGIN   ~0xfffUL/* region/segment table origin      */
+#define _REGION_ENTRY_RO       0x200   /* region protection bit            */
 #define _REGION_ENTRY_INV      0x20    /* invalid region table entry       */
 #define _REGION_ENTRY_TYPE_MASK        0x0c    /* region/segment table type mask   */
 #define _REGION_ENTRY_TYPE_R1  0x0c    /* region first table type          */
  */
 #define PAGE_NONE      __pgprot(_PAGE_TYPE_NONE)
 #define PAGE_RO                __pgprot(_PAGE_TYPE_RO)
-#define PAGE_RW                __pgprot(_PAGE_TYPE_RW)
+#define PAGE_RW                __pgprot(_PAGE_TYPE_RO | _PAGE_SWW)
+#define PAGE_RWC       __pgprot(_PAGE_TYPE_RW | _PAGE_SWW | _PAGE_SWC)
 
-#define PAGE_KERNEL    PAGE_RW
+#define PAGE_KERNEL    PAGE_RWC
 #define PAGE_SHARED    PAGE_KERNEL
 #define PAGE_COPY      PAGE_RO
 
        bits = skey & (_PAGE_CHANGED | _PAGE_REFERENCED);
        /* Clear page changed & referenced bit in the storage key */
        if (bits & _PAGE_CHANGED)
-               page_set_storage_key(address, skey ^ bits, 1);
+               page_set_storage_key(address, skey ^ bits, 0);
        else if (bits)
                page_reset_referenced(address);
        /* Transfer page changed & referenced bit to guest bits in pgste */
        pgste_val(pgste) |= bits << 48;         /* RCP_GR_BIT & RCP_GC_BIT */
        /* Get host changed & referenced bits from pgste */
        bits |= (pgste_val(pgste) & (RCP_HR_BIT | RCP_HC_BIT)) >> 52;
-       /* Clear host bits in pgste. */
+       /* Transfer page changed & referenced bit to kvm user bits */
+       pgste_val(pgste) |= bits << 45;         /* KVM_UR_BIT & KVM_UC_BIT */
+       /* Clear relevant host bits in pgste. */
        pgste_val(pgste) &= ~(RCP_HR_BIT | RCP_HC_BIT);
        pgste_val(pgste) &= ~(RCP_ACC_BITS | RCP_FP_BIT);
        /* Copy page access key and fetch protection bit to pgste */
        pgste_val(pgste) |=
                (unsigned long) (skey & (_PAGE_ACC_BITS | _PAGE_FP_BIT)) << 56;
-       /* Transfer changed and referenced to kvm user bits */
-       pgste_val(pgste) |= bits << 45;         /* KVM_UR_BIT & KVM_UC_BIT */
-       /* Transfer changed & referenced to pte sofware bits */
-       pte_val(*ptep) |= bits << 1;            /* _PAGE_SWR & _PAGE_SWC */
+       /* Transfer referenced bit to pte */
+       pte_val(*ptep) |= (bits & _PAGE_REFERENCED) << 1;
 #endif
        return pgste;
 
 
        if (!pte_present(*ptep))
                return pgste;
+       /* Get referenced bit from storage key */
        young = page_reset_referenced(pte_val(*ptep) & PAGE_MASK);
-       /* Transfer page referenced bit to pte software bit (host view) */
-       if (young || (pgste_val(pgste) & RCP_HR_BIT))
+       if (young)
+               pgste_val(pgste) |= RCP_GR_BIT;
+       /* Get host referenced bit from pgste */
+       if (pgste_val(pgste) & RCP_HR_BIT) {
+               pgste_val(pgste) &= ~RCP_HR_BIT;
+               young = 1;
+       }
+       /* Transfer referenced bit to kvm user bits and pte */
+       if (young) {
+               pgste_val(pgste) |= KVM_UR_BIT;
                pte_val(*ptep) |= _PAGE_SWR;
-       /* Clear host referenced bit in pgste. */
-       pgste_val(pgste) &= ~RCP_HR_BIT;
-       /* Transfer page referenced bit to guest bit in pgste */
-       pgste_val(pgste) |= (unsigned long) young << 50; /* set RCP_GR_BIT */
+       }
 #endif
        return pgste;
-
 }
 
-static inline void pgste_set_pte(pte_t *ptep, pgste_t pgste, pte_t entry)
+static inline void pgste_set_key(pte_t *ptep, pgste_t pgste, pte_t entry)
 {
 #ifdef CONFIG_PGSTE
        unsigned long address;
        /* Set page access key and fetch protection bit from pgste */
        nkey |= (pgste_val(pgste) & (RCP_ACC_BITS | RCP_FP_BIT)) >> 56;
        if (okey != nkey)
-               page_set_storage_key(address, nkey, 1);
+               page_set_storage_key(address, nkey, 0);
 #endif
 }
 
+static inline void pgste_set_pte(pte_t *ptep, pte_t entry)
+{
+       if (!MACHINE_HAS_ESOP && (pte_val(entry) & _PAGE_SWW)) {
+               /*
+                * Without enhanced suppression-on-protection force
+                * the dirty bit on for all writable ptes.
+                */
+               pte_val(entry) |= _PAGE_SWC;
+               pte_val(entry) &= ~_PAGE_RO;
+       }
+       *ptep = entry;
+}
+
 /**
  * struct gmap_struct - guest address space
  * @mm: pointer to the parent mm_struct
 
        if (mm_has_pgste(mm)) {
                pgste = pgste_get_lock(ptep);
-               pgste_set_pte(ptep, pgste, entry);
-               *ptep = entry;
+               pgste_set_key(ptep, pgste, entry);
+               pgste_set_pte(ptep, entry);
                pgste_set_unlock(ptep, pgste);
-       } else
+       } else {
+               if (!(pte_val(entry) & _PAGE_INVALID) && MACHINE_HAS_EDAT1)
+                       pte_val(entry) |= _PAGE_CO;
                *ptep = entry;
+       }
 }
 
 /*
  */
 static inline int pte_write(pte_t pte)
 {
-       return (pte_val(pte) & _PAGE_RO) == 0;
+       return (pte_val(pte) & _PAGE_SWW) != 0;
 }
 
 static inline int pte_dirty(pte_t pte)
 {
-#ifdef CONFIG_PGSTE
-       if (pte_val(pte) & _PAGE_SWC)
-               return 1;
-#endif
-       return 0;
+       return (pte_val(pte) & _PAGE_SWC) != 0;
 }
 
 static inline int pte_young(pte_t pte)
 {
        pte_val(pte) &= _PAGE_CHG_MASK;
        pte_val(pte) |= pgprot_val(newprot);
+       if ((pte_val(pte) & _PAGE_SWC) && (pte_val(pte) & _PAGE_SWW))
+               pte_val(pte) &= ~_PAGE_RO;
        return pte;
 }
 
 static inline pte_t pte_wrprotect(pte_t pte)
 {
+       pte_val(pte) &= ~_PAGE_SWW;
        /* Do not clobber _PAGE_TYPE_NONE pages!  */
        if (!(pte_val(pte) & _PAGE_INVALID))
                pte_val(pte) |= _PAGE_RO;
 
 static inline pte_t pte_mkwrite(pte_t pte)
 {
-       pte_val(pte) &= ~_PAGE_RO;
+       pte_val(pte) |= _PAGE_SWW;
+       if (pte_val(pte) & _PAGE_SWC)
+               pte_val(pte) &= ~_PAGE_RO;
        return pte;
 }
 
 static inline pte_t pte_mkclean(pte_t pte)
 {
-#ifdef CONFIG_PGSTE
        pte_val(pte) &= ~_PAGE_SWC;
-#endif
+       /* Do not clobber _PAGE_TYPE_NONE pages!  */
+       if (!(pte_val(pte) & _PAGE_INVALID))
+               pte_val(pte) |= _PAGE_RO;
        return pte;
 }
 
 static inline pte_t pte_mkdirty(pte_t pte)
 {
+       pte_val(pte) |= _PAGE_SWC;
+       if (pte_val(pte) & _PAGE_SWW)
+               pte_val(pte) &= ~_PAGE_RO;
        return pte;
 }
 
                pte_val(pte) |= _SEGMENT_ENTRY_INV;
        }
        /*
-        * Clear SW pte bits SWT and SWX, there are no SW bits in a segment
-        * table entry.
+        * Clear SW pte bits, there are no SW bits in a segment table entry.
         */
-       pte_val(pte) &= ~(_PAGE_SWT | _PAGE_SWX);
+       pte_val(pte) &= ~(_PAGE_SWT | _PAGE_SWX | _PAGE_SWC |
+                         _PAGE_SWR | _PAGE_SWW);
        /*
         * Also set the change-override bit because we don't need dirty bit
         * tracking for hugetlbfs pages.
                                           unsigned long address,
                                           pte_t *ptep, pte_t pte)
 {
-       *ptep = pte;
-       if (mm_has_pgste(mm))
+       if (mm_has_pgste(mm)) {
+               pgste_set_pte(ptep, pte);
                pgste_set_unlock(ptep, *(pgste_t *)(ptep + PTRS_PER_PTE));
+       } else
+               *ptep = pte;
 }
 
 #define __HAVE_ARCH_PTEP_CLEAR_FLUSH
 
                if (!mm_exclusive(mm))
                        __ptep_ipte(address, ptep);
-               *ptep = pte_wrprotect(pte);
+               pte = pte_wrprotect(pte);
 
-               if (mm_has_pgste(mm))
+               if (mm_has_pgste(mm)) {
+                       pgste_set_pte(ptep, pte);
                        pgste_set_unlock(ptep, pgste);
+               } else
+                       *ptep = pte;
        }
        return pte;
 }
                pgste = pgste_get_lock(ptep);
 
        __ptep_ipte(address, ptep);
-       *ptep = entry;
 
-       if (mm_has_pgste(vma->vm_mm))
+       if (mm_has_pgste(vma->vm_mm)) {
+               pgste_set_pte(ptep, entry);
                pgste_set_unlock(ptep, pgste);
+       } else
+               *ptep = entry;
        return 1;
 }
 
 static inline pte_t mk_pte(struct page *page, pgprot_t pgprot)
 {
        unsigned long physpage = page_to_phys(page);
+       pte_t __pte = mk_pte_phys(physpage, pgprot);
 
-       return mk_pte_phys(physpage, pgprot);
+       if ((pte_val(__pte) & _PAGE_SWW) && PageDirty(page)) {
+               pte_val(__pte) |= _PAGE_SWC;
+               pte_val(__pte) &= ~_PAGE_RO;
+       }
+       return __pte;
 }
 
 #define pgd_index(address) (((address) >> PGDIR_SHIFT) & (PTRS_PER_PGD-1))
 static inline void set_pmd_at(struct mm_struct *mm, unsigned long addr,
                              pmd_t *pmdp, pmd_t entry)
 {
+       if (!(pmd_val(entry) & _SEGMENT_ENTRY_INV) && MACHINE_HAS_EDAT1)
+               pmd_val(entry) |= _SEGMENT_ENTRY_CO;
        *pmdp = entry;
 }
 
 
        pud_t *pu_dir;
        pmd_t *pm_dir;
        pte_t *pt_dir;
-       pte_t  pte;
        int ret = -ENOMEM;
 
        while (address < end) {
-               pte = mk_pte_phys(address, __pgprot(ro ? _PAGE_RO : 0));
                pg_dir = pgd_offset_k(address);
                if (pgd_none(*pg_dir)) {
                        pu_dir = vmem_pud_alloc();
 #if defined(CONFIG_64BIT) && !defined(CONFIG_DEBUG_PAGEALLOC)
                if (MACHINE_HAS_EDAT2 && pud_none(*pu_dir) && address &&
                    !(address & ~PUD_MASK) && (address + PUD_SIZE <= end)) {
-                       pte_val(pte) |= _REGION3_ENTRY_LARGE;
-                       pte_val(pte) |= _REGION_ENTRY_TYPE_R3;
-                       pud_val(*pu_dir) = pte_val(pte);
+                       pud_val(*pu_dir) = __pa(address) |
+                               _REGION_ENTRY_TYPE_R3 | _REGION3_ENTRY_LARGE |
+                               (ro ? _REGION_ENTRY_RO : 0);
                        address += PUD_SIZE;
                        continue;
                }
 #if defined(CONFIG_64BIT) && !defined(CONFIG_DEBUG_PAGEALLOC)
                if (MACHINE_HAS_EDAT1 && pmd_none(*pm_dir) && address &&
                    !(address & ~PMD_MASK) && (address + PMD_SIZE <= end)) {
-                       pte_val(pte) |= _SEGMENT_ENTRY_LARGE;
-                       pmd_val(*pm_dir) = pte_val(pte);
+                       pmd_val(*pm_dir) = __pa(address) |
+                               _SEGMENT_ENTRY | _SEGMENT_ENTRY_LARGE |
+                               (ro ? _SEGMENT_ENTRY_RO : 0);
                        address += PMD_SIZE;
                        continue;
                }
                }
 
                pt_dir = pte_offset_kernel(pm_dir, address);
-               *pt_dir = pte;
+               pte_val(*pt_dir) = __pa(address) | (ro ? _PAGE_RO : 0);
                address += PAGE_SIZE;
        }
        ret = 0;
        pud_t *pu_dir;
        pmd_t *pm_dir;
        pte_t *pt_dir;
-       pte_t  pte;
        int ret = -ENOMEM;
 
        start_addr = (unsigned long) start;
                                new_page = vmemmap_alloc_block(PMD_SIZE, node);
                                if (!new_page)
                                        goto out;
-                               pte = mk_pte_phys(__pa(new_page), PAGE_RW);
-                               pte_val(pte) |= _SEGMENT_ENTRY_LARGE;
-                               pmd_val(*pm_dir) = pte_val(pte);
+                               pmd_val(*pm_dir) = __pa(new_page) |
+                                       _SEGMENT_ENTRY | _SEGMENT_ENTRY_LARGE;
                                address = (address + PMD_SIZE) & PMD_MASK;
                                continue;
                        }
                        new_page =__pa(vmem_alloc_pages(0));
                        if (!new_page)
                                goto out;
-                       pte = pfn_pte(new_page >> PAGE_SHIFT, PAGE_KERNEL);
-                       *pt_dir = pte;
+                       pte_val(*pt_dir) = __pa(new_page);
                }
                address += PAGE_SIZE;
        }