/*
         * Set the page table
         */
-       void (*switch_mm)(unsigned long pgd_phys, struct mm_struct *mm);
+       void (*switch_mm)(phys_addr_t pgd_phys, struct mm_struct *mm);
        /*
         * Set a possibly extended PTE.  Non-extended PTEs should
         * ignore 'ext'.
 extern void cpu_proc_fin(void);
 extern int cpu_do_idle(void);
 extern void cpu_dcache_clean_area(void *, int);
-extern void cpu_do_switch_mm(unsigned long pgd_phys, struct mm_struct *mm);
+extern void cpu_do_switch_mm(phys_addr_t pgd_phys, struct mm_struct *mm);
 #ifdef CONFIG_ARM_LPAE
 extern void cpu_set_pte_ext(pte_t *ptep, pte_t pte);
 #else
 
 #define TTB_FLAGS_SMP  (TTB_IRGN_WBWA|TTB_S|TTB_RGN_OC_WBWA)
 #define PMD_FLAGS_SMP  (PMD_SECT_WBWA|PMD_SECT_S)
 
+#ifndef __ARMEB__
+#  define rpgdl        r0
+#  define rpgdh        r1
+#else
+#  define rpgdl        r1
+#  define rpgdh        r0
+#endif
+
 /*
  * cpu_v7_switch_mm(pgd_phys, tsk)
  *
  */
 ENTRY(cpu_v7_switch_mm)
 #ifdef CONFIG_MMU
-       mmid    r1, r1                          @ get mm->context.id
-       asid    r3, r1
-       mov     r3, r3, lsl #(48 - 32)          @ ASID
-       mcrr    p15, 0, r0, r3, c2              @ set TTB 0
+       mmid    r2, r2
+       asid    r2, r2
+       orr     rpgdh, rpgdh, r2, lsl #(48 - 32)        @ upper 32-bits of pgd
+       mcrr    p15, 0, rpgdl, rpgdh, c2                @ set TTB 0
        isb
 #endif
        mov     pc, lr