Add implementation of preserve_iommu_context to preserve the root_table
and the associated context tables. Also mark the iommu unit as
preserved.

Signed-off-by: Samiullah Khawaja <skhawaja@google.com>
---
 drivers/iommu/intel/iommu.c      |  2 -
 drivers/iommu/intel/iommu.h      |  1 +
 drivers/iommu/intel/liveupdate.c | 80 +++++++++++++++++++++++++++++++-
 3 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/drivers/iommu/intel/iommu.c b/drivers/iommu/intel/iommu.c
index 7035ffca020f..caac4fd9a51e 100644
--- a/drivers/iommu/intel/iommu.c
+++ b/drivers/iommu/intel/iommu.c
@@ -67,8 +67,6 @@ static int force_on = 0;
 static int intel_iommu_tboot_noforce;
 static int no_platform_optin;
 
-#define ROOT_ENTRY_NR (VTD_PAGE_SIZE/sizeof(struct root_entry))
-
 /*
  * Take a root_entry and return the Lower Context Table Pointer (LCTP)
  * if marked present.
diff --git a/drivers/iommu/intel/iommu.h b/drivers/iommu/intel/iommu.h
index 93ac55eb49f0..273d40812d09 100644
--- a/drivers/iommu/intel/iommu.h
+++ b/drivers/iommu/intel/iommu.h
@@ -556,6 +556,8 @@ struct root_entry {
 	u64     lo;
 	u64     hi;
 };
+
+#define ROOT_ENTRY_NR (VTD_PAGE_SIZE / sizeof(struct root_entry))
 
 /*
  * low 64 bits:
diff --git a/drivers/iommu/intel/liveupdate.c b/drivers/iommu/intel/liveupdate.c
index 94aabf025a60..fb214736aa3c 100644
--- a/drivers/iommu/intel/liveupdate.c
+++ b/drivers/iommu/intel/liveupdate.c
@@ -59,11 +59,87 @@ static int preserve_device_state(struct pci_dev *dev, struct device_ser *ser)
 	return 0;
 }
 
+static int unpreserve_iommu_context(struct intel_iommu *iommu, int end)
+{
+	struct context_entry *context;
+	int i;
+
+	if (end < 0)
+		end = ROOT_ENTRY_NR;
+
+	for (i = 0; i < end; i++) {
+		context = iommu_context_addr(iommu, i, 0, 0);
+		if (context)
+			WARN_ON_ONCE(kho_unpreserve_folio(virt_to_folio(context)));
+
+		if (!sm_supported(iommu))
+			continue;
+
+		context = iommu_context_addr(iommu, i, 0x80, 0);
+		if (context)
+			WARN_ON_ONCE(kho_unpreserve_folio(virt_to_folio(context)));
+	}
+
+	return 0;
+}
+
+static int preserve_iommu_context(struct intel_iommu *iommu)
+{
+	struct context_entry *context;
+	int ret;
+	int i;
+
+	for (i = 0; i < ROOT_ENTRY_NR; i++) {
+		context = iommu_context_addr(iommu, i, 0, 0);
+		if (context) {
+			ret = kho_preserve_folio(virt_to_folio(context));
+			if (ret)
+				goto error;
+		}
+
+		if (!sm_supported(iommu))
+			continue;
+
+		context = iommu_context_addr(iommu, i, 0x80, 0);
+		if (context) {
+			ret = kho_preserve_folio(virt_to_folio(context));
+			if (ret)
+				goto error_sm;
+		}
+	}
+
+	return 0;
+
+error_sm:
+	context = iommu_context_addr(iommu, i, 0, 0);
+	WARN_ON_ONCE(kho_unpreserve_folio(virt_to_folio(context)));
+error:
+	WARN_ON_ONCE(unpreserve_iommu_context(iommu, i));
+	return ret;
+}
+
 static int preserve_iommu_state(struct intel_iommu *iommu,
 				struct iommu_unit_ser *ser)
 {
-	pr_warn("Not implemented\n");
-	return 0;
+	int ret;
+
+	spin_lock(&iommu->lock);
+	ret = preserve_iommu_context(iommu);
+	if (ret)
+		goto error;
+
+	ret = kho_preserve_folio(virt_to_folio(iommu->root_entry));
+	if (ret) {
+		unpreserve_iommu_context(iommu, -1);
+		goto error;
+	}
+
+	ser->phys_addr = iommu->reg_phys;
+	ser->root_table = __pa(iommu->root_entry);
+	atomic_set(&iommu->preserved, 1);
+error:
+	spin_unlock(&iommu->lock);
+	return ret;
 }
 
 static void unpreserve_state(struct iommu_ser *ser)
-- 
2.51.0.536.g15c5d4f767-goog

On Sun, Sep 28, 2025 at 07:06:16PM +0000, Samiullah Khawaja wrote:
> +static int unpreserve_iommu_context(struct intel_iommu *iommu, int end)
> +{
> +	struct context_entry *context;
> +	int i;
> +
> +	if (end < 0)
> +		end = ROOT_ENTRY_NR;
> +
> +	for (i = 0; i < end; i++) {
> +		context = iommu_context_addr(iommu, i, 0, 0);
> +		if (context)
> +			WARN_ON_ONCE(kho_unpreserve_folio(virt_to_folio(context)));

Wrong function. IIRC all of these allocations came from iommu-pages.c
and have struct page metadata.

iommu-pages needs to participate in restoring them and put back it's
struct page information.

>  static int preserve_iommu_state(struct intel_iommu *iommu,
>  				struct iommu_unit_ser *ser)
>  {
> -	pr_warn("Not implemented\n");
> -	return 0;
> +	int ret;
> +
> +	spin_lock(&iommu->lock);
> +	ret = preserve_iommu_context(iommu);
> +	if (ret)
> +		goto error;
> +
> +	ret = kho_preserve_folio(virt_to_folio(iommu->root_entry));
> +	if (ret) {
> +		unpreserve_iommu_context(iommu, -1);
> +		goto error;
> +	}
> +
> +	ser->phys_addr = iommu->reg_phys;
> +	ser->root_table = __pa(iommu->root_entry);
> +	atomic_set(&iommu->preserved, 1);

Why all these atomics??

Also most probably this should all be flowing through the core code as
I think the core code has to genericall prevent attach/detach/probe
from happing once serialization has started.

Jason