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

[Vipin Sharma <vipinsh@google.com>]

On Tue, Mar 17, 2026 at 01:06:34AM +0000, Samiullah Khawaja wrote:
> On Mon, Mar 16, 2026 at 03:54:50PM -0700, Vipin Sharma wrote:
> > On Tue, Feb 03, 2026 at 10:09:35PM +0000, Samiullah Khawaja wrote:
> > > +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.
> > > +
> > 
> > Do we need a separate config? Can't we just use CONFIG_LIVEUPDATE?
> 
> We have a separate CONFIG here so that the phase 1/2 split for iommu
> preservation doesn't break the vfio preservation. See following
> discussion in the RFCv2:
> 
> https://lore.kernel.org/all/aYEpHBYxlQxhXrwl@google.com/

Sounds good. 

> > > +static void iommu_liveupdate_free_objs(u64 next, bool incoming)
> > > +{
> > > +	struct iommu_objs_ser *objs;
> > > +
> > > +	while (next) {
> > > +		objs = __va(next);
> > > +		next = objs->next_objs;
> > > +
> > > +		if (!incoming)
> > > +			kho_unpreserve_free(objs);
> > > +		else
> > > +			folio_put(virt_to_folio(objs));
> > > +	}
> > > +}
> > 
> > Instead of passing boolean, and calling with different arguments, I
> > think it will be simpler to just have two functions
> > 
> > - iommu_liveupdate_unpreserve()
> > - iommu_liveupdate_folio_put()
> 
> This is a helper function to free the serialized state without
> duplicating multiple checks for various type of state (iommu,
> iommu_domain and devices).
> 
> Do you think maybe I should add these two functions and make it call the
> helper?

Read the next response.

