unsigned long fault_status)
 {
        int ret;
-       bool write_fault, exec_fault, writable, hugetlb = false, force_pte = false;
+       bool write_fault, exec_fault, writable, force_pte = false;
        unsigned long mmu_seq;
        gfn_t gfn = fault_ipa >> PAGE_SHIFT;
        struct kvm *kvm = vcpu->kvm;
        kvm_pfn_t pfn;
        pgprot_t mem_type = PAGE_S2;
        bool logging_active = memslot_is_logging(memslot);
-       unsigned long flags = 0;
+       unsigned long vma_pagesize, flags = 0;
 
        write_fault = kvm_is_write_fault(vcpu);
        exec_fault = kvm_vcpu_trap_is_iabt(vcpu);
                return -EFAULT;
        }
 
-       if (vma_kernel_pagesize(vma) == PMD_SIZE && !logging_active) {
-               hugetlb = true;
+       vma_pagesize = vma_kernel_pagesize(vma);
+       if (vma_pagesize == PMD_SIZE && !logging_active) {
                gfn = (fault_ipa & PMD_MASK) >> PAGE_SHIFT;
        } else {
+               /*
+                * Fallback to PTE if it's not one of the Stage 2
+                * supported hugepage sizes
+                */
+               vma_pagesize = PAGE_SIZE;
+
                /*
                 * Pages belonging to memslots that don't have the same
                 * alignment for userspace and IPA cannot be mapped using
        if (mmu_notifier_retry(kvm, mmu_seq))
                goto out_unlock;
 
-       if (!hugetlb && !force_pte)
-               hugetlb = transparent_hugepage_adjust(&pfn, &fault_ipa);
+       if (vma_pagesize == PAGE_SIZE && !force_pte) {
+               /*
+                * Only PMD_SIZE transparent hugepages(THP) are
+                * currently supported. This code will need to be
+                * updated to support other THP sizes.
+                */
+               if (transparent_hugepage_adjust(&pfn, &fault_ipa))
+                       vma_pagesize = PMD_SIZE;
+       }
+
+       if (writable)
+               kvm_set_pfn_dirty(pfn);
 
-       if (hugetlb) {
+       if (fault_status != FSC_PERM)
+               clean_dcache_guest_page(pfn, vma_pagesize);
+
+       if (exec_fault)
+               invalidate_icache_guest_page(pfn, vma_pagesize);
+
+       if (vma_pagesize == PMD_SIZE) {
                pmd_t new_pmd = pfn_pmd(pfn, mem_type);
                new_pmd = pmd_mkhuge(new_pmd);
-               if (writable) {
+               if (writable)
                        new_pmd = kvm_s2pmd_mkwrite(new_pmd);
-                       kvm_set_pfn_dirty(pfn);
-               }
-
-               if (fault_status != FSC_PERM)
-                       clean_dcache_guest_page(pfn, PMD_SIZE);
 
                if (exec_fault) {
                        new_pmd = kvm_s2pmd_mkexec(new_pmd);
-                       invalidate_icache_guest_page(pfn, PMD_SIZE);
                } else if (fault_status == FSC_PERM) {
                        /* Preserve execute if XN was already cleared */
                        if (stage2_is_exec(kvm, fault_ipa))
 
                if (writable) {
                        new_pte = kvm_s2pte_mkwrite(new_pte);
-                       kvm_set_pfn_dirty(pfn);
                        mark_page_dirty(kvm, gfn);
                }
 
-               if (fault_status != FSC_PERM)
-                       clean_dcache_guest_page(pfn, PAGE_SIZE);
-
                if (exec_fault) {
                        new_pte = kvm_s2pte_mkexec(new_pte);
-                       invalidate_icache_guest_page(pfn, PAGE_SIZE);
                } else if (fault_status == FSC_PERM) {
                        /* Preserve execute if XN was already cleared */
                        if (stage2_is_exec(kvm, fault_ipa))