[PATCH RFCv1 03/14] iommufd: Prepare for viommu structures and functions

[Nicolin Chen]

The following changes will introduce a new iommufd_viommu and its related
structures and functions. These new core-level structures will be embedded
in the driver-level structures. So a helper to allocate the bundle of core
and driver strucutres will be nice to have.

As a preparatory change, introduce a viommu.c file and put a new structure
allocator in it. Since some of the object allocation will be similar, also
move the common part to a new level-3 allocator ___iommufd_object_alloc().

Signed-off-by: Nicolin Chen <nicolinc at nvidia.com>
---
 drivers/iommu/iommufd/Makefile          |  3 ++-
 drivers/iommu/iommufd/iommufd_private.h | 20 ++++++++++++++++++++
 drivers/iommu/iommufd/main.c            |  9 +++------
 drivers/iommu/iommufd/viommu.c          | 19 +++++++++++++++++++
 4 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 drivers/iommu/iommufd/viommu.c

diff --git a/drivers/iommu/iommufd/Makefile b/drivers/iommu/iommufd/Makefile
index 34b446146961..76c2bc41af40 100644
--- a/drivers/iommu/iommufd/Makefile
+++ b/drivers/iommu/iommufd/Makefile
@@ -6,7 +6,8 @@ iommufd-y := \
 	ioas.o \
 	main.o \
 	pages.o \
-	vfio_compat.o
+	vfio_compat.o \
+	viommu.o
 
 iommufd-$(CONFIG_IOMMUFD_TEST) += selftest.o
 
diff --git a/drivers/iommu/iommufd/iommufd_private.h b/drivers/iommu/iommufd/iommufd_private.h
index 3acbc67dd5f0..eccc565ed38e 100644
--- a/drivers/iommu/iommufd/iommufd_private.h
+++ b/drivers/iommu/iommufd/iommufd_private.h
@@ -156,6 +156,26 @@ void iommufd_object_abort_and_destroy(struct iommufd_ctx *ictx,
 void iommufd_object_finalize(struct iommufd_ctx *ictx,
 			     struct iommufd_object *obj);
 
+static inline struct iommufd_object *___iommufd_object_alloc(size_t size)
+{
+	struct iommufd_object *obj;
+
+	obj = kzalloc(size, GFP_KERNEL_ACCOUNT);
+	if (!obj)
+		return ERR_PTR(-ENOMEM);
+
+	/* Starts out bias'd by 1 until it is removed from the xarray */
+	refcount_set(&obj->shortterm_users, 1);
+	refcount_set(&obj->users, 1);
+
+	/*
+	 * The allocation of an obj->id needs an ictx, so it has to be done
+	 * after this ___iommufd_object_alloc() callback.
+	 */
+
+	return obj;
+}
+
 enum {
 	REMOVE_WAIT_SHORTTERM = 1,
 };
diff --git a/drivers/iommu/iommufd/main.c b/drivers/iommu/iommufd/main.c
index a51ab766e183..5187942b375d 100644
--- a/drivers/iommu/iommufd/main.c
+++ b/drivers/iommu/iommufd/main.c
@@ -36,13 +36,10 @@ struct iommufd_object *__iommufd_object_alloc(struct iommufd_ctx *ictx,
 	struct iommufd_object *obj;
 	int rc;
 
-	obj = kzalloc(size, GFP_KERNEL_ACCOUNT);
-	if (!obj)
-		return ERR_PTR(-ENOMEM);
+	obj = ___iommufd_object_alloc(size);
+	if (IS_ERR(obj))
+		return obj;
 	obj->type = type;
-	/* Starts out bias'd by 1 until it is removed from the xarray */
-	refcount_set(&obj->shortterm_users, 1);
-	refcount_set(&obj->users, 1);
 
 	/*
 	 * Reserve an ID in the xarray but do not publish the pointer yet since
diff --git a/drivers/iommu/iommufd/viommu.c b/drivers/iommu/iommufd/viommu.c
new file mode 100644
index 000000000000..f77d6972d552
--- /dev/null
+++ b/drivers/iommu/iommufd/viommu.c
@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: GPL-2.0-only
+/* Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES
+ */
+
+#include <linux/iommufd.h>
+
+#include "iommufd_private.h"
+
+#define viommu_struct_alloc(name)                                              \
+	struct iommufd_##name *_iommufd_##name##_alloc(size_t size)            \
+	{                                                                      \
+		struct iommufd_object *obj;                                    \
+		if (WARN_ON(size < sizeof(struct iommufd_##name)))             \
+			return NULL;                                           \
+		obj = ___iommufd_object_alloc(size);                           \
+		if (IS_ERR(obj))                                               \
+			return NULL;                                           \
+		return container_of(obj, struct iommufd_##name, obj);          \
+	}
-- 
2.43.0