Merge changes from topic "rmm-planes"

* changes:
  fix(rme): RMM is not taking PSTATE into account on PN entry/exit
  feat(rme): update tests to alp14
  feat(rme): uplift RSI_SYSREG_READ/WRITE tests to alp13
diff --git a/include/lib/aarch64/arch.h b/include/lib/aarch64/arch.h
index e0e1a9d..e408afa 100644
--- a/include/lib/aarch64/arch.h
+++ b/include/lib/aarch64/arch.h
@@ -824,6 +824,14 @@
 #define CPTR_EL2_RESET_VAL	CPTR_EL2_RES1
 
 /* CPSR/SPSR definitions */
+#define SPSR_M3_0_SHIFT		U(0)
+#define SPSR_M3_0_MASK		U(0xf)
+#define SPSR_M3_0_EL0		U(0)
+#define SPSR_M3_0_EL1_SP_EL0	U(4)
+#define SPSR_M3_0_EL1_SP_EL1	U(5)
+#define SPSR_M3_0_EL2_SP_EL0	U(8)
+#define SPSR_M3_0_EL2_SP_EL2	U(9)
+
 #define DAIF_FIQ_BIT		(U(1) << 0)
 #define DAIF_IRQ_BIT		(U(1) << 1)
 #define DAIF_ABT_BIT		(U(1) << 2)
@@ -1768,69 +1776,73 @@
  * Definitions of system register identifiers
  *****************************************************************************/
 
-#define ESR_EL2_SYSREG_TRAP_OP0_SHIFT	20
-#define ESR_EL2_SYSREG_TRAP_OP0_WIDTH	U(2)
+#define SYSREG_ID_OP2_SHIFT		0
+#define SYSREG_ID_OP2_WIDTH		U(3)
 
-#define ESR_EL2_SYSREG_TRAP_OP2_SHIFT	17
-#define ESR_EL2_SYSREG_TRAP_OP2_WIDTH	U(3)
+#define SYSREG_ID_CRM_SHIFT		3
+#define SYSREG_ID_CRM_WIDTH		U(4)
 
-#define ESR_EL2_SYSREG_TRAP_OP1_SHIFT	14
-#define ESR_EL2_SYSREG_TRAP_OP1_WIDTH	U(3)
+#define SYSREG_ID_CRN_SHIFT		7
+#define SYSREG_ID_CRN_WIDTH		U(4)
 
-#define ESR_EL2_SYSREG_TRAP_CRN_SHIFT	10
-#define ESR_EL2_SYSREG_TRAP_CRN_WIDTH	U(4)
+#define SYSREG_ID_OP1_SHIFT		11
+#define SYSREG_ID_OP1_WIDTH		U(3)
 
-#define ESR_EL2_SYSREG_TRAP_CRM_SHIFT	1
-#define ESR_EL2_SYSREG_TRAP_CRM_WIDTH	U(4)
+#define SYSREG_ID_OP0_SHIFT		14
+#define SYSREG_ID_OP0_WIDTH		U(2)
 
-#define SYSREG_ESR(op0, op1, crn, crm, op2) \
-		((UL(op0) << ESR_EL2_SYSREG_TRAP_OP0_SHIFT) | \
-		 (UL(op1) << ESR_EL2_SYSREG_TRAP_OP1_SHIFT) | \
-		 (UL(crn) << ESR_EL2_SYSREG_TRAP_CRN_SHIFT) | \
-		 (UL(crm) << ESR_EL2_SYSREG_TRAP_CRM_SHIFT) | \
-		 (UL(op2) << ESR_EL2_SYSREG_TRAP_OP2_SHIFT))
+#define SYSREG_ID_D128_SHIFT		16
+#define SYSREG_ID_D128_WIDTH		U(1)
 
