armv8: fsl-lsch3: Rewrite MMU translation table entries

This patch rewrites MMU translation table entries. To start, all table
entries are written as "invalid", then "device-ngnrnr" and "normal" are
written to the entries to enable access to specific addresses.

Signed-off-by: Alison Wang <alison.wang@freescale.com>
Signed-off-by: York Sun <yorksun@freescale.com>
diff --git a/arch/arm/cpu/armv8/cache_v8.c b/arch/arm/cpu/armv8/cache_v8.c
index 835f6a6..6bde1cf 100644
--- a/arch/arm/cpu/armv8/cache_v8.c
+++ b/arch/arm/cpu/armv8/cache_v8.c
@@ -12,16 +12,25 @@
 DECLARE_GLOBAL_DATA_PTR;
 
 #ifndef CONFIG_SYS_DCACHE_OFF
-void set_pgtable_section(u64 *page_table, u64 index, u64 section,
-			 u64 memory_type)
+inline void set_pgtable_section(u64 *page_table, u64 index, u64 section,
+			 u64 memory_type, u64 share)
 {
 	u64 value;
 
 	value = section | PMD_TYPE_SECT | PMD_SECT_AF;
 	value |= PMD_ATTRINDX(memory_type);
+	value |= share;
 	page_table[index] = value;
 }
 
+inline void set_pgtable_table(u64 *page_table, u64 index, u64 *table_addr)
+{
+	u64 value;
+
+	value = (u64)table_addr | PMD_TYPE_TABLE;
+	page_table[index] = value;
+}
+
 /* to activate the MMU we need to set up virtual memory */
 static void mmu_setup(void)
 {
@@ -32,7 +41,7 @@
 	/* Setup an identity-mapping for all spaces */
 	for (i = 0; i < (PGTABLE_SIZE >> 3); i++) {
 		set_pgtable_section(page_table, i, i << SECTION_SHIFT,
-				    MT_DEVICE_NGNRNE);
+				    MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE);
 	}
 
 	/* Setup an identity-mapping for all RAM space */
@@ -42,7 +51,7 @@
 		for (j = start >> SECTION_SHIFT;
 		     j < end >> SECTION_SHIFT; j++) {
 			set_pgtable_section(page_table, j, j << SECTION_SHIFT,
-					    MT_NORMAL);
+					    MT_NORMAL, PMD_SECT_NON_SHARE);
 		}
 	}
 
diff --git a/arch/arm/cpu/armv8/fsl-lsch3/README b/arch/arm/cpu/armv8/fsl-lsch3/README
index 3c15479..08da7e4 100644
--- a/arch/arm/cpu/armv8/fsl-lsch3/README
+++ b/arch/arm/cpu/armv8/fsl-lsch3/README
@@ -171,3 +171,74 @@
 
 Notice the difference from QDS is SRC, SRC_ADDR and the offset of u-boot image
 to match board NAND device with 4KB/page, block size 512KB.
