return 0;
 }
 
-static bool stage2_is_exec(struct kvm *kvm, phys_addr_t addr)
+/*
+ * stage2_get_leaf_entry - walk the stage2 VM page tables and return
+ * true if a valid and present leaf-entry is found. A pointer to the
+ * leaf-entry is returned in the appropriate level variable - pudpp,
+ * pmdpp, ptepp.
+ */
+static bool stage2_get_leaf_entry(struct kvm *kvm, phys_addr_t addr,
+                                 pud_t **pudpp, pmd_t **pmdpp, pte_t **ptepp)
 {
+       pud_t *pudp;
        pmd_t *pmdp;
        pte_t *ptep;
 
-       pmdp = stage2_get_pmd(kvm, NULL, addr);
+       *pudpp = NULL;
+       *pmdpp = NULL;
+       *ptepp = NULL;
+
+       pudp = stage2_get_pud(kvm, NULL, addr);
+       if (!pudp || stage2_pud_none(kvm, *pudp) || !stage2_pud_present(kvm, *pudp))
+               return false;
+
+       if (stage2_pud_huge(kvm, *pudp)) {
+               *pudpp = pudp;
+               return true;
+       }
+
+       pmdp = stage2_pmd_offset(kvm, pudp, addr);
        if (!pmdp || pmd_none(*pmdp) || !pmd_present(*pmdp))
                return false;
 
-       if (pmd_thp_or_huge(*pmdp))
-               return kvm_s2pmd_exec(pmdp);
+       if (pmd_thp_or_huge(*pmdp)) {
+               *pmdpp = pmdp;
+               return true;
+       }
 
        ptep = pte_offset_kernel(pmdp, addr);
        if (!ptep || pte_none(*ptep) || !pte_present(*ptep))
                return false;
 
-       return kvm_s2pte_exec(ptep);
+       *ptepp = ptep;
+       return true;
+}
+
+static bool stage2_is_exec(struct kvm *kvm, phys_addr_t addr)
+{
+       pud_t *pudp;
+       pmd_t *pmdp;
+       pte_t *ptep;
+       bool found;
+
+       found = stage2_get_leaf_entry(kvm, addr, &pudp, &pmdp, &ptep);
+       if (!found)
+               return false;
+
+       if (pudp)
+               return kvm_s2pud_exec(pudp);
+       else if (pmdp)
+               return kvm_s2pmd_exec(pmdp);
+       else
+               return kvm_s2pte_exec(ptep);
 }
 
 static int stage2_set_pte(struct kvm *kvm, struct kvm_mmu_memory_cache *cache,