[PATCH 17/30] iommu: Add iommu_fw_alloc_per_device_ids()

[Jason Gunthorpe <jgg@nvidia.com>]

This helper can be called after iommu_of_get_single_iommu() to
automatically extract the raw list of IDs. It can be used by drivers that
expect to have a list of simple u32 IDs for a single IOMMU instance.

The driver must provide a per-driver structure with a trailing flex array
to hold the IDs. This helper will allocate the structure, size the ids
flex array, and fill it in with data.

Driver's should follow a typical pattern in their probe_device:

  struct tegra_smmu_device {
	  struct tegra_smmu *smmu;
	  unsigned int num_ids;
	  u32 ids[] __counted_by(num_ids);
  };

	smmu = iommu_of_get_single_iommu(pinf, &tegra_smmu_ops, 1,
					 struct tegra_smmu, iommu);
	if (IS_ERR(smmu)) return ERR_CAST(smmu);

	smmu_device = iommu_fw_alloc_per_device_ids(pinf, smmu_device);
	if (IS_ERR(smmu_device)) return ERR_CAST(smmu_device);
	[..]
	dev_iommu_priv_set(dev, data);
        return &iommu->iommu;

Signed-off-by: Jason Gunthorpe <jgg@nvidia.com>
---
 drivers/iommu/iommu.c        | 21 ++++++++++++++
 drivers/iommu/of_iommu.c     | 21 +++++++++++++-
 include/linux/iommu-driver.h | 56 ++++++++++++++++++++++++++++++++++++
 3 files changed, 97 insertions(+), 1 deletion(-)

diff --git a/drivers/iommu/iommu.c b/drivers/iommu/iommu.c
index ca411ad14c1182..caf14a53ed1952 100644
--- a/drivers/iommu/iommu.c
+++ b/drivers/iommu/iommu.c
@@ -3067,6 +3067,27 @@ struct iommu_device *iommu_fw_finish_get_single(struct iommu_probe_info *pinf)
 	return pinf->cached_iommu;
 }
 
+/*
+ * This function shouldn't be called directly without a pretty good reason,
+ * prefer to structure the driver to use iommu_fw_alloc_per_device_ids()
+ * instead.
+ */
+int iommu_fw_get_u32_ids(struct iommu_probe_info *pinf, u32 *ids)
+{
+	if (WARN_ON(!pinf->get_u32_ids))
+		return -EINVAL;
+
+	/*
+	 * We pre-parse a small number if IDs and keep it on the stack. If that
+	 * isn't enough then just reparse again.
+	 */
+	if (pinf->num_ids > ARRAY_SIZE(pinf->cached_ids))
+		return pinf->get_u32_ids(pinf, ids);
+	memcpy(ids, pinf->cached_ids, pinf->num_ids * sizeof(*ids));
+	return 0;
+}
+EXPORT_SYMBOL_GPL(iommu_fw_get_u32_ids);
+
 int iommu_fwspec_init(struct device *dev, struct fwnode_handle *iommu_fwnode,
 		      const struct iommu_ops *ops)
 {
diff --git a/drivers/iommu/of_iommu.c b/drivers/iommu/of_iommu.c
index 8d5495f03dbbcb..6f6e442f899ded 100644
--- a/drivers/iommu/of_iommu.c
+++ b/drivers/iommu/of_iommu.c
@@ -332,10 +332,28 @@ static int parse_single_iommu(struct of_phandle_args *iommu_spec, void *_info)
 	iommu = parse_iommu(info, iommu_spec);
 	if (IS_ERR(iommu))
 		return PTR_ERR(iommu);
-	info->pinf->num_ids++;
+	iommu_fw_cache_id(info->pinf, iommu_spec->args[0]);
 	return 0;
 }
 