+
+MMU Translation Tables
+======================
+
+(1) Early MMU Tables:
+
+     Level 0                   Level 1                   Level 2
+------------------        ------------------        ------------------
+| 0x00_0000_0000 | -----> | 0x00_0000_0000 | -----> | 0x00_0000_0000 |
+------------------        ------------------        ------------------
+| 0x80_0000_0000 | --|    | 0x00_4000_0000 |        | 0x00_0020_0000 |
+------------------   |    ------------------        ------------------
+|    invalid     |   |    | 0x00_8000_0000 |        | 0x00_0040_0000 |
+------------------   |    ------------------        ------------------
+                     |    | 0x00_c000_0000 |        | 0x00_0060_0000 |
+                     |    ------------------        ------------------
+                     |    | 0x01_0000_0000 |        | 0x00_0080_0000 |
+                     |    ------------------        ------------------
+                     |            ...                      ...
+                     |    ------------------
+                     |    | 0x05_8000_0000 |  --|
+                     |    ------------------    |
+                     |    | 0x05_c000_0000 |    |
+                     |    ------------------    |
+                     |            ...           |
+                     |    ------------------    |   ------------------
+                     |--> | 0x80_0000_0000 |    |-> | 0x00_3000_0000 |
+                          ------------------        ------------------
+                          | 0x80_4000_0000 |        | 0x00_3020_0000 |
+                          ------------------        ------------------
+                          | 0x80_8000_0000 |        | 0x00_3040_0000 |
+                          ------------------        ------------------
+                          | 0x80_c000_0000 |        | 0x00_3060_0000 |
+                          ------------------        ------------------
+                          | 0x81_0000_0000 |        | 0x00_3080_0000 |
+                          ------------------        ------------------
+			         ...	                   ...
+
+(2) Final MMU Tables:
+
+     Level 0                   Level 1                   Level 2
+------------------        ------------------        ------------------
+| 0x00_0000_0000 | -----> | 0x00_0000_0000 | -----> | 0x00_0000_0000 |
+------------------        ------------------        ------------------
+| 0x80_0000_0000 | --|    | 0x00_4000_0000 |        | 0x00_0020_0000 |
+------------------   |    ------------------        ------------------
+|    invalid     |   |    | 0x00_8000_0000 |        | 0x00_0040_0000 |
+------------------   |    ------------------        ------------------
+                     |    | 0x00_c000_0000 |        | 0x00_0060_0000 |
+                     |    ------------------        ------------------
+                     |    | 0x01_0000_0000 |        | 0x00_0080_0000 |
+                     |    ------------------        ------------------
+                     |            ...                      ...
+                     |    ------------------
+                     |    | 0x08_0000_0000 | --|
+                     |    ------------------   |
+                     |    | 0x08_4000_0000 |   |
+                     |    ------------------   |
+                     |            ...          |
+                     |    ------------------   |    ------------------
+                     |--> | 0x80_0000_0000 |   |--> | 0x08_0000_0000 |
+                          ------------------        ------------------
+                          | 0x80_4000_0000 |        | 0x08_0020_0000 |
+                          ------------------        ------------------
+                          | 0x80_8000_0000 |        | 0x08_0040_0000 |
+                          ------------------        ------------------
+                          | 0x80_c000_0000 |        | 0x08_0060_0000 |
+                          ------------------        ------------------
+                          | 0x81_0000_0000 |        | 0x08_0080_0000 |
+                          ------------------        ------------------
+			         ...	                   ...
diff --git a/arch/arm/cpu/armv8/fsl-lsch3/cpu.c b/arch/arm/cpu/armv8/fsl-lsch3/cpu.c
index d02c0be..eb1213e 100644
--- a/arch/arm/cpu/armv8/fsl-lsch3/cpu.c
+++ b/arch/arm/cpu/armv8/fsl-lsch3/cpu.c
@@ -6,6 +6,7 @@
 
 #include <common.h>
 #include <asm/io.h>
+#include <asm/errno.h>
 #include <asm/system.h>
 #include <asm/armv8/mmu.h>
 #include <asm/io.h>
@@ -53,27 +54,16 @@
 }
 
 #ifndef CONFIG_SYS_DCACHE_OFF
-/*
- * To start MMU before DDR is available, we create MMU table in SRAM.
- * The base address of SRAM is CONFIG_SYS_FSL_OCRAM_BASE. We use three
- * levels of translation tables here to cover 40-bit address space.
- * We use 4KB granule size, with 40 bits physical address, T0SZ=24
- * Level 0 IA[39], table address @0
- * Level 1 IA[31:30], table address @0x1000, 0x2000
- * Level 2 IA[29:21], table address @0x3000, 0x4000
- * Address above 0x5000 is free for other purpose.
- */
+
+#define SECTION_SHIFT_L0		39UL
+#define SECTION_SHIFT_L1		30UL
+#define SECTION_SHIFT_L2		21UL
+#define BLOCK_SIZE_L0			0x8000000000
+#define BLOCK_SIZE_L1			0x40000000
+#define BLOCK_SIZE_L2			0x200000
 
-#define SECTION_SHIFT_L0	39UL
-#define SECTION_SHIFT_L1	30UL
-#define SECTION_SHIFT_L2	21UL
-#define BLOCK_SIZE_L0		0x8000000000UL
-#define BLOCK_SIZE_L1		(1 << SECTION_SHIFT_L1)
-#define BLOCK_SIZE_L2		(1 << SECTION_SHIFT_L2)
-#define CONFIG_SYS_IFC_BASE	0x30000000
-#define CONFIG_SYS_IFC_SIZE	0x10000000
-#define CONFIG_SYS_IFC_BASE2	0x500000000
-#define CONFIG_SYS_IFC_SIZE2	0x100000000
+#define NUM_OF_ENTRY		512
+
 #define TCR_EL2_PS_40BIT	(2 << 16)
 #define LSCH3_VA_BITS		(40)
 #define LSCH3_TCR	(TCR_TG0_4K		| \
