diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index 626f8c708dc9..69465035cff9 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -381,6 +381,7 @@ enum kvm_iommu_driver {
 };
 
 enum pkvm_iommu_driver_id {
+	PKVM_IOMMU_DRIVER_S2MPU,
 	PKVM_IOMMU_NR_DRIVERS,
 };
 
@@ -395,6 +396,8 @@ int pkvm_iommu_register(struct device *dev, enum pkvm_iommu_driver_id drv_id,
 int pkvm_iommu_suspend(struct device *dev);
 int pkvm_iommu_resume(struct device *dev);
 
+int pkvm_iommu_s2mpu_register(struct device *dev, phys_addr_t pa);
+
 struct vcpu_reset_state {
 	unsigned long	pc;
 	unsigned long	r0;
diff --git a/arch/arm64/include/asm/kvm_s2mpu.h b/arch/arm64/include/asm/kvm_s2mpu.h
index 3c61842a25ee..0e322ede8421 100644
--- a/arch/arm64/include/asm/kvm_s2mpu.h
+++ b/arch/arm64/include/asm/kvm_s2mpu.h
@@ -123,8 +123,9 @@ static_assert(SMPT_GRAN <= PAGE_SIZE);
 #define SMPT_ELEMS_PER_WORD			(SMPT_WORD_SIZE * SMPT_ELEMS_PER_BYTE)
 #define SMPT_WORD_BYTE_RANGE			(SMPT_GRAN * SMPT_ELEMS_PER_WORD)
 #define SMPT_NUM_ELEMS				(SZ_1G / SMPT_GRAN)
-#define SMPT_NUM_WORDS				(SMPT_SIZE / SMPT_WORD_SIZE)
 #define SMPT_SIZE				(SMPT_NUM_ELEMS / SMPT_ELEMS_PER_BYTE)
+#define SMPT_NUM_WORDS				(SMPT_SIZE / SMPT_WORD_SIZE)
+#define SMPT_NUM_PAGES				(SMPT_SIZE / PAGE_SIZE)
 #define SMPT_ORDER				get_order(SMPT_SIZE)
 
 /*
diff --git a/arch/arm64/kvm/hyp/include/nvhe/iommu.h b/arch/arm64/kvm/hyp/include/nvhe/iommu.h
index 75a4790a8a5d..c8907c911673 100644
--- a/arch/arm64/kvm/hyp/include/nvhe/iommu.h
+++ b/arch/arm64/kvm/hyp/include/nvhe/iommu.h
@@ -80,6 +80,8 @@ bool pkvm_iommu_host_dabt_handler(struct kvm_cpu_context *host_ctxt, u32 esr,
 void pkvm_iommu_host_stage2_idmap(phys_addr_t start, phys_addr_t end,
 				  enum kvm_pgtable_prot prot);
 
+extern const struct pkvm_iommu_ops pkvm_s2mpu_ops;
+
 struct kvm_iommu_ops {
 	int (*init)(void);
 	bool (*host_smc_handler)(struct kvm_cpu_context *host_ctxt);
diff --git a/arch/arm64/kvm/hyp/nvhe/iommu.c b/arch/arm64/kvm/hyp/nvhe/iommu.c
index 80d85b442233..9b659abc7957 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu.c
@@ -63,6 +63,8 @@ static inline struct pkvm_iommu_driver *get_driver(enum pkvm_iommu_driver_id id)
 static const struct pkvm_iommu_ops *get_driver_ops(enum pkvm_iommu_driver_id id)
 {
 	switch (id) {
+	case PKVM_IOMMU_DRIVER_S2MPU:
+		return IS_ENABLED(CONFIG_KVM_S2MPU) ? &pkvm_s2mpu_ops : NULL;
 	default:
 		return NULL;
 	}
diff --git a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
index afa17a80d592..6bcea2c5fe6b 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
@@ -472,50 +472,72 @@ static bool s2mpu_host_mmio_dabt_handler(struct kvm_cpu_context *host_ctxt,
 	return true;
 }
 
-static int s2mpu_init(void)
+static int s2mpu_init(void *data, size_t size)
 {
-	struct pkvm_iommu *dev;
+	struct mpt in_mpt;
+	u32 *smpt;
+	phys_addr_t pa;
 	unsigned int gb;
-	int ret;
+	int ret = 0;
 
-	/* Map data structures in EL2 stage-1. */
-	ret = pkvm_create_mappings(s2mpus,
-				   s2mpus + nr_s2mpus,
-				   PAGE_HYP);
-	if (ret)
-		return ret;
+	if (size != sizeof(in_mpt))
+		return -EINVAL;
 
+	/* The host can concurrently modify 'data'. Copy it to avoid TOCTOU. */
+	memcpy(&in_mpt, data, sizeof(in_mpt));
+
+	/* Take ownership of all SMPT buffers. This will also map them in. */
 	for_each_gb(gb) {
-		ret = pkvm_create_mappings(
-			host_mpt.fmpt[gb].smpt,
-			host_mpt.fmpt[gb].smpt + SMPT_NUM_WORDS,
-			PAGE_HYP);
+		smpt = kern_hyp_va(in_mpt.fmpt[gb].smpt);
+		pa = __hyp_pa(smpt);
+
+		if (!IS_ALIGNED(pa, SMPT_SIZE)) {
+			ret = -EINVAL;
+			break;
+		}
+
+		ret = __pkvm_host_donate_hyp(pa >> PAGE_SHIFT, SMPT_NUM_PAGES);
 		if (ret)
-			return ret;
+			break;
+
+		host_mpt.fmpt[gb] = (struct fmpt){
+			.smpt = smpt,
+			.gran_1g = true,
+			.prot = MPT_PROT_NONE,
+		};
 	}
 
-	/* Map S2MPU MMIO regions in EL2 stage-1. */
-	for_each_s2mpu(dev) {
-		dev->va = (void __iomem *)__pkvm_create_private_mapping(
-			dev->pa, S2MPU_MMIO_SIZE, PAGE_HYP_DEVICE);
-		if (IS_ERR_OR_NULL(dev->va))
-			return PTR_ERR(dev->va);
+	/* Try to return memory back if there was an error. */
+	if (ret) {
+		for_each_gb(gb) {
+			smpt = host_mpt.fmpt[gb].smpt;
+			if (!smpt)
+				break;
+
+			WARN_ON(__pkvm_hyp_donate_host(__hyp_pa(smpt) >> PAGE_SHIFT,
+						       SMPT_NUM_PAGES));
+		}
+		memset(&host_mpt, 0, sizeof(host_mpt));
 	}
 
-	/*
-	 * Program all S2MPUs powered on at boot. Note that they may not be in
-	 * the blocking reset state as the bootloader may have programmed them.
-	 */
-	for_each_powered_s2mpu(dev) {
-		ret = initialize_with_mpt(dev, &host_mpt);
-		if (ret)
-			return ret;
-	}
+	return ret;
+}
+
+static int s2mpu_validate(phys_addr_t pa, size_t size)
+{
+	if (size != S2MPU_MMIO_SIZE)
+		return -EINVAL;
+
 	return 0;
 }
 
-const struct kvm_iommu_ops kvm_s2mpu_ops = (struct kvm_iommu_ops){
+const struct pkvm_iommu_ops pkvm_s2mpu_ops = (struct pkvm_iommu_ops){
 	.init = s2mpu_init,
+	.validate = s2mpu_validate,
+	.data_size = sizeof(struct s2mpu_drv_data),
+};
+
+const struct kvm_iommu_ops kvm_s2mpu_ops = (struct kvm_iommu_ops){
 	.host_smc_handler = s2mpu_host_smc_handler,
 	.host_mmio_dabt_handler = s2mpu_host_mmio_dabt_handler,
 	.host_stage2_set_owner = s2mpu_host_stage2_set_owner,
diff --git a/arch/arm64/kvm/iommu/s2mpu.c b/arch/arm64/kvm/iommu/s2mpu.c
index 496e75f76e56..be2b1ad09480 100644
--- a/arch/arm64/kvm/iommu/s2mpu.c
+++ b/arch/arm64/kvm/iommu/s2mpu.c
@@ -6,3 +6,81 @@
 
 #include <linux/kvm_host.h>
 #include <asm/kvm_s2mpu.h>
+
+static int init_s2mpu_driver(void)
+{
+	static DEFINE_MUTEX(lock);
+	static bool init_done;
+
+	struct mpt *mpt;
+	unsigned int gb;
+	unsigned long addr;
+	u64 pfn;
+	int ret = 0;
+
+	mutex_lock(&lock);
+	if (init_done)
+		goto out;
+
+	/* Allocate a page for driver data. Must fit MPT descriptor. */
+	BUILD_BUG_ON(sizeof(*mpt) > PAGE_SIZE);
+	addr = __get_free_page(GFP_KERNEL);
+	if (!addr) {
+		ret = -ENOMEM;
+		goto out;
+	}
+
+	mpt = (struct mpt *)addr;
+
+	/* Allocate SMPT buffers. */
+	for_each_gb(gb) {
+		addr = __get_free_pages(GFP_KERNEL, SMPT_ORDER);
+		if (!addr) {
+			ret = -ENOMEM;
+			goto out_free;
+		}
+		mpt->fmpt[gb].smpt = (u32 *)addr;
+	}
+
+	/* Share MPT descriptor with hyp. */
+	pfn = __pa(mpt) >> PAGE_SHIFT;
+	ret = kvm_call_hyp_nvhe(__pkvm_host_share_hyp, pfn);
+	if (ret)
+		goto out_free;
+
+	/* Hypercall to initialize EL2 driver. */
+	ret = pkvm_iommu_driver_init(PKVM_IOMMU_DRIVER_S2MPU, mpt, sizeof(*mpt));
+	if (ret)
+		goto out_unshare;
+
+	init_done = true;
+
+out_unshare:
+	WARN_ON(kvm_call_hyp_nvhe(__pkvm_host_unshare_hyp, pfn));
+out_free:
+	/* TODO - will driver return the memory? */
+	if (ret) {
+		for_each_gb(gb)
+			free_pages((unsigned long)mpt->fmpt[gb].smpt, SMPT_ORDER);
+		free_page((unsigned long)mpt);
+	}
+out:
+	mutex_unlock(&lock);
+	return ret;
+}
+
+int pkvm_iommu_s2mpu_register(struct device *dev, phys_addr_t addr)
+{
+	int ret;
+
+	if (!is_protected_kvm_enabled())
+		return -ENODEV;
+
+	ret = init_s2mpu_driver();
+	if (ret)
+		return ret;
+
+	return pkvm_iommu_register(dev, PKVM_IOMMU_DRIVER_S2MPU,
+				   addr, S2MPU_MMIO_SIZE);
+}
+EXPORT_SYMBOL_GPL(pkvm_iommu_s2mpu_register);