diff --git a/arch/arm64/include/asm/kvm_s2mpu.h b/arch/arm64/include/asm/kvm_s2mpu.h
index 4d517bc1d0eb..0e322ede8421 100644
--- a/arch/arm64/include/asm/kvm_s2mpu.h
+++ b/arch/arm64/include/asm/kvm_s2mpu.h
@@ -160,22 +160,21 @@ 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)
@@ -255,28 +254,24 @@ 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 void __set_fmpt_range(struct fmpt *fmpt, size_t start_gb_byte,
-				    size_t end_gb_byte, enum mpt_prot prot)
+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)
 {
 	if (start_gb_byte == 0 && end_gb_byte >= SZ_1G) {
 		/* Update covers the entire GB region. */
-		if (fmpt->gran_1g && fmpt->prot == prot) {
-			fmpt->flags = 0;
-			return;
-		}
+		if (fmpt->gran_1g && fmpt->prot == prot)
+			return 0;
 
 		fmpt->gran_1g = true;
 		fmpt->prot = prot;
-		fmpt->flags = MPT_UPDATE_L1;
-		return;
+		return MPT_UPDATE_L1;
 	}
 
 	if (fmpt->gran_1g) {
 		/* GB region currently uses 1G mapping. */
-		if (fmpt->prot == prot) {
-			fmpt->flags = 0;
-			return;
-		}
+		if (fmpt->prot == prot)
+			return 0;
 
 		/*
 		 * Range has different mapping than the rest of the GB.
@@ -286,22 +281,19 @@ static inline void __set_fmpt_range(struct fmpt *fmpt, size_t start_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);
-		fmpt->flags = MPT_UPDATE_L1 | MPT_UPDATE_L2;
-		return;
+		return MPT_UPDATE_L1 | MPT_UPDATE_L2;
 	}
 
 	/* 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)) {
-		fmpt->flags = MPT_UPDATE_L2;
-		return;
-	}
+	if (!__is_smpt_uniform(fmpt->smpt, prot))
+		return MPT_UPDATE_L2;
 
 	fmpt->gran_1g = true;
 	fmpt->prot = prot;
-	fmpt->flags = MPT_UPDATE_L1;
+	return 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 a1001f8f9aaf..cd621190120f 100644
--- a/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
+++ b/arch/arm64/kvm/hyp/nvhe/iommu/s2mpu.c
@@ -285,18 +285,19 @@ 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;
 
-		__set_fmpt_range(fmpt, start_gb_byte, end_gb_byte, prot);
+		flags = __set_fmpt_range(fmpt, start_gb_byte, end_gb_byte, prot);
 
-		if (fmpt->flags & MPT_UPDATE_L2)
+		if (flags & MPT_UPDATE_L2)
 			kvm_flush_dcache_to_poc(fmpt->smpt, SMPT_SIZE);
 
-		if (fmpt->flags & MPT_UPDATE_L1) {
+		if (flags & MPT_UPDATE_L1) {
 			for_each_powered_s2mpu(dev) {
 				for_each_vid(vid)
 					__set_l1entry_attr_with_fmpt(dev, gb, vid, fmpt);
diff --git a/arch/arm64/kvm/iommu/test_kvm_s2mpu.c b/arch/arm64/kvm/iommu/test_kvm_s2mpu.c
index fbc500de588e..67cc37f00dee 100644
--- a/arch/arm64/kvm/iommu/test_kvm_s2mpu.c
+++ b/arch/arm64/kvm/iommu/test_kvm_s2mpu.c
@@ -14,6 +14,10 @@
 
 KSTM_MODULE_GLOBALS();
 
+/*
+ * Kernel module for testing the foobinator
+ */
+
 #define ASSERT(cond)							\
 	do {								\
 		if (!(cond)) {						\
@@ -73,9 +77,11 @@ static bool __init check_smpt(size_t start_byte, size_t end_byte,
 /* Start with 1G granule, overwrite the whole 1G. */
 static int __init test_set_fmpt__fmpt_to_fmpt_whole(void)
 {
+	enum mpt_update_flags flags;
+
 	init_fmpt(MPT_PROT_NONE, /*gran_1g*/ true);
-	__set_fmpt_range(&g_fmpt, 0, SZ_1G, MPT_PROT_R);
-	ASSERT(g_fmpt.flags == MPT_UPDATE_L1);
+	flags = __set_fmpt_range(&g_fmpt, 0, SZ_1G, MPT_PROT_R);
+	ASSERT(flags == MPT_UPDATE_L1);
 	ASSERT(g_fmpt.gran_1g);
 	ASSERT(g_fmpt.prot == MPT_PROT_R);
 	return 0;
@@ -84,9 +90,11 @@ static int __init test_set_fmpt__fmpt_to_fmpt_whole(void)
 /* Start with 1G granule, overwrite the whole 1G with the same prot. */
 static int __init test_set_fmpt__fmpt_no_change_whole(void)
 {
+	enum mpt_update_flags flags;
+
 	init_fmpt(MPT_PROT_R, /*gran_1g*/ true);
-	__set_fmpt_range(&g_fmpt, 0, SZ_1G, MPT_PROT_R);
-	ASSERT(g_fmpt.flags == 0);
+	flags = __set_fmpt_range(&g_fmpt, 0, SZ_1G, MPT_PROT_R);
+	ASSERT(flags == 0);
 	ASSERT(g_fmpt.gran_1g);
 	ASSERT(g_fmpt.prot == MPT_PROT_R);
 	return 0;
@@ -95,9 +103,11 @@ static int __init test_set_fmpt__fmpt_no_change_whole(void)
 /* Start with 1G granule, partially overwrite with the same prot. */
 static int __init test_set_fmpt__fmpt_no_change_partial(void)
 {
+	enum mpt_update_flags flags;
+
 	init_fmpt(MPT_PROT_R, /*gran_1g*/ true);
-	__set_fmpt_range(&g_fmpt, 0, PAGE_SIZE, MPT_PROT_R);
-	ASSERT(g_fmpt.flags == 0);
+	flags = __set_fmpt_range(&g_fmpt, 0, PAGE_SIZE, MPT_PROT_R);
+	ASSERT(flags == 0);
 	ASSERT(g_fmpt.gran_1g);
 	ASSERT(g_fmpt.prot == MPT_PROT_R);
 	return 0;
@@ -108,10 +118,11 @@ static int __init test_set_fmpt__fmpt_to_smpt(void)
 {
 	size_t start = 5 * SMPT_WORD_BYTE_RANGE / 2;
 	size_t end = 20 * SMPT_WORD_BYTE_RANGE;
+	enum mpt_update_flags flags;
 
 	init_fmpt(MPT_PROT_R, /*gran_1g*/ true);
-	__set_fmpt_range(&g_fmpt, start, end, MPT_PROT_RW);
-	ASSERT(g_fmpt.flags == (MPT_UPDATE_L1 | MPT_UPDATE_L2));
+	flags = __set_fmpt_range(&g_fmpt, start, end, MPT_PROT_RW);
+	ASSERT(flags == (MPT_UPDATE_L1 | MPT_UPDATE_L2));
 	ASSERT(!g_fmpt.gran_1g);
 	return check_smpt(start, end, MPT_PROT_R, MPT_PROT_RW) ? 0 : 1;
 }
@@ -119,9 +130,11 @@ static int __init test_set_fmpt__fmpt_to_smpt(void)
 /* Convert from PAGE_SIZE to 1G granule by overwriting the whole 1G. */
 static int __init test_set_fmpt__smpt_to_fmpt_whole(void)
 {
+	enum mpt_update_flags flags;
+
 	init_fmpt(MPT_PROT_NONE, /*gran_1g*/ false);
-	__set_fmpt_range(&g_fmpt, 0, SZ_1G, MPT_PROT_R);
-	ASSERT(g_fmpt.flags == MPT_UPDATE_L1);
+	flags = __set_fmpt_range(&g_fmpt, 0, SZ_1G, MPT_PROT_R);
+	ASSERT(flags == MPT_UPDATE_L1);
 	ASSERT(g_fmpt.gran_1g);
 	ASSERT(g_fmpt.prot == MPT_PROT_R);
 	return 0;
@@ -132,14 +145,15 @@ static int __init test_set_fmpt__smpt_to_fmpt_partial(void)
 {
 	size_t start = 5 * SMPT_WORD_BYTE_RANGE / 2;
 	size_t end = 20 * SMPT_WORD_BYTE_RANGE;
+	enum mpt_update_flags flags;
 
 	/* Create SMPT with all PROT_W except a small subrange. */
 	init_fmpt(MPT_PROT_W, /*gran_1g*/ false);
 	__set_smpt_range(g_smpt, start, end, MPT_PROT_RW);
 
 	/* Fill the subrange with PROT_W to make the SMPT uniform. */
-	__set_fmpt_range(&g_fmpt, start, end, MPT_PROT_W);
-	ASSERT(g_fmpt.flags == MPT_UPDATE_L1);
+	flags = __set_fmpt_range(&g_fmpt, start, end, MPT_PROT_W);
+	ASSERT(flags == MPT_UPDATE_L1);
 	ASSERT(g_fmpt.gran_1g);
 	ASSERT(g_fmpt.prot == MPT_PROT_W);
 	return 0;
@@ -150,13 +164,14 @@ static int __init test_set_fmpt__smpt_to_smpt(void)
 {
 	size_t start = SZ_1G - SMPT_GRAN;
 	size_t end = SZ_1G;
+	enum mpt_update_flags flags;
 
 	init_fmpt(MPT_PROT_NONE, /*gran_1g*/ false);
 	ASSERT(__is_smpt_uniform(g_smpt, MPT_PROT_NONE));
 
 	/* Fill the subrange with PROT_W to make the SMPT uniform. */
-	__set_fmpt_range(&g_fmpt, start, end, MPT_PROT_RW);
-	ASSERT(g_fmpt.flags == MPT_UPDATE_L2);
+	flags = __set_fmpt_range(&g_fmpt, start, end, MPT_PROT_RW);
+	ASSERT(flags == MPT_UPDATE_L2);
 	ASSERT(!g_fmpt.gran_1g);
 	ASSERT(!__is_smpt_uniform(g_smpt, MPT_PROT_NONE));
 	return 0;