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

[Nicolin Chen]

On Wed, May 14, 2025 at 02:06:37PM -0300, Jason Gunthorpe wrote:
> On Thu, May 08, 2025 at 08:02:26PM -0700, Nicolin Chen wrote:
> > When an IOMMU driver calls iommufd_viommu_alloc(), it must pass in an ictx
> > pointer as the underlying _iommufd_object_alloc() helper function requires
> > that to allocate a new object. However, neither the iommufd_viommu_alloc()
> > nor its underlying _iommufd_object_alloc() saves the ictx in the allocated
> > viommu object, although viommu could hold an ictx pointer.
> > 
> > When the IOMMU driver wants to use another iommufd function passing in the
> > allocated viommu, it could have avoided passing in the ictx pointer again,
> > if viommu->ictx is valid.
> > 
> > Save ictx to viommu->ictx in the iommufd_viommu_alloc(), in order to ease
> > a new vIOMMU-based helper that would then get the ictx from viommu->ictx.
> > 
> > Signed-off-by: Nicolin Chen <nicolinc at nvidia.com>
> > ---
> >  include/linux/iommufd.h | 4 +++-
> >  1 file changed, 3 insertions(+), 1 deletion(-)
> 
> It is OK, but please think carefully if the ictx is actually
> needed. The reason most objects don't have an ictx in them is because
> their contexts are always inside an ioctl so they get the ictx from
> there.
> 
> We don't have a lot of viommu allocations so it isn't such a big deal,
> but just generally that is how it was built that the ictx comes from
> the ioctl not the object.

It's more of providing cleaner for-driver APIs.

Especially with the new ucmd changes, it can be just:
	// in my_viommu_alloc()
	my_viommu = iommufd_viommu_alloc(ucmd, ...);
	iommufd_viommu_alloc_mmap(my_viommu, ...);
	iommufd_viommu_destroy_mmap(my_viommu, ...);
	// in my_viommu_destory()
	iommufd_viommu_destroy_mmap(my_viommu, ...);

Otherwise, they would be something like:
	// in my_viommu_alloc()
	my_viommu = iommufd_viommu_alloc(ucmd, ...);
	iommufd_viommu_alloc_mmap(ucmd->ictx, my_viommu, ...);
	iommufd_viommu_destroy_mmap(ucmd->ictx, my_viommu, ...);
	// in my_viommu_destory()
	iommufd_viommu_destroy_mmap(my_viommu->ictx, my_viommu ...);

That being said, I ended up doing something more in this patch
for the ucmd rework and Kevin's comments:

diff --git a/drivers/iommu/iommufd/iommufd_private.h b/drivers/iommu/iommufd/iommufd_private.h
index 98fe5a49c9606..ddd144031af5d 100644
--- a/drivers/iommu/iommufd/iommufd_private.h
+++ b/drivers/iommu/iommufd/iommufd_private.h
@@ -136,14 +136,6 @@ int iopt_pin_pages(struct io_pagetable *iopt, unsigned long iova,
 void iopt_unpin_pages(struct io_pagetable *iopt, unsigned long iova,
 		      unsigned long length);
 
-struct iommufd_ucmd {
-	struct iommufd_ctx *ictx;
-	void __user *ubuffer;
-	u32 user_size;
-	void *cmd;
-	struct iommufd_object *alloced_obj;
-};
-
 int iommufd_vfio_ioctl(struct iommufd_ctx *ictx, unsigned int cmd,
 		       unsigned long arg);
 
diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
index 416cd281eb2b5..6c0a2a0885a32 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;
 	}
 
+	/* The iommufd_viommu_alloc helper saves ucmd->ictx in viommu->ictx */
+	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 a24b924cf6872..74d11c6779739 100644
--- a/include/linux/iommufd.h
+++ b/include/linux/iommufd.h
@@ -52,6 +52,14 @@ struct iommufd_object {
 	unsigned int id;
 };
 
+struct iommufd_ucmd {
+	struct iommufd_ctx *ictx;
+	void __user *ubuffer;
+	u32 user_size;
+	void *cmd;
+	struct iommufd_object *alloced_obj;
+};
+
 struct iommufd_device *iommufd_device_bind(struct iommufd_ctx *ictx,
 					   struct device *dev, u32 *id);
 void iommufd_device_unbind(struct iommufd_device *idev);
@@ -252,8 +260,10 @@ static inline int iommufd_viommu_report_event(struct iommufd_viommu *viommu,
 		static_assert(offsetof(drv_struct, member.obj) == 0);          \
 		ret = (drv_struct *)_iommufd_object_alloc_ucmd(                \
 			ucmd, sizeof(drv_struct), IOMMUFD_OBJ_VIOMMU);         \
-		if (!IS_ERR(ret))                                              \
+		if (!IS_ERR(ret)) {                                            \
 			ret->member.ops = viommu_ops;                          \
+			ret->member.ictx = ucmd->ictx;                         \
+		}                                                              \
 		ret;                                                           \
 	})
 #endif

Thanks
Nicolin