diff --git a/arch/arm64/include/asm/kvm_s2mpu.h b/arch/arm64/include/asm/kvm_s2mpu.h
index 3c61842a25ee..0e262b0c032d 100644
--- a/arch/arm64/include/asm/kvm_s2mpu.h
+++ b/arch/arm64/include/asm/kvm_s2mpu.h
@@ -144,6 +144,21 @@ enum s2mpu_version {
 	S2MPU_VERSION_9 = 0x20000000,
 };
 
+enum s2mpu_power_state {
+	S2MPU_POWER_ALWAYS_ON = 0,
+	S2MPU_POWER_ON,
+	S2MPU_POWER_OFF,
+};
+
+struct s2mpu {
+	phys_addr_t pa;
+	void __iomem *va;
+	u32 version;
+	enum s2mpu_power_state power_state;
+	u32 power_domain_id;
+	u32 context_cfg_valid_vid;
+};
+
 enum mpt_prot {
 	MPT_PROT_NONE	= 0,
 	MPT_PROT_R	= BIT(0),
diff --git a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
index 202440d8d5b7..3ffbc4578e6d 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
@@ -34,44 +34,52 @@
 	(CONTEXT_CFG_VALID_VID_CTX_VID(ctxid, vid) \
 	 | (((ctxid) < (nr_ctx)) ? CONTEXT_CFG_VALID_VID_CTX_VALID(ctxid) : 0))
 
-static size_t __ro_after_init			nr_s2mpus;
-static struct pkvm_iommu __ro_after_init	*s2mpus;
-static struct mpt				host_mpt;
-static hyp_spinlock_t				s2mpu_lock;
+static size_t __ro_after_init		nr_s2mpus;
+static struct s2mpu __ro_after_init	*s2mpus;
+static struct mpt			host_mpt;
+static hyp_spinlock_t			s2mpu_lock;
 
-struct s2mpu_drv_data {
-	u32 version;
-	u32 context_cfg_valid_vid;
-};
-
-static bool is_version(struct pkvm_iommu *dev, u32 version)
+static bool is_version(struct s2mpu *dev, u32 version)
 {
-	struct s2mpu_drv_data *data = (struct s2mpu_drv_data *)dev->data;
-
-	return (data->version & VERSION_CHECK_MASK) == version;
+	return (dev->version & VERSION_CHECK_MASK) == version;
 }
 
-static bool is_powered_on(struct pkvm_iommu *dev)
+static bool is_powered_on(struct s2mpu *dev)
 {
-	return dev->powered;
+	switch (dev->power_state) {
+	case S2MPU_POWER_ALWAYS_ON:
+	case S2MPU_POWER_ON:
+		return true;
+	case S2MPU_POWER_OFF:
+		return false;
+	default:
+		BUG();
+	}
 }
 
-static bool is_in_power_domain(struct pkvm_iommu *dev, u64 power_domain_id)
+static bool is_in_power_domain(struct s2mpu *dev, u64 power_domain_id)
 {
-	return false;
+	switch (dev->power_state) {
+	case S2MPU_POWER_ALWAYS_ON:
+		return false;
+	case S2MPU_POWER_ON:
+	case S2MPU_POWER_OFF:
+		return dev->power_domain_id == power_domain_id;
+	default:
+		BUG();
+	}
 }
 
-static u32 __context_cfg_valid_vid(struct pkvm_iommu *dev, u32 vid_bmap)
+static u32 __context_cfg_valid_vid(struct s2mpu *dev, u32 vid_bmap)
 {
-	struct s2mpu_drv_data *data = (struct s2mpu_drv_data *)dev->data;
 	u8 ctx_vid[NR_CTX_IDS] = { 0 };
 	unsigned int vid, ctx = 0;
 	unsigned int num_ctx;
 	u32 res;
 
 	/* Only initialize once. */
-	if (data->context_cfg_valid_vid)
-		return data->context_cfg_valid_vid;
+	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) {
@@ -95,11 +103,11 @@ static u32 __context_cfg_valid_vid(struct pkvm_iommu *dev, u32 vid_bmap)
 	    | CTX_CFG_ENTRY(6, ctx, ctx_vid[6])
 	    | CTX_CFG_ENTRY(7, ctx, ctx_vid[7]);
 
-	data->context_cfg_valid_vid = res;
+	dev->context_cfg_valid_vid = res;
 	return res;
 }
 
-static int __initialize_v9(struct pkvm_iommu *dev)
+static int __initialize_v9(struct s2mpu *dev)
 {
 	u32 ssmt_valid_vid_bmap, ctx_cfg;
 
@@ -118,14 +126,12 @@ static int __initialize_v9(struct pkvm_iommu *dev)
 	return 0;
 }
 
-static int __initialize(struct pkvm_iommu *dev)
+static int __initialize(struct s2mpu *dev)
 {
-	struct s2mpu_drv_data *data = (struct s2mpu_drv_data *)dev->data;
+	if (!dev->version)
+		dev->version = readl_relaxed(dev->va + REG_NS_VERSION);
 
-	if (!data->version)
-		data->version = readl_relaxed(dev->va + REG_NS_VERSION);
-
-	switch (data->version & VERSION_CHECK_MASK) {
+	switch (dev->version & VERSION_CHECK_MASK) {
 	case S2MPU_VERSION_8:
 		return 0;
 	case S2MPU_VERSION_9:
@@ -135,7 +141,7 @@ static int __initialize(struct pkvm_iommu *dev)
 	}
 }
 
-static void __set_control_regs(struct pkvm_iommu *dev)
+static void __set_control_regs(struct s2mpu *dev)
 {
 	u32 ctrl0 = 0, irq_vids;
 
@@ -175,7 +181,7 @@ static void __wait_while(void __iomem *addr, u32 mask)
 		continue;
 }
 
-static void __wait_for_invalidation_complete(struct pkvm_iommu *dev)
+static void __wait_for_invalidation_complete(struct s2mpu *dev)
 {
 	/* Must not access SFRs while S2MPU is busy invalidating (v9 only). */
 	if (is_version(dev, S2MPU_VERSION_9)) {
@@ -184,13 +190,13 @@ static void __wait_for_invalidation_complete(struct pkvm_iommu *dev)
 	}
 }
 
-static void __all_invalidation(struct pkvm_iommu *dev)
+static void __all_invalidation(struct s2mpu *dev)
 {
 	writel_relaxed(INVALIDATION_INVALIDATE, dev->va + REG_NS_ALL_INVALIDATION);
 	__wait_for_invalidation_complete(dev);
 }
 
-static void __range_invalidation(struct pkvm_iommu *dev, phys_addr_t first_byte,
+static void __range_invalidation(struct s2mpu *dev, phys_addr_t first_byte,
 				 phys_addr_t last_byte)
 {
 	u32 start_ppn = first_byte >> RANGE_INVALIDATION_PPN_SHIFT;
@@ -202,14 +208,14 @@ static void __range_invalidation(struct pkvm_iommu *dev, phys_addr_t first_byte,
 	__wait_for_invalidation_complete(dev);
 }
 
-static void __set_l1entry_attr_with_prot(struct pkvm_iommu *dev, unsigned int gb,
+static void __set_l1entry_attr_with_prot(struct s2mpu *dev, unsigned int gb,
 					 unsigned int vid, enum mpt_prot prot)
 {
 	writel_relaxed(L1ENTRY_ATTR_1G(prot),
 		       dev->va + REG_NS_L1ENTRY_ATTR(vid, gb));
 }
 
-static void __set_l1entry_attr_with_fmpt(struct pkvm_iommu *dev, unsigned int gb,
+static void __set_l1entry_attr_with_fmpt(struct s2mpu *dev, unsigned int gb,
 					 unsigned int vid, struct fmpt *fmpt)
 {
 	if (fmpt->gran_1g) {
@@ -221,7 +227,7 @@ static void __set_l1entry_attr_with_fmpt(struct pkvm_iommu *dev, unsigned int gb
 	}
 }
 
-static void __set_l1entry_l2table_addr(struct pkvm_iommu *dev, unsigned int gb,
+static void __set_l1entry_l2table_addr(struct s2mpu *dev, unsigned int gb,
 				       unsigned int vid, phys_addr_t addr)
 {
 	/* Order against writes to the SMPT. */
@@ -233,7 +239,7 @@ static void __set_l1entry_l2table_addr(struct pkvm_iommu *dev, unsigned int gb,
  * Initialize S2MPU device and set all GB regions to 1G granularity with
  * given protection bits.
  */
-static int initialize_with_prot(struct pkvm_iommu *dev, enum mpt_prot prot)
+static int initialize_with_prot(struct s2mpu *dev, enum mpt_prot prot)
 {
 	unsigned int gb, vid;
 	int ret;
@@ -255,7 +261,7 @@ static int initialize_with_prot(struct pkvm_iommu *dev, enum mpt_prot prot)
  * Initialize S2MPU device, set L2 table addresses and configure L1TABLE_ATTR
  * registers according to the given MPT struct.
  */
-static int initialize_with_mpt(struct pkvm_iommu *dev, struct mpt *mpt)
+static int initialize_with_mpt(struct s2mpu *dev, struct mpt *mpt)
 {
 	unsigned int gb, vid;
 	struct fmpt *fmpt;
@@ -288,7 +294,7 @@ static void set_mpt_range_locked(struct mpt *mpt, phys_addr_t first_byte,
 	unsigned int last_gb = last_byte / SZ_1G;
 	size_t start_gb_byte, end_gb_byte;
 	unsigned int gb, vid;
-	struct pkvm_iommu *dev;
+	struct s2mpu *dev;
 	struct fmpt *fmpt;
 	enum mpt_update_flags flags;
 
@@ -354,7 +360,7 @@ static void s2mpu_host_stage2_set_owner(phys_addr_t addr, size_t size,
 static int s2mpu_host_stage2_adjust_mmio_range(phys_addr_t addr, phys_addr_t *start,
 					       phys_addr_t *end)
 {
-	struct pkvm_iommu *dev;
+	struct s2mpu *dev;
 	phys_addr_t dev_start, dev_end, int_start, int_end;
 
 	/* Find the PA interval in the non-empty, sorted list of S2MPUs. */
@@ -387,7 +393,7 @@ static bool s2mpu_host_smc_handler(struct kvm_cpu_context *host_ctxt)
 	DECLARE_REG(u64, group, host_ctxt, 3);
 
 	struct arm_smccc_res res;
-	struct pkvm_iommu *dev;
+	struct s2mpu *dev;
 	int ret;
 
 	if (fn != SMC_CMD_PREPARE_PD_ONOFF)
@@ -421,11 +427,11 @@ static bool s2mpu_host_smc_handler(struct kvm_cpu_context *host_ctxt)
 				continue;
 
 			if (mode == SMC_MODE_POWER_UP) {
-				dev->powered = true;
+				dev->power_state = S2MPU_POWER_ON;
 				ret = initialize_with_mpt(dev, &host_mpt);
 			} else {
 				ret = initialize_with_prot(dev, MPT_PROT_NONE);
-				dev->powered = false;
+				dev->power_state = S2MPU_POWER_OFF;
 			}
 		}
 	}
@@ -435,9 +441,9 @@ static bool s2mpu_host_smc_handler(struct kvm_cpu_context *host_ctxt)
 	return true;  /* SMC handled */
 }
 
-static struct pkvm_iommu *find_s2mpu_by_addr(phys_addr_t addr)
+static struct s2mpu *find_s2mpu_by_addr(phys_addr_t addr)
 {
-	struct pkvm_iommu *dev;
+	struct s2mpu *dev;
 
 	for_each_s2mpu(dev) {
 		if (dev->pa <= addr && addr < (dev->pa + S2MPU_MMIO_SIZE))
@@ -476,7 +482,7 @@ static bool s2mpu_host_mmio_dabt_handler(struct kvm_cpu_context *host_ctxt,
 					 phys_addr_t fault_pa, unsigned int len,
 					 bool is_write, int rd)
 {
-	struct pkvm_iommu *dev;
+	struct s2mpu *dev;
 	size_t off;
 	u32 mask;
 
@@ -502,7 +508,7 @@ static bool s2mpu_host_mmio_dabt_handler(struct kvm_cpu_context *host_ctxt,
 
 static int s2mpu_init(void)
 {
-	struct pkvm_iommu *dev;
+	struct s2mpu *dev;
 	unsigned int gb;
 	int ret;