#define VFIO_AP_MDEV_NAME_HWVIRT "VFIO AP Passthrough Device"
 
 static int vfio_ap_mdev_reset_queues(struct mdev_device *mdev);
+static struct vfio_ap_queue *vfio_ap_find_queue(int apqn);
 
 static int match_apqn(struct device *dev, const void *data)
 {
                                        int apqn)
 {
        struct vfio_ap_queue *q;
-       struct device *dev;
 
        if (!test_bit_inv(AP_QID_CARD(apqn), matrix_mdev->matrix.apm))
                return NULL;
        if (!test_bit_inv(AP_QID_QUEUE(apqn), matrix_mdev->matrix.aqm))
                return NULL;
 
-       dev = driver_find_device(&matrix_dev->vfio_ap_drv->driver, NULL,
-                                &apqn, match_apqn);
-       if (!dev)
-               return NULL;
-       q = dev_get_drvdata(dev);
-       q->matrix_mdev = matrix_mdev;
-       put_device(dev);
+       q = vfio_ap_find_queue(apqn);
+       if (q)
+               q->matrix_mdev = matrix_mdev;
 
        return q;
 }
  */
 static void vfio_ap_free_aqic_resources(struct vfio_ap_queue *q)
 {
-       if (q->saved_isc != VFIO_AP_ISC_INVALID && q->matrix_mdev)
+       if (!q)
+               return;
+       if (q->saved_isc != VFIO_AP_ISC_INVALID &&
+           !WARN_ON(!(q->matrix_mdev && q->matrix_mdev->kvm))) {
                kvm_s390_gisc_unregister(q->matrix_mdev->kvm, q->saved_isc);
-       if (q->saved_pfn && q->matrix_mdev)
+               q->saved_isc = VFIO_AP_ISC_INVALID;
+       }
+       if (q->saved_pfn && !WARN_ON(!q->matrix_mdev)) {
                vfio_unpin_pages(mdev_dev(q->matrix_mdev->mdev),
                                 &q->saved_pfn, 1);
-       q->saved_pfn = 0;
-       q->saved_isc = VFIO_AP_ISC_INVALID;
+               q->saved_pfn = 0;
+       }
 }
 
 /**
  * Returns if ap_aqic function failed with invalid, deconfigured or
  * checkstopped AP.
  */
-struct ap_queue_status vfio_ap_irq_disable(struct vfio_ap_queue *q)
+static struct ap_queue_status vfio_ap_irq_disable(struct vfio_ap_queue *q)
 {
        struct ap_qirq_ctrl aqic_gisa = {};
        struct ap_queue_status status;
        return notify_rc;
 }
 
-static void vfio_ap_irq_disable_apqn(int apqn)
+static struct vfio_ap_queue *vfio_ap_find_queue(int apqn)
 {
        struct device *dev;
-       struct vfio_ap_queue *q;
+       struct vfio_ap_queue *q = NULL;
 
        dev = driver_find_device(&matrix_dev->vfio_ap_drv->driver, NULL,
                                 &apqn, match_apqn);
        if (dev) {
                q = dev_get_drvdata(dev);
-               vfio_ap_irq_disable(q);
                put_device(dev);
        }
+
+       return q;
 }
 
-int vfio_ap_mdev_reset_queue(unsigned int apid, unsigned int apqi,
+int vfio_ap_mdev_reset_queue(struct vfio_ap_queue *q,
                             unsigned int retry)
 {
        struct ap_queue_status status;
+       int ret;
        int retry2 = 2;
-       int apqn = AP_MKQID(apid, apqi);
 
-       do {
-               status = ap_zapq(apqn);
-               switch (status.response_code) {
-               case AP_RESPONSE_NORMAL:
-                       while (!status.queue_empty && retry2--) {
-                               msleep(20);
-                               status = ap_tapq(apqn, NULL);
-                       }
-                       WARN_ON_ONCE(retry2 <= 0);
-                       return 0;
-               case AP_RESPONSE_RESET_IN_PROGRESS:
-               case AP_RESPONSE_BUSY:
+       if (!q)
+               return 0;
+
+retry_zapq:
+       status = ap_zapq(q->apqn);
+       switch (status.response_code) {
+       case AP_RESPONSE_NORMAL:
+               ret = 0;
+               break;
+       case AP_RESPONSE_RESET_IN_PROGRESS:
+               if (retry--) {
                        msleep(20);
-                       break;
-               default:
-                       /* things are really broken, give up */
-                       return -EIO;
+                       goto retry_zapq;
                }
-       } while (retry--);
+               ret = -EBUSY;
+               break;
+       case AP_RESPONSE_Q_NOT_AVAIL:
+       case AP_RESPONSE_DECONFIGURED:
+       case AP_RESPONSE_CHECKSTOPPED:
+               WARN_ON_ONCE(status.irq_enabled);
+               ret = -EBUSY;
+               goto free_resources;
+       default:
+               /* things are really broken, give up */
+               WARN(true, "PQAP/ZAPQ completed with invalid rc (%x)\n",
+                    status.response_code);
+               return -EIO;
+       }
+
+       /* wait for the reset to take effect */
+       while (retry2--) {
+               if (status.queue_empty && !status.irq_enabled)
+                       break;
+               msleep(20);
+               status = ap_tapq(q->apqn, NULL);
+       }
+       WARN_ON_ONCE(retry2 <= 0);
 
-       return -EBUSY;
+free_resources:
+       vfio_ap_free_aqic_resources(q);
+
+       return ret;
 }
 
 static int vfio_ap_mdev_reset_queues(struct mdev_device *mdev)
        int ret;
        int rc = 0;
        unsigned long apid, apqi;
+       struct vfio_ap_queue *q;
        struct ap_matrix_mdev *matrix_mdev = mdev_get_drvdata(mdev);
 
        for_each_set_bit_inv(apid, matrix_mdev->matrix.apm,
                             matrix_mdev->matrix.apm_max + 1) {
                for_each_set_bit_inv(apqi, matrix_mdev->matrix.aqm,
                                     matrix_mdev->matrix.aqm_max + 1) {
-                       ret = vfio_ap_mdev_reset_queue(apid, apqi, 1);
+                       q = vfio_ap_find_queue(AP_MKQID(apid, apqi));
+                       ret = vfio_ap_mdev_reset_queue(q, 1);
                        /*
                         * Regardless whether a queue turns out to be busy, or
                         * is not operational, we need to continue resetting
                         */
                        if (ret)
                                rc = ret;
-                       vfio_ap_irq_disable_apqn(AP_MKQID(apid, apqi));
                }
        }