#include <drm/drm_drv.h>
 #include <drm/drm_file.h>
 #include <drm/drm_ioctl.h>
+#include <drm/drm_managed.h>
 #include <drm/drm_prime.h>
 
 #include "vgem_drv.h"
 
        platform_device_unregister(vgem->platform);
        drm_dev_fini(&vgem->drm);
-
-       kfree(vgem);
 }
 
 static struct drm_driver vgem_driver = {
                           &vgem_device->platform->dev);
        if (ret)
                goto out_unregister;
+       drmm_add_final_kfree(&vgem_device->drm, vgem_device);
 
        /* Final step: expose the device/driver to userspace */
-       ret  = drm_dev_register(&vgem_device->drm, 0);
+       ret = drm_dev_register(&vgem_device->drm, 0);
        if (ret)
-               goto out_fini;
+               goto out_put;
 
        return 0;
 
-out_fini:
-       drm_dev_fini(&vgem_device->drm);
+out_put:
+       drm_dev_put(&vgem_device->drm);
+       return ret;
+
 out_unregister:
        platform_device_unregister(vgem_device->platform);
 out_free: