[PATCH v5 12/24] iommu/amd: Store per-segment iommu_dev_data in an xarray

[Suravee Suthikulpanit <suravee.suthikulpanit@amd.com>]

Replace the per-segment llist of iommu_dev_data with an xarray keyed
by devid so lookup is O(1). These objects stay immortal for the life
of the PCI segment: amd_iommu_release_device() keeps them for replug,
and IRQ/DTE paths look them up locklessly via xa_load().

Same-devid intern on xa_cmpxchg() covers replug and racing
find_dev_data(). Aliases still get separate keys. Intern remains
local to iommu.c; only the lookup table changes.

Signed-off-by: Suravee Suthikulpanit <suravee.suthikulpanit@amd.com>
---
 drivers/iommu/amd/amd_iommu_types.h |  6 ++---
 drivers/iommu/amd/init.c            |  4 +++-
 drivers/iommu/amd/iommu.c           | 34 +++++++++++++++--------------
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/drivers/iommu/amd/amd_iommu_types.h b/drivers/iommu/amd/amd_iommu_types.h
index 9a90b8fe0fe3..d02f49a39a90 100644
--- a/drivers/iommu/amd/amd_iommu_types.h
+++ b/drivers/iommu/amd/amd_iommu_types.h
@@ -22,6 +22,7 @@
 #include <linux/irqreturn.h>
 #include <linux/generic_pt/iommu.h>
 #include <linux/idr.h>
+#include <linux/xarray.h>
 
 #include <uapi/linux/iommufd.h>
 
@@ -605,8 +606,8 @@ struct amd_iommu_pci_seg {
 	/* List with all PCI segments in the system */
 	struct list_head list;
 
-	/* List of all available dev_data structures */
-	struct llist_head dev_data_list;
+	/* Immortal iommu_dev_data objects keyed by devid */
+	struct xarray dev_data_xa;
 
 	/* PCI segment number */
 	u16 id;
@@ -861,7 +862,6 @@ struct iommu_dev_data {
 	spinlock_t dte_lock;              /* DTE lock for 256-bit access */
 
 	struct list_head list;		  /* For domain->dev_list */
-	struct llist_node dev_data_list;  /* For global dev_data_list */
 	struct protection_domain *domain; /* Domain the device is bound to */
 	struct gcr3_tbl_info gcr3_info;   /* Per-device GCR3 table */
 	struct device *dev;
diff --git a/drivers/iommu/amd/init.c b/drivers/iommu/amd/init.c
index 17d321412dcc..1616fcfd361b 100644
--- a/drivers/iommu/amd/init.c
+++ b/drivers/iommu/amd/init.c
@@ -1737,7 +1737,7 @@ static struct amd_iommu_pci_seg *__init alloc_pci_segment(u16 id,
 		    SZ_4K);
 
 	pci_seg->id = id;
-	init_llist_head(&pci_seg->dev_data_list);
+	xa_init(&pci_seg->dev_data_xa);
 	INIT_LIST_HEAD(&pci_seg->unity_map);
 	list_add_tail(&pci_seg->list, &amd_iommu_pci_seg_list);
 
@@ -1756,6 +1756,7 @@ static struct amd_iommu_pci_seg *__init alloc_pci_segment(u16 id,
 	free_dev_table(pci_seg);
 err_free_pci_seg:
 	list_del(&pci_seg->list);
+	xa_destroy(&pci_seg->dev_data_xa);
 	kfree(pci_seg);
 	return NULL;
 }
@@ -1779,6 +1780,7 @@ static void __init free_pci_segments(void)
 
 	for_each_pci_segment_safe(pci_seg, next) {
 		list_del(&pci_seg->list);
+		xa_destroy(&pci_seg->dev_data_xa);
 		free_irq_lookup_table(pci_seg);
 		free_rlookup_table(pci_seg);
 		free_alias_table(pci_seg);
diff --git a/drivers/iommu/amd/iommu.c b/drivers/iommu/amd/iommu.c
index dd2cb79d506f..2b039f38f9ef 100644
--- a/drivers/iommu/amd/iommu.c
+++ b/drivers/iommu/amd/iommu.c
@@ -390,9 +390,14 @@ static struct amd_iommu *rlookup_amd_iommu(struct device *dev)
 	return __rlookup_amd_iommu(seg, PCI_SBDF_TO_DEVID(devid));
 }
 
+/*
+ * Allocate an immortal per-devid object stored in pci_seg->dev_data_xa.
+ * These are never erased: amd_iommu_release_device() keeps them for
+ * replug, and IRQ/DTE paths look them up locklessly via xa_load().
+ */
 static struct iommu_dev_data *alloc_dev_data(struct amd_iommu *iommu, u16 devid)
 {
-	struct iommu_dev_data *dev_data;
+	struct iommu_dev_data *dev_data, *old;
 	struct amd_iommu_pci_seg *pci_seg = iommu->pci_seg;
 
 	dev_data = kzalloc_obj(*dev_data);
@@ -404,26 +409,23 @@ static struct iommu_dev_data *alloc_dev_data(struct amd_iommu *iommu, u16 devid)
 	dev_data->devid = devid;
 	ratelimit_default_init(&dev_data->rs);
 
-	llist_add(&dev_data->dev_data_list, &pci_seg->dev_data_list);
+	old = xa_cmpxchg(&pci_seg->dev_data_xa, devid, NULL, dev_data,
+			 GFP_KERNEL);
+	if (xa_is_err(old)) {
+		kfree(dev_data);
+		return NULL;
+	}
+	if (old) {
+		kfree(dev_data);
+		return old;
+	}
+
 	return dev_data;
 }
 
 struct iommu_dev_data *search_dev_data(struct amd_iommu *iommu, u16 devid)
 {
-	struct iommu_dev_data *dev_data;
-	struct llist_node *node;
-	struct amd_iommu_pci_seg *pci_seg = iommu->pci_seg;
-
-	if (llist_empty(&pci_seg->dev_data_list))
-		return NULL;
-
-	node = pci_seg->dev_data_list.first;
-	llist_for_each_entry(dev_data, node, dev_data_list) {
-		if (dev_data->devid == devid)
-			return dev_data;
-	}
-
-	return NULL;
+	return xa_load(&iommu->pci_seg->dev_data_xa, devid);
 }
 
 static int clone_alias(struct pci_dev *pdev_origin, u16 alias, void *data)
-- 
2.34.1