return 0;
 }
 
+static struct hisi_acc_vf_core_device *hssi_acc_drvdata(struct pci_dev *pdev)
+{
+       struct vfio_pci_core_device *core_device = dev_get_drvdata(&pdev->dev);
+
+       return container_of(core_device, struct hisi_acc_vf_core_device,
+                           core_device);
+}
+
 static void vf_qm_fun_reset(struct hisi_acc_vf_core_device *hisi_acc_vdev,
                            struct hisi_qm *qm)
 {
 
 static void hisi_acc_vf_pci_aer_reset_done(struct pci_dev *pdev)
 {
-       struct hisi_acc_vf_core_device *hisi_acc_vdev = dev_get_drvdata(&pdev->dev);
+       struct hisi_acc_vf_core_device *hisi_acc_vdev = hssi_acc_drvdata(pdev);
 
        if (hisi_acc_vdev->core_device.vdev.migration_flags !=
                                VFIO_MIGRATION_STOP_COPY)
                                          &hisi_acc_vfio_pci_ops);
        }
 
+       dev_set_drvdata(&pdev->dev, &hisi_acc_vdev->core_device);
        ret = vfio_pci_core_register_device(&hisi_acc_vdev->core_device);
        if (ret)
                goto out_free;
-
-       dev_set_drvdata(&pdev->dev, hisi_acc_vdev);
        return 0;
 
 out_free:
 
 static void hisi_acc_vfio_pci_remove(struct pci_dev *pdev)
 {
-       struct hisi_acc_vf_core_device *hisi_acc_vdev = dev_get_drvdata(&pdev->dev);
+       struct hisi_acc_vf_core_device *hisi_acc_vdev = hssi_acc_drvdata(pdev);
 
        vfio_pci_core_unregister_device(&hisi_acc_vdev->core_device);
        vfio_pci_core_uninit_device(&hisi_acc_vdev->core_device);
 
 /* Arbitrary to prevent userspace from consuming endless memory */
 #define MAX_MIGRATION_SIZE (512*1024*1024)
 
+static struct mlx5vf_pci_core_device *mlx5vf_drvdata(struct pci_dev *pdev)
+{
+       struct vfio_pci_core_device *core_device = dev_get_drvdata(&pdev->dev);
+
+       return container_of(core_device, struct mlx5vf_pci_core_device,
+                           core_device);
+}
+
 static struct page *
 mlx5vf_get_migration_page(struct mlx5_vf_migration_file *migf,
                          unsigned long offset)
 
 static void mlx5vf_pci_aer_reset_done(struct pci_dev *pdev)
 {
-       struct mlx5vf_pci_core_device *mvdev = dev_get_drvdata(&pdev->dev);
+       struct mlx5vf_pci_core_device *mvdev = mlx5vf_drvdata(pdev);
 
        if (!mvdev->migrate_cap)
                return;
                return -ENOMEM;
        vfio_pci_core_init_device(&mvdev->core_device, pdev, &mlx5vf_pci_ops);
        mlx5vf_cmd_set_migratable(mvdev);
+       dev_set_drvdata(&pdev->dev, &mvdev->core_device);
        ret = vfio_pci_core_register_device(&mvdev->core_device);
        if (ret)
                goto out_free;
-
-       dev_set_drvdata(&pdev->dev, mvdev);
        return 0;
 
 out_free:
 
 static void mlx5vf_pci_remove(struct pci_dev *pdev)
 {
-       struct mlx5vf_pci_core_device *mvdev = dev_get_drvdata(&pdev->dev);
+       struct mlx5vf_pci_core_device *mvdev = mlx5vf_drvdata(pdev);
 
        vfio_pci_core_unregister_device(&mvdev->core_device);
        mlx5vf_cmd_remove_migratable(mvdev);
 
                return -ENOMEM;
        vfio_pci_core_init_device(vdev, pdev, &vfio_pci_ops);
 
+       dev_set_drvdata(&pdev->dev, vdev);
        ret = vfio_pci_core_register_device(vdev);
        if (ret)
                goto out_free;
-       dev_set_drvdata(&pdev->dev, vdev);
        return 0;
 
 out_free:
 
        struct pci_dev *pdev = vdev->pdev;
        int ret;
 
+       /* Drivers must set the vfio_pci_core_device to their drvdata */
+       if (WARN_ON(vdev != dev_get_drvdata(&vdev->pdev->dev)))
+               return -EINVAL;
+
        if (pdev->hdr_type != PCI_HEADER_TYPE_NORMAL)
                return -EINVAL;