+static int parse_read_ids(struct of_phandle_args *iommu_spec, void *_info)
+{
+	struct parse_info *info = _info;
+	u32 *ids = info->priv;
+
+	*ids = iommu_spec->args[0];
+	info->priv = ids + 1;
+	return 0;
+}
+
+static int iommu_of_get_u32_ids(struct iommu_probe_info *pinf, u32 *ids)
+{
+	struct parse_info info = { .pinf = pinf, .priv = ids };
+
+	return of_iommu_for_each_id(pinf->dev, pinf->of_master_np,
+				    pinf->of_map_id, parse_read_ids, &info);
+}
+
 struct iommu_device *__iommu_of_get_single_iommu(struct iommu_probe_info *pinf,
 						 const struct iommu_ops *ops,
 						 int num_cells)
@@ -353,6 +371,7 @@ struct iommu_device *__iommu_of_get_single_iommu(struct iommu_probe_info *pinf,
 				   pinf->of_map_id, parse_single_iommu, &info);
 	if (err)
 		return ERR_PTR(err);
+	pinf->get_u32_ids = iommu_of_get_u32_ids;
 	return iommu_fw_finish_get_single(pinf);
 }
 EXPORT_SYMBOL_GPL(__iommu_of_get_single_iommu);
diff --git a/include/linux/iommu-driver.h b/include/linux/iommu-driver.h
index 622d6ad9056ce0..632c7b4a389abe 100644
--- a/include/linux/iommu-driver.h
+++ b/include/linux/iommu-driver.h
@@ -40,7 +40,9 @@ struct iommu_probe_info {
 	struct iommu_device *cached_iommu;
 	struct device_node *of_master_np;
 	const u32 *of_map_id;
+	int (*get_u32_ids)(struct iommu_probe_info *pinf, u32 *ids);
 	unsigned int num_ids;
+	u32 cached_ids[8];
 	bool defer_setup : 1;
 	bool is_dma_configure : 1;
 	bool cached_single_iommu : 1;
@@ -123,6 +125,60 @@ static inline unsigned int iommu_of_num_ids(struct iommu_probe_info *pinf)
 	return pinf->num_ids;
 }
 
+unsigned int iommu_of_num_ids(struct iommu_probe_info *pinf);
+int iommu_fw_get_u32_ids(struct iommu_probe_info *pinf, u32 *ids);
+
+static inline void iommu_fw_cache_id(struct iommu_probe_info *pinf, u32 id)
+{
+	if (pinf->num_ids < ARRAY_SIZE(pinf->cached_ids))
+		pinf->cached_ids[pinf->num_ids] = id;
+	pinf->num_ids++;
+}
+
+static inline void *
+__iommu_fw_alloc_per_device_ids(struct iommu_probe_info *pinf, void *mem,
+				unsigned int num_ids, unsigned int *num_ids_p,
+				u32 *ids_p)
+{
+	int ret;
+
+	if (!mem)
+		return ERR_PTR(-ENOMEM);
+
+	ret = iommu_fw_get_u32_ids(pinf, ids_p);
+	if (ret) {
+		kfree(mem);
+		return ERR_PTR(ret);
+	}
+
+	*num_ids_p = num_ids;
+	return mem;
+}
+
+/**
+ * iommu_fw_alloc_per_device_ids - Allocate a per-device struct with ids
+ * @pinf: The iommu_probe_info
+ * @drv_struct: Name of a variable to a pointer of the driver structure
+ *
+ * Called by a driver during probe this helper allocates and initializes the
+ * driver struct that embeds the ids array with the trailing members:
+ *
+ *	unsigned int num_ids;
+ *	u32 ids[] __counted_by(num_ids);
+ *
+ * The helper allocates the driver struct with the right size flex array,
+ * and initializes both members. Returns the driver struct or ERR_PTR.
+ */
+#define iommu_fw_alloc_per_device_ids(pinf, drv_struct)                    \
+	({                                                                 \
+		unsigned int num_ids = iommu_of_num_ids(pinf);             \
+		typeof(drv_struct) drv;                                    \
+                                                                           \
+		drv = kzalloc(struct_size(drv, ids, num_ids), GFP_KERNEL); \
+		drv = __iommu_fw_alloc_per_device_ids(                     \
+			pinf, drv, num_ids, &drv->num_ids, drv->ids);      \
+	})
+
 /*
  * Used temporarily to indicate drivers that have moved to the new probe method.
  */
-- 
2.42.0