#ifndef __ASM_PGTABLE_H
 #define __ASM_PGTABLE_H
 
+#include <asm/bug.h>
 #include <asm/proc-fns.h>
 
 #include <asm/memory.h>
 #define PTE_VALID              (_AT(pteval_t, 1) << 0)
 #define PTE_DIRTY              (_AT(pteval_t, 1) << 55)
 #define PTE_SPECIAL            (_AT(pteval_t, 1) << 56)
+#ifdef CONFIG_ARM64_HW_AFDBM
+#define PTE_WRITE              (PTE_DBM)                /* same as DBM */
+#else
 #define PTE_WRITE              (_AT(pteval_t, 1) << 57)
+#endif
 #define PTE_PROT_NONE          (_AT(pteval_t, 1) << 58) /* only when !PTE_VALID */
 
 /*
 #define FIRST_USER_ADDRESS     0UL
 
 #ifndef __ASSEMBLY__
+
+#include <linux/mmdebug.h>
+
 extern void __pte_error(const char *file, int line, unsigned long val);
 extern void __pmd_error(const char *file, int line, unsigned long val);
 extern void __pud_error(const char *file, int line, unsigned long val);
  * The following only work if pte_present(). Undefined behaviour otherwise.
  */
 #define pte_present(pte)       (!!(pte_val(pte) & (PTE_VALID | PTE_PROT_NONE)))
-#define pte_dirty(pte)         (!!(pte_val(pte) & PTE_DIRTY))
 #define pte_young(pte)         (!!(pte_val(pte) & PTE_AF))
 #define pte_special(pte)       (!!(pte_val(pte) & PTE_SPECIAL))
 #define pte_write(pte)         (!!(pte_val(pte) & PTE_WRITE))
 #define pte_exec(pte)          (!(pte_val(pte) & PTE_UXN))
 
+#ifdef CONFIG_ARM64_HW_AFDBM
+#define pte_hw_dirty(pte)      (!(pte_val(pte) & PTE_RDONLY))
+#else
+#define pte_hw_dirty(pte)      (0)
+#endif
+#define pte_sw_dirty(pte)      (!!(pte_val(pte) & PTE_DIRTY))
+#define pte_dirty(pte)         (pte_sw_dirty(pte) || pte_hw_dirty(pte))
+
+#define pte_valid(pte)         (!!(pte_val(pte) && PTE_VALID))
 #define pte_valid_user(pte) \
        ((pte_val(pte) & (PTE_VALID | PTE_USER)) == (PTE_VALID | PTE_USER))
 #define pte_valid_not_user(pte) \
        }
 }
 
+struct mm_struct;
+struct vm_area_struct;
+
 extern void __sync_icache_dcache(pte_t pteval, unsigned long addr);
 
+/*
+ * PTE bits configuration in the presence of hardware Dirty Bit Management
+ * (PTE_WRITE == PTE_DBM):
+ *
+ * Dirty  Writable | PTE_RDONLY  PTE_WRITE  PTE_DIRTY (sw)
+ *   0      0      |   1           0          0
+ *   0      1      |   1           1          0
+ *   1      0      |   1           0          1
+ *   1      1      |   0           1          x
+ *
+ * When hardware DBM is not present, the sofware PTE_DIRTY bit is updated via
+ * the page fault mechanism. Checking the dirty status of a pte becomes:
+ *
+ *   PTE_DIRTY || !PTE_RDONLY
+ */
 static inline void set_pte_at(struct mm_struct *mm, unsigned long addr,
                              pte_t *ptep, pte_t pte)
 {
        if (pte_valid_user(pte)) {
                if (!pte_special(pte) && pte_exec(pte))
                        __sync_icache_dcache(pte, addr);
-               if (pte_dirty(pte) && pte_write(pte))
+               if (pte_sw_dirty(pte) && pte_write(pte))
                        pte_val(pte) &= ~PTE_RDONLY;
                else
                        pte_val(pte) |= PTE_RDONLY;
        }
 
+       /*
+        * If the existing pte is valid, check for potential race with
+        * hardware updates of the pte (ptep_set_access_flags safely changes
+        * valid ptes without going through an invalid entry).
+        */
+       if (IS_ENABLED(CONFIG_DEBUG_VM) && IS_ENABLED(CONFIG_ARM64_HW_AFDBM) &&
+           pte_valid(*ptep)) {
+               BUG_ON(!pte_young(pte));
+               BUG_ON(pte_write(*ptep) && !pte_dirty(pte));
+       }
+
        set_pte(ptep, pte);
 }
 
 {
        const pteval_t mask = PTE_USER | PTE_PXN | PTE_UXN | PTE_RDONLY |
                              PTE_PROT_NONE | PTE_WRITE | PTE_TYPE_MASK;
+       /* preserve the hardware dirty information */
+       if (pte_hw_dirty(pte))
+               newprot |= PTE_DIRTY;
        pte_val(pte) = (pte_val(pte) & ~mask) | (pgprot_val(newprot) & mask);
        return pte;
 }
        return pte_pmd(pte_modify(pmd_pte(pmd), newprot));
 }
 
