: "=r" (reg));
 }
 
-static void __iomem *imx3_ioremap(unsigned long phys_addr, size_t size,
-                                 unsigned int mtype)
+static void __iomem *imx3_ioremap_caller(unsigned long phys_addr, size_t size,
+                                        unsigned int mtype, void *caller)
 {
        if (mtype == MT_DEVICE) {
                /*
                        mtype = MT_DEVICE_NONSHARED;
        }
 
-       return __arm_ioremap(phys_addr, size, mtype);
+       return __arm_ioremap_caller(phys_addr, size, mtype, caller);
 }
 
 void imx3_init_l2x0(void)
 {
        mxc_set_cpu_type(MXC_CPU_MX31);
        mxc_arch_reset_init(MX31_IO_ADDRESS(MX31_WDOG_BASE_ADDR));
-       imx_ioremap = imx3_ioremap;
+       arch_ioremap_caller = imx3_ioremap_caller;
        arm_pm_idle = imx3_idle;
 }
 
        mxc_iomux_v3_init(MX35_IO_ADDRESS(MX35_IOMUXC_BASE_ADDR));
        mxc_arch_reset_init(MX35_IO_ADDRESS(MX35_WDOG_BASE_ADDR));
        arm_pm_idle = imx3_idle;
-       imx_ioremap = imx3_ioremap;
+       arch_ioremap_caller = imx3_ioremap_caller;
 }
 
 void __init mx35_init_irq(void)
 
 #define IOMEM(addr)    ((void __force __iomem *)(addr))
 #endif
 
+#define addr_in_module(addr, mod) \
+       ((unsigned long)(addr) - mod ## _BASE_ADDR < mod ## _SIZE)
+
 #define IMX_IO_P2V_MODULE(addr, module)                                        \
        (((addr) - module ## _BASE_ADDR) < module ## _SIZE ?            \
         (addr) - (module ## _BASE_ADDR) + (module ## _BASE_ADDR_VIRT) : 0)
 
 /* Allow IO space to be anywhere in the memory */
 #define IO_SPACE_LIMIT 0xffffffff
 
-#define __arch_ioremap __imx_ioremap
-#define __arch_iounmap __iounmap
-
-#define addr_in_module(addr, mod) \
-       ((unsigned long)(addr) - mod ## _BASE_ADDR < mod ## _SIZE)
-
-extern void __iomem *(*imx_ioremap)(unsigned long, size_t, unsigned int);
-
-static inline void __iomem *
-__imx_ioremap(unsigned long phys_addr, size_t size, unsigned int mtype)
-{
-       if (imx_ioremap != NULL)
-               return imx_ioremap(phys_addr, size, mtype);
-       else
-               return __arm_ioremap(phys_addr, size, mtype);
-}
-
 /* io address mapping macro */
 #define __io(a)                __typesafe_io(a)