@@ -89,95 +79,265 @@
 			TCR_IRGN_WBWA		| \
 			TCR_T0SZ(LSCH3_VA_BITS))
 
+#define CONFIG_SYS_FSL_CCSR_BASE	0x00000000
+#define CONFIG_SYS_FSL_CCSR_SIZE	0x10000000
+#define CONFIG_SYS_FSL_QSPI_BASE1	0x20000000
+#define CONFIG_SYS_FSL_QSPI_SIZE1	0x10000000
+#define CONFIG_SYS_FSL_IFC_BASE1	0x30000000
+#define CONFIG_SYS_FSL_IFC_SIZE1	0x10000000
+#define CONFIG_SYS_FSL_IFC_SIZE1_1	0x400000
+#define CONFIG_SYS_FSL_DRAM_BASE1	0x80000000
+#define CONFIG_SYS_FSL_DRAM_SIZE1	0x80000000
+#define CONFIG_SYS_FSL_QSPI_BASE2	0x400000000
+#define CONFIG_SYS_FSL_QSPI_SIZE2	0x100000000
+#define CONFIG_SYS_FSL_IFC_BASE2	0x500000000
+#define CONFIG_SYS_FSL_IFC_SIZE2	0x100000000
+#define CONFIG_SYS_FSL_DCSR_BASE	0x700000000
+#define CONFIG_SYS_FSL_DCSR_SIZE	0x40000000
+#define CONFIG_SYS_FSL_MC_BASE		0x80c000000
+#define CONFIG_SYS_FSL_MC_SIZE		0x4000000
+#define CONFIG_SYS_FSL_NI_BASE		0x810000000
+#define CONFIG_SYS_FSL_NI_SIZE		0x8000000
+#define CONFIG_SYS_FSL_QBMAN_BASE	0x818000000
+#define CONFIG_SYS_FSL_QBMAN_SIZE	0x8000000
+#define CONFIG_SYS_FSL_QBMAN_SIZE_1	0x4000000
+#define CONFIG_SYS_PCIE1_PHYS_SIZE	0x200000000
+#define CONFIG_SYS_PCIE2_PHYS_SIZE	0x200000000
+#define CONFIG_SYS_PCIE3_PHYS_SIZE	0x200000000
+#define CONFIG_SYS_PCIE4_PHYS_SIZE	0x200000000
+#define CONFIG_SYS_FSL_WRIOP1_BASE	0x4300000000
+#define CONFIG_SYS_FSL_WRIOP1_SIZE	0x100000000
+#define CONFIG_SYS_FSL_AIOP1_BASE	0x4b00000000
+#define CONFIG_SYS_FSL_AIOP1_SIZE	0x100000000
+#define CONFIG_SYS_FSL_PEBUF_BASE	0x4c00000000
+#define CONFIG_SYS_FSL_PEBUF_SIZE	0x400000000
+#define CONFIG_SYS_FSL_DRAM_BASE2	0x8080000000
+#define CONFIG_SYS_FSL_DRAM_SIZE2	0x7F80000000
+
+struct sys_mmu_table {
+	u64 virt_addr;
+	u64 phys_addr;
+	u64 size;
+	u64 memory_type;
+	u64 share;
+};
+
+static const struct sys_mmu_table lsch3_early_mmu_table[] = {
+	{ CONFIG_SYS_FSL_CCSR_BASE, CONFIG_SYS_FSL_CCSR_BASE,
+	  CONFIG_SYS_FSL_CCSR_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_OCRAM_BASE, CONFIG_SYS_FSL_OCRAM_BASE,
+	  CONFIG_SYS_FSL_OCRAM_SIZE, MT_NORMAL, PMD_SECT_NON_SHARE },
+	/* For IFC Region #1, only the first 4MB is cache-enabled */
+	{ CONFIG_SYS_FSL_IFC_BASE1, CONFIG_SYS_FSL_IFC_BASE1,
+	  CONFIG_SYS_FSL_IFC_SIZE1_1, MT_NORMAL, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_IFC_BASE1 + CONFIG_SYS_FSL_IFC_SIZE1_1,
+	  CONFIG_SYS_FSL_IFC_BASE1 + CONFIG_SYS_FSL_IFC_SIZE1_1,
+	  CONFIG_SYS_FSL_IFC_SIZE1 - CONFIG_SYS_FSL_IFC_SIZE1_1,
+	  MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FLASH_BASE, CONFIG_SYS_FSL_IFC_BASE1,
+	  CONFIG_SYS_FSL_IFC_SIZE1, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_DRAM_BASE1, CONFIG_SYS_FSL_DRAM_BASE1,
+	  CONFIG_SYS_FSL_DRAM_SIZE1, MT_NORMAL, PMD_SECT_OUTER_SHARE },
+	{ CONFIG_SYS_FSL_DCSR_BASE, CONFIG_SYS_FSL_DCSR_BASE,
+	  CONFIG_SYS_FSL_DCSR_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_DRAM_BASE2, CONFIG_SYS_FSL_DRAM_BASE2,
+	  CONFIG_SYS_FSL_DRAM_SIZE2, MT_NORMAL, PMD_SECT_OUTER_SHARE },
+};
+
+static const struct sys_mmu_table lsch3_final_mmu_table[] = {
+	{ CONFIG_SYS_FSL_CCSR_BASE, CONFIG_SYS_FSL_CCSR_BASE,
+	  CONFIG_SYS_FSL_CCSR_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_OCRAM_BASE, CONFIG_SYS_FSL_OCRAM_BASE,
+	  CONFIG_SYS_FSL_OCRAM_SIZE, MT_NORMAL, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_DRAM_BASE1, CONFIG_SYS_FSL_DRAM_BASE1,
+	  CONFIG_SYS_FSL_DRAM_SIZE1, MT_NORMAL, PMD_SECT_OUTER_SHARE },
+	{ CONFIG_SYS_FSL_QSPI_BASE2, CONFIG_SYS_FSL_QSPI_BASE2,
+	  CONFIG_SYS_FSL_QSPI_SIZE2, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_IFC_BASE2, CONFIG_SYS_FSL_IFC_BASE2,
+	  CONFIG_SYS_FSL_IFC_SIZE2, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_DCSR_BASE, CONFIG_SYS_FSL_DCSR_BASE,
+	  CONFIG_SYS_FSL_DCSR_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_MC_BASE, CONFIG_SYS_FSL_MC_BASE,
+	  CONFIG_SYS_FSL_MC_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_NI_BASE, CONFIG_SYS_FSL_NI_BASE,
+	  CONFIG_SYS_FSL_NI_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	/* For QBMAN portal, only the first 64MB is cache-enabled */
+	{ CONFIG_SYS_FSL_QBMAN_BASE, CONFIG_SYS_FSL_QBMAN_BASE,
+	  CONFIG_SYS_FSL_QBMAN_SIZE_1, MT_NORMAL, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_QBMAN_BASE + CONFIG_SYS_FSL_QBMAN_SIZE_1,
+	  CONFIG_SYS_FSL_QBMAN_BASE + CONFIG_SYS_FSL_QBMAN_SIZE_1,
+	  CONFIG_SYS_FSL_QBMAN_SIZE - CONFIG_SYS_FSL_QBMAN_SIZE_1,
+	  MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_PCIE1_PHYS_ADDR, CONFIG_SYS_PCIE1_PHYS_ADDR,
+	  CONFIG_SYS_PCIE1_PHYS_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_PCIE2_PHYS_ADDR, CONFIG_SYS_PCIE2_PHYS_ADDR,
+	  CONFIG_SYS_PCIE2_PHYS_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_PCIE3_PHYS_ADDR, CONFIG_SYS_PCIE3_PHYS_ADDR,
+	  CONFIG_SYS_PCIE3_PHYS_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+#ifdef CONFIG_LS2085A
+	{ CONFIG_SYS_PCIE4_PHYS_ADDR, CONFIG_SYS_PCIE4_PHYS_ADDR,
+	  CONFIG_SYS_PCIE4_PHYS_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+#endif
+	{ CONFIG_SYS_FSL_WRIOP1_BASE, CONFIG_SYS_FSL_WRIOP1_BASE,
+	  CONFIG_SYS_FSL_WRIOP1_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_AIOP1_BASE, CONFIG_SYS_FSL_AIOP1_BASE,
+	  CONFIG_SYS_FSL_AIOP1_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_PEBUF_BASE, CONFIG_SYS_FSL_PEBUF_BASE,
+	  CONFIG_SYS_FSL_PEBUF_SIZE, MT_DEVICE_NGNRNE, PMD_SECT_NON_SHARE },
+	{ CONFIG_SYS_FSL_DRAM_BASE2, CONFIG_SYS_FSL_DRAM_BASE2,
+	  CONFIG_SYS_FSL_DRAM_SIZE2, MT_NORMAL, PMD_SECT_OUTER_SHARE },
+};
+
+struct table_info {
+	u64 *ptr;
+	u64 table_base;
+	u64 entry_size;
+};
+
 /*
- * Final MMU
- * Let's start from the same layout as early MMU and modify as needed.
- * IFC regions will be cache-inhibit.
+ * Set the block entries according to the information of the table.
  */
-#define FINAL_QBMAN_CACHED_MEM	0x818000000UL
-#define FINAL_QBMAN_CACHED_SIZE	0x4000000
+static int set_block_entry(const struct sys_mmu_table *list,
+			   struct table_info *table)
+{
+	u64 block_size = 0, block_shift = 0;
+	u64 block_addr, index;
+	int j;
 
+	if (table->entry_size == BLOCK_SIZE_L1) {
+		block_size = BLOCK_SIZE_L1;
+		block_shift = SECTION_SHIFT_L1;
+	} else if (table->entry_size == BLOCK_SIZE_L2) {
+		block_size = BLOCK_SIZE_L2;
+		block_shift = SECTION_SHIFT_L2;
+	} else {
+		return -EINVAL;
+	}
 
-static inline void early_mmu_setup(void)
+	block_addr = list->phys_addr;
+	index = (list->virt_addr - table->table_base) >> block_shift;
+
+	for (j = 0; j < (list->size >> block_shift); j++) {
+		set_pgtable_section(table->ptr,
+				    index,
+				    block_addr,
+				    list->memory_type,
+				    list->share);
+		block_addr += block_size;
+		index++;
+	}
+
+	return 0;
+}
+
+/*
+ * Find the corresponding table entry for the list.
+ */
+static int find_table(const struct sys_mmu_table *list,
+		      struct table_info *table, u64 *level0_table)
 {
-	int el;
-	u64 i;
-	u64 section_l1t0, section_l1t1, section_l2t0, section_l2t1;
-	u64 *level0_table = (u64 *)CONFIG_SYS_FSL_OCRAM_BASE;
-	u64 *level1_table_0 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x1000);
-	u64 *level1_table_1 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x2000);
-	u64 *level2_table_0 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x3000);
-	u64 *level2_table_1 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x4000);
+	u64 index = 0, level = 0;
+	u64 *level_table = level0_table;
+	u64 temp_base = 0, block_size = 0, block_shift = 0;
+
+	while (level < 3) {
+		if (level == 0) {
+			block_size = BLOCK_SIZE_L0;
+			block_shift = SECTION_SHIFT_L0;
+		} else if (level == 1) {
+			block_size = BLOCK_SIZE_L1;
+			block_shift = SECTION_SHIFT_L1;
+		} else if (level == 2) {
+			block_size = BLOCK_SIZE_L2;
+			block_shift = SECTION_SHIFT_L2;
+		}
+
+		index = 0;
+		while (list->virt_addr >= temp_base) {
+			index++;
+			temp_base += block_size;
+		}
+
+		temp_base -= block_size;
+
+		if ((level_table[index - 1] & PMD_TYPE_MASK) ==
+		    PMD_TYPE_TABLE) {
+			level_table = (u64 *)(level_table[index - 1] &
+				      ~PMD_TYPE_MASK);
+			level++;
+			continue;
+		} else {
+			if (level == 0)
+				return -EINVAL;
 
-	level0_table[0] =
-		(u64)level1_table_0 | PMD_TYPE_TABLE;
-	level0_table[1] =
-		(u64)level1_table_1 | PMD_TYPE_TABLE;
+			if ((list->phys_addr + list->size) >
+			    (temp_base + block_size * NUM_OF_ENTRY))
+				return -EINVAL;
 
-	/*
-	 * set level 1 table 0 to cache_inhibit, covering 0 to 512GB
-	 * set level 1 table 1 to cache enabled, covering 512GB to 1TB
-	 * set level 2 table to cache-inhibit, covering 0 to 1GB
-	 */
-	section_l1t0 = 0;
-	section_l1t1 = BLOCK_SIZE_L0;
-	section_l2t0 = 0;
-	section_l2t1 = CONFIG_SYS_FLASH_BASE;
-	for (i = 0; i < 512; i++) {
-		set_pgtable_section(level1_table_0, i, section_l1t0,
-				    MT_DEVICE_NGNRNE);
-		set_pgtable_section(level1_table_1, i, section_l1t1,
-				    MT_NORMAL);
-		set_pgtable_section(level2_table_0, i, section_l2t0,
-				    MT_DEVICE_NGNRNE);
-		set_pgtable_section(level2_table_1, i, section_l2t1,
-				    MT_DEVICE_NGNRNE);
-		section_l1t0 += BLOCK_SIZE_L1;
-		section_l1t1 += BLOCK_SIZE_L1;
-		section_l2t0 += BLOCK_SIZE_L2;
-		section_l2t1 += BLOCK_SIZE_L2;
+			/*
+			 * Check the address and size of the list member is
+			 * aligned with the block size.
+			 */
+			if (((list->phys_addr & (block_size - 1)) != 0) ||
+			    ((list->size & (block_size - 1)) != 0))
+				return -EINVAL;
+
+			table->ptr = level_table;
+			table->table_base = temp_base -
+					    ((index - 1) << block_shift);
+			table->entry_size = block_size;
+
+			return 0;
+		}
 	}
+	return -EINVAL;
+}
 
