[U-Boot] arm64: zynqmp: Setup MMU map for DDR at run time

diff --git a/arch/arm/cpu/armv8/zynqmp/cpu.c b/arch/arm/cpu/armv8/zynqmp/cpu.c
index 14e7d4006494..5db5fafc6f65 100644
--- a/arch/arm/cpu/armv8/zynqmp/cpu.c
+++ b/arch/arm/cpu/armv8/zynqmp/cpu.c
@@ -16,16 +16,28 @@ 
 
 DECLARE_GLOBAL_DATA_PTR;
 
-static struct mm_region zynqmp_mem_map[] = {
+/*
+ * Number of filled static entries and also the first empty
+ * slot in zynqmp_mem_map.
+ */
+#define ZYNQMP_MEM_MAP_USED	4
+
 #if !defined(CONFIG_ZYNQMP_NO_DDR)
-	{
-		.virt = 0x0UL,
-		.phys = 0x0UL,
-		.size = 0x80000000UL,
-		.attrs = PTE_BLOCK_MEMTYPE(MT_NORMAL) |
-			 PTE_BLOCK_INNER_SHARE
-	},
+#define DRAM_BANKS CONFIG_NR_DRAM_BANKS
+#else
+#define DRAM_BANKS 0
 #endif
+
+#if defined(CONFIG_DEFINE_TCM_OCM_MMAP)
+#define TCM_MAP 1
+#else
+#define TCM_MAP 0
+#endif
+
+/* +1 is end of list which needs to be empty */
+#define ZYNQMP_MEM_MAP_MAX (ZYNQMP_MEM_MAP_USED + DRAM_BANKS + TCM_MAP + 1)
+
+static struct mm_region zynqmp_mem_map[ZYNQMP_MEM_MAP_MAX] = {
 	{
 		.virt = 0x80000000UL,
 		.phys = 0x80000000UL,
@@ -33,8 +45,7 @@  static struct mm_region zynqmp_mem_map[] = {
 		.attrs = PTE_BLOCK_MEMTYPE(MT_DEVICE_NGNRNE) |
 			 PTE_BLOCK_NON_SHARE |
 			 PTE_BLOCK_PXN | PTE_BLOCK_UXN
-	},
-	{
+	}, {
 		.virt = 0xf8000000UL,
 		.phys = 0xf8000000UL,
 		.size = 0x07e00000UL,
@@ -42,42 +53,51 @@  static struct mm_region zynqmp_mem_map[] = {
 			 PTE_BLOCK_NON_SHARE |
 			 PTE_BLOCK_PXN | PTE_BLOCK_UXN
 	}, {
-#if defined(CONFIG_DEFINE_TCM_OCM_MMAP)
-		.virt = 0xffe00000UL,
-		.phys = 0xffe00000UL,
-		.size = 0x00200000UL,
-		.attrs = PTE_BLOCK_MEMTYPE(MT_NORMAL) |
-			 PTE_BLOCK_INNER_SHARE
-	}, {
-#endif
 		.virt = 0x400000000UL,
 		.phys = 0x400000000UL,
 		.size = 0x400000000UL,
 		.attrs = PTE_BLOCK_MEMTYPE(MT_DEVICE_NGNRNE) |
 			 PTE_BLOCK_NON_SHARE |
 			 PTE_BLOCK_PXN | PTE_BLOCK_UXN
-	},
-#if !defined(CONFIG_ZYNQMP_NO_DDR)
-	{
-		.virt = 0x800000000UL,
-		.phys = 0x800000000UL,
-		.size = 0x800000000UL,
-		.attrs = PTE_BLOCK_MEMTYPE(MT_NORMAL) |
-			 PTE_BLOCK_INNER_SHARE
-	},
-#endif
-	{
+	}, {
 		.virt = 0x1000000000UL,
 		.phys = 0x1000000000UL,
 		.size = 0xf000000000UL,
 		.attrs = PTE_BLOCK_MEMTYPE(MT_DEVICE_NGNRNE) |
 			 PTE_BLOCK_NON_SHARE |
 			 PTE_BLOCK_PXN | PTE_BLOCK_UXN
-	}, {
-		/* List terminator */
-		0,
 	}
 };
+
+void mem_map_fill(void)
+{
+	int banks = ZYNQMP_MEM_MAP_USED;
+
+#if defined(CONFIG_DEFINE_TCM_OCM_MMAP)
+	zynqmp_mem_map[banks].virt = 0xffe00000UL;
+	zynqmp_mem_map[banks].phys = 0xffe00000UL;
+	zynqmp_mem_map[banks].size = 0x00200000UL;
+	zynqmp_mem_map[banks].attrs = PTE_BLOCK_MEMTYPE(MT_NORMAL) |
+				      PTE_BLOCK_INNER_SHARE;
+	banks = banks + 1;
+#endif
+
+#if !defined(CONFIG_ZYNQMP_NO_DDR)
+	for (int i = 0; i < CONFIG_NR_DRAM_BANKS; i++) {
+		/* Skip empty banks */
+		if (!gd->bd->bi_dram[i].size)
+			break;
+
+		zynqmp_mem_map[banks].virt = gd->bd->bi_dram[i].start;
+		zynqmp_mem_map[banks].phys = gd->bd->bi_dram[i].start;
+		zynqmp_mem_map[banks].size = gd->bd->bi_dram[i].size;
+		zynqmp_mem_map[banks].attrs = PTE_BLOCK_MEMTYPE(MT_NORMAL) |
+					      PTE_BLOCK_INNER_SHARE;
+		banks = banks + 1;
+	}
+#endif
+}
+
 struct mm_region *mem_map = zynqmp_mem_map;
 
 u64 get_page_table_size(void)
diff --git a/arch/arm/include/asm/arch-zynqmp/sys_proto.h b/arch/arm/include/asm/arch-zynqmp/sys_proto.h
index 3daf0e81d80c..d11c0782b8dd 100644
--- a/arch/arm/include/asm/arch-zynqmp/sys_proto.h
+++ b/arch/arm/include/asm/arch-zynqmp/sys_proto.h
@@ -46,7 +46,7 @@  int invoke_smc(u32 pm_api_id, u32 arg0, u32 arg1, u32 arg2, u32 arg3,
 	       u32 *ret_payload);
 
 void initialize_tcm(bool mode);
-
+void mem_map_fill(void);
 int chip_id(unsigned char id);
 
 #endif /* _ASM_ARCH_SYS_PROTO_H */
diff --git a/board/xilinx/zynqmp/zynqmp.c b/board/xilinx/zynqmp/zynqmp.c
index aeef0b4cdd7e..6d09a4c73139 100644
--- a/board/xilinx/zynqmp/zynqmp.c
+++ b/board/xilinx/zynqmp/zynqmp.c
@@ -406,7 +406,15 @@  unsigned long do_go_exec(ulong (*entry)(int, char * const []), int argc,
 #if !defined(CONFIG_SYS_SDRAM_BASE) && !defined(CONFIG_SYS_SDRAM_SIZE)
 int dram_init_banksize(void)
 {
-	return fdtdec_setup_memory_banksize();
+	int ret;
+
+	ret = fdtdec_setup_memory_banksize();
+	if (ret)
+		return ret;
+
+	mem_map_fill();
+
+	return 0;
 }
 
 int dram_init(void)
@@ -417,6 +425,18 @@  int dram_init(void)
 	return 0;
 }
 #else
+int dram_init_banksize(void)
+{
+#if defined(CONFIG_NR_DRAM_BANKS)
+	gd->bd->bi_dram[0].start = CONFIG_SYS_SDRAM_BASE;
+	gd->bd->bi_dram[0].size = get_effective_memsize();
+#endif
+
+	mem_map_fill();
+
+	return 0;
+}
+
 int dram_init(void)
 {
 	gd->ram_size = get_ram_size((void *)CONFIG_SYS_SDRAM_BASE,