TMP Add set_rid_user

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

drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.c

@@ -2964,11 +2964,50 @@ arm_smmu_domain_alloc_user(struct device *dev, struct iommu_domain *parent,
 	return __arm_smmu_domain_alloc(type, s2, master, user_cfg);
 }
 
+static int arm_smmu_set_rid_user(struct device *dev, u32 rid, u32 rid_base)
+{
+	struct arm_smmu_master *master = dev_iommu_priv_get(dev);
+	struct arm_smmu_device *smmu = master->smmu;
+	u32 sid_user = rid_base | rid;
+	int ret = 0;
+
+	if (!sid_user)
+		return -EINVAL;
+
+	ret = xa_alloc(&smmu->user_streams, &sid_user, &master->streams[0],
+		       XA_LIMIT(sid_user, sid_user), GFP_KERNEL_ACCOUNT);
+	if (ret)
+		return ret;
+	return 0;
+}
+
+static int arm_smmu_unset_rid_user(struct device *dev, u32 rid, u32 rid_base)
+{
+	struct arm_smmu_master *master = dev_iommu_priv_get(dev);
+	struct arm_smmu_device *smmu = master->smmu;
+	struct arm_smmu_stream *stream;
+	u32 sid_user = rid_base | rid;
+
+	xa_lock(&smmu->user_streams);
+	stream = __xa_erase(&smmu->user_streams, sid_user);
+	if (stream != master->streams) {
+		WARN_ON(__xa_alloc(&smmu->user_streams, &sid_user, stream,
+				   XA_LIMIT(sid_user, sid_user),
+				   GFP_KERNEL_ACCOUNT));
+		xa_unlock(&smmu->user_streams);
+		return -EINVAL;
+	}
+	xa_unlock(&smmu->user_streams);
+	return 0;
+}
+
 static struct iommu_ops arm_smmu_ops = {
 	.capable		= arm_smmu_capable,
 	.hw_info		= arm_smmu_hw_info,
 	.domain_alloc		= arm_smmu_domain_alloc,
 	.domain_alloc_user	= arm_smmu_domain_alloc_user,
+	.set_rid_user		= arm_smmu_set_rid_user,
+	.unset_rid_user		= arm_smmu_unset_rid_user,
 	.probe_device		= arm_smmu_probe_device,
 	.release_device		= arm_smmu_release_device,
 	.device_group		= arm_smmu_device_group,
@@ -3205,6 +3244,7 @@ static int arm_smmu_init_structures(struct arm_smmu_device *smmu)
 
 	mutex_init(&smmu->streams_mutex);
 	smmu->streams = RB_ROOT;
+	xa_init_flags(&smmu->user_streams, XA_FLAGS_ALLOC1 | XA_FLAGS_ACCOUNT);
 
 	ret = arm_smmu_init_queues(smmu);
 	if (ret)

drivers/iommu/arm/arm-smmu-v3/arm-smmu-v3.h

@@ -680,6 +680,7 @@ struct arm_smmu_device {
 
 	struct rb_root			streams;
 	struct mutex			streams_mutex;
+	struct xarray			user_streams;
 };
 
 struct arm_smmu_stream {

drivers/iommu/iommufd/device.c

@@ -738,6 +738,45 @@ void iommufd_device_detach(struct iommufd_device *idev)
 }
 EXPORT_SYMBOL_NS_GPL(iommufd_device_detach, IOMMUFD);
 
+int iommufd_device_set_rid(struct iommufd_ucmd *ucmd)
+{
+	struct iommufd_device_set_rid *cmd = ucmd->cmd;
+	struct iommufd_device *idev;
+	const struct iommu_ops *ops;
+	int rc = 0;
+
+	if (!cmd->rid)
+		return -EINVAL;
+
+	idev = iommufd_get_device(ucmd, cmd->dev_id);
+	if (IS_ERR(idev))
+		return PTR_ERR(idev);
+
+	ops = dev_iommu_ops(idev->dev);
+	if (ops && ops->set_rid_user)
+		rc = ops->set_rid_user(idev->dev, cmd->rid, cmd->rid_base);
+	iommufd_put_object(&idev->obj);
+	return rc;
+}
+
+int iommufd_device_unset_rid(struct iommufd_ucmd *ucmd)
+{
+	struct iommufd_device_unset_rid *cmd = ucmd->cmd;
+	struct iommufd_device *idev;
+	const struct iommu_ops *ops;
+	int rc = 0;
+
+	idev = iommufd_get_device(ucmd, cmd->dev_id);
+	if (IS_ERR(idev))
+		return PTR_ERR(idev);
+
+	ops = dev_iommu_ops(idev->dev);
+	if (ops && ops->unset_rid_user)
+		rc = ops->unset_rid_user(idev->dev, 0, 0);
+	iommufd_put_object(&idev->obj);
+	return rc;
+}
+
 void iommufd_access_destroy_object(struct iommufd_object *obj)
 {
 	struct iommufd_access *access =

drivers/iommu/iommufd/iommufd_private.h

@@ -320,6 +320,8 @@ iommufd_get_device(struct iommufd_ucmd *ucmd, u32 id)
 
 void iommufd_device_destroy(struct iommufd_object *obj);
 int iommufd_device_get_hw_info(struct iommufd_ucmd *ucmd);
+int iommufd_device_set_rid(struct iommufd_ucmd *ucmd);
+int iommufd_device_unset_rid(struct iommufd_ucmd *ucmd);
 
 extern const u64 iommufd_hwpt_type_bitmaps[];
 

drivers/iommu/iommufd/main.c

@@ -323,6 +323,10 @@ static const struct iommufd_ioctl_op iommufd_ioctl_ops[] = {
 		 val64),
 	IOCTL_OP(IOMMU_VFIO_IOAS, iommufd_vfio_ioas, struct iommu_vfio_ioas,
 		 __reserved),
+	IOCTL_OP(IOMMU_DEVICE_SET_RID, iommufd_device_set_rid,
+		 struct iommufd_device_set_rid, rid_base),
+	IOCTL_OP(IOMMU_DEVICE_UNSET_RID, iommufd_device_unset_rid,
+		 struct iommufd_device_unset_rid, dev_id),
 #ifdef CONFIG_IOMMUFD_TEST
 	IOCTL_OP(IOMMU_TEST_CMD, iommufd_test, struct iommu_test_cmd, last),
 #endif

include/linux/iommu.h

@@ -274,6 +274,8 @@ struct iommu_ops {
 	struct iommu_domain *(*domain_alloc_user)(struct device *dev,
 						  struct iommu_domain *parent,
 						  const void *user_data);
+	int (*set_rid_user)(struct device *dev, u32 rid, u32 rid_base);
+	int (*unset_rid_user)(struct device *dev, u32 rid, u32 rid_base);
 
 	struct iommu_device *(*probe_device)(struct device *dev);
 	void (*release_device)(struct device *dev);

include/uapi/linux/iommufd.h

@@ -48,6 +48,8 @@ enum {
 	IOMMUFD_CMD_HWPT_ALLOC,
 	IOMMUFD_CMD_DEVICE_GET_HW_INFO,
 	IOMMUFD_CMD_HWPT_INVALIDATE,
+	IOMMUFD_CMD_DEVICE_SET_RID,
+	IOMMUFD_CMD_DEVICE_UNSET_RID,
 };
 
 /**
@@ -682,4 +684,33 @@ struct iommu_hwpt_invalidate {
 	__aligned_u64 data_uptr;
 };
 #define IOMMU_HWPT_INVALIDATE _IO(IOMMUFD_TYPE, IOMMUFD_CMD_HWPT_INVALIDATE)
+
+/**
+ * struct iommufd_device_set_rid - ioctl(IOMMU_DEVICE_SET_RID)
+ * @size: sizeof(struct iommufd_device_set_rid)
+ * @dev_id: The device to set a Request ID
+ * @rid: PCI Request ID
+ *       Bits [15:8] are the Bus number.
+ *       Bits [7:3] are the Device number.
+ *       Bits [2:0] are the Function number.
+ * @rid_base: Base ID to map PCI Request ID
+ */
+struct iommufd_device_set_rid {
+	__u32 size;
+	__u32 dev_id;
+	__u32 rid;
+	__u32 rid_base;
+};
+#define IOMMU_DEVICE_SET_RID _IO(IOMMUFD_TYPE, IOMMUFD_CMD_DEVICE_SET_RID)
+
+/**
+ * struct iommufd_device_unset_rid - ioctl(IOMMU_DEVICE_UNSET_RID)
+ * @size: sizeof(struct iommufd_device_unset_rid)
+ * @dev_id: The device to unset a Request ID
+ */
+struct iommufd_device_unset_rid {
+	__u32 size;
+	__u32 dev_id;
+};
+#define IOMMU_DEVICE_UNSET_RID _IO(IOMMUFD_TYPE, IOMMUFD_CMD_DEVICE_UNSET_RID)
 #endif