clean_pte_table(pte);
 }
 
-static inline void kvm_set_s2pte_writable(pte_t *pte)
+static inline pte_t kvm_s2pte_mkwrite(pte_t pte)
 {
-       pte_val(*pte) |= L_PTE_S2_RDWR;
+       pte_val(pte) |= L_PTE_S2_RDWR;
+       return pte;
 }
 
-static inline void kvm_set_s2pmd_writable(pmd_t *pmd)
+static inline pmd_t kvm_s2pmd_mkwrite(pmd_t pmd)
 {
-       pmd_val(*pmd) |= L_PMD_S2_RDWR;
+       pmd_val(pmd) |= L_PMD_S2_RDWR;
+       return pmd;
 }
 
 static inline void kvm_set_s2pte_readonly(pte_t *pte)
 
        return 0;
 }
 
+#ifndef __HAVE_ARCH_PTEP_TEST_AND_CLEAR_YOUNG
+static int stage2_ptep_test_and_clear_young(pte_t *pte)
+{
+       if (pte_young(*pte)) {
+               *pte = pte_mkold(*pte);
+               return 1;
+       }
+       return 0;
+}
+#else
+static int stage2_ptep_test_and_clear_young(pte_t *pte)
+{
+       return __ptep_test_and_clear_young(pte);
+}
+#endif
+
+static int stage2_pmdp_test_and_clear_young(pmd_t *pmd)
+{
+       return stage2_ptep_test_and_clear_young((pte_t *)pmd);
+}
+
 /**
  * kvm_phys_addr_ioremap - map a device range to guest IPA
  *
                pte_t pte = pfn_pte(pfn, PAGE_S2_DEVICE);
 
                if (writable)
-                       kvm_set_s2pte_writable(&pte);
+                       pte = kvm_s2pte_mkwrite(pte);
 
                ret = mmu_topup_memory_cache(&cache, KVM_MMU_CACHE_MIN_PAGES,
                                                KVM_NR_MEM_OBJS);
                pmd_t new_pmd = pfn_pmd(pfn, mem_type);
                new_pmd = pmd_mkhuge(new_pmd);
                if (writable) {
-                       kvm_set_s2pmd_writable(&new_pmd);
+                       new_pmd = kvm_s2pmd_mkwrite(new_pmd);
                        kvm_set_pfn_dirty(pfn);
                }
                coherent_cache_guest_page(vcpu, pfn, PMD_SIZE, fault_ipa_uncached);
                pte_t new_pte = pfn_pte(pfn, mem_type);
 
                if (writable) {
-                       kvm_set_s2pte_writable(&new_pte);
+                       new_pte = kvm_s2pte_mkwrite(new_pte);
                        kvm_set_pfn_dirty(pfn);
                        mark_page_dirty(kvm, gfn);
                }
  * Resolve the access fault by making the page young again.
  * Note that because the faulting entry is guaranteed not to be
  * cached in the TLB, we don't need to invalidate anything.
+ * Only the HW Access Flag updates are supported for Stage 2 (no DBM),
+ * so there is no need for atomic (pte|pmd)_mkyoung operations.
  */
 static void handle_access_fault(struct kvm_vcpu *vcpu, phys_addr_t fault_ipa)
 {
        if (!pmd || pmd_none(*pmd))     /* Nothing there */
                return 0;
 
-       if (pmd_thp_or_huge(*pmd)) {    /* THP, HugeTLB */
-               if (pmd_young(*pmd)) {
-                       *pmd = pmd_mkold(*pmd);
-                       return 1;
-               }
-
-               return 0;
-       }
+       if (pmd_thp_or_huge(*pmd))      /* THP, HugeTLB */
+               return stage2_pmdp_test_and_clear_young(pmd);
 
        pte = pte_offset_kernel(pmd, gpa);
        if (pte_none(*pte))
                return 0;
 
-       if (pte_young(*pte)) {
-               *pte = pte_mkold(*pte); /* Just a page... */
-               return 1;
-       }
-
-       return 0;
+       return stage2_ptep_test_and_clear_young(pte);
 }
 
 static int kvm_test_age_hva_handler(struct kvm *kvm, gpa_t gpa, void *data)
 
 
 /* VTCR_EL2 Registers bits */
 #define VTCR_EL2_RES1          (1 << 31)
