[PATCH v5 22/24] iommu/amd: Add translate-device-id alloc/free with vIOMMU owner

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

Extend the per-segment translate-device-id pool to record the owning
struct amd_iommu_viommu in the xarray.  Add amd_iommu_trans_devid_alloc()
and amd_iommu_trans_devid_free() to allocate ids from last_bdf down
to 0 so each id indexes the device and alias tables, then return them
on teardown.  Free loads the xarray slot and erases it only when the
owner matches @aviommu.

Signed-off-by: Suravee Suthikulpanit <suravee.suthikulpanit@amd.com>
---
 drivers/iommu/amd/amd_viommu.h  |   4 ++
 drivers/iommu/amd/trans_devid.c | 105 ++++++++++++++++++++++++++++++--
 2 files changed, 105 insertions(+), 4 deletions(-)

diff --git a/drivers/iommu/amd/amd_viommu.h b/drivers/iommu/amd/amd_viommu.h
index c91ba854847d..0c77c025f8b5 100644
--- a/drivers/iommu/amd/amd_viommu.h
+++ b/drivers/iommu/amd/amd_viommu.h
@@ -36,6 +36,10 @@ void amd_iommu_pci_seg_trans_devid_fini(struct amd_iommu_pci_seg *pci_seg);
 int amd_iommu_trans_devid_reserve(struct amd_iommu_pci_seg *pci_seg, u16 id);
 int amd_iommu_trans_devid_reserve_pci_aliases(struct amd_iommu *iommu,
 					      struct device *dev);
+int amd_iommu_trans_devid_alloc(struct amd_iommu_pci_seg *pci_seg,
+				struct amd_iommu_viommu *aviommu);
+void amd_iommu_trans_devid_free(struct amd_iommu_pci_seg *pci_seg, u16 id,
+				struct amd_iommu_viommu *aviommu);
 #else
 
 /*
diff --git a/drivers/iommu/amd/trans_devid.c b/drivers/iommu/amd/trans_devid.c
index 19a219519c92..3a4e09fe39d5 100644
--- a/drivers/iommu/amd/trans_devid.c
+++ b/drivers/iommu/amd/trans_devid.c
@@ -28,24 +28,55 @@ static inline bool trans_devid_xa_is_reserved(void *entry)
 	       xa_to_value(entry) == TRANS_DEVID_RESERVED;
 }
 
+static inline struct amd_iommu_viommu *trans_devid_xa_owner(void *entry)
+{
+	if (!entry || xa_is_value(entry))
+		return NULL;
+	return entry;
+}
+
 static inline void *trans_devid_xa_mk_reserved(void)
 {
 	return xa_mk_value(TRANS_DEVID_RESERVED);
 }
 
-static int trans_devid_xa_install_reserved_locked(struct amd_iommu_pci_seg *pci_seg,
-						  u16 id)
+static int trans_devid_find_free_locked(struct amd_iommu_pci_seg *pci_seg)
+{
+	int id;
+
+	for (id = pci_seg->last_bdf; id >= 0; id--) {
+		if (!xa_load(&pci_seg->trans_devid_xa, id))
+			return id;
+	}
+	return -ENOSPC;
+}
+
+static int trans_devid_xa_install_locked(struct amd_iommu_pci_seg *pci_seg,
+					 u16 id, void *entry)
 {
 	void *old;
 
-	old = xa_store(&pci_seg->trans_devid_xa, id,
-		       trans_devid_xa_mk_reserved(), GFP_KERNEL);
+	old = xa_store(&pci_seg->trans_devid_xa, id, entry, GFP_KERNEL);
 	if (xa_is_err(old))
 		return xa_err(old);
 	WARN_ON_ONCE(old);
 	return 0;
 }
 
+static int trans_devid_xa_install_allocated_locked(struct amd_iommu_pci_seg *pci_seg,
+						   u16 id,
+						   struct amd_iommu_viommu *aviommu)
+{
+	return trans_devid_xa_install_locked(pci_seg, id, aviommu);
+}
+
+static int trans_devid_xa_install_reserved_locked(struct amd_iommu_pci_seg *pci_seg,
+						u16 id)
+{
+	return trans_devid_xa_install_locked(pci_seg, id,
+					      trans_devid_xa_mk_reserved());
+}
+
 void amd_iommu_pci_seg_trans_devid_init(struct amd_iommu_pci_seg *pci_seg)
 {
 	mutex_init(&pci_seg->trans_devid_mutex);
@@ -130,3 +161,69 @@ int amd_iommu_trans_devid_reserve_pci_aliases(struct amd_iommu *iommu,
 	return pci_for_each_dma_alias(pdev, reserve_trans_devid_each_dma_alias,
 				      pci_seg);
 }
+
+/**
+ * amd_iommu_trans_devid_alloc - allocate a translate-device-id for @pci_seg
+ *
+ * The trans_devid is allocated from pci_seg->last_bdf down to 0 so the
+ * id is always a valid device-table and alias-table index.  PCI devices
+ * are generally enumerated from the low end of the bus range, so high
+ * ids are likely unused.
+ *
+ * Each vIOMMU receives its own translate-device-id from the per-segment pool.
+ * @aviommu is stored in the xarray as the slot owner.
+ *
+ * Return: allocated id on success, negative errno on failure.
+ */
+int amd_iommu_trans_devid_alloc(struct amd_iommu_pci_seg *pci_seg,
+				struct amd_iommu_viommu *aviommu)
+{
+	int id, ret;
+
+	mutex_lock(&pci_seg->trans_devid_mutex);
+	id = trans_devid_find_free_locked(pci_seg);
+	if (id < 0) {
+		ret = id;
+		goto unlock;
+	}
+
+	ret = trans_devid_xa_install_allocated_locked(pci_seg, id, aviommu);
+	if (ret)
+		goto unlock;
+
+	mutex_unlock(&pci_seg->trans_devid_mutex);
+	pr_debug("%s: Allocated trans_devid %#x (seg %#x)\n", __func__, id,
+		 pci_seg->id);
+	return id;
+
+unlock:
+	mutex_unlock(&pci_seg->trans_devid_mutex);
+	if (ret == -ENOSPC)
+		pr_err("%s: No free trans_devid found (seg %#x)\n", __func__,
+		       pci_seg->id);
+	return ret;
+}
+
+/**
+ * amd_iommu_trans_devid_free - return @id to the per-segment pool
+ *
+ * Erase @id only if the xarray owner is @aviommu. A reserved or foreign
+ * slot is left unchanged.
+ *
+ * Caller must hold @aviommu->trans_devid_lock if racing with relocation.
+ */
+void amd_iommu_trans_devid_free(struct amd_iommu_pci_seg *pci_seg, u16 id,
+				struct amd_iommu_viommu *aviommu)
+{
+	void *entry;
+
+	mutex_lock(&pci_seg->trans_devid_mutex);
+	entry = xa_load(&pci_seg->trans_devid_xa, id);
+	if (WARN_ON_ONCE(trans_devid_xa_owner(entry) != aviommu))
+		goto out;
+	xa_erase(&pci_seg->trans_devid_xa, id);
+	pr_debug("%s: Freed trans_devid %#x (seg %#x)\n", __func__, id,
+		 pci_seg->id);
+out:
+	mutex_unlock(&pci_seg->trans_devid_mutex);
+}
-- 
2.34.1