4 files changed, 22 insertions, 9 deletions

diff --git a/Documentation/arm64/memory.txt b/Documentation/arm64/memory.txt
index 8845d0847a66..344e85cc7323 100644
--- a/Documentation/arm64/memory.txt
+++ b/Documentation/arm64/memory.txt
@@ -45,6 +45,14 @@ Start			End			Size		Use
 fffffc0000000000        ffffffffffffffff           4TB          kernel
 
 
+AArch64 Linux memory layout with 64KB pages + 3 levels:
+
+Start                   End                     Size            Use
+-----------------------------------------------------------------------
+0000000000000000        0000ffffffffffff         256TB          user
+ffff000000000000        ffffffffffffffff         256TB          kernel
+
+
 For details of the virtual kernel memory layout please see the kernel
 booting log.
 
diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 8cdb7f36bfc8..f4615ba88a77 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -210,6 +210,7 @@ config ARM64_VA_BITS
 config ARM64_PGTABLE_LEVELS
         int
         default 2 if ARM64_64K_PAGES && ARM64_VA_BITS_42
+        default 3 if ARM64_64K_PAGES && ARM64_VA_BITS_48
         default 3 if ARM64_4K_PAGES && ARM64_VA_BITS_39
         default 4 if ARM64_4K_PAGES && ARM64_VA_BITS_48
 
diff --git a/arch/arm64/include/asm/page.h b/arch/arm64/include/asm/page.h
index 2502754d1921..7a3f462133b0 100644
--- a/arch/arm64/include/asm/page.h
+++ b/arch/arm64/include/asm/page.h
@@ -34,17 +34,19 @@
 /*
  * The idmap and swapper page tables need some space reserved in the kernel
  * image. Both require pgd, pud (4 levels only) and pmd tables to (section)
- * map the kernel. The swapper also maps the FDT (see __create_page_tables for
- * more information).
+ * map the kernel. With the 64K page configuration, swapper and idmap need to
+ * map to pte level. The swapper also maps the FDT (see __create_page_tables
+ * for more information).
  */
-#if CONFIG_ARM64_PGTABLE_LEVELS == 4
-#define SWAPPER_DIR_SIZE        (3 * PAGE_SIZE)
-#define IDMAP_DIR_SIZE          (3 * PAGE_SIZE)
+#ifdef CONFIG_ARM64_64K_PAGES
+#define SWAPPER_PGTABLE_LEVELS  (CONFIG_ARM64_PGTABLE_LEVELS)
 #else
-#define SWAPPER_DIR_SIZE        (2 * PAGE_SIZE)
-#define IDMAP_DIR_SIZE          (2 * PAGE_SIZE)
+#define SWAPPER_PGTABLE_LEVELS  (CONFIG_ARM64_PGTABLE_LEVELS - 1)
 #endif
 
+#define SWAPPER_DIR_SIZE        (SWAPPER_PGTABLE_LEVELS * PAGE_SIZE)
+#define IDMAP_DIR_SIZE          (SWAPPER_DIR_SIZE)
+
 #ifndef __ASSEMBLY__
 
 #include <asm/pgtable-types.h>
diff --git a/arch/arm64/kernel/head.S b/arch/arm64/kernel/head.S
index a6db505411bc..0bce493495e9 100644
--- a/arch/arm64/kernel/head.S
+++ b/arch/arm64/kernel/head.S
@@ -55,9 +55,11 @@
 #ifdef CONFIG_ARM64_64K_PAGES
 #define BLOCK_SHIFT     PAGE_SHIFT
 #define BLOCK_SIZE      PAGE_SIZE
+#define TABLE_SHIFT     PMD_SHIFT
 #else
 #define BLOCK_SHIFT     SECTION_SHIFT
 #define BLOCK_SIZE      SECTION_SIZE
+#define TABLE_SHIFT     PUD_SHIFT
 #endif
 
 #define KERNEL_START    KERNEL_RAM_VADDR
@@ -505,8 +507,8 @@ ENDPROC(__calc_phys_offset)
  */
         .macro  create_pgd_entry, tbl, virt, tmp1, tmp2
         create_table_entry \tbl, \virt, PGDIR_SHIFT, PTRS_PER_PGD, \tmp1, \tmp2
-#if CONFIG_ARM64_PGTABLE_LEVELS == 4
-        create_table_entry \tbl, \virt, PUD_SHIFT, PTRS_PER_PUD, \tmp1, \tmp2
+#if SWAPPER_PGTABLE_LEVELS == 3
+        create_table_entry \tbl, \virt, TABLE_SHIFT, PTRS_PER_PTE, \tmp1, \tmp2
 #endif
         .endm