struct page *base)
 {
        pte_t *pbase = (pte_t *)page_address(base);
-       unsigned long pfn, pfninc = 1;
+       unsigned long ref_pfn, pfn, pfninc = 1;
        unsigned int i, level;
        pte_t *tmp;
        pgprot_t ref_prot;
        }
 
        paravirt_alloc_pte(&init_mm, page_to_pfn(base));
-       ref_prot = pte_pgprot(pte_clrhuge(*kpte));
 
-       /* promote PAT bit to correct position */
-       if (level == PG_LEVEL_2M)
+       switch (level) {
+       case PG_LEVEL_2M:
+               ref_prot = pmd_pgprot(*(pmd_t *)kpte);
+               /* clear PSE and promote PAT bit to correct position */
                ref_prot = pgprot_large_2_4k(ref_prot);
+               ref_pfn = pmd_pfn(*(pmd_t *)kpte);
+               break;
 
-#ifdef CONFIG_X86_64
-       if (level == PG_LEVEL_1G) {
+       case PG_LEVEL_1G:
+               ref_prot = pud_pgprot(*(pud_t *)kpte);
+               ref_pfn = pud_pfn(*(pud_t *)kpte);
                pfninc = PMD_PAGE_SIZE >> PAGE_SHIFT;
+
                /*
-                * Set the PSE flags only if the PRESENT flag is set
+                * Clear the PSE flags if the PRESENT flag is not set
                 * otherwise pmd_present/pmd_huge will return true
                 * even on a non present pmd.
                 */
-               if (pgprot_val(ref_prot) & _PAGE_PRESENT)
-                       pgprot_val(ref_prot) |= _PAGE_PSE;
-               else
+               if (!(pgprot_val(ref_prot) & _PAGE_PRESENT))
                        pgprot_val(ref_prot) &= ~_PAGE_PSE;
+               break;
+
+       default:
+               spin_unlock(&pgd_lock);
+               return 1;
        }
-#endif
 
        /*
         * Set the GLOBAL flags only if the PRESENT flag is set
        /*
         * Get the target pfn from the original entry:
         */
-       pfn = pte_pfn(*kpte);
+       pfn = ref_pfn;
        for (i = 0; i < PTRS_PER_PTE; i++, pfn += pfninc)
                set_pte(&pbase[i], pfn_pte(pfn, canon_pgprot(ref_prot)));