+#ifdef CONFIG_ARM64_HW_AFDBM
+/*
+ * Atomic pte/pmd modifications.
+ */
+#define __HAVE_ARCH_PTEP_TEST_AND_CLEAR_YOUNG
+static inline int ptep_test_and_clear_young(struct vm_area_struct *vma,
+                                           unsigned long address,
+                                           pte_t *ptep)
+{
+       pteval_t pteval;
+       unsigned int tmp, res;
+
+       asm volatile("//        ptep_test_and_clear_young\n"
+       "       prfm    pstl1strm, %2\n"
+       "1:     ldxr    %0, %2\n"
+       "       ubfx    %w3, %w0, %5, #1        // extract PTE_AF (young)\n"
+       "       and     %0, %0, %4              // clear PTE_AF\n"
+       "       stxr    %w1, %0, %2\n"
+       "       cbnz    %w1, 1b\n"
+       : "=&r" (pteval), "=&r" (tmp), "+Q" (pte_val(*ptep)), "=&r" (res)
+       : "L" (~PTE_AF), "I" (ilog2(PTE_AF)));
+
+       return res;
+}
+
+#ifdef CONFIG_TRANSPARENT_HUGEPAGE
+#define __HAVE_ARCH_PMDP_TEST_AND_CLEAR_YOUNG
+static inline int pmdp_test_and_clear_young(struct vm_area_struct *vma,
+                                           unsigned long address,
+                                           pmd_t *pmdp)
+{
+       return ptep_test_and_clear_young(vma, address, (pte_t *)pmdp);
+}
+#endif /* CONFIG_TRANSPARENT_HUGEPAGE */
+
+#define __HAVE_ARCH_PTEP_GET_AND_CLEAR
+static inline pte_t ptep_get_and_clear(struct mm_struct *mm,
+                                      unsigned long address, pte_t *ptep)
+{
+       pteval_t old_pteval;
+       unsigned int tmp;
+
+       asm volatile("//        ptep_get_and_clear\n"
+       "       prfm    pstl1strm, %2\n"
+       "1:     ldxr    %0, %2\n"
+       "       stxr    %w1, xzr, %2\n"
+       "       cbnz    %w1, 1b\n"
+       : "=&r" (old_pteval), "=&r" (tmp), "+Q" (pte_val(*ptep)));
+
+       return __pte(old_pteval);
+}
+
+#ifdef CONFIG_TRANSPARENT_HUGEPAGE
+#define __HAVE_ARCH_PMDP_GET_AND_CLEAR
+static inline pmd_t pmdp_get_and_clear(struct mm_struct *mm,
+                                      unsigned long address, pmd_t *pmdp)
+{
+       return pte_pmd(ptep_get_and_clear(mm, address, (pte_t *)pmdp));
+}
+#endif /* CONFIG_TRANSPARENT_HUGEPAGE */
+
+/*
+ * ptep_set_wrprotect - mark read-only while trasferring potential hardware
+ * dirty status (PTE_DBM && !PTE_RDONLY) to the software PTE_DIRTY bit.
+ */
+#define __HAVE_ARCH_PTEP_SET_WRPROTECT
+static inline void ptep_set_wrprotect(struct mm_struct *mm, unsigned long address, pte_t *ptep)
+{
+       pteval_t pteval;
+       unsigned long tmp;
+
+       asm volatile("//        ptep_set_wrprotect\n"
+       "       prfm    pstl1strm, %2\n"
+       "1:     ldxr    %0, %2\n"
+       "       tst     %0, %4                  // check for hw dirty (!PTE_RDONLY)\n"
+       "       csel    %1, %3, xzr, eq         // set PTE_DIRTY|PTE_RDONLY if dirty\n"
+       "       orr     %0, %0, %1              // if !dirty, PTE_RDONLY is already set\n"
+       "       and     %0, %0, %5              // clear PTE_WRITE/PTE_DBM\n"
+       "       stxr    %w1, %0, %2\n"
+       "       cbnz    %w1, 1b\n"
+       : "=&r" (pteval), "=&r" (tmp), "+Q" (pte_val(*ptep))
+       : "r" (PTE_DIRTY|PTE_RDONLY), "L" (PTE_RDONLY), "L" (~PTE_WRITE)
+       : "cc");
+}
+
+#ifdef CONFIG_TRANSPARENT_HUGEPAGE
+#define __HAVE_ARCH_PMDP_SET_WRPROTECT
+static inline void pmdp_set_wrprotect(struct mm_struct *mm,
+                                     unsigned long address, pmd_t *pmdp)
+{
+       ptep_set_wrprotect(mm, address, (pte_t *)pmdp);
+}
+#endif
+#endif /* CONFIG_ARM64_HW_AFDBM */
+
 extern pgd_t swapper_pg_dir[PTRS_PER_PGD];
 extern pgd_t idmap_pg_dir[PTRS_PER_PGD];