diff --git a/arch/arm/mach-imx/imx8m/soc.c b/arch/arm/mach-imx/imx8m/soc.c
index 5115471..dd393b5 100644
--- a/arch/arm/mach-imx/imx8m/soc.c
+++ b/arch/arm/mach-imx/imx8m/soc.c
@@ -112,16 +112,18 @@
 		/* DRAM1 */
 		.virt = 0x40000000UL,
 		.phys = 0x40000000UL,
-		.size = 0xC0000000UL,
+		.size = PHYS_SDRAM_SIZE,
 		.attrs = PTE_BLOCK_MEMTYPE(MT_NORMAL) |
 			 PTE_BLOCK_OUTER_SHARE
+#ifdef PHYS_SDRAM_2_SIZE
 	}, {
 		/* DRAM2 */
 		.virt = 0x100000000UL,
 		.phys = 0x100000000UL,
-		.size = 0x040000000UL,
+		.size = PHYS_SDRAM_2_SIZE,
 		.attrs = PTE_BLOCK_MEMTYPE(MT_NORMAL) |
 			 PTE_BLOCK_OUTER_SHARE
+#endif
 	}, {
 		/* List terminator */
 		0,
@@ -130,6 +132,20 @@
 
 struct mm_region *mem_map = imx8m_mem_map;
 
+void enable_caches(void)
+{
+	/*
+	 * If OPTEE runs, remove OPTEE memory from MMU table to
+	 * avoid speculative prefetch. OPTEE runs at the top of
+	 * the first memory bank
+	 */
+	if (rom_pointer[1])
+		imx8m_mem_map[5].size -= rom_pointer[1];
+
+	icache_enable();
+	dcache_enable();
+}
+
 static u32 get_cpu_variant_type(u32 type)
 {
 	struct ocotp_regs *ocotp = (struct ocotp_regs *)OCOTP_BASE_ADDR;