-	level1_table_0[0] =
-		(u64)level2_table_0 | PMD_TYPE_TABLE;
-	level1_table_0[1] =
-		0x40000000 | PMD_SECT_AF | PMD_TYPE_SECT |
-		PMD_ATTRINDX(MT_DEVICE_NGNRNE);
-	level1_table_0[2] =
-		0x80000000 | PMD_SECT_AF | PMD_TYPE_SECT |
-		PMD_ATTRINDX(MT_NORMAL);
-	level1_table_0[3] =
-		0xc0000000 | PMD_SECT_AF | PMD_TYPE_SECT |
-		PMD_ATTRINDX(MT_NORMAL);
+/*
+ * To start MMU before DDR is available, we create MMU table in SRAM.
+ * The base address of SRAM is CONFIG_SYS_FSL_OCRAM_BASE. We use three
+ * levels of translation tables here to cover 40-bit address space.
+ * We use 4KB granule size, with 40 bits physical address, T0SZ=24
+ * Level 0 IA[39], table address @0
+ * Level 1 IA[38:30], table address @0x1000, 0x2000
+ * Level 2 IA[29:21], table address @0x3000, 0x4000
+ * Address above 0x5000 is free for other purpose.
+ */
+static inline void early_mmu_setup(void)
+{
+	unsigned int el, i;
+	u64 *level0_table = (u64 *)CONFIG_SYS_FSL_OCRAM_BASE;
+	u64 *level1_table0 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x1000);
+	u64 *level1_table1 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x2000);
+	u64 *level2_table0 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x3000);
+	u64 *level2_table1 = (u64 *)(CONFIG_SYS_FSL_OCRAM_BASE + 0x4000);
+	struct table_info table = {level0_table, 0, BLOCK_SIZE_L0};
 