-#define SYSREG_ID_sp_el0			SYSREG_ESR(3, 0, 4, 1, 0)
-#define SYSREG_ID_sp_el1			SYSREG_ESR(3, 4, 4, 1, 0)
-#define SYSREG_ID_elr_el1			SYSREG_ESR(3, 0, 4, 0, 1)
-#define SYSREG_ID_spsr_el1			SYSREG_ESR(3, 0, 4, 0, 0)
-#define SYSREG_ID_pmcr_el0			SYSREG_ESR(3, 3, 9, 12, 0)
-#define SYSREG_ID_tpidrro_el0			SYSREG_ESR(3, 3, 13, 0, 3)
-#define SYSREG_ID_tpidr_el0			SYSREG_ESR(3, 3, 13, 0, 2)
-#define SYSREG_ID_csselr_el1			SYSREG_ESR(3, 2, 0, 0, 0)
-#define SYSREG_ID_sctlr_el1			SYSREG_ESR(3, 0, 1, 0, 0)
-#define SYSREG_ID_actlr_el1			SYSREG_ESR(3, 0, 1, 0, 1)
-#define SYSREG_ID_cpacr_el1			SYSREG_ESR(3, 0, 1, 0, 2)
-#define SYSREG_ID_zcr_el1			SYSREG_ESR(3, 0, 1, 2, 0)
-#define SYSREG_ID_ttbr0_el1			SYSREG_ESR(3, 0, 2, 0, 0)
-#define SYSREG_ID_ttbr1_el1			SYSREG_ESR(3, 0, 2, 0, 1)
-#define SYSREG_ID_tcr_el1			SYSREG_ESR(3, 0, 2, 0, 2)
-#define SYSREG_ID_esr_el1			SYSREG_ESR(3, 0, 5, 2, 0)
-#define SYSREG_ID_afsr0_el1			SYSREG_ESR(3, 0, 5, 1, 0)
-#define SYSREG_ID_afsr1_el1			SYSREG_ESR(3, 0, 5, 1, 1)
-#define SYSREG_ID_far_el1			SYSREG_ESR(3, 0, 6, 0, 0)
-#define SYSREG_ID_mair_el1			SYSREG_ESR(3, 0, 10, 2, 0)
-#define SYSREG_ID_vbar_el1			SYSREG_ESR(3, 0, 12, 0, 0)
-#define SYSREG_ID_contextidr_el1		SYSREG_ESR(3, 0, 13, 0, 1)
-#define SYSREG_ID_tpidr_el1			SYSREG_ESR(3, 0, 13, 0, 4)
-#define SYSREG_ID_amair_el1			SYSREG_ESR(3, 0, 10, 3, 0)
-#define SYSREG_ID_cntkctl_el1			SYSREG_ESR(3, 0, 14, 1, 0)
-#define SYSREG_ID_par_el1			SYSREG_ESR(3, 0, 7, 4, 0)
-#define SYSREG_ID_mdscr_el1			SYSREG_ESR(2, 0, 0, 2, 2)
-#define SYSREG_ID_mdccint_el1			SYSREG_ESR(2, 0, 0, 2, 0)
-#define SYSREG_ID_disr_el1			SYSREG_ESR(3, 0, 12, 1, 1)
-#define SYSREG_ID_mpam0_el1			SYSREG_ESR(3, 0, 10, 5, 1)
-#define SYSREG_ID_apiakeylo_el1			SYSREG_ESR(3, 0, 2,  1, 0)
-#define SYSREG_ID_apiakeyhi_el1			SYSREG_ESR(3, 0, 2,  1, 1)
-#define SYSREG_ID_apibkeylo_el1			SYSREG_ESR(3, 0, 2,  1, 2)
-#define SYSREG_ID_apibkeyhi_el1			SYSREG_ESR(3, 0, 2,  1, 3)
-#define SYSREG_ID_apdakeylo_el1			SYSREG_ESR(3, 0, 2,  2, 0)
-#define SYSREG_ID_apdakeyhi_el1			SYSREG_ESR(3, 0, 2,  2, 1)
-#define SYSREG_ID_apdbkeylo_el1			SYSREG_ESR(3, 0, 2,  2, 2)
-#define SYSREG_ID_apdbkeyhi_el1			SYSREG_ESR(3, 0, 2,  2, 3)
-#define SYSREG_ID_apgakeylo_el1			SYSREG_ESR(3, 0, 2,  3, 0)
-#define SYSREG_ID_apgakeyhi_el1			SYSREG_ESR(3, 0, 2,  3, 1)
-#define SYSREG_ID_mpamidr_el1			SYSREG_ESR(3, 0, 10, 4, 4)
+#define SYSREG_ID(d128, op0, op1, crn, crm, op2) \
+		((UL(op0) << SYSREG_ID_OP0_SHIFT) | \
+		 (UL(op1) << SYSREG_ID_OP1_SHIFT) | \
+		 (UL(crn) << SYSREG_ID_CRN_SHIFT) | \
+		 (UL(crm) << SYSREG_ID_CRM_SHIFT) | \
+		 (UL(op2) << SYSREG_ID_OP2_SHIFT) | \
+		 (UL(d128) << SYSREG_ID_D128_SHIFT))
+
+#define SYSREG_ID_sp_el0			SYSREG_ID(0, 3, 0, 4, 1, 0)
+#define SYSREG_ID_sp_el1			SYSREG_ID(0, 3, 4, 4, 1, 0)
+#define SYSREG_ID_elr_el1			SYSREG_ID(0, 3, 0, 4, 0, 1)
+#define SYSREG_ID_spsr_el1			SYSREG_ID(0, 3, 0, 4, 0, 0)
+#define SYSREG_ID_pmcr_el0			SYSREG_ID(0, 3, 3, 9, 12, 0)
+#define SYSREG_ID_tpidrro_el0			SYSREG_ID(0, 3, 3, 13, 0, 3)
+#define SYSREG_ID_tpidr_el0			SYSREG_ID(0, 3, 3, 13, 0, 2)
+#define SYSREG_ID_csselr_el1			SYSREG_ID(0, 3, 2, 0, 0, 0)
+#define SYSREG_ID_sctlr_el1			SYSREG_ID(0, 3, 0, 1, 0, 0)
+#define SYSREG_ID_actlr_el1			SYSREG_ID(0, 3, 0, 1, 0, 1)
+#define SYSREG_ID_cpacr_el1			SYSREG_ID(0, 3, 0, 1, 0, 2)
+#define SYSREG_ID_zcr_el1			SYSREG_ID(0, 3, 0, 1, 2, 0)
+#define SYSREG_ID_ttbr0_el1			SYSREG_ID(1, 3, 0, 2, 0, 0)
+#define SYSREG_ID_ttbr1_el1			SYSREG_ID(1, 3, 0, 2, 0, 1)
+#define SYSREG_ID_tcr_el1			SYSREG_ID(0, 3, 0, 2, 0, 2)
+#define SYSREG_ID_esr_el1			SYSREG_ID(0, 3, 0, 5, 2, 0)
+#define SYSREG_ID_afsr0_el1			SYSREG_ID(0, 3, 0, 5, 1, 0)
+#define SYSREG_ID_afsr1_el1			SYSREG_ID(0, 3, 0, 5, 1, 1)
+#define SYSREG_ID_far_el1			SYSREG_ID(0, 3, 0, 6, 0, 0)
+#define SYSREG_ID_mair_el1			SYSREG_ID(0, 3, 0, 10, 2, 0)
+#define SYSREG_ID_vbar_el1			SYSREG_ID(0, 3, 0, 12, 0, 0)
+#define SYSREG_ID_contextidr_el1		SYSREG_ID(0, 3, 0, 13, 0, 1)
+#define SYSREG_ID_tpidr_el1			SYSREG_ID(0, 3, 0, 13, 0, 4)
+#define SYSREG_ID_amair_el1			SYSREG_ID(0, 3, 0, 10, 3, 0)
+#define SYSREG_ID_cntkctl_el1			SYSREG_ID(0, 3, 0, 14, 1, 0)
+#define SYSREG_ID_par_el1			SYSREG_ID(1, 3, 0, 7, 4, 0)
+#define SYSREG_ID_mdscr_el1			SYSREG_ID(0, 2, 0, 0, 2, 2)
+#define SYSREG_ID_mdccint_el1			SYSREG_ID(0, 2, 0, 0, 2, 0)
+#define SYSREG_ID_disr_el1			SYSREG_ID(0, 3, 0, 12, 1, 1)
+#define SYSREG_ID_mpam0_el1			SYSREG_ID(0, 3, 0, 10, 5, 1)
+#define SYSREG_ID_apiakeylo_el1			SYSREG_ID(0, 3, 0, 2,  1, 0)
+#define SYSREG_ID_apiakeyhi_el1			SYSREG_ID(0, 3, 0, 2,  1, 1)
+#define SYSREG_ID_apibkeylo_el1			SYSREG_ID(0, 3, 0, 2,  1, 2)
+#define SYSREG_ID_apibkeyhi_el1			SYSREG_ID(0, 3, 0, 2,  1, 3)
+#define SYSREG_ID_apdakeylo_el1			SYSREG_ID(0, 3, 0, 2,  2, 0)
+#define SYSREG_ID_apdakeyhi_el1			SYSREG_ID(0, 3, 0, 2,  2, 1)
+#define SYSREG_ID_apdbkeylo_el1			SYSREG_ID(0, 3, 0, 2,  2, 2)
+#define SYSREG_ID_apdbkeyhi_el1			SYSREG_ID(0, 3, 0, 2,  2, 3)
+#define SYSREG_ID_apgakeylo_el1			SYSREG_ID(0, 3, 0, 2,  3, 0)
+#define SYSREG_ID_apgakeyhi_el1			SYSREG_ID(0, 3, 0, 2,  3, 1)
+#define SYSREG_ID_mpamidr_el1			SYSREG_ID(0, 3, 0, 10, 4, 4)
 
 /* RNDR definition */
 #define RNDR			S3_3_C2_C4_0
