Re: [PATCH v4 05/23] iommufd/driver: Let iommufd_viommu_alloc helper save ictx to viommu->ictx

[Nicolin Chen <nicolinc@nvidia.com>]

On Mon, May 26, 2025 at 10:30:46AM -0300, Jason Gunthorpe wrote:
> On Fri, May 16, 2025 at 01:56:26PM -0700, Nicolin Chen wrote:
> 
> > > You don't need to move this unless you are using inlines. Just use a
> > > forward declaration.
> > 
> > Since we forward ucmd now, ictx is in the ucmd so we need this
> > structure for:
> > 
> > -		if (!IS_ERR(ret))                                              \
> > +		if (!IS_ERR(ret)) {                                            \
> >  			ret->member.ops = viommu_ops;                          \
> > +			ret->member.ictx = ucmd->ictx;                         \
> > +		}                                                              \
> 
> De-inline more of that function probably..
> 
> Also seem my other remarks about not storing ictx so much..

I found that all other ictx pointers in vdev/hw_queue are unused,
as the core simply gets an ictx from their viommu pointers. This
means that only this viommu allocator here needs such a storing.

With that, how about a change like this v.s. inline:

[ Sol-1: Store ucmd ]
----------------------------------------------------------------------------
diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
index 2b30627d1d8ed..513af9ba04f17 100644
--- a/drivers/iommu/iommufd/viommu.c
+++ b/drivers/iommu/iommufd/viommu.c
@@ -60,9 +60,15 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
 		goto out_put_hwpt;
 	}
 
+	/* Verify if the driver correctly calls iommufd_viommu_alloc() */
+	static_assert(offsetof(typeof(*ucmd), ictx) == 0);
+	if (WARN_ON_ONCE(viommu->ictx != ucmd->ictx)) {
+		rc = -EINVAL;
+		goto out_put_hwpt;
+	}
+
 	xa_init(&viommu->vdevs);
 	viommu->type = cmd->type;
-	viommu->ictx = ucmd->ictx;
 	viommu->hwpt = hwpt_paging;
 	refcount_inc(&viommu->hwpt->common.obj.users);
 	INIT_LIST_HEAD(&viommu->veventqs);
diff --git a/include/linux/iommufd.h b/include/linux/iommufd.h
index d74c97feb9b52..6843c748a6f74 100644
--- a/include/linux/iommufd.h
+++ b/include/linux/iommufd.h
@@ -104,7 +104,10 @@ void iommufd_ctx_get(struct iommufd_ctx *ictx);
 
 struct iommufd_viommu {
 	struct iommufd_object obj;
-	struct iommufd_ctx *ictx;
+	union {
+		struct iommufd_ctx *ictx;
+		struct iommufd_ucmd *ucmd; /* for allocation only */
+	};
 	struct iommu_device *iommu_dev;
 	struct iommufd_hwpt_paging *hwpt;
 
@@ -262,8 +265,10 @@ static inline int iommufd_viommu_report_event(struct iommufd_viommu *viommu,
 									       \
 		ret = (drv_struct *)__iommufd_object_alloc_ucmd(               \
 			ucmd, ret, IOMMUFD_OBJ_VIOMMU, member.obj);            \
-		if (!IS_ERR(ret))                                              \
+		if (!IS_ERR(ret)) {                                            \
 			ret->member.ops = viommu_ops;                          \
+			ret->member.ucmd = ucmd;                               \
+		}                                                              \
 		ret;                                                           \
 	})
 #endif
----------------------------------------------------------------------------

[ Sol-2: De-inline ]
----------------------------------------------------------------------------
diff --git a/drivers/iommu/iommufd/driver.c b/drivers/iommu/iommufd/driver.c
index 28dec9e097204..0adb1a552e457 100644
--- a/drivers/iommu/iommufd/driver.c
+++ b/drivers/iommu/iommufd/driver.c
@@ -54,6 +54,23 @@ struct iommufd_object *_iommufd_object_alloc_ucmd(struct iommufd_ucmd *ucmd,
 }
 EXPORT_SYMBOL_NS_GPL(_iommufd_object_alloc_ucmd, "IOMMUFD");
 
