[PATCH v5 19/24] iommu/amd: Add per-segment translate device ID pool

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

A translate device ID (TransDevID) is a host device-table index used
by hardware vIOMMU. VFCTRL guest-misc names the slot; that DTE holds
the nest-parent v1 table for GPA->SPA walks of guest IOMMU traffic.

Reserve PCI requestor IDs in a per-segment pool so they are not
used as vIOMMU translation DTEs. The device table is per-segment,
so the pool is too.

Add trans_devid.c with pci_seg init/fini and
amd_iommu_trans_devid_reserve() from probe_device(). Keep
reservations for the pci_seg lifetime, including after
amd_iommu_release_device().

Signed-off-by: Suravee Suthikulpanit <suravee.suthikulpanit@amd.com>
---
 drivers/iommu/amd/Makefile          |  2 +-
 drivers/iommu/amd/amd_iommu_types.h | 10 ++++
 drivers/iommu/amd/amd_viommu.h      | 14 +++++
 drivers/iommu/amd/init.c            |  3 +
 drivers/iommu/amd/iommu.c           | 19 +++++++
 drivers/iommu/amd/trans_devid.c     | 87 +++++++++++++++++++++++++++++
 6 files changed, 134 insertions(+), 1 deletion(-)
 create mode 100644 drivers/iommu/amd/trans_devid.c

diff --git a/drivers/iommu/amd/Makefile b/drivers/iommu/amd/Makefile
index e1e824b9c7b0..12c3fe83e4ce 100644
--- a/drivers/iommu/amd/Makefile
+++ b/drivers/iommu/amd/Makefile
@@ -1,4 +1,4 @@
 # SPDX-License-Identifier: GPL-2.0-only
 obj-y += iommu.o init.o quirks.o ppr.o pasid.o
-obj-$(CONFIG_AMD_IOMMU_IOMMUFD) += iommufd.o nested.o viommu.o
+obj-$(CONFIG_AMD_IOMMU_IOMMUFD) += iommufd.o nested.o viommu.o trans_devid.o
 obj-$(CONFIG_AMD_IOMMU_DEBUGFS) += debugfs.o
diff --git a/drivers/iommu/amd/amd_iommu_types.h b/drivers/iommu/amd/amd_iommu_types.h
index b5d1b23791da..6b3a43952797 100644
--- a/drivers/iommu/amd/amd_iommu_types.h
+++ b/drivers/iommu/amd/amd_iommu_types.h
@@ -681,6 +681,16 @@ struct amd_iommu_pci_seg {
 	 * parsing time.
 	 */
 	struct list_head unity_map;
+
+#ifdef CONFIG_AMD_IOMMU_IOMMUFD
+	/*
+	 * Per-segment translate-device-id pool indexed by id.  Entries are:
+	 * absent (FREE), xa value TRANS_DEVID_RESERVED, or a pointer to the
+	 * owning struct amd_iommu_viommu (ALLOCATED).
+	 */
+	struct mutex trans_devid_mutex;
+	struct xarray trans_devid_xa;
+#endif
 };
 
 /*
diff --git a/drivers/iommu/amd/amd_viommu.h b/drivers/iommu/amd/amd_viommu.h
index f17ac13da213..b3717a006301 100644
--- a/drivers/iommu/amd/amd_viommu.h
+++ b/drivers/iommu/amd/amd_viommu.h
@@ -30,6 +30,10 @@ void amd_viommu_domain_id_update_locked(struct amd_iommu_viommu *aviommu,
 
 void amd_viommu_set_device_mapping(struct amd_iommu_viommu *aviommu,
 				   u16 hdev_id, u16 gdev_id);
+
+void amd_iommu_pci_seg_trans_devid_init(struct amd_iommu_pci_seg *pci_seg);
+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);
 #else
 
 /*
@@ -66,6 +70,16 @@ static inline void amd_viommu_domain_id_update_locked(struct amd_iommu_viommu *a
 {
 }
 
+static inline void
+amd_iommu_pci_seg_trans_devid_init(struct amd_iommu_pci_seg *pci_seg)
+{
+}
+
+static inline void
+amd_iommu_pci_seg_trans_devid_fini(struct amd_iommu_pci_seg *pci_seg)
+{
+}
+
 #endif /* CONFIG_AMD_IOMMU_IOMMUFD */
 
 #endif /* AMD_VIOMMU_H */
diff --git a/drivers/iommu/amd/init.c b/drivers/iommu/amd/init.c
index 1616fcfd361b..c646cf6bea6e 100644
--- a/drivers/iommu/amd/init.c
+++ b/drivers/iommu/amd/init.c
@@ -1740,6 +1740,7 @@ static struct amd_iommu_pci_seg *__init alloc_pci_segment(u16 id,
 	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);
+	amd_iommu_pci_seg_trans_devid_init(pci_seg);
 
 	if (alloc_dev_table(pci_seg))
 		goto err_free_pci_seg;
@@ -1756,6 +1757,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);
+	amd_iommu_pci_seg_trans_devid_fini(pci_seg);
 	xa_destroy(&pci_seg->dev_data_xa);
 	kfree(pci_seg);
 	return NULL;
