diff --git a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
index f10326b7bf8b..c138c48eb14f 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
@@ -30,6 +30,10 @@
 #define for_each_powered_s2mpu(i) \
 	for_each_s2mpu((i)) if (is_powered_on((i)))
 
+#define CTX_CFG_ENTRY(ctxid, nr_ctx, vid) \
+	(CONTEXT_CFG_VALID_VID_CTX_VID(ctxid, vid) \
+	 | (((ctxid) < (nr_ctx)) ? CONTEXT_CFG_VALID_VID_CTX_VALID(ctxid) : 0))
+
 size_t __ro_after_init		kvm_hyp_nr_s2mpus;
 struct s2mpu __ro_after_init	*kvm_hyp_s2mpus;
 struct mpt			kvm_hyp_host_mpt;
@@ -67,18 +71,75 @@ static bool is_in_power_domain(struct s2mpu *dev, u64 power_domain_id)
 	}
 }
 
-/*
- * Write CONTEXT_CFG_VALID_VID configuration before touching L1ENTRY* registers.
- * Writes to those registers are ignored unless there is a context ID allocated
- * to the corresponding VID (v9 only).
- */
-static void __set_context_ids(struct s2mpu *dev)
+static u32 __context_cfg_valid_vid(struct s2mpu *dev, u32 vid_bmap)
 {
-	if (!is_version(dev, S2MPU_VERSION_9))
-		return;
+	u8 ctx_vid[NR_CTX_IDS] = { 0 };
+	unsigned int vid, ctx = 0;
+	unsigned int num_ctx;
+	u32 res;
 
-	writel_relaxed(dev->context_cfg_valid_vid,
-		       dev->va + REG_NS_CONTEXT_CFG_VALID_VID);
+	/* Only initialize once. */
+	if (dev->context_cfg_valid_vid)
+		return dev->context_cfg_valid_vid;
+
+	num_ctx = readl_relaxed(dev->va + REG_NS_NUM_CONTEXT) & NUM_CONTEXT_MASK;
+	while (vid_bmap) {
+		/* Break if we cannot allocate more. */
+		if (ctx >= num_ctx)
+			break;
+
+		vid = __ffs(vid_bmap);
+		vid_bmap &= ~BIT(vid);
+		ctx_vid[ctx++] = vid;
+	}
+
+	/* The following loop was unrolled so bitmasks are constant. */
+	BUILD_BUG_ON(NR_CTX_IDS != 8);
+	res = CTX_CFG_ENTRY(0, ctx, ctx_vid[0])
+	    | CTX_CFG_ENTRY(1, ctx, ctx_vid[1])
+	    | CTX_CFG_ENTRY(2, ctx, ctx_vid[2])
+	    | CTX_CFG_ENTRY(3, ctx, ctx_vid[3])
+	    | CTX_CFG_ENTRY(4, ctx, ctx_vid[4])
+	    | CTX_CFG_ENTRY(5, ctx, ctx_vid[5])
+	    | CTX_CFG_ENTRY(6, ctx, ctx_vid[6])
+	    | CTX_CFG_ENTRY(7, ctx, ctx_vid[7]);
+
+	dev->context_cfg_valid_vid = res;
+	return res;
+}
+
+static int __initialize_v9(struct s2mpu *dev)
+{
+	u32 ssmt_valid_vid_bmap, ctx_cfg;
+
+	/* Assume all VIDs may be generated by the connected SSMTs for now. */
+	ssmt_valid_vid_bmap = ALL_VIDS_BITMAP;
+	ctx_cfg = __context_cfg_valid_vid(dev, ssmt_valid_vid_bmap);
+	if (!ctx_cfg)
+		return -EINVAL;
+
+	/*
+	 * Write CONTEXT_CFG_VALID_VID configuration before touching L1ENTRY*
+	 * registers. Writes to those registers are ignored unless there is
+	 * a context ID allocated to the corresponding VID (v9 only).
+	 */
+	writel_relaxed(ctx_cfg, dev->va + REG_NS_CONTEXT_CFG_VALID_VID);
+	return 0;
+}
+
+static int __initialize(struct s2mpu *dev)
+{
+	if (!dev->version)
+		dev->version = readl_relaxed(dev->va + REG_NS_VERSION);
+
+	switch (dev->version & VERSION_CHECK_MASK) {
+	case S2MPU_VERSION_8:
+		return 0;
+	case S2MPU_VERSION_9:
+		return __initialize_v9(dev);
+	default:
+		return -EINVAL;
+	}
 }
 
 static void __set_control_regs(struct s2mpu *dev)
