#define VPIF_CH2_MAX_MODES     15
 #define VPIF_CH3_MAX_MODES     2
 
+struct vpif_data {
+       struct platform_device *capture;
+       struct platform_device *display;
+};
+
 DEFINE_SPINLOCK(vpif_lock);
 EXPORT_SYMBOL_GPL(vpif_lock);
 
 }
 EXPORT_SYMBOL(vpif_channel_getfid);
 
+static void vpif_pdev_release(struct device *dev)
+{
+       struct platform_device *pdev = to_platform_device(dev);
+
+       kfree(pdev);
+}
+
 static int vpif_probe(struct platform_device *pdev)
 {
        static struct resource *res_irq;
        struct platform_device *pdev_capture, *pdev_display;
        struct device_node *endpoint = NULL;
+       struct vpif_data *data;
        int ret;
 
        vpif_base = devm_platform_ioremap_resource(pdev, 0);
        if (IS_ERR(vpif_base))
                return PTR_ERR(vpif_base);
 
+       data = kzalloc(sizeof(*data), GFP_KERNEL);
+       if (!data)
+               return -ENOMEM;
+
+       platform_set_drvdata(pdev, data);
+
        pm_runtime_enable(&pdev->dev);
        pm_runtime_get(&pdev->dev);
 
                goto err_put_rpm;
        }
 
-       pdev_capture = devm_kzalloc(&pdev->dev, sizeof(*pdev_capture),
-                                   GFP_KERNEL);
-       if (pdev_capture) {
-               pdev_capture->name = "vpif_capture";
-               pdev_capture->id = -1;
-               pdev_capture->resource = res_irq;
-               pdev_capture->num_resources = 1;
-               pdev_capture->dev.dma_mask = pdev->dev.dma_mask;
-               pdev_capture->dev.coherent_dma_mask = pdev->dev.coherent_dma_mask;
-               pdev_capture->dev.parent = &pdev->dev;
-               platform_device_register(pdev_capture);
-       } else {
-               dev_warn(&pdev->dev, "Unable to allocate memory for pdev_capture.\n");
+       pdev_capture = kzalloc(sizeof(*pdev_capture), GFP_KERNEL);
+       if (!pdev_capture) {
+               ret = -ENOMEM;
+               goto err_put_rpm;
        }
 
-       pdev_display = devm_kzalloc(&pdev->dev, sizeof(*pdev_display),
-                                   GFP_KERNEL);
-       if (pdev_display) {
-               pdev_display->name = "vpif_display";
-               pdev_display->id = -1;
-               pdev_display->resource = res_irq;
-               pdev_display->num_resources = 1;
-               pdev_display->dev.dma_mask = pdev->dev.dma_mask;
-               pdev_display->dev.coherent_dma_mask = pdev->dev.coherent_dma_mask;
-               pdev_display->dev.parent = &pdev->dev;
-               platform_device_register(pdev_display);
-       } else {
-               dev_warn(&pdev->dev, "Unable to allocate memory for pdev_display.\n");
+       pdev_capture->name = "vpif_capture";
+       pdev_capture->id = -1;
+       pdev_capture->resource = res_irq;
+       pdev_capture->num_resources = 1;
+       pdev_capture->dev.dma_mask = pdev->dev.dma_mask;
+       pdev_capture->dev.coherent_dma_mask = pdev->dev.coherent_dma_mask;
+       pdev_capture->dev.parent = &pdev->dev;
+       pdev_capture->dev.release = vpif_pdev_release;
+
+       ret = platform_device_register(pdev_capture);
+       if (ret)
+               goto err_put_pdev_capture;
+
+       pdev_display = kzalloc(sizeof(*pdev_display), GFP_KERNEL);
+       if (!pdev_display) {
+               ret = -ENOMEM;
+               goto err_put_pdev_capture;
        }
 
+       pdev_display->name = "vpif_display";
+       pdev_display->id = -1;
+       pdev_display->resource = res_irq;
+       pdev_display->num_resources = 1;
+       pdev_display->dev.dma_mask = pdev->dev.dma_mask;
+       pdev_display->dev.coherent_dma_mask = pdev->dev.coherent_dma_mask;
+       pdev_display->dev.parent = &pdev->dev;
+       pdev_display->dev.release = vpif_pdev_release;
+
+       ret = platform_device_register(pdev_display);
+       if (ret)
+               goto err_put_pdev_display;
+
+       data->capture = pdev_capture;
+       data->display = pdev_display;
+
        return 0;
 
+err_put_pdev_display:
+       platform_device_put(pdev_display);
+err_put_pdev_capture:
+       platform_device_put(pdev_capture);
 err_put_rpm:
        pm_runtime_put(&pdev->dev);
        pm_runtime_disable(&pdev->dev);
+       kfree(data);
 
        return ret;
 }
 
 static int vpif_remove(struct platform_device *pdev)
 {
+       struct vpif_data *data = platform_get_drvdata(pdev);
+
+       if (data->capture)
+               platform_device_unregister(data->capture);
+       if (data->display)
+               platform_device_unregister(data->display);
+
        pm_runtime_put(&pdev->dev);
        pm_runtime_disable(&pdev->dev);
+
+       kfree(data);
+
        return 0;
 }