3 files changed, 14 insertions, 6 deletions

diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 2f12e9fcf1a2..0a42d1521aa4 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -1260,6 +1260,7 @@ int do_huge_pmd_numa_page(struct mm_struct *mm, struct vm_area_struct *vma,
         int target_nid, last_cpupid = -1;
         bool page_locked;
         bool migrated = false;
+        bool was_writable;
         int flags = 0;
 
         /* A PROT_NONE fault should not end up here */
@@ -1354,7 +1355,10 @@ int do_huge_pmd_numa_page(struct mm_struct *mm, struct vm_area_struct *vma,
         goto out;
 clear_pmdnuma:
         BUG_ON(!PageLocked(page));
+        was_writable = pmd_write(pmd);
         pmd = pmd_modify(pmd, vma->vm_page_prot);
+        if (was_writable)
+                pmd = pmd_mkwrite(pmd);
         set_pmd_at(mm, haddr, pmdp, pmd);
         update_mmu_cache_pmd(vma, addr, pmdp);
         unlock_page(page);
@@ -1478,6 +1482,7 @@ int change_huge_pmd(struct vm_area_struct *vma, pmd_t *pmd,
 
         if (__pmd_trans_huge_lock(pmd, vma, &ptl) == 1) {
                 pmd_t entry;
+                bool preserve_write = prot_numa && pmd_write(*pmd);
                 ret = 1;
 
                 /*
@@ -1493,9 +1498,11 @@ int change_huge_pmd(struct vm_area_struct *vma, pmd_t *pmd,
                 if (!prot_numa || !pmd_protnone(*pmd)) {
                         entry = pmdp_get_and_clear_notify(mm, addr, pmd);
                         entry = pmd_modify(entry, newprot);
+                        if (preserve_write)
+                                entry = pmd_mkwrite(entry);
                         ret = HPAGE_PMD_NR;
                         set_pmd_at(mm, addr, pmd, entry);
-                        BUG_ON(pmd_write(entry));
+                        BUG_ON(!preserve_write && pmd_write(entry));
                 }
                 spin_unlock(ptl);
         }
diff --git a/mm/memory.c b/mm/memory.c
index 20beb6647dba..d20e12da3a3c 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -3035,6 +3035,7 @@ static int do_numa_page(struct mm_struct *mm, struct vm_area_struct *vma,
         int last_cpupid;
         int target_nid;
         bool migrated = false;
+        bool was_writable = pte_write(pte);
         int flags = 0;
 
         /* A PROT_NONE fault should not end up here */
@@ -3059,6 +3060,8 @@ static int do_numa_page(struct mm_struct *mm, struct vm_area_struct *vma,
         /* Make it present again */
         pte = pte_modify(pte, vma->vm_page_prot);
         pte = pte_mkyoung(pte);
+        if (was_writable)
+                pte = pte_mkwrite(pte);
         set_pte_at(mm, addr, ptep, pte);
         update_mmu_cache(vma, addr, ptep);
 
@@ -3075,11 +3078,6 @@ static int do_numa_page(struct mm_struct *mm, struct vm_area_struct *vma,
          * to it but pte_write gets cleared during protection updates and
          * pte_dirty has unpredictable behaviour between PTE scan updates,
          * background writeback, dirty balancing and application behaviour.
-         *
-         * TODO: Note that the ideal here would be to avoid a situation where a
-         * NUMA fault is taken immediately followed by a write fault in
-         * some cases which would have lower overhead overall but would be
-         * invasive as the fault paths would need to be unified.
          */
         if (!(vma->vm_flags & VM_WRITE))
                 flags |= TNF_NO_GROUP;
diff --git a/mm/mprotect.c b/mm/mprotect.c
index 44727811bf4c..88584838e704 100644
--- a/mm/mprotect.c
+++ b/mm/mprotect.c
@@ -75,6 +75,7 @@ static unsigned long change_pte_range(struct vm_area_struct *vma, pmd_t *pmd,
                 oldpte = *pte;
                 if (pte_present(oldpte)) {
                         pte_t ptent;
+                        bool preserve_write = prot_numa && pte_write(oldpte);
 
                         /*
                          * Avoid trapping faults against the zero or KSM
@@ -94,6 +95,8 @@ static unsigned long change_pte_range(struct vm_area_struct *vma, pmd_t *pmd,
 
                         ptent = ptep_modify_prot_start(mm, addr, pte);
                         ptent = pte_modify(ptent, newprot);
+                        if (preserve_write)
+                                ptent = pte_mkwrite(ptent);
 
                         /* Avoid taking write faults for known dirty pages */
                         if (dirty_accountable && pte_dirty(ptent) &&