[PATCH 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 hand out ids from the high end of
the range and return them on teardown.

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

diff --git a/drivers/iommu/amd/amd_iommu.h b/drivers/iommu/amd/amd_iommu.h
index 200640ed0ef6..94ad1902ce14 100644
--- a/drivers/iommu/amd/amd_iommu.h
+++ b/drivers/iommu/amd/amd_iommu.h
@@ -225,6 +225,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
 static inline void
 amd_iommu_pci_seg_trans_devid_init(struct amd_iommu_pci_seg *pci_seg) { }
diff --git a/drivers/iommu/amd/trans_devid.c b/drivers/iommu/amd/trans_devid.c
index a40686a576ec..9d7d93c5dcdb 100644
--- a/drivers/iommu/amd/trans_devid.c
+++ b/drivers/iommu/amd/trans_devid.c
@@ -22,24 +22,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 = U16_MAX; 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);
@@ -121,3 +152,64 @@ 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 the highest id to the lowest id.
+ * Generally, the PCI devices enumerated from the beginning of the bus range.
+ * Therefore, ids in the high range are likely to not be used.
+ *
+ * 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
+ *
+ * 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_erase(&pci_seg->trans_devid_xa, id);
+	if (WARN_ON_ONCE(!entry || trans_devid_xa_owner(entry) != aviommu))
+		goto out;
+	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