matrix_mdev = container_of(vcpu->kvm->arch.crypto.pqap_hook,
                                   struct ap_matrix_mdev, pqap_hook);
 
+       /*
+        * If the KVM pointer is in the process of being set, wait until the
+        * process has completed.
+        */
+       wait_event_cmd(matrix_mdev->wait_for_kvm,
+                      !matrix_mdev->kvm_busy,
+                      mutex_unlock(&matrix_dev->lock),
+                      mutex_lock(&matrix_dev->lock));
+
+       /* If the there is no guest using the mdev, there is nothing to do */
+       if (!matrix_mdev->kvm)
+               goto out_unlock;
+
        q = vfio_ap_get_queue(matrix_mdev, apqn);
        if (!q)
                goto out_unlock;
 
        matrix_mdev->mdev = mdev;
        vfio_ap_matrix_init(&matrix_dev->info, &matrix_mdev->matrix);
+       init_waitqueue_head(&matrix_mdev->wait_for_kvm);
        mdev_set_drvdata(mdev, matrix_mdev);
        matrix_mdev->pqap_hook.hook = handle_pqap;
        matrix_mdev->pqap_hook.owner = THIS_MODULE;
 {
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
-       if (matrix_mdev->kvm)
+       mutex_lock(&matrix_dev->lock);
+
+       /*
+        * If the KVM pointer is in flux or the guest is running, disallow
+        * un-assignment of control domain.
+        */
+       if (matrix_mdev->kvm_busy || matrix_mdev->kvm) {
+               mutex_unlock(&matrix_dev->lock);
                return -EBUSY;
+       }
 
-       mutex_lock(&matrix_dev->lock);
        vfio_ap_mdev_reset_queues(mdev);
        list_del(&matrix_mdev->node);
-       mutex_unlock(&matrix_dev->lock);
-
        kfree(matrix_mdev);
        mdev_set_drvdata(mdev, NULL);
        atomic_inc(&matrix_dev->available_instances);
+       mutex_unlock(&matrix_dev->lock);
 
        return 0;
 }
        struct mdev_device *mdev = mdev_from_dev(dev);
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
-       /* If the guest is running, disallow assignment of adapter */
-       if (matrix_mdev->kvm)
-               return -EBUSY;
+       mutex_lock(&matrix_dev->lock);
+
+       /*
+        * If the KVM pointer is in flux or the guest is running, disallow
+        * un-assignment of adapter
+        */
+       if (matrix_mdev->kvm_busy || matrix_mdev->kvm) {
+               ret = -EBUSY;
+               goto done;
+       }
 
        ret = kstrtoul(buf, 0, &apid);
        if (ret)
-               return ret;
+               goto done;
 
-       if (apid > matrix_mdev->matrix.apm_max)
-               return -ENODEV;
+       if (apid > matrix_mdev->matrix.apm_max) {
+               ret = -ENODEV;
+               goto done;
+       }
 
        /*
         * Set the bit in the AP mask (APM) corresponding to the AP adapter
         * number (APID). The bits in the mask, from most significant to least
         * significant bit, correspond to APIDs 0-255.
         */
-       mutex_lock(&matrix_dev->lock);
-
        ret = vfio_ap_mdev_verify_queues_reserved_for_apid(matrix_mdev, apid);
        if (ret)
                goto done;
        struct mdev_device *mdev = mdev_from_dev(dev);
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
-       /* If the guest is running, disallow un-assignment of adapter */
-       if (matrix_mdev->kvm)
-               return -EBUSY;
+       mutex_lock(&matrix_dev->lock);
+
+       /*
+        * If the KVM pointer is in flux or the guest is running, disallow
+        * un-assignment of adapter
+        */
+       if (matrix_mdev->kvm_busy || matrix_mdev->kvm) {
+               ret = -EBUSY;
+               goto done;
+       }
 
        ret = kstrtoul(buf, 0, &apid);
        if (ret)
-               return ret;
+               goto done;
 
-       if (apid > matrix_mdev->matrix.apm_max)
-               return -ENODEV;
+       if (apid > matrix_mdev->matrix.apm_max) {
+               ret = -ENODEV;
+               goto done;
+       }
 