@@ -179,12 +240,14 @@ static void __set_l1entry_l2table_addr(struct s2mpu *dev, unsigned int gb,
  * Initialize S2MPU device and set all GB regions to 1G granularity with
  * given protection bits.
  */
-static void initialize_with_prot(struct s2mpu *dev, enum mpt_prot prot)
+static int initialize_with_prot(struct s2mpu *dev, enum mpt_prot prot)
 {
 	unsigned int gb, vid;
+	int ret;
 
-	/* Must write CONTEXT_CFG_VALID_VID before setting L1ENTRY registers. */
-	__set_context_ids(dev);
+	ret = __initialize(dev);
+	if (ret)
+		return ret;
 
 	for_each_gb_and_vid(gb, vid)
 		__set_l1entry_attr_with_prot(dev, gb, vid, prot);
@@ -192,19 +255,22 @@ static void initialize_with_prot(struct s2mpu *dev, enum mpt_prot prot)
 
 	/* Set control registers, enable the S2MPU. */
 	__set_control_regs(dev);
+	return 0;
 }
 
 /*
  * Initialize S2MPU device, set L2 table addresses and configure L1TABLE_ATTR
  * registers according to the given MPT struct.
  */
-static void initialize_with_mpt(struct s2mpu *dev, struct mpt *mpt)
+static int initialize_with_mpt(struct s2mpu *dev, struct mpt *mpt)
 {
 	unsigned int gb, vid;
 	struct fmpt *fmpt;
+	int ret;
 
-	/* Must write CONTEXT_CFG_VALID_VID before setting L1ENTRY registers. */
-	__set_context_ids(dev);
+	ret = __initialize(dev);
+	if (ret)
+		return ret;
 
 	for_each_gb_and_vid(gb, vid) {
 		fmpt = &mpt->fmpt[gb];
@@ -215,6 +281,7 @@ static void initialize_with_mpt(struct s2mpu *dev, struct mpt *mpt)
 
 	/* Set control registers, enable the S2MPU. */
 	__set_control_regs(dev);
+	return 0;
 }
 
 /*
@@ -326,6 +393,7 @@ static bool s2mpu_host_smc_handler(struct kvm_cpu_context *host_ctxt)
 
 	struct arm_smccc_res res;
 	struct s2mpu *dev;
+	int ret;
 
 	if (fn != SMC_CMD_PREPARE_PD_ONOFF)
 		return false; /* SMC not handled */
@@ -350,23 +418,25 @@ static bool s2mpu_host_smc_handler(struct kvm_cpu_context *host_ctxt)
 
 	hyp_spin_lock(&s2mpu_lock);
 	arm_smccc_1_1_smc(fn, mode, domain_id, group, &res);
-	if (res.a0 == SMCCC_RET_SUCCESS) {
+	ret = res.a0;
+
+	if (ret == SMCCC_RET_SUCCESS) {
 		for_each_s2mpu(dev) {
 			if (!is_in_power_domain(dev, domain_id))
 				continue;
 
 			if (mode == SMC_MODE_POWER_UP) {
 				dev->power_state = S2MPU_POWER_ON;
-				initialize_with_mpt(dev, &kvm_hyp_host_mpt);
+				ret = initialize_with_mpt(dev, &kvm_hyp_host_mpt);
 			} else {
-				initialize_with_prot(dev, MPT_PROT_NONE);
+				ret = initialize_with_prot(dev, MPT_PROT_NONE);
 				dev->power_state = S2MPU_POWER_OFF;
 			}
 		}
 	}
 	hyp_spin_unlock(&s2mpu_lock);
 
-	cpu_reg(host_ctxt, 0) = res.a0;
+	cpu_reg(host_ctxt, 0) = ret;
 	return true;  /* SMC handled */
 }
 
@@ -469,8 +539,11 @@ static int s2mpu_init(void)
 	 * 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)
-		initialize_with_mpt(dev, &kvm_hyp_host_mpt);
+	for_each_powered_s2mpu(dev) {
+		ret = initialize_with_mpt(dev, &kvm_hyp_host_mpt);
+		if (ret)
+			return ret;
+	}
 	return 0;
 }
 
