return rc;
 }
 
-static void afu_unmap(struct kref *ref)
-{
-       struct afu *afu = container_of(ref, struct afu, mapcount);
-
-       if (likely(afu->afu_map)) {
-               cxl_psa_unmap((void __iomem *)afu->afu_map);
-               afu->afu_map = NULL;
-       }
-}
-
 /**
  * cxlflash_driver_info() - information handler for this host driver
  * @host:      SCSI host associated with device.
        ulong lock_flags;
        int nseg = 0;
        int rc = 0;
-       int kref_got = 0;
 
        dev_dbg_ratelimited(dev, "%s: (scp=%p) %d/%d/%d/%llu "
                            "cdb=(%08x-%08x-%08x-%08x)\n",
                break;
        }
 
-       kref_get(&cfg->afu->mapcount);
-       kref_got = 1;
-
        if (likely(sg)) {
                nseg = scsi_dma_map(scp);
                if (unlikely(nseg < 0)) {
        if (unlikely(rc))
                scsi_dma_unmap(scp);
 out:
-       if (kref_got)
-               kref_put(&afu->mapcount, afu_unmap);
        return rc;
 }
 
  *
  * Safe to call with AFU in a partially allocated/initialized state.
  *
- * Waits for any active internal AFU commands to timeout and then unmaps
- * the MMIO space.
+ * Cancels scheduled worker threads, waits for any active internal AFU
+ * commands to timeout and then unmaps the MMIO space.
  */
 static void stop_afu(struct cxlflash_cfg *cfg)
 {
        struct afu *afu = cfg->afu;
 
+       cancel_work_sync(&cfg->work_q);
+
        if (likely(afu)) {
                while (atomic_read(&afu->cmds_active))
                        ssleep(1);
                        cxl_psa_unmap((void __iomem *)afu->afu_map);
                        afu->afu_map = NULL;
                }
-               kref_put(&afu->mapcount, afu_unmap);
        }
 }
 
                scsi_remove_host(cfg->host);
                /* fall through */
        case INIT_STATE_AFU:
-               cancel_work_sync(&cfg->work_q);
                term_afu(cfg);
        case INIT_STATE_PCI:
                pci_disable_device(pdev);
                                __func__, port);
                        cfg->lr_state = LINK_RESET_REQUIRED;
                        cfg->lr_port = port;
-                       kref_get(&cfg->afu->mapcount);
                        schedule_work(&cfg->work_q);
                }
 
 
                if (info->action & SCAN_HOST) {
                        atomic_inc(&cfg->scan_host_needed);
-                       kref_get(&cfg->afu->mapcount);
                        schedule_work(&cfg->work_q);
                }
        }
                rc = -ENOMEM;
                goto err1;
        }
-       kref_init(&afu->mapcount);
 
        /* No byte reverse on reading afu_version or string will be backwards */
        reg = readq(&afu->afu_map->global.regs.afu_version);
                        "interface version %016llx\n", afu->version,
                       afu->interface_version);
                rc = -EINVAL;
-               goto err2;
+               goto err1;
        }
 
        if (afu_is_sq_cmd_mode(afu)) {
        rc = start_afu(cfg);
        if (rc) {
                dev_err(dev, "%s: start_afu failed, rc=%d\n", __func__, rc);
-               goto err2;
+               goto err1;
        }
 
        afu_err_intr_init(cfg->afu);
        dev_dbg(dev, "%s: returning rc=%d\n", __func__, rc);
        return rc;
 
-err2:
-       kref_put(&afu->mapcount, afu_unmap);
 err1:
        term_intr(cfg, UNMAP_THREE);
        term_mc(cfg);
 
        if (atomic_dec_if_positive(&cfg->scan_host_needed) >= 0)
                scsi_scan_host(cfg->host);
-       kref_put(&afu->mapcount, afu_unmap);
 }
 
 /**