-       mutex_lock(&matrix_dev->lock);
        clear_bit_inv((unsigned long)apid, matrix_mdev->matrix.apm);
+       ret = count;
+done:
        mutex_unlock(&matrix_dev->lock);
-
-       return count;
+       return ret;
 }
 static DEVICE_ATTR_WO(unassign_adapter);
 
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
        unsigned long max_apqi = matrix_mdev->matrix.aqm_max;
 
-       /* If the guest is running, disallow assignment of domain */
-       if (matrix_mdev->kvm)
-               return -EBUSY;
+       mutex_lock(&matrix_dev->lock);
+
+       /*
+        * If the KVM pointer is in flux or the guest is running, disallow
+        * assignment of domain
+        */
+       if (matrix_mdev->kvm_busy || matrix_mdev->kvm) {
+               ret = -EBUSY;
+               goto done;
+       }
 
        ret = kstrtoul(buf, 0, &apqi);
        if (ret)
-               return ret;
-       if (apqi > max_apqi)
-               return -ENODEV;
-
-       mutex_lock(&matrix_dev->lock);
+               goto done;
+       if (apqi > max_apqi) {
+               ret = -ENODEV;
+               goto done;
+       }
 
        ret = vfio_ap_mdev_verify_queues_reserved_for_apqi(matrix_mdev, apqi);
        if (ret)
        struct mdev_device *mdev = mdev_from_dev(dev);
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
-       /* If the guest is running, disallow un-assignment of domain */
-       if (matrix_mdev->kvm)
-               return -EBUSY;
+       mutex_lock(&matrix_dev->lock);
+
+       /*
+        * If the KVM pointer is in flux or the guest is running, disallow
+        * un-assignment of domain
+        */
+       if (matrix_mdev->kvm_busy || matrix_mdev->kvm) {
+               ret = -EBUSY;
+               goto done;
+       }
 
        ret = kstrtoul(buf, 0, &apqi);
        if (ret)
-               return ret;
+               goto done;
 
-       if (apqi > matrix_mdev->matrix.aqm_max)
-               return -ENODEV;
+       if (apqi > matrix_mdev->matrix.aqm_max) {
+               ret = -ENODEV;
+               goto done;
+       }
 
-       mutex_lock(&matrix_dev->lock);
        clear_bit_inv((unsigned long)apqi, matrix_mdev->matrix.aqm);
-       mutex_unlock(&matrix_dev->lock);
+       ret = count;
 
-       return count;
+done:
+       mutex_unlock(&matrix_dev->lock);
+       return ret;
 }
 static DEVICE_ATTR_WO(unassign_domain);
 
        struct mdev_device *mdev = mdev_from_dev(dev);
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
-       /* If the guest is running, disallow assignment of control domain */
-       if (matrix_mdev->kvm)
-               return -EBUSY;
+       mutex_lock(&matrix_dev->lock);
+
+       /*
+        * If the KVM pointer is in flux or the guest is running, disallow
+        * assignment of control domain.
+        */
+       if (matrix_mdev->kvm_busy || matrix_mdev->kvm) {
+               ret = -EBUSY;
+               goto done;
+       }
 
        ret = kstrtoul(buf, 0, &id);
        if (ret)
-               return ret;
+               goto done;
 
-       if (id > matrix_mdev->matrix.adm_max)
-               return -ENODEV;
+       if (id > matrix_mdev->matrix.adm_max) {
+               ret = -ENODEV;
+               goto done;
+       }
 
        /* Set the bit in the ADM (bitmask) corresponding to the AP control
         * domain number (id). The bits in the mask, from most significant to
         * least significant, correspond to IDs 0 up to the one less than the
         * number of control domains that can be assigned.
         */
-       mutex_lock(&matrix_dev->lock);
        set_bit_inv(id, matrix_mdev->matrix.adm);
+       ret = count;
+done:
        mutex_unlock(&matrix_dev->lock);
-
-       return count;
+       return ret;
 }
 static DEVICE_ATTR_WO(assign_control_domain);
 
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
        unsigned long max_domid =  matrix_mdev->matrix.adm_max;
 