diff --git a/arch/arm64/kvm/iommu/s2mpu.c b/arch/arm64/kvm/iommu/s2mpu.c
index 00111f237de3..79c33815267b 100644
--- a/arch/arm64/kvm/iommu/s2mpu.c
+++ b/arch/arm64/kvm/iommu/s2mpu.c
@@ -13,10 +13,6 @@
 #include <asm/kvm_mmu.h>
 #include <asm/kvm_s2mpu.h>
 
-#define CTX_CFG_ENTRY(ctxid, nr_ctx, vid) \
-	(CONTEXT_CFG_VALID_VID_CTX_VID(ctxid, vid) \
-	 | (((ctxid) < (nr_ctx)) ? CONTEXT_CFG_VALID_VID_CTX_VALID(ctxid) : 0))
-
 struct s2mpu_irq_info {
 	struct device *dev;
 	void __iomem *va;
@@ -81,59 +77,6 @@ static irqreturn_t s2mpu_irq_handler(int irq, void *data)
 	return ret;
 }
 
-static u32 gen_ctx_cfg_valid_vid(struct platform_device *pdev,
-				 unsigned int num_ctx, u32 vid_bmap)
-{
-	u8 ctx_vid[NR_CTX_IDS] = { 0 };
-	unsigned int vid, ctx = 0;
-
-	/* Check NUM_CONTEXT value is within bounds. This should not happen. */
-	if (WARN_ON(num_ctx > NR_CTX_IDS))
-		num_ctx = NR_CTX_IDS;
-
-	while (vid_bmap) {
-		/* Break if we cannot allocate more. */
-		if (ctx >= num_ctx) {
-			dev_warn(&pdev->dev,
-				 "could not allocate all context IDs, DMA may be blocked (VID bitmap: 0x%x)",
-				 vid_bmap);
-			break;
-		}
-
-		vid = __ffs(vid_bmap);
-		vid_bmap &= ~BIT(vid);
-		ctx_vid[ctx++] = vid;
-	}
-
-	/* The following loop was unrolled so bitmasks are constant. */
-	BUILD_BUG_ON(NR_CTX_IDS != 8);
-	return CTX_CFG_ENTRY(0, ctx, ctx_vid[0])
-	     | CTX_CFG_ENTRY(1, ctx, ctx_vid[1])
-	     | CTX_CFG_ENTRY(2, ctx, ctx_vid[2])
-	     | CTX_CFG_ENTRY(3, ctx, ctx_vid[3])
-	     | CTX_CFG_ENTRY(4, ctx, ctx_vid[4])
-	     | CTX_CFG_ENTRY(5, ctx, ctx_vid[5])
-	     | CTX_CFG_ENTRY(6, ctx, ctx_vid[6])
-	     | CTX_CFG_ENTRY(7, ctx, ctx_vid[7]);
-}
-
-static int s2mpu_probe_v9(struct platform_device *pdev, void __iomem *kaddr,
-			  struct s2mpu *info)
-{
-	unsigned int num_ctx;
-	u32 ssmt_valid_vid_bmap;
-
-	ssmt_valid_vid_bmap = ALL_VIDS_BITMAP;
-	num_ctx = readl_relaxed(kaddr + REG_NS_NUM_CONTEXT) & NUM_CONTEXT_MASK;
-	info->context_cfg_valid_vid = gen_ctx_cfg_valid_vid(pdev, num_ctx, ssmt_valid_vid_bmap);
-	if (!info->context_cfg_valid_vid) {
-		dev_err(&pdev->dev, "failed to allocate context IDs");
-		return -EINVAL;
-	}
-
-	return 0;
-}
-
 /*
  * Parse interrupt information from DT and if found, register IRQ handler.
  * This is considered optional and will not fail even if the initialization is
@@ -232,20 +175,6 @@ static int s2mpu_probe(struct platform_device *pdev)
 	 */
 	s2mpu_probe_irq(pdev, kaddr);
 
-	info->version = readl_relaxed(kaddr + REG_NS_VERSION);
-	switch (info->version & VERSION_CHECK_MASK) {
-	case S2MPU_VERSION_8:
-		break;
-	case S2MPU_VERSION_9:
-		ret = s2mpu_probe_v9(pdev, kaddr, info);
-		if (ret)
-			return ret;
-		break;
-	default:
-		dev_err(&pdev->dev, "unexpected version 0x%08x", info->version);
-		return -EINVAL;
-	}
-
 	/* Insert successfully parsed devices to a list later copied to hyp. */
 	list_add_tail(&entry->list, &s2mpu_list);
 	kvm_hyp_nr_s2mpus++;