#include <drm/drm_file.h>
 #include <drm/drm_gem_framebuffer_helper.h>
 #include <drm/drm_ioctl.h>
+#include <drm/drm_managed.h>
 #include <drm/drm_probe_helper.h>
 #include <drm/drm_vblank.h>
 
                           &vkms_device->platform->dev);
        if (ret)
                goto out_unregister;
+       drmm_add_final_kfree(&vkms_device->drm, vkms_device);
 
        ret = dma_coerce_mask_and_coherent(vkms_device->drm.dev,
                                           DMA_BIT_MASK(64));
 
        if (ret) {
                DRM_ERROR("Could not initialize DMA support\n");
-               goto out_fini;
+               goto out_put;
        }
 
        vkms_device->drm.irq_enabled = true;
        ret = drm_vblank_init(&vkms_device->drm, 1);
        if (ret) {
                DRM_ERROR("Failed to vblank\n");
-               goto out_fini;
+               goto out_put;
        }
 
        ret = vkms_modeset_init(vkms_device);
        if (ret)
-               goto out_fini;
+               goto out_put;
 
        ret = drm_dev_register(&vkms_device->drm, 0);
        if (ret)
-               goto out_fini;
+               goto out_put;
 
        return 0;
 
-out_fini:
-       drm_dev_fini(&vkms_device->drm);
+out_put:
+       drm_dev_put(&vkms_device->drm);
+       return ret;
 
 out_unregister:
        platform_device_unregister(vkms_device->platform);
-
 out_free:
        kfree(vkms_device);
        return ret;
 
        drm_dev_unregister(&vkms_device->drm);
        drm_dev_put(&vkms_device->drm);
-
-       kfree(vkms_device);
 }
 
 module_init(vkms_init);