diff --git a/arch/x86/include/asm/pgtable.h b/arch/x86/include/asm/pgtable.h
index 0996f8a6979a..1053b940b35c 100644
--- a/arch/x86/include/asm/pgtable.h
+++ b/arch/x86/include/asm/pgtable.h
@@ -1148,11 +1148,27 @@ static inline pte_t ptep_get_and_clear_full(struct mm_struct *mm,
 	return pte;
 }
 
-#define __HAVE_ARCH_PTEP_SET_WRPROTECT
-static inline void ptep_set_wrprotect(struct mm_struct *mm,
-				      unsigned long addr, pte_t *ptep)
-{
-	clear_bit(_PAGE_BIT_RW, (unsigned long *)&ptep->pte);
+#define __HAVE_ARCH_PTEP_SET_WRPROTECT_FLUSH
+extern pte_t ptep_clear_flush(struct vm_area_struct *vma,
+			      unsigned long address,
+			      pte_t *ptep);
+static inline void ptep_set_wrprotect_flush(struct vm_area_struct *vma,
+					    unsigned long addr, pte_t *ptep)
+{
+	bool rw;
+
+	rw = test_and_clear_bit(_PAGE_BIT_RW, (unsigned long *)&ptep->pte);
+	if (IS_ENABLED(CONFIG_X86_INTEL_SHADOW_STACK_USER)) {
+		struct mm_struct *mm = vma->vm_mm;
+		pte_t pte;
+
+		if (rw && (atomic_read(&mm->mm_users) > 1))
+			pte = ptep_clear_flush(vma, addr, ptep);
+		else
+			pte = *ptep;
+		pte = pte_move_flags(pte, _PAGE_DIRTY_HW, _PAGE_DIRTY_SW);
+		set_pte_at(mm, addr, ptep, pte);
+	}
 }
 
 #define flush_tlb_fix_spurious_fault(vma, address) do { } while (0)
@@ -1198,11 +1214,33 @@ static inline pud_t pudp_huge_get_and_clear(struct mm_struct *mm,
 	return native_pudp_get_and_clear(pudp);
 }
 
-#define __HAVE_ARCH_PMDP_SET_WRPROTECT
-static inline void pmdp_set_wrprotect(struct mm_struct *mm,
-				      unsigned long addr, pmd_t *pmdp)
+#define __HAVE_ARCH_HUGE_PTEP_SET_WRPROTECT_FLUSH
+static inline void huge_ptep_set_wrprotect_flush(struct vm_area_struct *vma,
+						 unsigned long addr, pte_t *ptep)
 {
-	clear_bit(_PAGE_BIT_RW, (unsigned long *)pmdp);
+	ptep_set_wrprotect_flush(vma, addr, ptep);
+}
+
+#define __HAVE_ARCH_PMDP_SET_WRPROTECT_FLUSH
+extern pmd_t pmdp_huge_clear_flush(struct vm_area_struct *vma,
+				   unsigned long address,
+				   pmd_t *pmdp);
+static inline void pmdp_set_wrprotect_flush(struct vm_area_struct *vma,
+					    unsigned long addr, pmd_t *pmdp)
+{	bool rw;
+
+	rw = test_and_clear_bit(_PAGE_BIT_RW, (unsigned long *)&pmdp);
+	if (IS_ENABLED(CONFIG_X86_INTEL_SHADOW_STACK_USER)) {
+		struct mm_struct *mm = vma->vm_mm;
+		pmd_t pmd;
+
+		if (rw && (atomic_read(&mm->mm_users) > 1))
+			pmd = pmdp_huge_clear_flush(vma, addr, pmdp);
+		else
+			pmd = *pmdp;
+		pmd = pmd_move_flags(pmd, _PAGE_DIRTY_HW, _PAGE_DIRTY_SW);
+		set_pmd_at(mm, addr, pmdp, pmd);
+	}
 }
 
 #define pud_write pud_write
diff --git a/include/asm-generic/pgtable.h b/include/asm-generic/pgtable.h
index 3f6f998509f0..9bcfdfc045bb 100644
--- a/include/asm-generic/pgtable.h
+++ b/include/asm-generic/pgtable.h
@@ -121,6 +121,15 @@ static inline int pmdp_clear_flush_young(struct vm_area_struct *vma,
 #endif /* CONFIG_TRANSPARENT_HUGEPAGE */
 #endif
 
+#ifndef __HAVE_ARCH_HUGE_PTEP_SET_WRPROTECT_FLUSH
+static inline void huge_ptep_set_wrorptect_flush(struct vm_area_struct *vma,
+						 unsigned long addr,
+						 pte_t *ptep)
+{
+	huge_ptep_set_wrprotect(vma->vm_mm, addr, ptep);
+}
+#endif
+
 #ifndef __HAVE_ARCH_PTEP_GET_AND_CLEAR
 static inline pte_t ptep_get_and_clear(struct mm_struct *mm,
 				       unsigned long address,
@@ -226,6 +235,15 @@ static inline void ptep_set_wrprotect(struct mm_struct *mm, unsigned long addres
 }
 #endif
 
+#ifndef __HAVE_ARCH_PTEP_SET_WRPROTECT_FLUSH
+static inline void ptep_set_wrprotect_flush(struct vm_area_struct *vma,
+					    unsigned long address,
+					    pte_t *ptep)
+{
+	ptep_set_wrprotect(vma->vm_mm, address, ptep);
+}
+#endif
+
 #ifndef pte_savedwrite
 #define pte_savedwrite pte_write
 #endif
@@ -266,6 +284,14 @@ static inline void pmdp_set_wrprotect(struct mm_struct *mm,
 }
 #endif /* CONFIG_TRANSPARENT_HUGEPAGE */
 #endif
+#ifndef __HAVE_ARCH_PMDP_SET_WRPROTECT_FLUSH
+static inline void pmdp_set_wrprotect_flush(struct vm_area_struct *vma,
+					    unsigned long address,
+					    pmd_t *pmdp)
+{
+	pmdp_set_wrprotect(vma->vm_mm, address, pmdp);
+}
+#endif
 #ifndef __HAVE_ARCH_PUDP_SET_WRPROTECT
 #ifdef CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD
 static inline void pudp_set_wrprotect(struct mm_struct *mm,
