struct vfio_pci_irq_ctx *vfio_irq_ctx_get(struct vfio_pci_core_device *vdev,
                                          unsigned long index)
 {
-       if (index >= vdev->num_ctx)
-               return NULL;
-       return &vdev->ctx[index];
+       return xa_load(&vdev->ctx, index);
 }
 
-static void vfio_irq_ctx_free_all(struct vfio_pci_core_device *vdev)
+static void vfio_irq_ctx_free(struct vfio_pci_core_device *vdev,
+                             struct vfio_pci_irq_ctx *ctx, unsigned long index)
 {
-       kfree(vdev->ctx);
+       xa_erase(&vdev->ctx, index);
+       kfree(ctx);
 }
 
-static int vfio_irq_ctx_alloc_num(struct vfio_pci_core_device *vdev,
-                                 unsigned long num)
+static struct vfio_pci_irq_ctx *
+vfio_irq_ctx_alloc(struct vfio_pci_core_device *vdev, unsigned long index)
 {
-       vdev->ctx = kcalloc(num, sizeof(struct vfio_pci_irq_ctx),
-                           GFP_KERNEL_ACCOUNT);
-       if (!vdev->ctx)
-               return -ENOMEM;
+       struct vfio_pci_irq_ctx *ctx;
+       int ret;
 
-       return 0;
+       ctx = kzalloc(sizeof(*ctx), GFP_KERNEL_ACCOUNT);
+       if (!ctx)
+               return NULL;
+
+       ret = xa_insert(&vdev->ctx, index, ctx, GFP_KERNEL_ACCOUNT);
+       if (ret) {
+               kfree(ctx);
+               return NULL;
+       }
+
+       return ctx;
 }
 
 /*
 static int vfio_intx_enable(struct vfio_pci_core_device *vdev)
 {
        struct vfio_pci_irq_ctx *ctx;
-       int ret;
 
        if (!is_irq_none(vdev))
                return -EINVAL;
        if (!vdev->pdev->irq)
                return -ENODEV;
 
-       ret = vfio_irq_ctx_alloc_num(vdev, 1);
-       if (ret)
-               return ret;
-
-       ctx = vfio_irq_ctx_get(vdev, 0);
-       if (!ctx) {
-               vfio_irq_ctx_free_all(vdev);
-               return -EINVAL;
-       }
+       ctx = vfio_irq_ctx_alloc(vdev, 0);
+       if (!ctx)
+               return -ENOMEM;
 
        vdev->num_ctx = 1;
 
        vfio_intx_set_signal(vdev, -1);
        vdev->irq_type = VFIO_PCI_NUM_IRQS;
        vdev->num_ctx = 0;
-       vfio_irq_ctx_free_all(vdev);
+       vfio_irq_ctx_free(vdev, ctx, 0);
 }
 
 /*
        if (!is_irq_none(vdev))
                return -EINVAL;
 
-       ret = vfio_irq_ctx_alloc_num(vdev, nvec);
-       if (ret)
-               return ret;
-
        /* return the number of supported vectors if we can't get all: */
        cmd = vfio_pci_memory_lock_and_enable(vdev);
        ret = pci_alloc_irq_vectors(pdev, 1, nvec, flag);
                if (ret > 0)
                        pci_free_irq_vectors(pdev);
                vfio_pci_memory_unlock_and_restore(vdev, cmd);
-               vfio_irq_ctx_free_all(vdev);
                return ret;
        }
        vfio_pci_memory_unlock_and_restore(vdev, cmd);
        if (vector >= vdev->num_ctx)
                return -EINVAL;
 
-       ctx = vfio_irq_ctx_get(vdev, vector);
-       if (!ctx)
-               return -EINVAL;
        irq = pci_irq_vector(pdev, vector);
+       if (irq < 0)
+               return -EINVAL;
 
-       if (ctx->trigger) {
+       ctx = vfio_irq_ctx_get(vdev, vector);
+
+       if (ctx) {
                irq_bypass_unregister_producer(&ctx->producer);
 
                cmd = vfio_pci_memory_lock_and_enable(vdev);
                vfio_pci_memory_unlock_and_restore(vdev, cmd);
                kfree(ctx->name);
                eventfd_ctx_put(ctx->trigger);
-               ctx->trigger = NULL;
+               vfio_irq_ctx_free(vdev, ctx, vector);
        }
 
        if (fd < 0)
                return 0;
 
+       ctx = vfio_irq_ctx_alloc(vdev, vector);
+       if (!ctx)
+               return -ENOMEM;
+
        ctx->name = kasprintf(GFP_KERNEL_ACCOUNT, "vfio-msi%s[%d](%s)",
                              msix ? "x" : "", vector, pci_name(pdev));
-       if (!ctx->name)
-               return -ENOMEM;
+       if (!ctx->name) {
+               ret = -ENOMEM;
+               goto out_free_ctx;
+       }
 
        trigger = eventfd_ctx_fdget(fd);
        if (IS_ERR(trigger)) {
        eventfd_ctx_put(trigger);
 out_free_name:
        kfree(ctx->name);
+out_free_ctx:
+       vfio_irq_ctx_free(vdev, ctx, vector);
        return ret;
 }
 
 {
        struct pci_dev *pdev = vdev->pdev;
        struct vfio_pci_irq_ctx *ctx;
-       unsigned int i;
+       unsigned long i;
        u16 cmd;
 
-       for (i = 0; i < vdev->num_ctx; i++) {
-               ctx = vfio_irq_ctx_get(vdev, i);
-               if (ctx) {
-                       vfio_virqfd_disable(&ctx->unmask);
-                       vfio_virqfd_disable(&ctx->mask);
-                       vfio_msi_set_vector_signal(vdev, i, -1, msix);
-               }
+       xa_for_each(&vdev->ctx, i, ctx) {
+               vfio_virqfd_disable(&ctx->unmask);
+               vfio_virqfd_disable(&ctx->mask);
+               vfio_msi_set_vector_signal(vdev, i, -1, msix);
        }
 
        cmd = vfio_pci_memory_lock_and_enable(vdev);
 
        vdev->irq_type = VFIO_PCI_NUM_IRQS;
        vdev->num_ctx = 0;
-       vfio_irq_ctx_free_all(vdev);
 }
 
 /*
 
        for (i = start; i < start + count; i++) {
                ctx = vfio_irq_ctx_get(vdev, i);
-               if (!ctx || !ctx->trigger)
+               if (!ctx)
                        continue;
                if (flags & VFIO_IRQ_SET_DATA_NONE) {
                        eventfd_signal(ctx->trigger, 1);