> > 
> > > +
> > > +static void iommu_liveupdate_flb_free(struct iommu_lu_flb_obj *obj)
> > > +{
> > > +	if (obj->iommu_domains)
> > > +		iommu_liveupdate_free_objs(obj->ser->iommu_domains_phys, false);
> > > +
> > > +	if (obj->devices)
> > > +		iommu_liveupdate_free_objs(obj->ser->devices_phys, false);
> > > +
> > > +	if (obj->iommus)
> > > +		iommu_liveupdate_free_objs(obj->ser->iommus_phys, false);
> > > +
> > > +	kho_unpreserve_free(obj->ser);
> > > +	kfree(obj);
> > > +}
> > > +
> > > +static int iommu_liveupdate_flb_preserve(struct liveupdate_flb_op_args *argp)
> > > +{
> > > +	struct iommu_lu_flb_obj *obj;
> > > +	struct iommu_lu_flb_ser *ser;
> > > +	void *mem;
> > > +
> > > +	obj = kzalloc(sizeof(*obj), GFP_KERNEL);
> > > +	if (!obj)
> > > +		return -ENOMEM;
> > > +
> > > +	mutex_init(&obj->lock);
> > > +	mem = kho_alloc_preserve(sizeof(*ser));
> > > +	if (IS_ERR(mem))
> > > +		goto err_free;
> > > +
> > > +	ser = mem;
> > > +	obj->ser = ser;
> > > +
> > > +	mem = kho_alloc_preserve(PAGE_SIZE);
> > > +	if (IS_ERR(mem))
> > > +		goto err_free;
> > > +
> > > +	obj->iommu_domains = mem;
> > > +	ser->iommu_domains_phys = virt_to_phys(obj->iommu_domains);
> > > +
> > > +	mem = kho_alloc_preserve(PAGE_SIZE);
> > > +	if (IS_ERR(mem))
> > > +		goto err_free;
> > > +
> > > +	obj->devices = mem;
> > > +	ser->devices_phys = virt_to_phys(obj->devices);
> > > +
> > > +	mem = kho_alloc_preserve(PAGE_SIZE);
> > > +	if (IS_ERR(mem))
> > > +		goto err_free;
> > > +
> > > +	obj->iommus = mem;
> > > +	ser->iommus_phys = virt_to_phys(obj->iommus);
> > > +
> > > +	argp->obj = obj;
> > > +	argp->data = virt_to_phys(ser);
> > > +	return 0;
> > > +
> > > +err_free:
> > > +	iommu_liveupdate_flb_free(obj);
> > 
> > Generally, I have seen in the function goto will call corresponding
> > error tags, and free corresponding allocations and all the one which
> > happend before. It is easier to read code that way. I know you are
> > combining the free call from iommu_liveupdate_flb_unpreserve() also.
> > IMHO, code readability will be better this way.
> 
> I had that originally when I was writing this function, but it gets
> really cluttered :(. Instead it is more clean without code duplication
> using this one cleanup function here to free the state on error and also
> when doing unpreserve. Please consider this a "destroy" function of obj
> and it can be called from 2 places,
> 
> - Error during allocation of internal state.
> - During unpreserve.

It is removing code duplication in

 - iommu_liveupdate_flb_preserve()
 - iommu_liveupdate_flb_unpreserve()

However, there is still duplicate code in iommu_liveupdate_flb_finish().
Another thing is iommu_liveupdate_free_objs() is doing two different
things based on current liveupdate state (before or after kexec) passed by a
bool argument. IMO, it is cleaner if we explicitly write whether we are
doing unpreserve or just folio put.

I meant something like:

static void iommu_liveupdate_unpreserve_free(u64 next)
{
       while (next) {
               struct iommu_objs_ser *objs = __va(next);
               next = objs->next_objs;
               kho_unpreserve_free(objs);
       }
}

static void iommu_liveupdate_folio_put(u64 next)
{
       while (next) {
               struct iommu_objs_ser *objs = __va(next);
               next = objs->next_objs;
               folio_put(virt_to_folio(objs));
       }
}

static int iommu_liveupdate_flb_preserve(struct liveupdate_flb_op_args *argp)
{

...

err_free_devices:
	iommu_liveupdate_unpreserve_free(obj->ser->devices_phys);
err_free_iommu_domains:
	iommu_liveupdate_unpreserve_free(obj->ser->iommu_domains_phys);
err_free_ser:
	kho_unpreserve_free(obj->ser);
err_free_obj:
	kfree(obj);
	return PTR_ERR(mem);
}

static void iommu_liveupdate_flb_unpreserve(struct liveupdate_flb_op_args *argp)
{
	struct iommu_lu_flb_obj *obj = argp->obj;

	iommu_liveupdate_unpreserve_free(obj->ser->iommus_phys);
	iommu_liveupdate_unpreserve_free(obj->ser->devices_phys);
	iommu_liveupdate_unpreserve_free(obj->ser->iommu_domains_phys);
	kho_unpreserve_free(obj->ser);
	kfree(obj);
	argp->obj = NULL;
}

static void iommu_liveupdate_flb_finish(struct liveupdate_flb_op_args *argp)
{
	struct iommu_lu_flb_obj *obj = argp->obj;

	iommu_liveupdate_folio_put(obj->ser->iommus_phys);
	iommu_liveupdate_folio_put(obj->ser->devices_phys);
	iommu_liveupdate_folio_put(obj->ser->iommu_domains_phys);
	folio_put(virt_to_folio(obj->ser));
	kfree(obj);
	argp->obj = NULL
}

This way code is pretty explicit and clear what is happening. Let me
know if you meant something else by cluttered code.
> > 
> > > +	return PTR_ERR(mem);
> > > +}
> > > +
> > > +static void iommu_liveupdate_flb_unpreserve(struct liveupdate_flb_op_args *argp)
> > > +{
> > > +	iommu_liveupdate_flb_free(argp->obj);
> > > +}
> > > +
> > > +static void iommu_liveupdate_flb_finish(struct liveupdate_flb_op_args *argp)
> > > +{
> > > +	struct iommu_lu_flb_obj *obj = argp->obj;
> > > +
> > > +	if (obj->iommu_domains)
> > > +		iommu_liveupdate_free_objs(obj->ser->iommu_domains_phys, true);
> > 
> > Can there be the case where obj->iommu_domains is NULL but
> > obj->ser->iommu_domains_phys is not? If that is not possible, I will
> > just simplify the patch and unconditionally call
> > iommu_liveupdate_free_objs()?
> 
> Are you suggesting that on flb_finish() the obj->iommu_domains should be
> non-NULL as flb_retrieve() succeeded? If yes, then that is correct. I
> will update this to call the free_objs() without checking
> obj->iommu_domains. I will do same for other types.

Yes.