-       /* If the guest is running, disallow un-assignment of control domain */
-       if (matrix_mdev->kvm)
-               return -EBUSY;
+       mutex_lock(&matrix_dev->lock);
+
+       /*
+        * If the KVM pointer is in flux or the guest is running, disallow
+        * un-assignment of control domain.
+        */
+       if (matrix_mdev->kvm_busy || matrix_mdev->kvm) {
+               ret = -EBUSY;
+               goto done;
+       }
 
        ret = kstrtoul(buf, 0, &domid);
        if (ret)
-               return ret;
-       if (domid > max_domid)
-               return -ENODEV;
+               goto done;
+       if (domid > max_domid) {
+               ret = -ENODEV;
+               goto done;
+       }
 
-       mutex_lock(&matrix_dev->lock);
        clear_bit_inv(domid, matrix_mdev->matrix.adm);
+       ret = count;
+done:
        mutex_unlock(&matrix_dev->lock);
-
-       return count;
+       return ret;
 }
 static DEVICE_ATTR_WO(unassign_control_domain);
 
  * @matrix_mdev: a mediated matrix device
  * @kvm: reference to KVM instance
  *
- * Verifies no other mediated matrix device has @kvm and sets a reference to
- * it in @matrix_mdev->kvm.
+ * Sets all data for @matrix_mdev that are needed to manage AP resources
+ * for the guest whose state is represented by @kvm.
+ *
+ * Note: The matrix_dev->lock must be taken prior to calling
+ * this function; however, the lock will be temporarily released while the
+ * guest's AP configuration is set to avoid a potential lockdep splat.
+ * The kvm->lock is taken to set the guest's AP configuration which, under
+ * certain circumstances, will result in a circular lock dependency if this is
+ * done under the @matrix_mdev->lock.
  *
  * Return 0 if no other mediated matrix device has a reference to @kvm;
  * otherwise, returns an -EPERM.
 {
        struct ap_matrix_mdev *m;
 
-       list_for_each_entry(m, &matrix_dev->mdev_list, node) {
-               if ((m != matrix_mdev) && (m->kvm == kvm))
-                       return -EPERM;
-       }
+       if (kvm->arch.crypto.crycbd) {
+               list_for_each_entry(m, &matrix_dev->mdev_list, node) {
+                       if (m != matrix_mdev && m->kvm == kvm)
+                               return -EPERM;
+               }
 
-       matrix_mdev->kvm = kvm;
-       kvm_get_kvm(kvm);
-       kvm->arch.crypto.pqap_hook = &matrix_mdev->pqap_hook;
+               kvm_get_kvm(kvm);
+               matrix_mdev->kvm_busy = true;
+               mutex_unlock(&matrix_dev->lock);
+               kvm_arch_crypto_set_masks(kvm,
+                                         matrix_mdev->matrix.apm,
+                                         matrix_mdev->matrix.aqm,
+                                         matrix_mdev->matrix.adm);
+               mutex_lock(&matrix_dev->lock);
+               kvm->arch.crypto.pqap_hook = &matrix_mdev->pqap_hook;
+               matrix_mdev->kvm = kvm;
+               matrix_mdev->kvm_busy = false;
+               wake_up_all(&matrix_mdev->wait_for_kvm);
+       }
 
        return 0;
 }
        return NOTIFY_DONE;
 }
 