-	/* Rewerite table to enable cache for OCRAM */
-	set_pgtable_section(level2_table_0,
-			    CONFIG_SYS_FSL_OCRAM_BASE >> SECTION_SHIFT_L2,
-			    CONFIG_SYS_FSL_OCRAM_BASE,
-			    MT_NORMAL);
+	/* Invalidate all table entries */
+	memset(level0_table, 0, 0x5000);
 
-#if defined(CONFIG_SYS_NOR0_CSPR_EARLY) && defined(CONFIG_SYS_NOR_AMASK_EARLY)
-	/* Rewrite table to enable cache for two entries (4MB) */
-	section_l2t1 = CONFIG_SYS_IFC_BASE;
-	set_pgtable_section(level2_table_0,
-			    section_l2t1 >> SECTION_SHIFT_L2,
-			    section_l2t1,
-			    MT_NORMAL);
-	section_l2t1 += BLOCK_SIZE_L2;
-	set_pgtable_section(level2_table_0,
-			    section_l2t1 >> SECTION_SHIFT_L2,
-			    section_l2t1,
-			    MT_NORMAL);
-#endif
+	/* Fill in the table entries */
+	set_pgtable_table(level0_table, 0, level1_table0);
+	set_pgtable_table(level0_table, 1, level1_table1);
+	set_pgtable_table(level1_table0, 0, level2_table0);
+	set_pgtable_table(level1_table0,
+			  CONFIG_SYS_FLASH_BASE >> SECTION_SHIFT_L1,
+			  level2_table1);
 