@@ -1780,6 +1782,7 @@ static void __init free_pci_segments(void)
 
 	for_each_pci_segment_safe(pci_seg, next) {
 		list_del(&pci_seg->list);
+		amd_iommu_pci_seg_trans_devid_fini(pci_seg);
 		xa_destroy(&pci_seg->dev_data_xa);
 		free_irq_lookup_table(pci_seg);
 		free_rlookup_table(pci_seg);
diff --git a/drivers/iommu/amd/iommu.c b/drivers/iommu/amd/iommu.c
index d620586d5734..34915e5ff156 100644
--- a/drivers/iommu/amd/iommu.c
+++ b/drivers/iommu/amd/iommu.c
@@ -2731,6 +2731,22 @@ static struct iommu_device *amd_iommu_probe_device(struct device *dev)
 	if (dev_is_pci(dev))
 		pci_prepare_ats(to_pci_dev(dev), PAGE_SHIFT);
 
+#if IS_ENABLED(CONFIG_AMD_IOMMU_IOMMUFD)
+	/*
+	 * Occupy this RID in the per-segment translate-device-id pool
+	 * so vIOMMU alloc cannot reuse it. Probe is the RID lifetime;
+	 * it is not released on unplug. If reserve fails, the IOMMU
+	 * core frees dev->iommu and a later re-probe retries.
+	 */
+	ret = amd_iommu_trans_devid_reserve(iommu->pci_seg, dev_data->devid);
+	if (ret) {
+		pr_err("%s: Failed to reserve device id %#x\n", __func__,
+		       dev_data->devid);
+		iommu_dev = ERR_PTR(ret);
+		goto out_err;
+	}
+#endif
+
 out_err:
 	return iommu_dev;
 }
@@ -2744,6 +2760,9 @@ static void amd_iommu_release_device(struct device *dev)
 	/*
 	 * We keep dev_data around for unplugged devices and reuse it when the
 	 * device is re-plugged - not doing so would introduce a ton of races.
+	 * Translate-device-id reservations stay as well: the RID remains
+	 * occupied for the pci_seg lifetime so later alloc cannot reuse a
+	 * DTE index that still belongs to this function.
 	 */
 }
 
diff --git a/drivers/iommu/amd/trans_devid.c b/drivers/iommu/amd/trans_devid.c
new file mode 100644
index 000000000000..195551fcfc1a
--- /dev/null
+++ b/drivers/iommu/amd/trans_devid.c
@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/*
+ * Copyright (C) 2026 Advanced Micro Devices, Inc.
+ *
+ * AMD vIOMMU translate-device-id management.
+ *
+ * The pool is per PCI segment because the AMD IOMMU device table is
+ * per-segment.  Each id must be allocated from unused slots in that
+ * segment.  It is used to program the vIOMMU VF Control register to
+ * specify the DTE used to contain the GPA->SPA mapping (v1 page table).
+ */
+
+#include <linux/kernel.h>
+#include <linux/xarray.h>
+
+#include "amd_iommu.h"
+#include "amd_viommu.h"
+
+enum trans_devid_state {
+	TRANS_DEVID_FREE = 0,
+	TRANS_DEVID_RESERVED,
+};
+
+static inline bool trans_devid_xa_is_reserved(void *entry)
+{
+	return entry && xa_is_value(entry) &&
+	       xa_to_value(entry) == TRANS_DEVID_RESERVED;
+}
+
+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)
+{
+	void *old;
+
+	old = xa_store(&pci_seg->trans_devid_xa, id,
+		       trans_devid_xa_mk_reserved(), GFP_KERNEL);
+	if (xa_is_err(old))
+		return xa_err(old);
+	WARN_ON_ONCE(old);
+	return 0;
+}
+
+void amd_iommu_pci_seg_trans_devid_init(struct amd_iommu_pci_seg *pci_seg)
+{
+	mutex_init(&pci_seg->trans_devid_mutex);
+	xa_init(&pci_seg->trans_devid_xa);
+}
+
+void amd_iommu_pci_seg_trans_devid_fini(struct amd_iommu_pci_seg *pci_seg)
+{
+	xa_destroy(&pci_seg->trans_devid_xa);
+	mutex_destroy(&pci_seg->trans_devid_mutex);
+}
+
+/*
+ * amd_iommu_trans_devid_reserve - occupy @id so it is never returned by alloc
+ *
+ * Reservation is done when probing the device (see amd_iommu_probe_device()).
+ * It is not released from amd_iommu_release_device(); the slot stays
+ * reserved for the pci_seg lifetime so replug cannot race later alloc.
+ *
+ * Return: 0 on success.  A second reserve of an already-reserved @id succeeds.
+ */
+int amd_iommu_trans_devid_reserve(struct amd_iommu_pci_seg *pci_seg, u16 id)
+{
+	void *entry;
+	int ret = 0;
+
+	mutex_lock(&pci_seg->trans_devid_mutex);
+	entry = xa_load(&pci_seg->trans_devid_xa, id);
+	if (trans_devid_xa_is_reserved(entry))
+		goto unlock;
+
+	ret = trans_devid_xa_install_reserved_locked(pci_seg, id);
+unlock:
+	mutex_unlock(&pci_seg->trans_devid_mutex);
+
+	if (!ret)
+		pr_debug("%s: Reserved trans_devid %#x (seg %#x)\n", __func__, id,
+			 pci_seg->id);
+	return ret;
+}
-- 
2.34.1