8 files changed, 84 insertions, 9 deletions

diff --git a/arch/x86/include/asm/pgtable.h b/arch/x86/include/asm/pgtable.h
index 7dc305a46058..bd0518a7f197 100644
--- a/arch/x86/include/asm/pgtable.h
+++ b/arch/x86/include/asm/pgtable.h
@@ -314,6 +314,21 @@ static inline pmd_t pmd_mksoft_dirty(pmd_t pmd)
         return pmd_set_flags(pmd, _PAGE_SOFT_DIRTY);
 }
 
+static inline pte_t pte_swp_mksoft_dirty(pte_t pte)
+{
+        return pte_set_flags(pte, _PAGE_SWP_SOFT_DIRTY);
+}
+
+static inline int pte_swp_soft_dirty(pte_t pte)
+{
+        return pte_flags(pte) & _PAGE_SWP_SOFT_DIRTY;
+}
+
+static inline pte_t pte_swp_clear_soft_dirty(pte_t pte)
+{
+        return pte_clear_flags(pte, _PAGE_SWP_SOFT_DIRTY);
+}
+
 /*
  * Mask out unsupported bits in a present pgprot.  Non-present pgprots
  * can use those bits for other purposes, so leave them be.
diff --git a/arch/x86/include/asm/pgtable_types.h b/arch/x86/include/asm/pgtable_types.h
index c98ac63aae48..5e8442f178f9 100644
--- a/arch/x86/include/asm/pgtable_types.h
+++ b/arch/x86/include/asm/pgtable_types.h
@@ -67,6 +67,19 @@
 #define _PAGE_SOFT_DIRTY        (_AT(pteval_t, 0))
 #endif
 
+/*
+ * Tracking soft dirty bit when a page goes to a swap is tricky.
+ * We need a bit which can be stored in pte _and_ not conflict
+ * with swap entry format. On x86 bits 6 and 7 are *not* involved
+ * into swap entry computation, but bit 6 is used for nonlinear
+ * file mapping, so we borrow bit 7 for soft dirty tracking.
+ */
+#ifdef CONFIG_MEM_SOFT_DIRTY
+#define _PAGE_SWP_SOFT_DIRTY    _PAGE_PSE
+#else
+#define _PAGE_SWP_SOFT_DIRTY    (_AT(pteval_t, 0))
+#endif
+
 #if defined(CONFIG_X86_64) || defined(CONFIG_X86_PAE)
 #define _PAGE_NX        (_AT(pteval_t, 1) << _PAGE_BIT_NX)
 #else
diff --git a/fs/proc/task_mmu.c b/fs/proc/task_mmu.c
index dbf61f6174f0..e2d9bdce5e7e 100644
--- a/fs/proc/task_mmu.c
+++ b/fs/proc/task_mmu.c
@@ -730,8 +730,14 @@ static inline void clear_soft_dirty(struct vm_area_struct *vma,
          * of how soft-dirty works.
          */
         pte_t ptent = *pte;
-        ptent = pte_wrprotect(ptent);
-        ptent = pte_clear_flags(ptent, _PAGE_SOFT_DIRTY);
+
+        if (pte_present(ptent)) {
+                ptent = pte_wrprotect(ptent);
+                ptent = pte_clear_flags(ptent, _PAGE_SOFT_DIRTY);
+        } else if (is_swap_pte(ptent)) {
+                ptent = pte_swp_clear_soft_dirty(ptent);
+        }
+
         set_pte_at(vma->vm_mm, addr, pte, ptent);
 #endif
 }