diff --git a/include/lib/utils_def.h b/include/lib/utils_def.h
index 0128277..c9bd527 100644
--- a/include/lib/utils_def.h
+++ b/include/lib/utils_def.h
@@ -172,6 +172,15 @@
 	(((reg) & MASK(regfield)) >> (regfield##_SHIFT))
 
 /*
+ * Generates an unsigned long long (64-bit) value where the bits @_msb
+ * through @_lsb (inclusive) are set to one and all other bits are zero.  The
+ * parameters can hold values from 0 through 63 and if _msb == _lsb a single
+ * bit is set at that location.
+ */
+#define BIT_MASK_ULL(_msb, _lsb) \
+	((~ULL(0) >> (63UL - (_msb))) & (~ULL(0) << (_lsb)))
+
+/*
  * Defines member of structure and reserves space
  * for the next member with specified offset.
  */
diff --git a/include/runtime_services/host_realm_managment/realm_def.h b/include/runtime_services/host_realm_managment/realm_def.h
index 6fca05c..5f662e2 100644
--- a/include/runtime_services/host_realm_managment/realm_def.h
+++ b/include/runtime_services/host_realm_managment/realm_def.h
@@ -29,11 +29,33 @@
 #if (PAGE_SIZE == PAGE_SIZE_4KB)
 #define PAGE_ALIGNMENT			PAGE_SIZE_4KB
 #define TCR_TG0				TCR_TG0_4K
+#define GRANULE_SHIFT			U(12)
 #else
 #error "Undefined value for PAGE_SIZE"
 #endif
 
 /*
+ * TTE_STRIDE: The number of bits resolved in a single level of translation
+ * walk (except for the starting level which may resolve more or fewer bits)
+ */
+#define MAX_IPA_BITS			U(48)
+#define TTE_STRIDE			(GRANULE_SHIFT - 3U)
+
+/*
+ * TTE_STRIDE_LM1: The number of bits resolved al level -1 when FEAT_LPA2
+ * is enabled. This value is equal to
+ * MAX_IPA_BITS_LPA2 - ((4 * S2TTE_STRIDE) + GRANULE_SHIFT)
+ * as Level -1 only has 4 bits for the index (bits [51:48])
+ */
+#define MAX_IPA_BITS_LPA2		U(52)
+#define TTE_STRIDE_LM1			U(4)
+
+#define TT_PAGE_LEVEL			U(3)
+#define tte_map_size(level)						\
+	(1ULL << (unsigned int)(((TT_PAGE_LEVEL - (level)) *		\
+				 (int)TTE_STRIDE) + (int)GRANULE_SHIFT))
+
+/*
  * 'MPIDR_EL1_AFF<n>_VAL_SHIFT' constants specify the left shift
  * for affinity <n> value that gives the field's actual value.
  */
diff --git a/realm/include/realm_rsi.h b/realm/include/realm_rsi.h
index a6260e4..38793e9 100644
--- a/realm/include/realm_rsi.h
+++ b/realm/include/realm_rsi.h
@@ -359,12 +359,12 @@
 /*
  * FID: 0xC40001AE
  */
-#define SMC_RSI_PLANE_REG_READ		SMC64_RSI_FID(U(0x1E))
+#define SMC_RSI_PLANE_SYSREG_READ	SMC64_RSI_FID(U(0x1E))
 
 /*
  * FID: 0xC40001AF
  */
-#define SMC_RSI_PLANE_REG_WRITE		SMC64_RSI_FID(U(0x1F))
+#define SMC_RSI_PLANE_SYSREG_WRITE	SMC64_RSI_FID(U(0x1F))
 
 typedef enum {
 	RSI_EMPTY = 0U,
@@ -560,9 +560,11 @@
 /* Data structure used to pass values from P0 to the RMM on Plane entry */
 struct rsi_plane_entry {
 	/* Flags */
-	SET_MEMBER(u_register_t flags, 0, 0x8);	/* Offset 0 */
+	SET_MEMBER(u_register_t flags, 0, 0x8);		/* Offset 0 */
 	/* PC */
-	SET_MEMBER(u_register_t pc, 0x8, 0x100);	/* Offset 0x8 */
+	SET_MEMBER(u_register_t pc, 0x8, 0x10);		/* Offset 0x8 */
+	/* PSTATE */
+	SET_MEMBER(u_register_t pstate, 0x10, 0x100);	/* Offset 0x10 */
 	/* General-purpose registers */
 	SET_MEMBER(u_register_t gprs[RSI_PLANE_NR_GPRS], 0x100, 0x200);	/* 0x100 */
 	/* EL1 system registers */
@@ -584,9 +586,11 @@
 		/* Exception Syndrome Register */
 		u_register_t esr;				/* 0x108 */
 		/* Fault Address Register */
-		u_register_t far;				/* 0x108 */
+		u_register_t far;				/* 0x110 */
 		/* Hypervisor IPA Fault Address register */
-		u_register_t hpfar;				/* 0x110 */
+		u_register_t hpfar;				/* 0x118 */
+		/* Hypervisor PSTATE */
+		u_register_t pstate;				/* 0x120 */
 	}, 0x100, 0x200);
 	/* General-purpose registers */
 	SET_MEMBER(u_register_t gprs[RSI_PLANE_NR_GPRS], 0x200, 0x300); /* 0x200 */
@@ -609,37 +613,11 @@
 	SET_MEMBER(struct rsi_plane_exit exit, 0x800, 0x1000);/* 0x800 */
 } rsi_plane_run;
 
-/*
- * arg1 == plane index
- * arg2 == run pointer
- *
- * ret0 == status
- */
-#define RSI_PLANE_ENTER		SMC_RSI_FID(0x13U)
+u_register_t rsi_plane_sysreg_read(u_register_t plane_index, u_register_t register_encoding,
+		u_register_t *value_lo, u_register_t *value_hi);
 
-/*
- * arg1 == plane index
- * arg2 == register encoding
- *
- * ret0 == status
- * ret1 = register value
- */
-#define RSI_PLANE_REG_READ	SMC_RSI_FID(0x1EU)
-
-u_register_t rsi_plane_reg_read(u_register_t plane_index, u_register_t register_encoding,
-		u_register_t *value);
-
-u_register_t rsi_plane_reg_write(u_register_t plane_index, u_register_t register_encoding,
-		u_register_t value);
-
-/*
- * arg1 == plane index
- * arg2 == register encoding
- * arg3 == register value
- *
- * ret0 == status
- */
-#define RSI_PLANE_REG_WRITE	SMC_RSI_FID(0x1FU)
+u_register_t rsi_plane_sysreg_write(u_register_t plane_index, u_register_t register_encoding,
+		u_register_t value_lo, u_register_t value_hi);
 
 /*
  * Function to set overlay permission value for a specified
diff --git a/realm/realm_multiple_rec.c b/realm/realm_multiple_rec.c
index 95bb6f1..2b364c5 100644
--- a/realm/realm_multiple_rec.c
+++ b/realm/realm_multiple_rec.c
@@ -48,6 +48,11 @@
 	/* enter plane */
 	u_register_t flags = 0U;
 
+	/* Setup the initial PSTATE for the plane */
+	run[rec].enter.pstate = ((SPSR_M3_0_EL1_SP_EL1 << SPSR_M3_0_SHIFT) |
+				 (SPSR_M_AARCH64 << SPSR_M_SHIFT) |
+				 (0xf << SPSR_DAIF_SHIFT));
+
 	/* Use Base adr, plane_index, perm_index programmed by P0 rec0 */
 	run[rec].enter.pc = base;
 	realm_printf("Entering plane %ld, ep=0x%lx rec=0x%lx\n", plane_index, base, rec);
diff --git a/realm/realm_payload_main.c b/realm/realm_payload_main.c
index 50b78c8..c891746 100644
--- a/realm/realm_payload_main.c
+++ b/realm/realm_payload_main.c
@@ -92,41 +92,85 @@
 			realm_printf("P0 read 0x%lx 0x%lx\n", reg1, reg2);
 
 			/* read pauth register for plane1 */
-			ret = rsi_plane_reg_read(plane_index, SYSREG_ID_apiakeylo_el1, &reg3);
+			ret = rsi_plane_sysreg_read(plane_index,
+						    SYSREG_ID_apiakeylo_el1,
+						    &reg3, NULL);
 			if ((ret != RSI_SUCCESS) || (reg1 != reg3)) {
-				realm_printf("pauth register mismatch 0x%lx 0x%lx\n", reg1, reg3);
+				realm_printf("pauth register mismatch 0x%lx 0x%lx\n",
+					     reg1, reg3);
 				return false;
 			}
 
 			/* read sctlr register for plane1 */
-			ret = rsi_plane_reg_read(plane_index, SYSREG_ID_sctlr_el1, &reg4);
+			ret = rsi_plane_sysreg_read(plane_index,
+						    SYSREG_ID_sctlr_el1, &reg4, NULL);
 			if ((ret != RSI_SUCCESS) || (reg2 != reg4)) {
-				realm_printf("sctlr register mismatch 0x%lx 0x%lx\n", reg2, reg4);
+				realm_printf("sctlr register mismatch 0x%lx 0x%lx\n",
+					     reg2, reg4);
 				return false;
 			}
 
 			/* write pauth register and verify it is same after exiting plane n */
-			ret = rsi_plane_reg_write(plane_index, SYSREG_ID_apibkeylo_el1, 0xABCD);
+			ret = rsi_plane_sysreg_write(plane_index,
+						     SYSREG_ID_apibkeylo_el1,
+						     0xABCD, 0UL);
 			if (ret != RSI_SUCCESS) {
 				realm_printf("pauth register write failed\n");
 				return false;
 			}
 
+			/*
+			 * write ttbr0_el1 register and verify it is same after
+			 * exiting plane n to validate access to 128-bit registers.
+			 */
+			ret = rsi_plane_sysreg_write(plane_index,
+						     SYSREG_ID_ttbr0_el1,
+						     0x12340000, 0xA0000);
+			if (ret != RSI_SUCCESS) {
+				realm_printf("ttbr0_el1 register write failed\n");
+				return false;
+			}
+
 			/* enter plane n */
 			ret = realm_plane_enter(plane_index, perm_index, base, flags, &run);
 			if (ret) {
 				/* read pauth register for plane1 */
-				ret = rsi_plane_reg_read(plane_index, SYSREG_ID_apibkeylo_el1,
-						&reg3);
+				ret = rsi_plane_sysreg_read(plane_index, SYSREG_ID_apibkeylo_el1,
+						&reg3, NULL);
 
 				if ((ret != RSI_SUCCESS) || (reg3 != 0xABCD)) {
 					realm_printf("reg mismatch after write 0x%lx\n", reg3);
 					return false;
 				}
+
+				/*
+				 * read ttbr0_el1 register for plane1 to validate
+				 * access to 128-bit registers.
+				 */
+				ret = rsi_plane_sysreg_read(plane_index, SYSREG_ID_ttbr0_el1,
+						&reg3, &reg4);
+
+				if (!is_feat_d128_supported()) {
+					/*
+					 * As FEAT_D128 is not supported, manually write
+					 * the expected value into reg4.
+					 */
+					reg4 = 0xA0000;
+				}
+
+				if ((ret != RSI_SUCCESS) ||
+				    (reg3 != 0x12340000) ||
+				    (reg4 != 0xA0000)) {
+					realm_printf("reg mismatch after write 0x%lx - 0x%lx\n",
+						     reg3, reg4);
+					return false;
+				}
 			}
 
 			/* read sysreg not supported by rmm, expect error */
-			ret = rsi_plane_reg_read(plane_index, SYSREG_ID_mpamidr_el1, &reg3);
+			ret = rsi_plane_sysreg_read(plane_index,
+						    SYSREG_ID_mpamidr_el1,
+						    &reg3, NULL);
 			if (ret == RSI_SUCCESS) {
 				realm_printf("reg read should have failed\n");
 				return false;
diff --git a/realm/realm_plane.c b/realm/realm_plane.c
index ff42ba9..dba6cff 100644
--- a/realm/realm_plane.c
+++ b/realm/realm_plane.c
@@ -172,6 +172,11 @@
 	memset(run, 0, sizeof(rsi_plane_run));
 	run->enter.pc = base;
 
+	/* Setup the initial PSTATE for the plane */
+	run->enter.pstate = ((SPSR_M3_0_EL1_SP_EL1 << SPSR_M3_0_SHIFT) |
+			     (SPSR_M_AARCH64 << SPSR_M_SHIFT) |
+			     (0xf << SPSR_DAIF_SHIFT));
+
 	/* Perm init */
 	if (plane_init[plane_index]) {
 		return true;
@@ -198,6 +203,10 @@
 
 	while (true) {
 		ret = rsi_plane_enter(plane_index, (u_register_t)run);
+
+		/* Restore PSTATE with the value on run->exit for the next entry */
+		run->enter.pstate = run->exit.pstate;
+
 		if (ret != RSI_SUCCESS) {
 			ERROR("Plane %u enter failed ret= 0x%lx\n", plane_index, ret);
 			return false;
diff --git a/realm/realm_rsi.c b/realm/realm_rsi.c
index 163443e..a22ddc2 100644
--- a/realm/realm_rsi.c
+++ b/realm/realm_rsi.c
@@ -218,30 +218,44 @@
 	return res.ret0;
 }
 
-u_register_t rsi_plane_reg_read(u_register_t plane_index,
+/*
+ * Call RSI_PLANE_SYSREG_READ.
+ *
+ * Note that if @value_hi is not NULL and FEAT_D128 is disabled, its return
+ * value will be undefined.
+ */
+u_register_t rsi_plane_sysreg_read(u_register_t plane_index,
 				u_register_t register_encoding,
-				u_register_t *value)
+				u_register_t *value_lo, u_register_t *value_hi)
 {
 	smc_ret_values res = {};
 
 	res = tftf_smc(&(smc_args)
-			{SMC_RSI_PLANE_REG_READ, plane_index,
+			{SMC_RSI_PLANE_SYSREG_READ, plane_index,
 			 register_encoding});
 	if (res.ret0 == RSI_SUCCESS) {
-		*value = res.ret1;
+		*value_lo = res.ret1;
+		if (value_hi != NULL) {
+			*value_hi = res.ret2;
+		}
 	}
 	return res.ret0;
 }
 
-u_register_t rsi_plane_reg_write(u_register_t plane_index,
+/*
+ * Call RSI_PLANE_SYSREG_WRITE.
+ *
+ * Note that @value_hi will be ignored if FEAT_D128 is disabled.
+ */
+u_register_t rsi_plane_sysreg_write(u_register_t plane_index,
 				 u_register_t register_encoding,
-				 u_register_t value)
+				 u_register_t value_lo, u_register_t value_hi)
 {
 	smc_ret_values res = {};
 
 	res = tftf_smc(&(smc_args)
-			{SMC_RSI_PLANE_REG_WRITE, plane_index,
-			 register_encoding, value});
+			{SMC_RSI_PLANE_SYSREG_WRITE, plane_index,
+			 register_encoding, value_lo, value_hi});
 	return res.ret0;
 }
 
diff --git a/tftf/tests/runtime_services/host_realm_managment/host_realm_rmi.c b/tftf/tests/runtime_services/host_realm_managment/host_realm_rmi.c
index 6fb4405..21e6c59 100644
--- a/tftf/tests/runtime_services/host_realm_managment/host_realm_rmi.c
+++ b/tftf/tests/runtime_services/host_realm_managment/host_realm_rmi.c
@@ -40,6 +40,21 @@
 static spinlock_t pool_lock;
 static unsigned int pool_counter;
 
+/*
+ * Return an IPA mask @level
+ */
+static unsigned long tte_ipa_lvl_mask(long level, bool lpa2)
+{
+	unsigned long mask;
+	unsigned long levels = (unsigned long)(TT_PAGE_LEVEL - level);
+	unsigned long lsb = (levels * TTE_STRIDE) + GRANULE_SHIFT;
+	unsigned long oa_bits = (lpa2 == true) ? MAX_IPA_BITS_LPA2 : MAX_IPA_BITS;
+
+	mask = BIT_MASK_ULL((oa_bits - 1U), lsb);
+
+	return mask;
+}
+
 static smc_ret_values host_rmi_handler(smc_args *args, unsigned int in_reg)
 {
 	u_register_t regs[8];
@@ -325,18 +340,12 @@
 
 u_register_t host_rmi_rtt_aux_map_unprotected(u_register_t rd,
 					      u_register_t map_addr,
-					      u_register_t tree_index,
-					      u_register_t *fail_index,
-					      u_register_t *level_pri,
-					      u_register_t *state)
+					      u_register_t tree_index)
 {
 	smc_ret_values rets;
 
 	rets = host_rmi_handler(&(smc_args) {SMC_RMI_RTT_AUX_MAP_UNPROTECTED,
 				rd, map_addr, tree_index}, 4U);
-	*fail_index = rets.ret1;
-	*level_pri = rets.ret2;
-	*state = rets.ret3;
 	return rets.ret0;
 }
 
@@ -391,15 +400,13 @@
 u_register_t host_rmi_rtt_aux_unmap_unprotected(u_register_t rd,
 					  u_register_t map_addr,
 					  u_register_t tree_index,
-					  u_register_t *top,
-					  u_register_t *level)
+					  u_register_t *top)
 {
 	smc_ret_values rets;
 
 	rets = host_rmi_handler(&(smc_args) {SMC_RMI_RTT_AUX_UNMAP_UNPROTECTED,
 					rd, map_addr, tree_index}, 4U);
 	*top = rets.ret1;
-	*level = rets.ret2;
 	return rets.ret0;
 }
 
@@ -612,6 +619,7 @@
 		}
 		map_addr += PAGE_SIZE;
 	}
+
 	return REALM_SUCCESS;
 }
 
