#ifndef CONFIG_TRANSPARENT_HUGEPAGE
 #define pmd_large(pmd)         0
 #endif
-pte_t *__find_linux_pte_or_hugepte(pgd_t *pgdir, unsigned long ea,
-                                  bool *is_thp, unsigned *shift);
-static inline pte_t *find_linux_pte_or_hugepte(pgd_t *pgdir, unsigned long ea,
-                                              bool *is_thp, unsigned *shift)
-{
-       VM_WARN(!arch_irqs_disabled(),
-               "%s called with irq enabled\n", __func__);
-       return __find_linux_pte_or_hugepte(pgdir, ea, is_thp, shift);
-}
 
+/* can we use this in kvm */
 unsigned long vmalloc_to_phys(void *vmalloc_addr);
 
 void pgtable_cache_add(unsigned shift, void (*ctor)(void *));
 
--- /dev/null
+#ifndef _ASM_POWERPC_PTE_WALK_H
+#define _ASM_POWERPC_PTE_WALK_H
+
+#include <linux/sched.h>
+
+/* Don't use this directly */
+extern pte_t *__find_linux_pte(pgd_t *pgdir, unsigned long ea,
+                              bool *is_thp, unsigned *hshift);
+
+static inline pte_t *find_linux_pte(pgd_t *pgdir, unsigned long ea,
+                                   bool *is_thp, unsigned *hshift)
+{
+       VM_WARN(!arch_irqs_disabled(), "%s called with irq enabled\n", __func__);
+       return __find_linux_pte(pgdir, ea, is_thp, hshift);
+}
+
+static inline pte_t *find_init_mm_pte(unsigned long ea, unsigned *hshift)
+{
+       pgd_t *pgdir = init_mm.pgd;
+       return __find_linux_pte(pgdir, ea, NULL, hshift);
+}
+/*
+ * This is what we should always use. Any other lockless page table lookup needs
+ * careful audit against THP split.
+ */
+static inline pte_t *find_current_mm_pte(pgd_t *pgdir, unsigned long ea,
+                                        bool *is_thp, unsigned *hshift)
+{
+       VM_WARN(!arch_irqs_disabled(), "%s called with irq enabled\n", __func__);
+       VM_WARN(pgdir != current->mm->pgd,
+               "%s lock less page table lookup called on wrong mm\n", __func__);
+       return __find_linux_pte(pgdir, ea, is_thp, hshift);
+}
+
+#endif /* _ASM_POWERPC_PTE_WALK_H */
 
 #include <asm/machdep.h>
 #include <asm/ppc-pci.h>
 #include <asm/rtas.h>
+#include <asm/pte-walk.h>
 
 
 /** Overview:
         * worried about _PAGE_SPLITTING/collapse. Also we will not hit
         * page table free, because of init_mm.
         */
-       ptep = __find_linux_pte_or_hugepte(init_mm.pgd, token,
-                                          NULL, &hugepage_shift);
+       ptep = find_init_mm_pte(token, &hugepage_shift);
        if (!ptep)
                return token;
        WARN_ON(hugepage_shift);
 
 #include <asm/pgtable.h>
 #include <asm/ppc-pci.h>
 #include <asm/io-workarounds.h>
+#include <asm/pte-walk.h>
+
 
 #define IOWA_MAX_BUS   8
 
                 * We won't find huge pages here (iomem). Also can't hit
                 * a page table free due to init_mm
                 */
-               ptep = __find_linux_pte_or_hugepte(init_mm.pgd, vaddr,
-                                                  NULL, &hugepage_shift);
+               ptep = find_init_mm_pte(vaddr, &hugepage_shift);
                if (ptep == NULL)
                        paddr = 0;
                else {
 
 #include <asm/synch.h>
 #include <asm/ppc-opcode.h>
 #include <asm/cputable.h>
+#include <asm/pte-walk.h>
 
 #include "trace_hv.h"
 
                         * hugepage split and collapse.
                         */
                        local_irq_save(flags);
-                       ptep = find_linux_pte_or_hugepte(current->mm->pgd,
-                                                        hva, NULL, NULL);
+                       ptep = find_current_mm_pte(current->mm->pgd,
+                                                  hva, NULL, NULL);
                        if (ptep) {
                                pte = kvmppc_read_update_linux_pte(ptep, 1);
                                if (__pte_write(pte))
 
 #include <asm/mmu.h>
 #include <asm/pgtable.h>
 #include <asm/pgalloc.h>
+#include <asm/pte-walk.h>
 
 /*
  * Supported radix tree geometry.
                if (writing)
                        pgflags |= _PAGE_DIRTY;
                local_irq_save(flags);
-               ptep = __find_linux_pte_or_hugepte(current->mm->pgd, hva,
-                                                  NULL, NULL);
+               ptep = find_current_mm_pte(current->mm->pgd, hva, NULL, NULL);
                if (ptep) {
                        pte = READ_ONCE(*ptep);
                        if (pte_present(pte) &&
                                spin_unlock(&kvm->mmu_lock);
                                return RESUME_GUEST;
                        }
-                       ptep = __find_linux_pte_or_hugepte(kvm->arch.pgtable,
-                                                       gpa, NULL, &shift);
+                       /*
+                        * We are walking the secondary page table here. We can do this
+                        * without disabling irq.
+                        */
+                       ptep = __find_linux_pte(kvm->arch.pgtable,
+                                               gpa, NULL, &shift);
                        if (ptep && pte_present(*ptep)) {
                                kvmppc_radix_update_pte(kvm, ptep, 0, pgflags,
                                                        gpa, shift);
                        pgflags |= _PAGE_WRITE;
                } else {
                        local_irq_save(flags);
-                       ptep = __find_linux_pte_or_hugepte(current->mm->pgd,
-                                                       hva, NULL, NULL);
+                       ptep = find_current_mm_pte(current->mm->pgd,
+                                                  hva, NULL, NULL);
                        if (ptep && pte_write(*ptep) && pte_dirty(*ptep))
                                pgflags |= _PAGE_WRITE;
                        local_irq_restore(flags);
        unsigned int shift;
        unsigned long old;
 
-       ptep = __find_linux_pte_or_hugepte(kvm->arch.pgtable, gpa,
-                                          NULL, &shift);
+       ptep = __find_linux_pte(kvm->arch.pgtable, gpa, NULL, &shift);
        if (ptep && pte_present(*ptep)) {
                old = kvmppc_radix_update_pte(kvm, ptep, _PAGE_PRESENT, 0,
                                              gpa, shift);
        unsigned int shift;
        int ref = 0;
 
-       ptep = __find_linux_pte_or_hugepte(kvm->arch.pgtable, gpa,
-                                          NULL, &shift);
+       ptep = __find_linux_pte(kvm->arch.pgtable, gpa, NULL, &shift);
        if (ptep && pte_present(*ptep) && pte_young(*ptep)) {
                kvmppc_radix_update_pte(kvm, ptep, _PAGE_ACCESSED, 0,
                                        gpa, shift);
        unsigned int shift;
        int ref = 0;
 
-       ptep = __find_linux_pte_or_hugepte(kvm->arch.pgtable, gpa,
-                                          NULL, &shift);
+       ptep = __find_linux_pte(kvm->arch.pgtable, gpa, NULL, &shift);
        if (ptep && pte_present(*ptep) && pte_young(*ptep))
                ref = 1;
        return ref;
        unsigned int shift;
        int ret = 0;
 
-       ptep = __find_linux_pte_or_hugepte(kvm->arch.pgtable, gpa,
-                                          NULL, &shift);
+       ptep = __find_linux_pte(kvm->arch.pgtable, gpa, NULL, &shift);
        if (ptep && pte_present(*ptep) && pte_dirty(*ptep)) {
                ret = 1;
                if (shift)
 
 #include <asm/udbg.h>
 #include <asm/iommu.h>
 #include <asm/tce.h>
+#include <asm/pte-walk.h>
 
 #ifdef CONFIG_BUG
 
        pte_t *ptep, pte;
        unsigned shift = 0;
 
-       ptep = __find_linux_pte_or_hugepte(vcpu->arch.pgdir, ua, NULL, &shift);
+       /*
+        * Called in real mode with MSR_EE = 0. We are safe here.
+        * It is ok to do the lookup with arch.pgdir here, because
+        * we are doing this on secondary cpus and current task there
+        * is not the hypervisor. Also this is safe against THP in the
+        * host, because an IPI to primary thread will wait for the secondary
+        * to exit which will agains result in the below page table walk
+        * to finish.
+        */
+       ptep = __find_linux_pte(vcpu->arch.pgdir, ua, NULL, &shift);
        if (!ptep || !pte_present(*ptep))
                return -ENXIO;
        pte = *ptep;
 
 #include <asm/hvcall.h>
 #include <asm/synch.h>
 #include <asm/ppc-opcode.h>
+#include <asm/pte-walk.h>
 
 /* Translate address of a vmalloc'd thing to a linear map address */
 static void *real_vmalloc_addr(void *x)
        /*
         * assume we don't have huge pages in vmalloc space...
         * So don't worry about THP collapse/split. Called
-        * Only in realmode, hence won't need irq_save/restore.
+        * Only in realmode with MSR_EE = 0, hence won't need irq_save/restore.
         */
-       p = __find_linux_pte_or_hugepte(swapper_pg_dir, addr, NULL, NULL);
+       p = find_init_mm_pte(addr, NULL);
        if (!p || !pte_present(*p))
                return NULL;
        addr = (pte_pfn(*p) << PAGE_SHIFT) | (addr & ~PAGE_MASK);
         * If we had a page table table change after lookup, we would
         * retry via mmu_notifier_retry.
         */
-       if (realmode)
-               ptep = __find_linux_pte_or_hugepte(pgdir, hva, NULL,
-                                                  &hpage_shift);
-       else {
+       if (!realmode)
                local_irq_save(irq_flags);
-               ptep = find_linux_pte_or_hugepte(pgdir, hva, NULL,
-                                                &hpage_shift);
-       }
+       /*
+        * If called in real mode we have MSR_EE = 0. Otherwise
+        * we disable irq above.
+        */
+       ptep = __find_linux_pte(pgdir, hva, NULL, &hpage_shift);
        if (ptep) {
                pte_t pte;
                unsigned int host_pte_size;
 
 #include <linux/vmalloc.h>
 #include <linux/hugetlb.h>
 #include <asm/kvm_ppc.h>
+#include <asm/pte-walk.h>
 
 #include "e500.h"
 #include "timing.h"
         * can't run hence pfn won't change.
         */
        local_irq_save(flags);
-       ptep = find_linux_pte_or_hugepte(pgdir, hva, NULL, NULL);
+       ptep = find_linux_pte(pgdir, hva, NULL, NULL);
        if (ptep) {
                pte_t pte = READ_ONCE(*ptep);
 
 
 #include <asm/tm.h>
 #include <asm/trace.h>
 #include <asm/ps3.h>
+#include <asm/pte-walk.h>
 
 #ifdef DEBUG
 #define DBG(fmt...) udbg_printf(fmt)
 #endif /* CONFIG_PPC_64K_PAGES */
 
        /* Get PTE and page size from page tables */
-       ptep = __find_linux_pte_or_hugepte(pgdir, ea, &is_thp, &hugeshift);
+       ptep = find_linux_pte(pgdir, ea, &is_thp, &hugeshift);
        if (ptep == NULL || !pte_present(*ptep)) {
                DBG_LOW(" no PTE !\n");
                rc = 1;
         * THP pages use update_mmu_cache_pmd. We don't do
         * hash preload there. Hence can ignore THP here
         */
-       ptep = find_linux_pte_or_hugepte(pgdir, ea, NULL, &hugepage_shift);
+       ptep = find_current_mm_pte(pgdir, ea, NULL, &hugepage_shift);
        if (!ptep)
                goto out_exit;
 
 
 #include <asm/tlb.h>
 #include <asm/setup.h>
 #include <asm/hugetlb.h>
+#include <asm/pte-walk.h>
+
 
 #ifdef CONFIG_HUGETLB_PAGE
 
 
 pte_t *huge_pte_offset(struct mm_struct *mm, unsigned long addr, unsigned long sz)
 {
-       /* Only called for hugetlbfs pages, hence can ignore THP */
-       return __find_linux_pte_or_hugepte(mm->pgd, addr, NULL, NULL);
+       /*
+        * Only called for hugetlbfs pages, hence can ignore THP and the
+        * irq disabled walk.
+        */
+       return __find_linux_pte(mm->pgd, addr, NULL, NULL);
 }
 
 static int __hugepte_alloc(struct mm_struct *mm, hugepd_t *hpdp,
  * This function need to be called with interrupts disabled. We use this variant
  * when we have MSR[EE] = 0 but the paca->soft_enabled = 1
  */
-
-pte_t *__find_linux_pte_or_hugepte(pgd_t *pgdir, unsigned long ea,
-                                  bool *is_thp, unsigned *shift)
+pte_t *__find_linux_pte(pgd_t *pgdir, unsigned long ea,
+                       bool *is_thp, unsigned *hpage_shift)
 {
        pgd_t pgd, *pgdp;
        pud_t pud, *pudp;
        hugepd_t *hpdp = NULL;
        unsigned pdshift = PGDIR_SHIFT;
 
-       if (shift)
-               *shift = 0;
+       if (hpage_shift)
+               *hpage_shift = 0;
 
        if (is_thp)
                *is_thp = false;
        ret_pte = hugepte_offset(*hpdp, ea, pdshift);
        pdshift = hugepd_shift(*hpdp);
 out:
-       if (shift)
-               *shift = pdshift;
+       if (hpage_shift)
+               *hpage_shift = pdshift;
        return ret_pte;
 }
-EXPORT_SYMBOL_GPL(__find_linux_pte_or_hugepte);
+EXPORT_SYMBOL_GPL(__find_linux_pte);
 
 int gup_hugepte(pte_t *ptep, unsigned long sz, unsigned long addr,
                unsigned long end, int write, struct page **pages, int *nr)
 
 #include <asm/tlbflush.h>
 #include <asm/tlb.h>
 #include <asm/bug.h>
+#include <asm/pte-walk.h>
+
 
 #include <trace/events/thp.h>
 
        local_irq_save(flags);
        arch_enter_lazy_mmu_mode();
        for (; start < end; start += PAGE_SIZE) {
-               pte_t *ptep = find_linux_pte_or_hugepte(mm->pgd, start, &is_thp,
-                                                       &hugepage_shift);
+               pte_t *ptep = find_current_mm_pte(mm->pgd, start, &is_thp,
+                                                 &hugepage_shift);
                unsigned long pte;
 
                if (ptep == NULL)
 
 #ifdef CONFIG_PPC64
 #include "../kernel/ppc32.h"
 #endif
+#include <asm/pte-walk.h>
 
 
 /*
                return -EFAULT;
 
        local_irq_save(flags);
-       ptep = find_linux_pte_or_hugepte(pgdir, addr, NULL, &shift);
+       ptep = find_current_mm_pte(pgdir, addr, NULL, &shift);
        if (!ptep)
                goto err_out;
        if (!shift)