@@ -752,14 +758,15 @@ static int clear_refs_pte_range(pmd_t *pmd, unsigned long addr,
         pte = pte_offset_map_lock(vma->vm_mm, pmd, addr, &ptl);
         for (; addr != end; pte++, addr += PAGE_SIZE) {
                 ptent = *pte;
-                if (!pte_present(ptent))
-                        continue;
 
                 if (cp->type == CLEAR_REFS_SOFT_DIRTY) {
                         clear_soft_dirty(vma, addr, pte);
                         continue;
                 }
 
+                if (!pte_present(ptent))
+                        continue;
+
                 page = vm_normal_page(vma, addr, ptent);
                 if (!page)
                         continue;
@@ -930,8 +937,10 @@ static void pte_to_pagemap_entry(pagemap_entry_t *pme, struct pagemapread *pm,
                 flags = PM_PRESENT;
                 page = vm_normal_page(vma, addr, pte);
         } else if (is_swap_pte(pte)) {
-                swp_entry_t entry = pte_to_swp_entry(pte);
-
+                swp_entry_t entry;
+                if (pte_swp_soft_dirty(pte))
+                        flags2 |= __PM_SOFT_DIRTY;
+                entry = pte_to_swp_entry(pte);
                 frame = swp_type(entry) |
                         (swp_offset(entry) << MAX_SWAPFILES_SHIFT);
                 flags = PM_SWAP;
diff --git a/include/asm-generic/pgtable.h b/include/asm-generic/pgtable.h
index 2f47ade1b567..2a7e0d10ad9a 100644
--- a/include/asm-generic/pgtable.h
+++ b/include/asm-generic/pgtable.h
@@ -417,6 +417,21 @@ static inline pmd_t pmd_mksoft_dirty(pmd_t pmd)
 {
         return pmd;
 }
+
+static inline pte_t pte_swp_mksoft_dirty(pte_t pte)
+{
+        return pte;
+}
+
+static inline int pte_swp_soft_dirty(pte_t pte)
+{
+        return 0;
+}
+
+static inline pte_t pte_swp_clear_soft_dirty(pte_t pte)
+{
+        return pte;
+}
 #endif
 
 #ifndef __HAVE_PFNMAP_TRACKING
diff --git a/include/linux/swapops.h b/include/linux/swapops.h
index c5fd30d2a415..8d4fa82bfb91 100644
--- a/include/linux/swapops.h
+++ b/include/linux/swapops.h
@@ -67,6 +67,8 @@ static inline swp_entry_t pte_to_swp_entry(pte_t pte)
         swp_entry_t arch_entry;
 
         BUG_ON(pte_file(pte));
+        if (pte_swp_soft_dirty(pte))
+                pte = pte_swp_clear_soft_dirty(pte);
         arch_entry = __pte_to_swp_entry(pte);
         return swp_entry(__swp_type(arch_entry), __swp_offset(arch_entry));
 }
diff --git a/mm/memory.c b/mm/memory.c
index 1ce2e2a734fc..e98ecad2b9c8 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -3115,6 +3115,8 @@ static int do_swap_page(struct mm_struct *mm, struct vm_area_struct *vma,
                 exclusive = 1;
         }
         flush_icache_page(vma, page);
+        if (pte_swp_soft_dirty(orig_pte))
+                pte = pte_mksoft_dirty(pte);
         set_pte_at(mm, address, page_table, pte);
         if (page == swapcache)
                 do_page_add_anon_rmap(page, vma, address, exclusive);
diff --git a/mm/rmap.c b/mm/rmap.c
index cd356df4f71a..83325b80142b 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1236,6 +1236,7 @@ int try_to_unmap_one(struct page *page, struct vm_area_struct *vma,
                            swp_entry_to_pte(make_hwpoison_entry(page)));
         } else if (PageAnon(page)) {
                 swp_entry_t entry = { .val = page_private(page) };
+                pte_t swp_pte;
 
                 if (PageSwapCache(page)) {
                         /*
@@ -1264,7 +1265,10 @@ int try_to_unmap_one(struct page *page, struct vm_area_struct *vma,
                         BUG_ON(TTU_ACTION(flags) != TTU_MIGRATION);
                         entry = make_migration_entry(page, pte_write(pteval));
                 }
-                set_pte_at(mm, address, pte, swp_entry_to_pte(entry));
+                swp_pte = swp_entry_to_pte(entry);
+                if (pte_soft_dirty(pteval))
+                        swp_pte = pte_swp_mksoft_dirty(swp_pte);
+                set_pte_at(mm, address, pte, swp_pte);
                 BUG_ON(pte_file(*pte));
         } else if (IS_ENABLED(CONFIG_MIGRATION) &&
                    (TTU_ACTION(flags) == TTU_MIGRATION)) {
diff --git a/mm/swapfile.c b/mm/swapfile.c
index 36af6eeaa67e..6cf2e60983b7 100644
--- a/mm/swapfile.c
+++ b/mm/swapfile.c
@@ -866,6 +866,21 @@ unsigned int count_swap_pages(int type, int free)
 }
 #endif /* CONFIG_HIBERNATION */
 
+static inline int maybe_same_pte(pte_t pte, pte_t swp_pte)
+{
+#ifdef CONFIG_MEM_SOFT_DIRTY
+        /*
+         * When pte keeps soft dirty bit the pte generated
+         * from swap entry does not has it, still it's same
+         * pte from logical point of view.
+         */
+        pte_t swp_pte_dirty = pte_swp_mksoft_dirty(swp_pte);
+        return pte_same(pte, swp_pte) || pte_same(pte, swp_pte_dirty);
+#else
+        return pte_same(pte, swp_pte);
+#endif
+}
+
 /*
  * No need to decide whether this PTE shares the swap entry with others,
  * just let do_wp_page work it out if a write is requested later - to
@@ -892,7 +907,7 @@ static int unuse_pte(struct vm_area_struct *vma, pmd_t *pmd,
         }
 
         pte = pte_offset_map_lock(vma->vm_mm, pmd, addr, &ptl);
-        if (unlikely(!pte_same(*pte, swp_entry_to_pte(entry)))) {
+        if (unlikely(!maybe_same_pte(*pte, swp_entry_to_pte(entry)))) {
                 mem_cgroup_cancel_charge_swapin(memcg);
                 ret = 0;
                 goto out;
@@ -947,7 +962,7 @@ static int unuse_pte_range(struct vm_area_struct *vma, pmd_t *pmd,
                  * swapoff spends a _lot_ of time in this loop!
                  * Test inline before going to call unuse_pte.
                  */
-                if (unlikely(pte_same(*pte, swp_pte))) {
+                if (unlikely(maybe_same_pte(*pte, swp_pte))) {
                         pte_unmap(pte);
                         ret = unuse_pte(vma, pmd, addr, entry, page);
                         if (ret)