+struct iommufd_viommu *
+_iommufd_viommu_alloc_ucmd(struct iommufd_ucmd *ucmd, size_t size,
+			   const struct iommufd_viommu_ops *viommu_ops)
+{
+	struct iommufd_viommu *viommu;
+
+	static_assert(offsetof(typeof(*viommu), obj) == 0);
+	viommu = (struct iommufd_viommu *)_iommufd_object_alloc_ucmd(
+		ucmd, size, IOMMUFD_OBJ_VIOMMU);
+	if (IS_ERR(viommu))
+		return viommu;
+	viommu->ops = viommu_ops;
+	viommu->ictx = ucmd->ictx;
+	return viommu;
+}
+EXPORT_SYMBOL_NS_GPL(_iommufd_viommu_alloc_ucmd, "IOMMUFD");
+
 /* Caller should xa_lock(&viommu->vdevs) to protect the return value */
 struct device *iommufd_viommu_find_dev(struct iommufd_viommu *viommu,
 				       unsigned long vdev_id)
diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
index 2b30627d1d8ed..259fcfd85b1bd 100644
--- a/drivers/iommu/iommufd/viommu.c
+++ b/drivers/iommu/iommufd/viommu.c
@@ -60,9 +60,14 @@ int iommufd_viommu_alloc_ioctl(struct iommufd_ucmd *ucmd)
 		goto out_put_hwpt;
 	}
 
+	/* Verify if the driver correctly calls iommufd_viommu_alloc() */
+	if (WARN_ON_ONCE(viommu->ictx != ucmd->ictx)) {
+		rc = -EINVAL;
+		goto out_put_hwpt;
+	}
+
 	xa_init(&viommu->vdevs);
 	viommu->type = cmd->type;
-	viommu->ictx = ucmd->ictx;
 	viommu->hwpt = hwpt_paging;
 	refcount_inc(&viommu->hwpt->common.obj.users);
 	INIT_LIST_HEAD(&viommu->veventqs);
diff --git a/include/linux/iommufd.h b/include/linux/iommufd.h
index d74c97feb9b52..50ac3a5703520 100644
--- a/include/linux/iommufd.h
+++ b/include/linux/iommufd.h
@@ -209,6 +209,9 @@ struct iommufd_object *_iommufd_object_alloc(struct iommufd_ctx *ictx,
 struct iommufd_object *
 _iommufd_object_alloc_ucmd(struct iommufd_ucmd *ucmd, size_t size,
 			   enum iommufd_object_type type);
+struct iommufd_viommu *
+_iommufd_viommu_alloc_ucmd(struct iommufd_ucmd *ucmd, size_t size,
+			   const struct iommufd_viommu_ops *viommu_ops);
 struct device *iommufd_viommu_find_dev(struct iommufd_viommu *viommu,
 				       unsigned long vdev_id);
 int iommufd_viommu_get_vdev_id(struct iommufd_viommu *viommu,
@@ -231,6 +234,13 @@ _iommufd_object_alloc_ucmd(struct iommufd_ucmd *ucmd, size_t size,
 	return ERR_PTR(-EOPNOTSUPP);
 }
 
+static inline struct iommufd_viommu *
+_iommufd_viommu_alloc_ucmd(struct iommufd_ucmd *ucmd, size_t size,
+			   const struct iommufd_viommu_ops *viommu_ops)
+{
+	return ERR_PTR(-EOPNOTSUPP);
+}
+
 static inline struct device *
 iommufd_viommu_find_dev(struct iommufd_viommu *viommu, unsigned long vdev_id)
 {
@@ -257,13 +267,7 @@ static inline int iommufd_viommu_report_event(struct iommufd_viommu *viommu,
  * the iommufd core. The free op will be called prior to freeing the memory.
  */
 #define iommufd_viommu_alloc(ucmd, drv_struct, member, viommu_ops)             \
-	({                                                                     \
-		drv_struct *ret;                                               \
-									       \
-		ret = (drv_struct *)__iommufd_object_alloc_ucmd(               \
-			ucmd, ret, IOMMUFD_OBJ_VIOMMU, member.obj);            \
-		if (!IS_ERR(ret))                                              \
-			ret->member.ops = viommu_ops;                          \
-		ret;                                                           \
-	})
+	container_of(_iommufd_viommu_alloc_ucmd(                               \
+			     ucmd, sizeof(drv_struct), viommu_ops),            \
+		     typeof(drv_struct), member)
 #endif
----------------------------------------------------------------------------

Thanks
Nicolin