-	/* Create a mapping for 256MB IFC region to final flash location */
-	level1_table_0[CONFIG_SYS_FLASH_BASE >> SECTION_SHIFT_L1] =
-		(u64)level2_table_1 | PMD_TYPE_TABLE;
-	section_l2t1 = CONFIG_SYS_IFC_BASE;
-	for (i = 0; i < 0x10000000 >> SECTION_SHIFT_L2; i++) {
-		set_pgtable_section(level2_table_1, i,
-				    section_l2t1, MT_DEVICE_NGNRNE);
-		section_l2t1 += BLOCK_SIZE_L2;
+	/* Find the table and fill in the block entries */
+	for (i = 0; i < ARRAY_SIZE(lsch3_early_mmu_table); i++) {
+		if (find_table(&lsch3_early_mmu_table[i],
+			       &table, level0_table) == 0) {
+			/*
+			 * If find_table() returns error, it cannot be dealt
+			 * with here. Breakpoint can be added for debugging.
+			 */
+			set_block_entry(&lsch3_early_mmu_table[i], &table);
+			/*
+			 * If set_block_entry() returns error, it cannot be
+			 * dealt with here too.
+			 */
+		}
 	}
 
 	el = current_el();
@@ -186,89 +346,55 @@
 }
 
 /*
- * This final tale looks similar to early table, but different in detail.
- * These tables are in regular memory. Cache on IFC is disabled. One sub table
- * is added to enable cache for QBMan.
+ * The final tables look similar to early tables, but different in detail.
+ * These tables are in DRAM. Sub tables are added to enable cache for
+ * QBMan and OCRAM.
+ *
+ * Level 1 table 0 contains 512 entries for each 1GB from 0 to 512GB.
+ * Level 1 table 1 contains 512 entries for each 1GB from 512GB to 1TB.
+ * Level 2 table 0 contains 512 entries for each 2MB from 0 to 1GB.
+ * Level 2 table 1 contains 512 entries for each 2MB from 32GB to 33GB.
  */
 static inline void final_mmu_setup(void)
 {
-	int el;
-	u64 i, tbl_base, tbl_limit, section_base;
-	u64 section_l1t0, section_l1t1, section_l2;
+	unsigned int el, i;
 	u64 *level0_table = (u64 *)gd->arch.tlb_addr;
-	u64 *level1_table_0 = (u64 *)(gd->arch.tlb_addr + 0x1000);
-	u64 *level1_table_1 = (u64 *)(gd->arch.tlb_addr + 0x2000);
-	u64 *level2_table_0 = (u64 *)(gd->arch.tlb_addr + 0x3000);
-	u64 *level2_table_1 = (u64 *)(gd->arch.tlb_addr + 0x4000);
+	u64 *level1_table0 = (u64 *)(gd->arch.tlb_addr + 0x1000);
+	u64 *level1_table1 = (u64 *)(gd->arch.tlb_addr + 0x2000);
+	u64 *level2_table0 = (u64 *)(gd->arch.tlb_addr + 0x3000);
+	u64 *level2_table1 = (u64 *)(gd->arch.tlb_addr + 0x4000);
+	struct table_info table = {level0_table, 0, BLOCK_SIZE_L0};
 
+	/* Invalidate all table entries */
+	memset(level0_table, 0, PGTABLE_SIZE);
 
-	level0_table[0] =
-		(u64)level1_table_0 | PMD_TYPE_TABLE;
-	level0_table[1] =
-		(u64)level1_table_1 | PMD_TYPE_TABLE;
+	/* Fill in the table entries */
+	set_pgtable_table(level0_table, 0, level1_table0);
+	set_pgtable_table(level0_table, 1, level1_table1);
+	set_pgtable_table(level1_table0, 0, level2_table0);
+	set_pgtable_table(level1_table0,
+			  CONFIG_SYS_FSL_QBMAN_BASE >> SECTION_SHIFT_L1,
+			  level2_table1);
 
-	/*
-	 * set level 1 table 0 to cache_inhibit, covering 0 to 512GB
-	 * set level 1 table 1 to cache enabled, covering 512GB to 1TB
-	 * set level 2 table 0 to cache-inhibit, covering 0 to 1GB
-	 */
-	section_l1t0 = 0;
-	section_l1t1 = BLOCK_SIZE_L0 | PMD_SECT_OUTER_SHARE;
-	section_l2 = 0;
-	for (i = 0; i < 512; i++) {
-		set_pgtable_section(level1_table_0, i, section_l1t0,
-				    MT_DEVICE_NGNRNE);
-		set_pgtable_section(level1_table_1, i, section_l1t1,
-				    MT_NORMAL);
-		set_pgtable_section(level2_table_0, i, section_l2,
-				    MT_DEVICE_NGNRNE);
-		section_l1t0 += BLOCK_SIZE_L1;
-		section_l1t1 += BLOCK_SIZE_L1;
-		section_l2 += BLOCK_SIZE_L2;
-	}
-
-	level1_table_0[0] =
-		(u64)level2_table_0 | PMD_TYPE_TABLE;
-	level1_table_0[2] =
-		0x80000000 | PMD_SECT_AF | PMD_TYPE_SECT |
-		PMD_SECT_OUTER_SHARE | PMD_ATTRINDX(MT_NORMAL);
-	level1_table_0[3] =
-		0xc0000000 | PMD_SECT_AF | PMD_TYPE_SECT |
-		PMD_SECT_OUTER_SHARE | PMD_ATTRINDX(MT_NORMAL);
+	/* Find the table and fill in the block entries */
+	for (i = 0; i < ARRAY_SIZE(lsch3_final_mmu_table); i++) {
+		if (find_table(&lsch3_final_mmu_table[i],
+			       &table, level0_table) == 0) {
+			if (set_block_entry(&lsch3_final_mmu_table[i],
+					    &table) != 0) {
+				printf("MMU error: could not set block entry for %p\n",
+				       &lsch3_final_mmu_table[i]);
+			}
 
-	/* Rewrite table to enable cache */
-	set_pgtable_section(level2_table_0,
-			    CONFIG_SYS_FSL_OCRAM_BASE >> SECTION_SHIFT_L2,
-			    CONFIG_SYS_FSL_OCRAM_BASE,
-			    MT_NORMAL);
-
-	/*
-	 * Fill in other part of tables if cache is needed
-	 * If finer granularity than 1GB is needed, sub table
-	 * should be created.
-	 */
-	section_base = FINAL_QBMAN_CACHED_MEM & ~(BLOCK_SIZE_L1 - 1);
-	i = section_base >> SECTION_SHIFT_L1;
-	level1_table_0[i] = (u64)level2_table_1 | PMD_TYPE_TABLE;
-	section_l2 = section_base;
-	for (i = 0; i < 512; i++) {
-		set_pgtable_section(level2_table_1, i, section_l2,
-				    MT_DEVICE_NGNRNE);
-		section_l2 += BLOCK_SIZE_L2;
-	}
-	tbl_base = FINAL_QBMAN_CACHED_MEM & (BLOCK_SIZE_L1 - 1);
-	tbl_limit = (FINAL_QBMAN_CACHED_MEM + FINAL_QBMAN_CACHED_SIZE) &
-		    (BLOCK_SIZE_L1 - 1);
-	for (i = tbl_base >> SECTION_SHIFT_L2;
-	     i < tbl_limit >> SECTION_SHIFT_L2; i++) {
-		section_l2 = section_base + (i << SECTION_SHIFT_L2);
-		set_pgtable_section(level2_table_1, i,
-				    section_l2, MT_NORMAL);
+		} else {
+			printf("MMU error: could not find the table for %p\n",
+			       &lsch3_final_mmu_table[i]);
+		}
 	}
 
 	/* flush new MMU table */
 	flush_dcache_range(gd->arch.tlb_addr,
-			   gd->arch.tlb_addr +  gd->arch.tlb_size);
+			   gd->arch.tlb_addr + gd->arch.tlb_size);
 
 	/* point TTBR to the new table */
 	el = current_el();