+#define VTCR_EL2_HD            (1 << 22)
+#define VTCR_EL2_HA            (1 << 21)
 #define VTCR_EL2_PS_MASK       TCR_EL2_PS_MASK
 #define VTCR_EL2_TG0_MASK      TCR_TG0_MASK
 #define VTCR_EL2_TG0_4K                TCR_TG0_4K
 
 static inline void kvm_clean_pte(pte_t *pte) {}
 static inline void kvm_clean_pte_entry(pte_t *pte) {}
 
-static inline void kvm_set_s2pte_writable(pte_t *pte)
+static inline pte_t kvm_s2pte_mkwrite(pte_t pte)
 {
-       pte_val(*pte) |= PTE_S2_RDWR;
+       pte_val(pte) |= PTE_S2_RDWR;
+       return pte;
 }
 
-static inline void kvm_set_s2pmd_writable(pmd_t *pmd)
+static inline pmd_t kvm_s2pmd_mkwrite(pmd_t pmd)
 {
-       pmd_val(*pmd) |= PMD_S2_RDWR;
+       pmd_val(pmd) |= PMD_S2_RDWR;
+       return pmd;
 }
 
 static inline void kvm_set_s2pte_readonly(pte_t *pte)
 {
-       pte_val(*pte) = (pte_val(*pte) & ~PTE_S2_RDWR) | PTE_S2_RDONLY;
+       pteval_t pteval;
+       unsigned long tmp;
+
+       asm volatile("//        kvm_set_s2pte_readonly\n"
+       "       prfm    pstl1strm, %2\n"
+       "1:     ldxr    %0, %2\n"
+       "       and     %0, %0, %3              // clear PTE_S2_RDWR\n"
+       "       orr     %0, %0, %4              // set PTE_S2_RDONLY\n"
+       "       stxr    %w1, %0, %2\n"
+       "       cbnz    %w1, 1b\n"
+       : "=&r" (pteval), "=&r" (tmp), "+Q" (pte_val(*pte))
+       : "L" (~PTE_S2_RDWR), "L" (PTE_S2_RDONLY));
 }
 
 static inline bool kvm_s2pte_readonly(pte_t *pte)
 
 static inline void kvm_set_s2pmd_readonly(pmd_t *pmd)
 {
-       pmd_val(*pmd) = (pmd_val(*pmd) & ~PMD_S2_RDWR) | PMD_S2_RDONLY;
+       kvm_set_s2pte_readonly((pte_t *)pmd);
 }
 
 static inline bool kvm_s2pmd_readonly(pmd_t *pmd)
 {
-       return (pmd_val(*pmd) & PMD_S2_RDWR) == PMD_S2_RDONLY;
+       return kvm_s2pte_readonly((pte_t *)pmd);
 }
 
 static inline bool kvm_page_empty(void *ptr)
 
  * 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)
+static inline int __ptep_test_and_clear_young(pte_t *ptep)
 {
        pteval_t pteval;
        unsigned int tmp, res;
 
-       asm volatile("//        ptep_test_and_clear_young\n"
+       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"
        return res;
 }
 
+static inline int ptep_test_and_clear_young(struct vm_area_struct *vma,
+                                           unsigned long address,
+                                           pte_t *ptep)
+{
+       return __ptep_test_and_clear_young(ptep);
+}
+
 #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,
 
         */
        val |= 64 - (parange > 40 ? 40 : parange);
 
+       /*
+        * Check the availability of Hardware Access Flag / Dirty Bit
+        * Management in ID_AA64MMFR1_EL1 and enable the feature in VTCR_EL2.
+        */
+       tmp = (read_sysreg(id_aa64mmfr1_el1) >> ID_AA64MMFR1_HADBS_SHIFT) & 0xf;
+       if (IS_ENABLED(CONFIG_ARM64_HW_AFDBM) && tmp)
+               val |= VTCR_EL2_HA;
+
        /*
         * Read the VMIDBits bits from ID_AA64MMFR1_EL1 and set the VS
         * bit in VTCR_EL2.