/* DMA */
 int zpci_dma_init(void);
 void zpci_dma_exit(void);
+int zpci_dma_init_device(struct zpci_dev *zdev);
+int zpci_dma_exit_device(struct zpci_dev *zdev);
 
 /* IRQ */
 int __init zpci_irq_init(void);
 
 }
 
 /* Prototypes */
-int zpci_dma_init_device(struct zpci_dev *);
-void zpci_dma_exit_device(struct zpci_dev *);
 void dma_free_seg_table(unsigned long);
 unsigned long *dma_alloc_cpu_table(void);
 void dma_cleanup_tables(unsigned long *);
 
 {
        u64 req = ZPCI_CREATE_REQ(zdev->fh, dmaas, ZPCI_MOD_FC_REG_IOAT);
        struct zpci_fib fib = {0};
-       u8 status;
+       u8 cc, status;
 
        WARN_ON_ONCE(iota & 0x3fff);
        fib.pba = base;
        fib.pal = limit;
        fib.iota = iota | ZPCI_IOTA_RTTO_FLAG;
-       return zpci_mod_fc(req, &fib, &status) ? -EIO : 0;
+       cc = zpci_mod_fc(req, &fib, &status);
+       if (cc)
+               zpci_dbg(3, "reg ioat fid:%x, cc:%d, status:%d\n", zdev->fid, cc, status);
+       return cc;
 }
 
 /* Modify PCI: Unregister I/O address translation parameters */
        u8 cc, status;
 
        cc = zpci_mod_fc(req, &fib, &status);
-       if (cc == 3) /* Function already gone. */
-               cc = 0;
-       return cc ? -EIO : 0;
+       if (cc)
+               zpci_dbg(3, "unreg ioat fid:%x, cc:%d, status:%d\n", zdev->fid, cc, status);
+       return cc;
 }
 
 /* Modify PCI: Set PCI function measurement parameters */
 int zpci_enable_device(struct zpci_dev *zdev)
 {
        u32 fh = zdev->fh;
-       int rc;
+       int rc = 0;
 
-       if (clp_enable_fh(zdev, &fh, ZPCI_NR_DMA_SPACES)) {
+       if (clp_enable_fh(zdev, &fh, ZPCI_NR_DMA_SPACES))
                rc = -EIO;
-               goto out;
-       }
-       zdev->fh = fh;
-
-       rc = zpci_dma_init_device(zdev);
-       if (rc)
-               goto out_dma;
-
-       return 0;
-
-out_dma:
-       clp_disable_fh(zdev, &fh);
-out:
-       zdev->fh = fh;
+       else
+               zdev->fh = fh;
        return rc;
 }
 
        u32 fh = zdev->fh;
        int cc, rc = 0;
 
-       zpci_dma_exit_device(zdev);
-       if (!zdev_enabled(zdev))
-               return 0;
        cc = clp_disable_fh(zdev, &fh);
        if (!cc) {
                zdev->fh = fh;
        if (zdev->zbus->bus)
                zpci_bus_remove_device(zdev, false);
 
+       if (zdev->dma_table) {
+               rc = zpci_dma_exit_device(zdev);
+               if (rc)
+                       return rc;
+       }
        if (zdev_enabled(zdev)) {
                rc = zpci_disable_device(zdev);
                if (rc)
        if (zdev->zbus->bus)
                zpci_bus_remove_device(zdev, false);
 
+       if (zdev->dma_table)
+               zpci_dma_exit_device(zdev);
        if (zdev_enabled(zdev))
                zpci_disable_device(zdev);
 
 
                rc = zpci_enable_device(zdev);
                if (rc)
                        return rc;
+               rc = zpci_dma_init_device(zdev);
+               if (rc) {
+                       zpci_disable_device(zdev);
+                       return rc;
+               }
        }
 
        if (!zdev->has_resources) {
 
                }
 
        }
