#include <mach/eac.h>
 #include <mach/mmc.h>
 
-#if defined(CONFIG_OMAP_DSP) || defined(CONFIG_OMAP_DSP_MODULE)
-#define OMAP2_MBOX_BASE                IO_ADDRESS(OMAP24XX_MAILBOX_BASE)
+#if defined(CONFIG_OMAP_MBOX_FWK) || defined(CONFIG_OMAP_MBOX_FWK_MODULE)
 
-static struct resource mbox_resources[] = {
+#define MBOX_REG_SIZE  0x120
+
+static struct resource omap2_mbox_resources[] = {
        {
-               .start          = OMAP2_MBOX_BASE,
-               .end            = OMAP2_MBOX_BASE + 0x11f,
+               .start          = OMAP24XX_MAILBOX_BASE,
+               .end            = OMAP24XX_MAILBOX_BASE + MBOX_REG_SIZE - 1,
                .flags          = IORESOURCE_MEM,
        },
        {
        },
 };
 
+static struct resource omap3_mbox_resources[] = {
+       {
+               .start          = OMAP34XX_MAILBOX_BASE,
+               .end            = OMAP34XX_MAILBOX_BASE + MBOX_REG_SIZE - 1,
+               .flags          = IORESOURCE_MEM,
+       },
+       {
+               .start          = INT_24XX_MAIL_U0_MPU,
+               .flags          = IORESOURCE_IRQ,
+       },
+};
+
 static struct platform_device mbox_device = {
        .name           = "mailbox",
        .id             = -1,
-       .num_resources  = ARRAY_SIZE(mbox_resources),
-       .resource       = mbox_resources,
 };
 
 static inline void omap_init_mbox(void)
 {
+       if (cpu_is_omap2420()) {
+               mbox_device.num_resources = ARRAY_SIZE(omap2_mbox_resources);
+               mbox_device.resource = omap2_mbox_resources;
+       } else if (cpu_is_omap3430()) {
+               mbox_device.num_resources = ARRAY_SIZE(omap3_mbox_resources);
+               mbox_device.resource = omap3_mbox_resources;
+       } else {
+               pr_err("%s: platform not supported\n", __func__);
+               return;
+       }
        platform_device_register(&mbox_device);
 }
 #else
 static inline void omap_init_mbox(void) { }
-#endif
+#endif /* CONFIG_OMAP_MBOX_FWK */
 
 #if defined(CONFIG_OMAP_STI)
 
 
 #define MAILBOX_IRQ_NEWMSG(u)          (1 << (2 * (u)))
 #define MAILBOX_IRQ_NOTFULL(u)         (1 << (2 * (u) + 1))
 
-static unsigned long mbox_base;
+static void __iomem *mbox_base;
 
 struct omap_mbox2_fifo {
        unsigned long msg;
 static void omap2_mbox_enable_irq(struct omap_mbox *mbox,
                                  omap_mbox_type_t irq);
 
-static inline unsigned int mbox_read_reg(unsigned int reg)
+static inline unsigned int mbox_read_reg(size_t ofs)
 {
-       return __raw_readl(mbox_base + reg);
+       return __raw_readl(mbox_base + ofs);
 }
 
-static inline void mbox_write_reg(unsigned int val, unsigned int reg)
+static inline void mbox_write_reg(u32 val, size_t ofs)
 {
-       __raw_writel(val, mbox_base + reg);
+       __raw_writel(val, mbox_base + ofs);
 }
 
 /* Mailbox H/W preparations */
 };
 EXPORT_SYMBOL(mbox_dsp_info);
 
-/* IVA */
+#if defined(CONFIG_ARCH_OMAP2420) /* IVA */
 static struct omap_mbox2_priv omap2_mbox_iva_priv = {
        .tx_fifo = {
                .msg            = MAILBOX_MESSAGE(2),
        .ops    = &omap2_mbox_ops,
        .priv   = &omap2_mbox_iva_priv,
 };
+#endif
 
 static int __init omap2_mbox_probe(struct platform_device *pdev)
 {
        struct resource *res;
-       int ret = 0;
-
-       if (pdev->num_resources != 3) {
-               dev_err(&pdev->dev, "invalid number of resources: %d\n",
-                       pdev->num_resources);
-               return -ENODEV;
-       }
+       int ret;
 
        /* MBOX base */
        res = platform_get_resource(pdev, IORESOURCE_MEM, 0);
                dev_err(&pdev->dev, "invalid mem resource\n");
                return -ENODEV;
        }
-       mbox_base = res->start;
+       mbox_base = ioremap(res->start, res->end - res->start);
+       if (!mbox_base)
+               return -ENOMEM;
 
-       /* DSP IRQ */
-       res = platform_get_resource(pdev, IORESOURCE_IRQ, 0);
-       if (unlikely(!res)) {
+       /* DSP or IVA2 IRQ */
+       mbox_dsp_info.irq = platform_get_irq(pdev, 0);
+       if (mbox_dsp_info.irq < 0) {
                dev_err(&pdev->dev, "invalid irq resource\n");
-               return -ENODEV;
+               ret = -ENODEV;
+               goto err_dsp;
        }
-       mbox_dsp_info.irq = res->start;
 
        ret = omap_mbox_register(&mbox_dsp_info);
-
-       /* IVA IRQ */
-       res = platform_get_resource(pdev, IORESOURCE_IRQ, 1);
-       if (unlikely(!res)) {
-               dev_err(&pdev->dev, "invalid irq resource\n");
-               return -ENODEV;
+       if (ret)
+               goto err_dsp;
+
+#if defined(CONFIG_ARCH_OMAP2420) /* IVA */
+       if (cpu_is_omap2420()) {
+               /* IVA IRQ */
+               res = platform_get_resource(pdev, IORESOURCE_IRQ, 1);
+               if (unlikely(!res)) {
+                       dev_err(&pdev->dev, "invalid irq resource\n");
+                       ret = -ENODEV;
+                       goto err_iva1;
+               }
+               mbox_iva_info.irq = res->start;
+               ret = omap_mbox_register(&mbox_iva_info);
+               if (ret)
+                       goto err_iva1;
        }
-       mbox_iva_info.irq = res->start;
-
-       ret = omap_mbox_register(&mbox_iva_info);
+#endif
+       return 0;
 
+err_iva1:
+       omap_mbox_unregister(&mbox_dsp_info);
+err_dsp:
+       iounmap(mbox_base);
        return ret;
 }
 
 static int omap2_mbox_remove(struct platform_device *pdev)
 {
+#if defined(CONFIG_ARCH_OMAP2420)
+       omap_mbox_unregister(&mbox_iva_info);
+#endif
        omap_mbox_unregister(&mbox_dsp_info);
+       iounmap(mbox_base);
        return 0;
 }