2 files changed, 43 insertions, 0 deletions
diff --git a/src/arch/arm64/armv8/mmu.c b/src/arch/arm64/armv8/mmu.c
index a24e7c6fdd..48f77ace83 100644
--- a/src/arch/arm64/armv8/mmu.c
+++ b/src/arch/arm64/armv8/mmu.c
@@ -274,6 +274,40 @@ void mmu_init(void)
 			  TCR_TBI_USED);
 }
 
+/* Func : mmu_save_context
+ * Desc : Save mmu context (registers and ttbr base).
+ */
+void mmu_save_context(struct mmu_context *mmu_context)
+{
+	assert(mmu_context);
+
+	/* Back-up MAIR_ATTRIBUTES */
+	mmu_context->mair = raw_read_mair_el3();
+
+	/* Back-up TCR value */
+	mmu_context->tcr = raw_read_tcr_el3();
+}
+
+/* Func : mmu_restore_context
+ * Desc : Restore mmu context using input backed-up context
+ */
+void mmu_restore_context(const struct mmu_context *mmu_context)
+{
+	assert(mmu_context);
+
+	/* Restore TTBR */
+	raw_write_ttbr0_el3((uintptr_t)_ttb);
+
+	/* Restore MAIR indices */
+	raw_write_mair_el3(mmu_context->mair);
+
+	/* Restore TCR flags */
+	raw_write_tcr_el3(mmu_context->tcr);
+
+	/* invalidate tlb since ttbr is updated. */
+	tlb_invalidate_all();
+}
+
 void mmu_enable(void)
 {
 	if (((get_pte(_ttb) >> BLOCK_INDEX_SHIFT) & BLOCK_INDEX_MASK)
diff --git a/src/arch/arm64/include/armv8/arch/mmu.h b/src/arch/arm64/include/armv8/arch/mmu.h
index f0e551e52d..6c83749e95 100644
--- a/src/arch/arm64/include/armv8/arch/mmu.h
+++ b/src/arch/arm64/include/armv8/arch/mmu.h
@@ -148,8 +148,17 @@
 #define TCR_TBI_USED               (0x0 << TCR_TBI_SHIFT)
 #define TCR_TBI_IGNORED            (0x1 << TCR_TBI_SHIFT)
 
+struct mmu_context {
+	uint64_t mair;         /* MAIR attributes */
+	uint64_t tcr;          /* TCR Attributes */
+};
+
 /* Initialize MMU registers and page table memory region. */
 void mmu_init(void);
+/* Desc : Save mmu context (registers and ttbr base/size). */
+void mmu_save_context(struct mmu_context *mmu_context);
+/* Desc : Restore mmu context using input backed-up context */
+void mmu_restore_context(const struct mmu_context *mmu_context);
 /* Change a memory type for a range of bytes at runtime. */
 void mmu_config_range(void *start, size_t size, uint64_t tag);
 /* Enable the MMU (need previous mmu_init() and configured ranges!). */