-       rc = zpci_register_ioat(zdev, 0, zdev->start_dma, zdev->end_dma,
-                               (u64) zdev->dma_table);
-       if (rc)
+       if (zpci_register_ioat(zdev, 0, zdev->start_dma, zdev->end_dma,
+                              (u64)zdev->dma_table)) {
+               rc = -EIO;
                goto free_bitmap;
+       }
 
        return 0;
 free_bitmap:
        return rc;
 }
 
-void zpci_dma_exit_device(struct zpci_dev *zdev)
+int zpci_dma_exit_device(struct zpci_dev *zdev)
 {
+       int cc = 0;
+
        /*
         * At this point, if the device is part of an IOMMU domain, this would
         * be a strong hint towards a bug in the IOMMU API (common) code and/or
         * simultaneous access via IOMMU and DMA API. So let's issue a warning.
         */
        WARN_ON(zdev->s390_domain);
-
-       if (zpci_unregister_ioat(zdev, 0))
-               return;
+       if (zdev_enabled(zdev))
+               cc = zpci_unregister_ioat(zdev, 0);
+       /*
+        * cc == 3 indicates the function is gone already. This can happen
+        * if the function was deconfigured/disabled suddenly and we have not
+        * received a new handle yet.
+        */
+       if (cc && cc != 3)
+               return -EIO;
 
        dma_cleanup_tables(zdev->dma_table);
        zdev->dma_table = NULL;
        zdev->iommu_bitmap = NULL;
        vfree(zdev->lazy_bitmap);
        zdev->lazy_bitmap = NULL;
-
        zdev->next_bit = 0;
+       return 0;
 }
 
 static int __init dma_alloc_cpu_table_caches(void)
 
        /* Even though the device is already gone we still
         * need to free zPCI resources as part of the disable.
         */
-       zpci_disable_device(zdev);
+       if (zdev->dma_table)
+               zpci_dma_exit_device(zdev);
+       if (zdev_enabled(zdev))
+               zpci_disable_device(zdev);
        zdev->state = ZPCI_FN_STATE_STANDBY;
 }
 
 
        pci_lock_rescan_remove();
        if (pci_dev_is_added(pdev)) {
                pci_stop_and_remove_bus_device(pdev);
-               ret = zpci_disable_device(zdev);
-               if (ret)
-                       goto out;
+               if (zdev->dma_table) {
+                       ret = zpci_dma_exit_device(zdev);
+                       if (ret)
+                               goto out;
+               }
+
+               if (zdev_enabled(zdev)) {
+                       ret = zpci_disable_device(zdev);
+                       if (ret)
+                               goto out;
+               }
 
                ret = zpci_enable_device(zdev);
                if (ret)
                        goto out;
+               ret = zpci_dma_init_device(zdev);
+               if (ret) {
+                       zpci_disable_device(zdev);
+                       goto out;
+               }
                pci_rescan_bus(zdev->zbus->bus);
        }
 out:
 
        struct zpci_dev *zdev = to_zpci_dev(dev);
        struct s390_domain_device *domain_device;
        unsigned long flags;
-       int rc;
+       int cc, rc;
 
        if (!zdev)
                return -ENODEV;
        if (!domain_device)
                return -ENOMEM;
 
-       if (zdev->dma_table)
-               zpci_dma_exit_device(zdev);
+       if (zdev->dma_table) {
+               cc = zpci_dma_exit_device(zdev);
+               if (cc) {
+                       rc = -EIO;
+                       goto out_free;
+               }
+       }
 
        zdev->dma_table = s390_domain->dma_table;
-       rc = zpci_register_ioat(zdev, 0, zdev->start_dma, zdev->end_dma,
+       cc = zpci_register_ioat(zdev, 0, zdev->start_dma, zdev->end_dma,
                                (u64) zdev->dma_table);
-       if (rc)
+       if (cc) {
+               rc = -EIO;
                goto out_restore;
+       }
 
        spin_lock_irqsave(&s390_domain->list_lock, flags);
        /* First device defines the DMA range limits */
 
 out_restore:
        zpci_dma_init_device(zdev);
+out_free:
        kfree(domain_device);
 
        return rc;