#include <drm/drm_atomic_helper.h>
 #include <drm/drm_ioctl.h>
 #include <drm/drm_irq.h>
+#include <drm/drm_managed.h>
 #include <drm/drm_probe_helper.h>
 
 #include "display/intel_acpi.h"
                return ERR_PTR(err);
        }
 
+       drmm_add_final_kfree(&i915->drm, i915);
+
        i915->drm.pdev = pdev;
        pci_set_drvdata(pdev, i915);
 
 static void i915_driver_destroy(struct drm_i915_private *i915)
 {
        drm_dev_fini(&i915->drm);
-       kfree(i915);
 }
 
 /**
 
        i915_welcome_messages(i915);
 
+       i915->do_release = true;
+
        return 0;
 
 out_cleanup_irq:
 out_fini:
        i915_probe_error(i915, "Device initialization failed (%d)\n", ret);
        i915_driver_destroy(i915);
+       drm_dev_put(&i915->drm);
        return ret;
 }
 
        struct drm_i915_private *dev_priv = to_i915(dev);
        struct intel_runtime_pm *rpm = &dev_priv->runtime_pm;
 
+       if (!dev_priv->do_release)
+               return;
+
        disable_rpm_wakeref_asserts(rpm);
 
        i915_gem_driver_release(dev_priv);
 
 #include <linux/pm_domain.h>
 #include <linux/pm_runtime.h>
 
+#include <drm/drm_managed.h>
+
 #include "gt/intel_gt.h"
 #include "gt/intel_gt_requests.h"
 #include "gt/mock_engine.h"
 {
        struct drm_i915_private *i915 = to_i915(dev);
 
+       if (!i915->do_release)
+               goto out;
+
        mock_device_flush(i915);
        intel_gt_driver_remove(&i915->gt);
 
        drm_mode_config_cleanup(&i915->drm);
 
        drm_dev_fini(&i915->drm);
+out:
        put_device(&i915->drm.pdev->dev);
+       i915->drm.pdev = NULL;
 }
 
 static struct drm_driver mock_driver = {
        struct pci_dev *pdev;
        int err;
 
-       pdev = kzalloc(sizeof(*pdev) + sizeof(*i915), GFP_KERNEL);
+       pdev = kzalloc(sizeof(*pdev), GFP_KERNEL);
        if (!pdev)
-               goto err;
+               return NULL;
+       i915 = kzalloc(sizeof(*i915), GFP_KERNEL);
+       if (!i915) {
+               kfree(pdev);
+               return NULL;
+       }
 
        device_initialize(&pdev->dev);
        pdev->class = PCI_BASE_CLASS_DISPLAY << 16;
        pdev->dev.archdata.iommu = (void *)-1;
 #endif
 
-       i915 = (struct drm_i915_private *)(pdev + 1);
        pci_set_drvdata(pdev, i915);
 
        dev_pm_domain_set(&pdev->dev, &pm_domain);
        err = drm_dev_init(&i915->drm, &mock_driver, &pdev->dev);
        if (err) {
                pr_err("Failed to initialise mock GEM device: err=%d\n", err);
-               goto put_device;
+               put_device(&pdev->dev);
+               kfree(i915);
+
+               return NULL;
        }
        i915->drm.pdev = pdev;
+       drmm_add_final_kfree(&i915->drm, i915);
 
        intel_runtime_pm_init_early(&i915->runtime_pm);
 
        __clear_bit(I915_WEDGED, &i915->gt.reset.flags);
        intel_engines_driver_register(i915);
 
+       i915->do_release = true;
+
        return i915;
 
 err_context:
        intel_memory_regions_driver_release(i915);
        drm_mode_config_cleanup(&i915->drm);
        drm_dev_fini(&i915->drm);
-put_device:
-       put_device(&pdev->dev);
-err:
+       drm_dev_put(&i915->drm);
+
        return NULL;
 }