b       dcache_inval_poc                // tail call
 SYM_CODE_END(preserve_boot_args)
 
-/*
- * Macro to create a table entry to the next page.
- *
- *     tbl:    page table address
- *     virt:   virtual address
- *     shift:  #imm page table shift
- *     ptrs:   #imm pointers per table page
- *
- * Preserves:  virt
- * Corrupts:   ptrs, tmp1, tmp2
- * Returns:    tbl -> next level table page address
- */
-       .macro  create_table_entry, tbl, virt, shift, ptrs, tmp1, tmp2
-       add     \tmp1, \tbl, #PAGE_SIZE
-       phys_to_pte \tmp2, \tmp1
-       orr     \tmp2, \tmp2, #PMD_TYPE_TABLE   // address of next table and entry type
-       lsr     \tmp1, \virt, #\shift
-       sub     \ptrs, \ptrs, #1
-       and     \tmp1, \tmp1, \ptrs             // table index
-       str     \tmp2, [\tbl, \tmp1, lsl #3]
-       add     \tbl, \tbl, #PAGE_SIZE          // next level table page
-       .endm
-
 /*
  * Macro to populate page table entries, these entries can be pointers to the next level
  * or last level entries pointing to physical memory.
  *     phys:   physical address corresponding to vstart - physical memory is contiguous
  *     order:  #imm 2log(number of entries in PGD table)
  *
+ * If extra_shift is set, an extra level will be populated if the end address does
+ * not fit in 'extra_shift' bits. This assumes vend is in the TTBR0 range.
+ *
  * Temporaries:        istart, iend, tmp, count, sv - these need to be different registers
  * Preserves:  vstart, flags
  * Corrupts:   tbl, rtbl, vend, istart, iend, tmp, count, sv
  */
-       .macro map_memory, tbl, rtbl, vstart, vend, flags, phys, order, istart, iend, tmp, count, sv
+       .macro map_memory, tbl, rtbl, vstart, vend, flags, phys, order, istart, iend, tmp, count, sv, extra_shift
        sub \vend, \vend, #1
        add \rtbl, \tbl, #PAGE_SIZE
        mov \count, #0
 
+       .ifnb   \extra_shift
+       tst     \vend, #~((1 << (\extra_shift)) - 1)
+       b.eq    .L_\@
+       compute_indices \vstart, \vend, #\extra_shift, #(PAGE_SHIFT - 3), \istart, \iend, \count
+       mov \sv, \rtbl
+       populate_entries \tbl, \rtbl, \istart, \iend, #PMD_TYPE_TABLE, #PAGE_SIZE, \tmp
+       mov \tbl, \sv
+       .endif
+.L_\@:
        compute_indices \vstart, \vend, #PGDIR_SHIFT, #\order, \istart, \iend, \count
        mov \sv, \rtbl
        populate_entries \tbl, \rtbl, \istart, \iend, #PMD_TYPE_TABLE, #PAGE_SIZE, \tmp
        adrp    x3, __idmap_text_start          // __pa(__idmap_text_start)
 
        /*
-        * VA_BITS may be too small to allow for an ID mapping to be created
-        * that covers system RAM if that is located sufficiently high in the
-        * physical address space. So for the ID map, use an extended virtual
-        * range in that case, and configure an additional translation level
-        * if needed.
+        * The ID map carries a 1:1 mapping of the physical address range
+        * covered by the loaded image, which could be anywhere in DRAM. This
+        * means that the required size of the VA (== PA) space is decided at
+        * boot time, and could be more than the configured size of the VA
+        * space for ordinary kernel and user space mappings.
+        *
+        * There are three cases to consider here:
+        * - 39 <= VA_BITS < 48, and the ID map needs up to 48 VA bits to cover
+        *   the placement of the image. In this case, we configure one extra
+        *   level of translation on the fly for the ID map only. (This case
+        *   also covers 42-bit VA/52-bit PA on 64k pages).
+        *
+        * - VA_BITS == 48, and the ID map needs more than 48 VA bits. This can
+        *   only happen when using 64k pages, in which case we need to extend
+        *   the root level table rather than add a level. Note that we can
+        *   treat this case as 'always extended' as long as we take care not
+        *   to program an unsupported T0SZ value into the TCR register.
+        *
+        * - Combinations that would require two additional levels of
+        *   translation are not supported, e.g., VA_BITS==36 on 16k pages, or
+        *   VA_BITS==39/4k pages with 5-level paging, where the input address
+        *   requires more than 47 or 48 bits, respectively.
         */
-       idmap_get_t0sz x5
-       cmp     x5, TCR_T0SZ(VA_BITS_MIN) // default T0SZ small enough?
-       b.ge    1f                      // .. then skip VA range extension
-
 #if (VA_BITS < 48)
 #define IDMAP_PGD_ORDER        (VA_BITS - PGDIR_SHIFT)
 #define EXTRA_SHIFT    (PGDIR_SHIFT + PAGE_SHIFT - 3)
-#define EXTRA_PTRS     (1 << (PHYS_MASK_SHIFT - EXTRA_SHIFT))
 
        /*
         * If VA_BITS < 48, we have to configure an additional table level.
 #if VA_BITS != EXTRA_SHIFT
 #error "Mismatch between VA_BITS and page size/number of translation levels"
 #endif
-
-       mov     x2, EXTRA_PTRS
-       create_table_entry x0, x3, EXTRA_SHIFT, x2, x5, x6
 #else
 #define IDMAP_PGD_ORDER        (PHYS_MASK_SHIFT - PGDIR_SHIFT)
+#define EXTRA_SHIFT
        /*
         * If VA_BITS == 48, we don't have to configure an additional
         * translation level, but the top-level table has more entries.
         */
 #endif
-1:
        adr_l   x6, __idmap_text_end            // __pa(__idmap_text_end)
 
-       map_memory x0, x1, x3, x6, x7, x3, IDMAP_PGD_ORDER, x10, x11, x12, x13, x14
+       map_memory x0, x1, x3, x6, x7, x3, IDMAP_PGD_ORDER, x10, x11, x12, x13, x14, EXTRA_SHIFT
 
        /*
         * Map the kernel image (starting with PHYS_OFFSET).