@@ -887,6 +895,7 @@
 					ipa, ret, tree_index);
 				}
 			}
+
 			/* Retry DATA_DESTROY */
 			continue;
 		}
@@ -913,6 +922,31 @@
 	return REALM_SUCCESS;
 }
 
+
+static u_register_t host_realm_aux_unmap_unprotected(struct realm *realm,
+						     u_register_t unmap_addr)
+{
+	u_register_t ret = RMI_SUCCESS;
+	__unused u_register_t top;
+
+	for (unsigned int tree_index = 1U;
+	     tree_index <= realm->num_aux_planes; tree_index++) {
+		bool lpa2 = (realm->rmm_feat_reg0 & RMI_FEATURE_REGISTER_0_LPA2);
+		unsigned long unmap_ns_addr = unmap_addr &
+						tte_ipa_lvl_mask(realm->start_level, lpa2);
+
+		ret = host_rmi_rtt_aux_unmap_unprotected(realm->rd,
+							 unmap_ns_addr,
+							 tree_index, &top);
+
+		if (ret != RMI_SUCCESS) {
+			return ret;
+		}
+	}
+
+	return ret;
+}
+
 static u_register_t host_realm_tear_down_rtt_range(struct realm *realm,
 						   long level,
 						   u_register_t start,
@@ -939,26 +973,17 @@
 		switch (rtt.state) {
 		case RMI_ASSIGNED:
 			if (host_ipa_is_ns(map_addr, realm->rmm_feat_reg0)) {
-
-				u_register_t level1;
-
 				/* Unmap from all Aux RTT */
 				if (!realm->rtt_s2ap_enc_indirect) {
-					for (unsigned int tree_index = 1U;
-						tree_index <= realm->num_aux_planes;
-						tree_index++) {
 
-						ret = host_rmi_rtt_aux_unmap_unprotected(
-								rd,
-								map_addr,
-								tree_index,
-								&top, &level1);
+					u_register_t ret =
+						host_realm_aux_unmap_unprotected(
+							realm, map_addr);
 
-						if (ret != RMI_SUCCESS) {
-							ERROR("%s() failed, addr=0x%lx ret=0x%lx\n",
+					if (ret != RMI_SUCCESS) {
+						ERROR("%s() failed, addr=0x%lx ret=0x%lx\n",
 							"host_rmi_rtt_unmap_unprotected",
 							map_addr, ret);
-						}
 					}
 				}
 
@@ -1398,47 +1423,31 @@
 	/* AUX MAP NS buffer for all RTTs */
 	if (!realm->rtt_tree_single) {
 		for (unsigned int j = 0U; j < realm->num_aux_planes; j++) {
-			for (unsigned int i = 0U; i < ns_shared_mem_size / PAGE_SIZE; i++) {
-				u_register_t fail_index, level_pri, state;
-
+			bool lpa2 = (realm->rmm_feat_reg0 & RMI_FEATURE_REGISTER_0_LPA2);
+			unsigned long sl_map_size = tte_map_size(realm->start_level);
+			unsigned long map_ns_addr = realm->ipa_ns_buffer &
+						tte_ipa_lvl_mask(realm->start_level, lpa2);
+			do {
 				ret = host_rmi_rtt_aux_map_unprotected(realm->rd,
-						realm->ipa_ns_buffer + (i * PAGE_SIZE),
-						j + 1, &fail_index, &level_pri, &state);
+								       map_ns_addr, j + 1U);
 
-				if (ret == RMI_SUCCESS) {
-					continue;
-				} else if (RMI_RETURN_STATUS(ret) == RMI_ERROR_RTT_AUX) {
-
-					/* Create Level and retry */
+				if (RMI_RETURN_STATUS(ret) == RMI_ERROR_RTT) {
 					int8_t ulevel = RMI_RETURN_INDEX(ret);
 					long level = (long)ulevel;
 
-					ret = host_realm_create_rtt_aux_levels(realm,
-							realm->ipa_ns_buffer +
-							(i * PAGE_SIZE),
-							level, 3, j + 1);
-
-					if (ret != RMI_SUCCESS) {
-						return REALM_ERROR;
-					}
-
-					ret = host_rmi_rtt_aux_map_unprotected(realm->rd,
-							realm->ipa_ns_buffer +
-							(i * PAGE_SIZE),
-							j + 1, &fail_index,
-							&level_pri, &state);
-
-					if (ret == RMI_SUCCESS) {
-						continue;
-					}
+					ERROR("%s() failed, walk_pri.level = %ld\n",
+					      "host_realm_map_unprotected", level);
+					return REALM_ERROR;
 				}
 
-				ERROR("%s() failed, par_base=0x%lx ret=0x%lx\n",
-						"host_realm_aux_map_unprotected",
-						(ns_shared_mem_adr), ret);
-				return REALM_ERROR;
+				if (ret != RMI_SUCCESS) {
+					ERROR("%s() failed\n", "host_realm_map_unprotected");
+					return REALM_ERROR;
+				}
 
-			}
+				map_ns_addr += sl_map_size;
+			} while (map_ns_addr < (realm->ipa_ns_buffer +
+						realm->ns_buffer_size));
 		}
 	}
 
diff --git a/tftf/tests/runtime_services/realm_payload/host_realm_payload_tests.c b/tftf/tests/runtime_services/realm_payload/host_realm_payload_tests.c
index 34ead9f..be999c3 100644
--- a/tftf/tests/runtime_services/realm_payload/host_realm_payload_tests.c
+++ b/tftf/tests/runtime_services/realm_payload/host_realm_payload_tests.c
@@ -40,7 +40,7 @@
 #endif
 
 /*
- * @Test_Aim@ Test RSI_PLANE_REG_READ/WRITE
+ * @Test_Aim@ Test RSI_PLANE_SYSREG_READ/WRITE
  */
 test_result_t host_test_realm_create_planes_register_rw(void)
 {