[PATCH v2 10/19] iommufd/viommu: Add vdev_id helpers for IOMMU drivers

[Nicolin Chen <nicolinc@xxxxxxxxxx>]

Driver can call the iommufd_viommu_find_device() to find a device pointer
using its per-viommu virtual ID. The returned device must be protected by
the pair of iommufd_viommu_lock/unlock_vdev_id() function.

Put these three functions into a new viommu_api file, to build it with the
IOMMUFD_DRIVER config.

Signed-off-by: Nicolin Chen <nicolinc@xxxxxxxxxx>
---
 drivers/iommu/iommufd/Makefile     |  2 +-
 drivers/iommu/iommufd/viommu_api.c | 39 ++++++++++++++++++++++++++++++
 include/linux/iommufd.h            | 16 ++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 drivers/iommu/iommufd/viommu_api.c

diff --git a/drivers/iommu/iommufd/Makefile b/drivers/iommu/iommufd/Makefile
index df490e836b30..288ef3e895e3 100644
--- a/drivers/iommu/iommufd/Makefile
+++ b/drivers/iommu/iommufd/Makefile
@@ -13,4 +13,4 @@ iommufd-y := \
 iommufd-$(CONFIG_IOMMUFD_TEST) += selftest.o
 
 obj-$(CONFIG_IOMMUFD) += iommufd.o
-obj-$(CONFIG_IOMMUFD_DRIVER) += iova_bitmap.o
+obj-$(CONFIG_IOMMUFD_DRIVER) += iova_bitmap.o viommu_api.o
diff --git a/drivers/iommu/iommufd/viommu_api.c b/drivers/iommu/iommufd/viommu_api.c
new file mode 100644
index 000000000000..e0ee592ce834
--- /dev/null
+++ b/drivers/iommu/iommufd/viommu_api.c
@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/* Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES
+ */
+
+#include "iommufd_private.h"
+
+void iommufd_viommu_lock_vdev_id(struct iommufd_viommu *viommu)
+{
+	down_read(&viommu->vdev_ids_rwsem);
+}
+EXPORT_SYMBOL_NS_GPL(iommufd_viommu_lock_vdev_id, IOMMUFD);
+
+void iommufd_viommu_unlock_vdev_id(struct iommufd_viommu *viommu)
+{
+	up_read(&viommu->vdev_ids_rwsem);
+}
+EXPORT_SYMBOL_NS_GPL(iommufd_viommu_unlock_vdev_id, IOMMUFD);
+
+/*
+ * Find a device attached to an VIOMMU object using a virtual device ID that was
+ * set via an IOMMUFD_CMD_VIOMMU_SET_VDEV_ID. Callers of this function must call
+ * iommufd_viommu_lock_vdev_id() prior and iommufd_viommu_unlock_vdev_id() after
+ *
+ * Return device or NULL.
+ */
+struct device *iommufd_viommu_find_device(struct iommufd_viommu *viommu, u64 id)
+{
+	struct iommufd_vdev_id *vdev_id;
+
+	lockdep_assert_held(&viommu->vdev_ids_rwsem);
+
+	xa_lock(&viommu->vdev_ids);
+	vdev_id = xa_load(&viommu->vdev_ids, (unsigned long)id);
+	xa_unlock(&viommu->vdev_ids);
+	if (!vdev_id || vdev_id->id != id)
+		return NULL;
+	return vdev_id->idev->dev;
+}
+EXPORT_SYMBOL_NS_GPL(iommufd_viommu_find_device, IOMMUFD);
diff --git a/include/linux/iommufd.h b/include/linux/iommufd.h
index 85291b346348..364f151d281d 100644
--- a/include/linux/iommufd.h
+++ b/include/linux/iommufd.h
@@ -89,6 +89,9 @@ int iommufd_access_rw(struct iommufd_access *access, unsigned long iova,
 int iommufd_vfio_compat_ioas_get_id(struct iommufd_ctx *ictx, u32 *out_ioas_id);
 int iommufd_vfio_compat_ioas_create(struct iommufd_ctx *ictx);
 int iommufd_vfio_compat_set_no_iommu(struct iommufd_ctx *ictx);
+void iommufd_viommu_lock_vdev_id(struct iommufd_viommu *viommu);
+void iommufd_viommu_unlock_vdev_id(struct iommufd_viommu *viommu);
+struct device *iommufd_viommu_find_device(struct iommufd_viommu *viommu, u64 id);
 #else /* !CONFIG_IOMMUFD */
 static inline struct iommufd_ctx *iommufd_ctx_from_file(struct file *file)
 {
@@ -129,5 +132,18 @@ static inline int iommufd_vfio_compat_set_no_iommu(struct iommufd_ctx *ictx)
 {
 	return -EOPNOTSUPP;
 }
+
+void iommufd_viommu_lock_vdev_id(struct iommufd_viommu *viommu)
+{
+}
+
+void iommufd_viommu_unlock_vdev_id(struct iommufd_viommu *viommu)
+{
+}
+
+struct device *iommufd_viommu_find_device(struct iommufd_viommu *viommu, u64 id)
+{
+	return NULL;
+}
 #endif /* CONFIG_IOMMUFD */
 #endif
-- 
2.43.0