+/**
+ * vfio_ap_mdev_unset_kvm
+ *
+ * @matrix_mdev: a matrix mediated device
+ *
+ * Performs clean-up of resources no longer needed by @matrix_mdev.
+ *
+ * Note: The matrix_dev->lock must be taken prior to calling
+ * this function; however, the lock will be temporarily released while the
+ * guest's AP configuration is cleared to avoid a potential lockdep splat.
+ * The kvm->lock is taken to clear the guest's AP configuration which, under
+ * certain circumstances, will result in a circular lock dependency if this is
+ * done under the @matrix_mdev->lock.
+ *
+ */
 static void vfio_ap_mdev_unset_kvm(struct ap_matrix_mdev *matrix_mdev)
 {
-       kvm_arch_crypto_clear_masks(matrix_mdev->kvm);
-       matrix_mdev->kvm->arch.crypto.pqap_hook = NULL;
-       vfio_ap_mdev_reset_queues(matrix_mdev->mdev);
-       kvm_put_kvm(matrix_mdev->kvm);
-       matrix_mdev->kvm = NULL;
+       /*
+        * If the KVM pointer is in the process of being set, wait until the
+        * process has completed.
+        */
+       wait_event_cmd(matrix_mdev->wait_for_kvm,
+                      !matrix_mdev->kvm_busy,
+                      mutex_unlock(&matrix_dev->lock),
+                      mutex_lock(&matrix_dev->lock));
+
+       if (matrix_mdev->kvm) {
+               matrix_mdev->kvm_busy = true;
+               mutex_unlock(&matrix_dev->lock);
+               kvm_arch_crypto_clear_masks(matrix_mdev->kvm);
+               mutex_lock(&matrix_dev->lock);
+               vfio_ap_mdev_reset_queues(matrix_mdev->mdev);
+               matrix_mdev->kvm->arch.crypto.pqap_hook = NULL;
+               kvm_put_kvm(matrix_mdev->kvm);
+               matrix_mdev->kvm = NULL;
+               matrix_mdev->kvm_busy = false;
+               wake_up_all(&matrix_mdev->wait_for_kvm);
+       }
 }
 
 static int vfio_ap_mdev_group_notifier(struct notifier_block *nb,
                                       unsigned long action, void *data)
 {
-       int ret, notify_rc = NOTIFY_OK;
+       int notify_rc = NOTIFY_OK;
        struct ap_matrix_mdev *matrix_mdev;
 
        if (action != VFIO_GROUP_NOTIFY_SET_KVM)
                return NOTIFY_OK;
 
-       matrix_mdev = container_of(nb, struct ap_matrix_mdev, group_notifier);
        mutex_lock(&matrix_dev->lock);
+       matrix_mdev = container_of(nb, struct ap_matrix_mdev, group_notifier);
 
-       if (!data) {
-               if (matrix_mdev->kvm)
-                       vfio_ap_mdev_unset_kvm(matrix_mdev);
-               goto notify_done;
-       }
-
-       ret = vfio_ap_mdev_set_kvm(matrix_mdev, data);
-       if (ret) {
-               notify_rc = NOTIFY_DONE;
-               goto notify_done;
-       }
-
-       /* If there is no CRYCB pointer, then we can't copy the masks */
-       if (!matrix_mdev->kvm->arch.crypto.crycbd) {
+       if (!data)
+               vfio_ap_mdev_unset_kvm(matrix_mdev);
+       else if (vfio_ap_mdev_set_kvm(matrix_mdev, data))
                notify_rc = NOTIFY_DONE;
-               goto notify_done;
-       }
-
-       kvm_arch_crypto_set_masks(matrix_mdev->kvm, matrix_mdev->matrix.apm,
-                                 matrix_mdev->matrix.aqm,
-                                 matrix_mdev->matrix.adm);
 
-notify_done:
        mutex_unlock(&matrix_dev->lock);
+
        return notify_rc;
 }
 
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
        mutex_lock(&matrix_dev->lock);
-       if (matrix_mdev->kvm)
-               vfio_ap_mdev_unset_kvm(matrix_mdev);
+       vfio_ap_mdev_unset_kvm(matrix_mdev);
        mutex_unlock(&matrix_dev->lock);
 
        vfio_unregister_notifier(mdev_dev(mdev), VFIO_IOMMU_NOTIFY,
                                    unsigned int cmd, unsigned long arg)
 {
        int ret;
+       struct ap_matrix_mdev *matrix_mdev;
 
        mutex_lock(&matrix_dev->lock);
        switch (cmd) {
                ret = vfio_ap_mdev_get_device_info(arg);
                break;
        case VFIO_DEVICE_RESET:
+               matrix_mdev = mdev_get_drvdata(mdev);
+               if (WARN(!matrix_mdev, "Driver data missing from mdev!!")) {
+                       ret = -EINVAL;
+                       break;
+               }
+
+               /*
+                * If the KVM pointer is in the process of being set, wait until
+                * the process has completed.
+                */
+               wait_event_cmd(matrix_mdev->wait_for_kvm,
+                              !matrix_mdev->kvm_busy,
+                              mutex_unlock(&matrix_dev->lock),
+                              mutex_lock(&matrix_dev->lock));
+
                ret = vfio_ap_mdev_reset_queues(mdev);
                break;
        default: