[RFC PATCHES 06/17] iommu: Add iommu page fault cookie helpers

[Lu Baolu]

Add an xarray in iommu_fault_param as place holder for per-{device, pasid}
fault cookie. The iommufd will use it to store the mapping of device object
ID and the device pointer. This allows the iommufd to quickly retrieve the
device object ID for a given {device, pasid} pair in the hot path of IO
page fault delivery.

Otherwise, the iommufd would have to maintain its own data structures to
map {device, pasid} pairs to object IDs, and then look up the object ID on
the critical path. This is not performance friendly.

The iommufd is supposed to set the cookie when a fault capable domain is
attached to the physical device or pasid, and clear the fault cookie when the
domain is removed.

Signed-off-by: Lu Baolu <baolu.lu@xxxxxxxxxxxxxxx>
---
 include/linux/iommu.h      |  2 ++
 drivers/iommu/iommu-priv.h |  3 +++
 drivers/iommu/iommu.c      | 45 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/include/linux/iommu.h b/include/linux/iommu.h
index c201704f9aea..9b0058ac971c 100644
--- a/include/linux/iommu.h
+++ b/include/linux/iommu.h
@@ -617,12 +617,14 @@ struct iommu_fault_event {
  * @data: handler private data
  * @faults: holds the pending faults which needs response
  * @lock: protect pending faults list
+ * @pasid_cookie: per-pasid fault cookie
  */
 struct iommu_fault_param {
 	iommu_dev_fault_handler_t handler;
 	void *data;
 	struct list_head faults;
 	struct mutex lock;
+	struct xarray pasid_cookie;
 };
 
 /**
diff --git a/drivers/iommu/iommu-priv.h b/drivers/iommu/iommu-priv.h
index a6e694f59f64..17ab989702a0 100644
--- a/drivers/iommu/iommu-priv.h
+++ b/drivers/iommu/iommu-priv.h
@@ -17,5 +17,8 @@ static inline const struct iommu_ops *dev_iommu_ops(struct device *dev)
 
 int iommu_group_replace_domain(struct iommu_group *group,
 			       struct iommu_domain *new_domain);
+void *iommu_set_device_fault_cookie(struct device *dev, ioasid_t pasid,
+				    void *cookie);
+void *iommu_get_device_fault_cookie(struct device *dev, ioasid_t pasid);
 
 #endif /* __LINUX_IOMMU_PRIV_H */
diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index cace57c066f4..2f81be7f3a90 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -1270,6 +1270,7 @@ int iommu_register_device_fault_handler(struct device *dev,
 	param->fault_param->data = data;
 	mutex_init(&param->fault_param->lock);
 	INIT_LIST_HEAD(&param->fault_param->faults);
+	xa_init(&param->fault_param->pasid_cookie);
 
 done_unlock:
 	mutex_unlock(&param->lock);
@@ -1435,6 +1436,50 @@ int iommu_page_response(struct device *dev,
 }
 EXPORT_SYMBOL_GPL(iommu_page_response);
 
+/**
+ * iommu_set_device_fault_cookie - Set a fault cookie for per-{device, pasid}
+ * @dev: the device to set the cookie
+ * @pasid: the pasid on this device
+ * @cookie: the opaque data
+ *
+ * Return the old cookie on success, or ERR_PTR(err#) on failure.
+ */
+void *iommu_set_device_fault_cookie(struct device *dev, ioasid_t pasid,
+				    void *cookie)
+{
+	struct iommu_fault_param *fault_param;
+	void *curr;
+
+	if (!dev->iommu || !dev->iommu->fault_param)
+		return ERR_PTR(-ENODEV);
+
+	fault_param = dev->iommu->fault_param;
+	curr = xa_store(&fault_param->pasid_cookie, pasid, cookie, GFP_KERNEL);
+
+	return xa_is_err(curr) ? ERR_PTR(xa_err(curr)) : curr;
+}
+EXPORT_SYMBOL_NS_GPL(iommu_set_device_fault_cookie, IOMMUFD_INTERNAL);
+
+/**
+ * iommu_get_device_fault_cookie - Get the fault cookie for {device, pasid}
+ * @dev: the device to set the cookie
+ * @pasid: the pasid on this device
+ *
+ * Return the cookie on success, or ERR_PTR(err#) on failure.
+ */
+void *iommu_get_device_fault_cookie(struct device *dev, ioasid_t pasid)
+{
+	struct iommu_fault_param *fault_param;
+
+	if (!dev->iommu || !dev->iommu->fault_param)
+		return ERR_PTR(-ENODEV);
+
+	fault_param = dev->iommu->fault_param;
+
+	return xa_load(&fault_param->pasid_cookie, pasid);
+}
+EXPORT_SYMBOL_NS_GPL(iommu_get_device_fault_cookie, IOMMUFD_INTERNAL);
+
 /**
  * iommu_group_id - Return ID for a group
  * @group: the group to ID
-- 
2.34.1