Re: [PATCH 01/14] iommu: Implement IOMMU LU FLB callbacks

[Pranjal Shrivastava <praan@google.com>]

On Thu, Mar 12, 2026 at 04:43:00PM +0000, Samiullah Khawaja wrote:
> On Wed, Mar 11, 2026 at 09:07:00PM +0000, Pranjal Shrivastava wrote:
> > On Tue, Feb 03, 2026 at 10:09:35PM +0000, Samiullah Khawaja wrote:
> > > Add liveupdate FLB for IOMMU state preservation. Use KHO preserve memory
> > > alloc/free helper functions to allocate memory for the IOMMU LU FLB
> > > object and the serialization structs for device, domain and iommu.
> > > 
> > > During retrieve, walk through the preserved objs nodes and restore each
> > > folio. Also recreate the FLB obj.
> > > 
> > > Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
> > > ---
> > >  drivers/iommu/Kconfig         |  11 +++
> > >  drivers/iommu/Makefile        |   1 +
> > >  drivers/iommu/liveupdate.c    | 177 ++++++++++++++++++++++++++++++++++
> > >  include/linux/iommu-lu.h      |  17 ++++
> > >  include/linux/kho/abi/iommu.h | 119 +++++++++++++++++++++++
> > >  5 files changed, 325 insertions(+)
> > >  create mode 100644 drivers/iommu/liveupdate.c
> > >  create mode 100644 include/linux/iommu-lu.h
> > >  create mode 100644 include/linux/kho/abi/iommu.h
> > > 
> > > diff --git a/drivers/iommu/Kconfig b/drivers/iommu/Kconfig
> > > index f86262b11416..fdcfbedee5ed 100644
> > > --- a/drivers/iommu/Kconfig
> > > +++ b/drivers/iommu/Kconfig
> > > @@ -11,6 +11,17 @@ config IOMMUFD_DRIVER
> > >  	bool
> > >  	default n
> > > 
> > > +config IOMMU_LIVEUPDATE
> > > +	bool "IOMMU live update state preservation support"
> > > +	depends on LIVEUPDATE && IOMMUFD
> > > +	help
> > > +	  Enable support for preserving IOMMU state across a kexec live update.
> > > +
> > > +	  This allows devices managed by iommufd to maintain their DMA mappings
> > > +	  during kexec base kernel update.
> > > +
> > > +	  If unsure, say N.
> > > +
> > 
> > I'm wondering if this should be under the if IOMMU_SUPPORT below? I
> > believe this was added here because IOMMUFD isn't under IOMMU_SUPPORT,
> > but it wouldn't make sense to "preserve" IOMMU across a liveupdate if
> > IOMMU_SUPPORT is disabled? Should we probably be move it inside the
> > if IOMMU_SUPPORT block for better organization, or at least have a depends
> > on IOMMU_SUPPORT added to it? The IOMMU_LUO still depends on the
> > IOMMU_SUPPORT infrastructure to actually function.. as we add calls
> > within core functions like dev_iommu_get etc.
> 
> Agreed. I will move it under IOMMU_SUPPORT and sort out any other
> dependencies.
> > 
> > >  menuconfig IOMMU_SUPPORT
> > >  	bool "IOMMU Hardware Support"
> > >  	depends on MMU
> > > diff --git a/drivers/iommu/Makefile b/drivers/iommu/Makefile
> > > index 0275821f4ef9..b3715c5a6b97 100644
> > > --- a/drivers/iommu/Makefile
> > > +++ b/drivers/iommu/Makefile
> > > @@ -15,6 +15,7 @@ obj-$(CONFIG_IOMMU_IO_PGTABLE_ARMV7S) += io-pgtable-arm-v7s.o
> > >  obj-$(CONFIG_IOMMU_IO_PGTABLE_LPAE) += io-pgtable-arm.o
> > >  obj-$(CONFIG_IOMMU_IO_PGTABLE_LPAE_KUNIT_TEST) += io-pgtable-arm-selftests.o
> > >  obj-$(CONFIG_IOMMU_IO_PGTABLE_DART) += io-pgtable-dart.o
> > > +obj-$(CONFIG_IOMMU_LIVEUPDATE) += liveupdate.o
> > >  obj-$(CONFIG_IOMMU_IOVA) += iova.o
> > >  obj-$(CONFIG_OF_IOMMU)	+= of_iommu.o
> > >  obj-$(CONFIG_MSM_IOMMU) += msm_iommu.o
> > > diff --git a/drivers/iommu/liveupdate.c b/drivers/iommu/liveupdate.c
> > > new file mode 100644
> > > index 000000000000..6189ba32ff2c
> > > --- /dev/null
> > > +++ b/drivers/iommu/liveupdate.c
> > > @@ -0,0 +1,177 @@
> > > +// SPDX-License-Identifier: GPL-2.0-only
> > > +
> > > +/*
> > > + * Copyright (C) 2025, Google LLC
> > 
> > Minor nit: 2026 OR 2025-26, here and everywhere else
> 
> Will fix in next revision.
> > 
> > > + * Author: Samiullah Khawaja <skhawaja@google.com>
> > > + */
> > > +
> > > +#define pr_fmt(fmt)    "iommu: liveupdate: " fmt
> > > +
> > > +#include <linux/kexec_handover.h>
> > > +#include <linux/liveupdate.h>
> > > +#include <linux/iommu-lu.h>
> > > +#include <linux/iommu.h>
> > > +#include <linux/errno.h>
> > > +
> > > +static void iommu_liveupdate_restore_objs(u64 next)
> > > +{
> > > +	struct iommu_objs_ser *objs;
> > > +
> > > +	while (next) {
> > > +		BUG_ON(!kho_restore_folio(next));
> > 
> > Same thing about BUG_ON [1] as mentioned below in the
> > iommu_liveupdate_flb_retrieve() function, can we consider returning an
> > error which can be checked in the caller and the error can be bubbled up
> > as -ENODATA?
> 
> Please see the explanation below on the BUG_ON.

[------- snip >8 --------]

> > > +
> > > +static int iommu_liveupdate_flb_retrieve(struct liveupdate_flb_op_args *argp)
> > > +{
> > > +	struct iommu_lu_flb_obj *obj;
> > > +	struct iommu_lu_flb_ser *ser;
> > > +
> > > +	obj = kzalloc(sizeof(*obj), GFP_ATOMIC);
> > 
> > Why does this have to be GFP_ATOMIC? IIUC, the retrieve path is
> > triggered by a userspace IOCTL in the new kernel. The system should be
> > able to sleep here? (unless we have a use-case to call this in IRQ-ctx?)
> > AFAICT, we call this under mutexes already, hence there's no situation
> > where we could sleep in a spinlock context?
> > 
> > GFP_ATOMIC creates a point of failure if the system is under memory
> > pressure. I believe we should be allowed to sleep for this allocation
> > because the "preserved" mappings still allow DMAs to go on and we're in
> > no hurry to restore the IOMMU state? I believe this could be GFP_KERNEL.
> > 

I guess we missed discussing this comment about s/GFP_ATOMIC/GFP_KERNEL?

> > > +	if (!obj)
> > > +		return -ENOMEM;
> > > +
> > > +	mutex_init(&obj->lock);
> > > +	BUG_ON(!kho_restore_folio(argp->data));
> > 
> > The use of BUG_ON in new code is heavily discouraged [1].
> > If KHO can't restore the folio for whatever reason, we can be treat it
> > as a corruption of the handover data. I believe crashing the kernel for
> > it would be an overkill?
> 
> The FLB restore is done during early boot and this has been discussed in
> past in KHO/LUO and IOMMU context also. But basically if this fails, the
> restoration of IOMMU state cannot be done, preserved devices would
> already have corrupted memory due to ongoing DMA as we didn't disable
> translation in the previous kernel. So logging an error and a BUG_ON at
> this point would be most appropriate.
> 
> Please see discussion here on this topic:
> https://lore.kernel.org/all/20251118153631.GB90703@nvidia.com/

I see.. so a failure here would mean an entire VM tear-down?

> > 
> > Can we consider returning a graceful failure like -ENODATA or something?
> > BUG_ON would instantly cause a kernel panic without providing no
> > opportunity for the system to log the failure or attempt a graceful
> > teardown of the 'preserved' mapping.
> 
> I will update this by adding a comment and also log an error.

Ack. Thanks

[ ---- snip >8 -----]

> > > +enum iommu_lu_type {
> > > +	IOMMU_INVALID,
> > > +	IOMMU_INTEL,
> > > +};
> > > +
> > > +struct iommu_obj_ser {
> > > +	u32 idx;
> > > +	u32 ref_count;
> > > +	u32 deleted:1;
> > > +	u32 incoming:1;
> > > +} __packed;
> > > +
> > > +struct iommu_domain_ser {
> > > +	struct iommu_obj_ser obj;
> > > +	u64 top_table;
> > > +	u64 top_level;
> > > +	struct iommu_domain *restored_domain;
> > > +} __packed;
> > > +
> > > +struct device_domain_iommu_ser {
> > > +	u32 did;

Nit: `did` sounds intel-specific, we can either call it something
generic or make it into a union when we support other archs in the
future.

> > > +	u64 domain_phys;
> > > +	u64 iommu_phys;
> > > +} __packed;
> > > +
> > > +struct device_ser {
> > > +	struct iommu_obj_ser obj;
> > > +	u64 token;
> > > +	u32 devid;
> > > +	u32 pci_domain;
> > > +	struct device_domain_iommu_ser domain_iommu_ser;
> > > +	enum iommu_lu_type type;
> > > +} __packed;
> > > +
> > > +struct iommu_intel_ser {
> > > +	u64 phys_addr;
> > > +	u64 root_table;
> > > +} __packed;
> > > +
> > > +struct iommu_ser {
> > > +	struct iommu_obj_ser obj;
> > > +	u64 token;
> > > +	enum iommu_lu_type type;
> > > +	union {
> > > +		struct iommu_intel_ser intel;
> > > +	};
> > > +} __packed;
> > > +
> > > +struct iommu_objs_ser {
> > > +	u64 next_objs;
> > > +	u64 nr_objs;
> > > +} __packed;
> > > +
> > > +struct iommus_ser {
> > > +	struct iommu_objs_ser objs;
> > > +	struct iommu_ser iommus[];
> > > +} __packed;
> > > +
> > > +struct iommu_domains_ser {
> > > +	struct iommu_objs_ser objs;
> > > +	struct iommu_domain_ser iommu_domains[];
> > > +} __packed;
> > > +
> > > +struct devices_ser {
> > > +	struct iommu_objs_ser objs;
> > > +	struct device_ser devices[];
> > > +} __packed;

I have a bone to pick here about the naming LOL, the names are kinda
confusing and make the code unreadable, I had to keep re-visitng this
patch while looking at the subsequent ones to understand what's going
on..

One suggestion is to add a graphic that could help understand the layout
like:

----------------------------------------------------------------------
[ PAGE START ]                                                       |
----------------------------------------------------------------------
| iommu_objs_ser (The Page Header)                                   |
|   - next_objs: 0x0000 (End of the page-chain)                      |
|   - nr_objs: 2                                                     |
----------------------------------------------------------------------
| ITEM 0: iommu_domain_ser                                           |
|   [ iommu_obj_ser (The entry header) ]                             |
|     - idx: 0                                                       |
|     - ref_count: 1                                                 |
|     - deleted: 0                                                   |
|   [ Domain Data ]                                                  |
----------------------------------------------------------------------
| ITEM 1: iommu_domain_ser                                           |
|   [ iommu_obj_ser (The Price Tag) ]                                |
|     - idx: 1                                                       |
|     - ref_count: 1                                                 |
|     - deleted: 0                                                   |
|   [ Domain Data ]                                                  |
----------------------------------------------------------------------
| ... (Empty space for more domains) ...                             |
|                                                                    |
----------------------------------------------------------------------
[ PAGE END ]                                                         |
----------------------------------------------------------------------

Additionally, a few naming suggestions here:

1. struct iommu_obj_ser -> struct iommu_ser_entry_hdr
2. struct iommu_objs_ser -> struct iommu_ser_page_hdr
3. struct iommu_domains_ser -> struct iommu_ser_domain_page

This makes things clearer:

struct iommu_ser_page_hdr {
    u64 next_page_phys;
    u64 entry_count;
} __packed;

/* The Container Page */
struct iommu_ser_domain_page {
    struct iommu_ser_page_hdr hdr;
    struct iommu_ser_domain domain_entries[];
} __packed;

Similarly, something like:

4. struct devices_ser -> struct iommu_ser_device_page
5. struct iommu_lu_flb_ser -> struct iommu_flb_metadata
6. struct iommu_lu_flb_obj -> struct iommu_flb_ctx

struct iommu_flb_ctx {
	struct mutex lock;
	struct iommu_flb_metadata *cookie;

	struct iommu_ser_domain_page    *curr_domains_page;
	struct iommu_ser_iommu_ctx_page *curr_iommu_ctx_page;
	struct iommu_ser_device_page    *curr_devices_page;
} __packed;

Makes things slightly more readable.

> > > +
> > > +#define MAX_IOMMU_SERS ((PAGE_SIZE - sizeof(struct iommus_ser)) / sizeof(struct iommu_ser))

Nit: For clarity, can we consider adding another set of braces:

+#define MAX_IOMMU_SERS ((PAGE_SIZE - (sizeof(struct iommus_ser)) / sizeof(struct iommu_ser)))

> > > +#define MAX_IOMMU_DOMAIN_SERS \
> > > +		((PAGE_SIZE - sizeof(struct iommu_domains_ser)) / sizeof(struct iommu_domain_ser))
> > > +#define MAX_DEVICE_SERS ((PAGE_SIZE - sizeof(struct devices_ser)) / sizeof(struct device_ser))
> > > +
> > > +struct iommu_lu_flb_ser {
> > > +	u64 iommus_phys;
> > > +	u64 nr_iommus;
> > > +	u64 iommu_domains_phys;
> > > +	u64 nr_domains;
> > > +	u64 devices_phys;
> > > +	u64 nr_devices;
> > > +} __packed;
> > > +
> > > +struct iommu_lu_flb_obj {
> > > +	struct mutex lock;
> > > +	struct iommu_lu_flb_ser *ser;
> > > +
> > > +	struct iommu_domains_ser *iommu_domains;
> > > +	struct iommus_ser *iommus;
> > > +	struct devices_ser *devices;
> > > +} __packed;
> > > +
> > 
> > Please let's add some comments describing the structs & their members
> > here like we have in memfd [2]. This should be descriptive for the user.
> > For example:
> 
> I agree. Will add comments for these and others.
> > 
> > +/**
> > + * struct iommu_lu_flb_ser - Main serialization header for IOMMU state.
> > + * @iommus_phys:        Physical address of the first page in the IOMMU unit chain.
> > + * @nr_iommus:          Total number of hardware IOMMU units preserved.
> > + * @iommu_domains_phys: [...]
> > + * @nr_domains:         [...]
> > + * @devices_phys:       [...]
> > + * @nr_devices:         [...]
> > + *
> > + * This structure acts as the root of the IOMMU state tree. It is hitching a ride
> > + * on the iommufd file descriptor's preservation flow.
> > + */
> > +struct iommu_lu_flb_ser {
> > +	u64 iommus_phys;
> > +	u64 nr_iommus;
> > +	u64 iommu_domains_phys;
> > +	u64 nr_domains;
> > +	u64 devices_phys;
> > +	u64 nr_devices;
> > +} __packed;
> > 
> > > +#endif /* _LINUX_KHO_ABI_IOMMU_H */
> > 
> > Thanks,
> > Praan
> > 
> > [1] https://docs.kernel.org/process/coding-style.html#use-warn-rather-than-bug
> > [2] https://elixir.bootlin.com/linux/v7.0-rc3/source/include/linux/kho/abi/memfd.h
> > 
> 
> Thanks for the review.
> 
> Sami

Thanks,
Praan