diff --git a/arch/arm64/include/asm/kvm_s2mpu.h b/arch/arm64/include/asm/kvm_s2mpu.h
index 0e322ede8421..4d517bc1d0eb 100644
--- a/arch/arm64/include/asm/kvm_s2mpu.h
+++ b/arch/arm64/include/asm/kvm_s2mpu.h
@@ -160,21 +160,22 @@ static const u64 mpt_prot_doubleword[] = {
 	[MPT_PROT_RW]   = 0xffffffffffffffff,
 };
 
+enum mpt_update_flags {
+	MPT_UPDATE_L1 = BIT(0),
+	MPT_UPDATE_L2 = BIT(1),
+};
+
 struct fmpt {
 	u32 *smpt;
 	bool gran_1g;
 	enum mpt_prot prot;
+	enum mpt_update_flags flags;
 };
 
 struct mpt {
 	struct fmpt fmpt[NR_GIGABYTES];
 };
 
-enum mpt_update_flags {
-	MPT_UPDATE_L1 = BIT(0),
-	MPT_UPDATE_L2 = BIT(1),
-};
-
 /* Set protection bits of SMPT in a given range without using memset. */
 static inline void __set_smpt_range_slow(u32 *smpt, size_t start_gb_byte,
 					 size_t end_gb_byte, enum mpt_prot prot)
@@ -254,24 +255,28 @@ static inline bool __is_smpt_uniform(u32 *smpt, enum mpt_prot prot)
  * Returns flags specifying whether L1/L2 changes need to be made visible
  * to the device.
  */
-static inline enum mpt_update_flags
-__set_fmpt_range(struct fmpt *fmpt, size_t start_gb_byte, size_t end_gb_byte,
-		 enum mpt_prot prot)
+static inline void __set_fmpt_range(struct fmpt *fmpt, size_t start_gb_byte,
+				    size_t end_gb_byte, enum mpt_prot prot)
 {
 	if (start_gb_byte == 0 && end_gb_byte >= SZ_1G) {
 		/* Update covers the entire GB region. */
-		if (fmpt->gran_1g && fmpt->prot == prot)
-			return 0;
+		if (fmpt->gran_1g && fmpt->prot == prot) {
+			fmpt->flags = 0;
+			return;
+		}
 
 		fmpt->gran_1g = true;
 		fmpt->prot = prot;
-		return MPT_UPDATE_L1;
+		fmpt->flags = MPT_UPDATE_L1;
+		return;
 	}
 
 	if (fmpt->gran_1g) {
 		/* GB region currently uses 1G mapping. */
-		if (fmpt->prot == prot)
-			return 0;
+		if (fmpt->prot == prot) {
+			fmpt->flags = 0;
+			return;
+		}
 
 		/*
 		 * Range has different mapping than the rest of the GB.
@@ -281,19 +286,22 @@ __set_fmpt_range(struct fmpt *fmpt, size_t start_gb_byte, size_t end_gb_byte,
 		__set_smpt_range(fmpt->smpt, 0, start_gb_byte, fmpt->prot);
 		__set_smpt_range(fmpt->smpt, start_gb_byte, end_gb_byte, prot);
 		__set_smpt_range(fmpt->smpt, end_gb_byte, SZ_1G, fmpt->prot);
-		return MPT_UPDATE_L1 | MPT_UPDATE_L2;
+		fmpt->flags = MPT_UPDATE_L1 | MPT_UPDATE_L2;
+		return;
 	}
 
 	/* GB region currently uses PAGE_SIZE mapping. */
 	__set_smpt_range(fmpt->smpt, start_gb_byte, end_gb_byte, prot);
 
 	/* Check if the entire GB region has the same prot bits. */
-	if (!__is_smpt_uniform(fmpt->smpt, prot))
-		return MPT_UPDATE_L2;
+	if (!__is_smpt_uniform(fmpt->smpt, prot)) {
+		fmpt->flags = MPT_UPDATE_L2;
+		return;
+	}
 
 	fmpt->gran_1g = true;
 	fmpt->prot = prot;
-	return MPT_UPDATE_L1;
+	fmpt->flags = MPT_UPDATE_L1;
 }
 
 #endif /* __ARM64_KVM_S2MPU_H__ */
diff --git a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
index 17c5f4548dc4..16cf1503005e 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
@@ -285,19 +285,18 @@ static void set_mpt_range_locked(struct mpt *mpt, phys_addr_t first_byte,
 	unsigned int gb, vid;
 	struct pkvm_iommu *dev;
 	struct fmpt *fmpt;
-	enum mpt_update_flags flags;
 
 	for_each_gb_in_range(gb, first_gb, last_gb) {
 		fmpt = &mpt->fmpt[gb];
 		start_gb_byte = (gb == first_gb) ? first_byte % SZ_1G : 0;
 		end_gb_byte = (gb == last_gb) ? (last_byte % SZ_1G) + 1 : SZ_1G;
 
-		flags = __set_fmpt_range(fmpt, start_gb_byte, end_gb_byte, prot);
+		__set_fmpt_range(fmpt, start_gb_byte, end_gb_byte, prot);
 
-		if (flags & MPT_UPDATE_L2)
+		if (fmpt->flags & MPT_UPDATE_L2)
 			kvm_flush_dcache_to_poc(fmpt->smpt, SMPT_SIZE);
 
-		if (flags & MPT_UPDATE_L1) {
+		if (fmpt->flags & MPT_UPDATE_L1) {
 			for_each_powered_s2mpu(dev) {
 				for_each_vid(vid)
 					__set_l1entry_attr_with_fmpt(dev, gb, vid, fmpt);