#include <linux/init.h>
 #include <linux/module.h>
 #include <linux/io.h>
+#include <linux/of.h>
+#include <linux/of_address.h>
 
 #include "hardware.h"
 #include "common.h"
 
 #define IIM_SREV 0x24
 
+static u32 imx5_read_srev_reg(const char *compat)
+{
+       void __iomem *iim_base;
+       struct device_node *np;
+       u32 srev;
+
+       np = of_find_compatible_node(NULL, NULL, compat);
+       iim_base = of_iomap(np, 0);
+       WARN_ON(!iim_base);
+
+       srev = readl(iim_base + IIM_SREV) & 0xff;
+
+       iounmap(iim_base);
+
+       return srev;
+}
+
 static int get_mx51_srev(void)
 {
-       void __iomem *iim_base = MX51_IO_ADDRESS(MX51_IIM_BASE_ADDR);
-       u32 rev = readl(iim_base + IIM_SREV) & 0xff;
+       u32 rev = imx5_read_srev_reg("fsl,imx51-iim");
 
        switch (rev) {
        case 0x0:
 
 static int get_mx53_srev(void)
 {
-       void __iomem *iim_base = MX51_IO_ADDRESS(MX53_IIM_BASE_ADDR);
-       u32 rev = readl(iim_base + IIM_SREV) & 0xff;
+       u32 rev = imx5_read_srev_reg("fsl,imx53-iim");
 
        switch (rev) {
        case 0x0: