/* VTCR_EL2 Registers bits */
 #define VTCR_EL2_PS_MASK       (7 << 16)
-#define VTCR_EL2_PS_40B                (2 << 16)
 #define VTCR_EL2_TG0_MASK      (1 << 14)
 #define VTCR_EL2_TG0_4K                (0 << 14)
 #define VTCR_EL2_TG0_64K       (1 << 14)
  * 64kB pages (TG0 = 1)
  * 2 level page tables (SL = 1)
  */
-#define VTCR_EL2_FLAGS         (VTCR_EL2_PS_40B | VTCR_EL2_TG0_64K | \
-                                VTCR_EL2_SH0_INNER | VTCR_EL2_ORGN0_WBWA | \
-                                VTCR_EL2_IRGN0_WBWA | VTCR_EL2_SL0_LVL1 | \
-                                VTCR_EL2_T0SZ_40B)
+#define VTCR_EL2_FLAGS         (VTCR_EL2_TG0_64K | VTCR_EL2_SH0_INNER | \
+                                VTCR_EL2_ORGN0_WBWA | VTCR_EL2_IRGN0_WBWA | \
+                                VTCR_EL2_SL0_LVL1 | VTCR_EL2_T0SZ_40B)
 #define VTTBR_X                (38 - VTCR_EL2_T0SZ_40B)
 #else
 /*
  * 4kB pages (TG0 = 0)
  * 3 level page tables (SL = 1)
  */
-#define VTCR_EL2_FLAGS         (VTCR_EL2_PS_40B | VTCR_EL2_TG0_4K | \
-                                VTCR_EL2_SH0_INNER | VTCR_EL2_ORGN0_WBWA | \
-                                VTCR_EL2_IRGN0_WBWA | VTCR_EL2_SL0_LVL1 | \
-                                VTCR_EL2_T0SZ_40B)
+#define VTCR_EL2_FLAGS         (VTCR_EL2_TG0_4K | VTCR_EL2_SH0_INNER | \
+                                VTCR_EL2_ORGN0_WBWA | VTCR_EL2_IRGN0_WBWA | \
+                                VTCR_EL2_SL0_LVL1 | VTCR_EL2_T0SZ_40B)
 #define VTTBR_X                (37 - VTCR_EL2_T0SZ_40B)
 #endif
 
 
 #define PTE_HYP                        PTE_USER
 
 /*
- * 40-bit physical address supported.
+ * Highest possible physical address supported.
  */
-#define PHYS_MASK_SHIFT                (40)
+#define PHYS_MASK_SHIFT                (48)
 #define PHYS_MASK              ((UL(1) << PHYS_MASK_SHIFT) - 1)
 
 /*
 #define TCR_SHARED             ((UL(3) << 12) | (UL(3) << 28))
 #define TCR_TG0_64K            (UL(1) << 14)
 #define TCR_TG1_64K            (UL(1) << 30)
-#define TCR_IPS_40BIT          (UL(2) << 32)
 #define TCR_ASID16             (UL(1) << 36)
 #define TCR_TBI0               (UL(1) << 37)
 
 
        msr     tcr_el2, x4
 
        ldr     x4, =VTCR_EL2_FLAGS
+       /*
+        * Read the PARange bits from ID_AA64MMFR0_EL1 and set the PS bits in
+        * VTCR_EL2.
+        */
+       mrs     x5, ID_AA64MMFR0_EL1
+       bfi     x4, x5, #16, #3
        msr     vtcr_el2, x4
 
        mrs     x4, mair_el1
 
         * Set/prepare TCR and TTBR. We use 512GB (39-bit) address range for
         * both user and kernel.
         */
-       ldr     x10, =TCR_TxSZ(VA_BITS) | TCR_FLAGS | TCR_IPS_40BIT | \
+       ldr     x10, =TCR_TxSZ(VA_BITS) | TCR_FLAGS | \
                      TCR_ASID16 | TCR_TBI0 | (1 << 31)
+       /*
+        * Read the PARange bits from ID_AA64MMFR0_EL1 and set the IPS bits in
+        * TCR_EL1.
+        */
+       mrs     x9, ID_AA64MMFR0_EL1
+       bfi     x10, x9, #32, #3
 #ifdef CONFIG_ARM64_64K_PAGES
        orr     x10, x10, TCR_TG0_64K
        orr     x10, x10, TCR_TG1_64K