#define kvm_pmd_mkhuge(pmd)    pmd_mkhuge(pmd)
 
+/*
+ * The following kvm_*pud*() functions are provided strictly to allow
+ * sharing code with arm64. They should never be called in practice.
+ */
+static inline void kvm_set_s2pud_readonly(pud_t *pud)
+{
+       WARN_ON(1);
+}
+
+static inline bool kvm_s2pud_readonly(pud_t *pud)
+{
+       WARN_ON(1);
+       return false;
+}
+
 static inline pte_t kvm_s2pte_mkwrite(pte_t pte)
 {
        pte_val(pte) |= L_PTE_S2_RDWR;
 
        return !(READ_ONCE(pmd_val(*pmdp)) & PMD_S2_XN);
 }
 
+static inline void kvm_set_s2pud_readonly(pud_t *pudp)
+{
+       kvm_set_s2pte_readonly((pte_t *)pudp);
+}
+
+static inline bool kvm_s2pud_readonly(pud_t *pudp)
+{
+       return kvm_s2pte_readonly((pte_t *)pudp);
+}
+
 #define hyp_pte_table_empty(ptep) kvm_page_empty(ptep)
 
 #ifdef __PAGETABLE_PMD_FOLDED
 
        do {
                next = stage2_pud_addr_end(kvm, addr, end);
                if (!stage2_pud_none(kvm, *pud)) {
-                       /* TODO:PUD not supported, revisit later if supported */
-                       BUG_ON(stage2_pud_huge(kvm, *pud));
-                       stage2_wp_pmds(kvm, pud, addr, next);
+                       if (stage2_pud_huge(kvm, *pud)) {
+                               if (!kvm_s2pud_readonly(pud))
+                                       kvm_set_s2pud_readonly(pud);
+                       } else {
+                               stage2_wp_pmds(kvm, pud, addr, next);
+                       }
                }
        } while (pud++, addr = next, addr != end);
 }
  *
  * Called to start logging dirty pages after memory region
  * KVM_MEM_LOG_DIRTY_PAGES operation is called. After this function returns
- * all present PMD and PTEs are write protected in the memory region.
+ * all present PUD, PMD and PTEs are write protected in the memory region.
  * Afterwards read of dirty page log can be called.
  *
  * Acquires kvm_mmu_lock. Called with kvm->slots_lock mutex acquired,