iommu: Add iommu->id and iommu_device_get_from_id

For nested translation cases in a multi-IOMMU environment, add iommu->id
to struct iommu_device so a driver can simply get the driver private data
via the returning iommu pointer.

This allows iommufd reporting to a user space VMM, via its hw_info ioctl,
the existence of multiple underlying pyhsical IOMMUs. VMM should use this
information to create a corresponding number of virtual IOMMUs for nested
translation support. This is helpful for VMM to redirect its invalidation
requests to physical IOMMUs, and very essential for HW features, such as
AMD's vIOMMU and SMMU's VCMDQ, to share HW resources with vIOMMUs.

Signed-off-by: Nicolin Chen <nicolinc@nvidia.com>

drivers/iommu/iommu-priv.h

@@ -27,4 +27,6 @@ void iommu_device_unregister_bus(struct iommu_device *iommu,
 				 struct bus_type *bus,
 				 struct notifier_block *nb);
 
+struct iommu_device *iommu_device_get_from_id(u32 iommu_id);
+
 #endif /* __LINUX_IOMMU_PRIV_H */

drivers/iommu/iommu.c

@@ -39,6 +39,7 @@
 #include "iommu-sva.h"
 
 static struct kset *iommu_group_kset;
+static DEFINE_XARRAY_ALLOC(iommu_ids);
 static DEFINE_IDA(iommu_group_ida);
 static DEFINE_IDA(iommu_global_pasid_ida);
 
@@ -258,6 +259,10 @@ int iommu_device_register(struct iommu_device *iommu,
 	if (WARN_ON(is_module_address((unsigned long)ops) && !ops->owner))
 		return -EINVAL;
 
+	err = xa_alloc(&iommu_ids, &iommu->id, iommu, xa_limit_32b, GFP_KERNEL);
+	if (err)
+		return err;
+
 	iommu->ops = ops;
 	if (hwdev)
 		iommu->fwnode = dev_fwnode(hwdev);
@@ -268,8 +273,10 @@ int iommu_device_register(struct iommu_device *iommu,
 
 	for (int i = 0; i < ARRAY_SIZE(iommu_buses) && !err; i++)
 		err = bus_iommu_probe(iommu_buses[i]);
-	if (err)
+	if (err) {
 		iommu_device_unregister(iommu);
+		WARN_ON(iommu != xa_erase(&iommu_ids, iommu->id));
+	}
 	return err;
 }
 EXPORT_SYMBOL_GPL(iommu_device_register);
@@ -286,9 +293,15 @@ void iommu_device_unregister(struct iommu_device *iommu)
 	/* Pairs with the alloc in generic_single_device_group() */
 	iommu_group_put(iommu->singleton_group);
 	iommu->singleton_group = NULL;
+	WARN_ON(iommu != xa_erase(&iommu_ids, iommu->id));
 }
 EXPORT_SYMBOL_GPL(iommu_device_unregister);
 
+struct iommu_device *iommu_device_get_from_id(u32 iommu_id)
+{
+	return xa_load(&iommu_ids, iommu_id);
+}
+
 #if IS_ENABLED(CONFIG_IOMMUFD_TEST)
 void iommu_device_unregister_bus(struct iommu_device *iommu,
 				 struct bus_type *bus,

include/linux/iommu.h

@@ -583,6 +583,7 @@ struct iommu_domain_ops {
  * @dev: struct device for sysfs handling
  * @singleton_group: Used internally for drivers that have only one group
  * @max_pasids: number of supported PASIDs
+ * @id: iommu instance id
  */
 struct iommu_device {
 	struct list_head list;
@@ -591,6 +592,7 @@ struct iommu_device {
 	struct device *dev;
 	struct iommu_group *singleton_group;
 	u32 max_pasids;
+	u32 id;
 };
 
 /**