{
        struct ap_matrix_mdev *m;
 
-       mutex_lock(&matrix_dev->lock);
-
        list_for_each_entry(m, &matrix_dev->mdev_list, node) {
-               if ((m != matrix_mdev) && (m->kvm == kvm)) {
-                       mutex_unlock(&matrix_dev->lock);
+               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;
-       mutex_unlock(&matrix_dev->lock);
 
        return 0;
 }
        return NOTIFY_DONE;
 }
 
+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;
+}
+
 static int vfio_ap_mdev_group_notifier(struct notifier_block *nb,
                                       unsigned long action, void *data)
 {
-       int ret;
+       int ret, 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);
 
        if (!data) {
-               matrix_mdev->kvm = NULL;
-               return NOTIFY_OK;
+               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)
-               return NOTIFY_DONE;
+       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)
-               return NOTIFY_DONE;
+       if (!matrix_mdev->kvm->arch.crypto.crycbd) {
+               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);
 
-       return NOTIFY_OK;
+notify_done:
+       mutex_unlock(&matrix_dev->lock);
+       return notify_rc;
 }
 
 static void vfio_ap_irq_disable_apqn(int apqn)
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
        mutex_lock(&matrix_dev->lock);
-       if (matrix_mdev->kvm) {
-               kvm_arch_crypto_clear_masks(matrix_mdev->kvm);
-               matrix_mdev->kvm->arch.crypto.pqap_hook = NULL;
-               vfio_ap_mdev_reset_queues(mdev);
-               kvm_put_kvm(matrix_mdev->kvm);
-               matrix_mdev->kvm = NULL;
-       }
+       if (matrix_mdev->kvm)
+               vfio_ap_mdev_unset_kvm(matrix_mdev);
        mutex_unlock(&matrix_dev->lock);
 
        vfio_unregister_notifier(mdev_dev(mdev), VFIO_IOMMU_NOTIFY,