diff --git a/xen/arch/arm/p2m.c b/xen/arch/arm/p2m.c
index c171568..1262ae9 100644
--- a/xen/arch/arm/p2m.c
+++ b/xen/arch/arm/p2m.c
@@ -34,7 +34,7 @@ static unsigned int __read_mostly max_vmid = MAX_VMID_8_BIT;
 
 #define P2M_ROOT_PAGES    (1<<P2M_ROOT_ORDER)
 
-unsigned int __read_mostly p2m_ipa_bits;
+unsigned int __read_mostly p2m_ipa_bits = 0;
 
 /* Helpers to lookup the properties of each level */
 static const paddr_t level_masks[] =
@@ -1981,10 +1981,25 @@ void __init setup_virt_paging(void)
         [7] = { 0 }  /* Invalid */
     };
 
-    unsigned int cpu;
+    unsigned int i, cpu;
     unsigned int pa_range = 0x10; /* Larger than any possible value */
     bool vmid_8_bit = false;
 
+    if ( iommu_enabled )
+    {
+        /* We expect "p2m_ipa_bits" to be updated by the IOMMU driver */
+        ASSERT(p2m_ipa_bits);
+
+        /* We need to restrict "pa_range" according to the IOMMU requirements */
+        for ( i = 0; i < ARRAY_SIZE(pa_range_info); i++ )
+        {
+            if ( p2m_ipa_bits >= 64 - pa_range_info[i].t0sz )
+                pa_range = i;
+            else
+                break;
+        }
+    }
+
     for_each_online_cpu ( cpu )
     {
         const struct cpuinfo_arm *info = &cpu_data[cpu];
diff --git a/xen/arch/arm/setup.c b/xen/arch/arm/setup.c
index 51a6677..01cd83d 100644
--- a/xen/arch/arm/setup.c
+++ b/xen/arch/arm/setup.c
@@ -936,12 +936,12 @@ void __init start_xen(unsigned long boot_phys_offset,
     printk("Brought up %ld CPUs\n", (long)num_online_cpus());
     /* TODO: smp_cpus_done(); */
 
-    setup_virt_paging();
-
     rc = iommu_setup();
     if ( !iommu_enabled && rc != -ENODEV )
         panic("Couldn't configure correctly all the IOMMUs.");
 
+    setup_virt_paging();
+
     do_initcalls();
 
     /*
diff --git a/xen/drivers/passthrough/arm/ipmmu-vmsa.c b/xen/drivers/passthrough/arm/ipmmu-vmsa.c
index ec543c3..0dc6351 100644
--- a/xen/drivers/passthrough/arm/ipmmu-vmsa.c
+++ b/xen/drivers/passthrough/arm/ipmmu-vmsa.c
@@ -526,6 +526,7 @@ static int ipmmu_domain_init_context(struct ipmmu_vmsa_domain *domain)
      * to TTBR0. Use 4KB page granule. Start page table walks at first level.
      * Always bypass stage 1 translation.
      */
+    ASSERT(p2m_ipa_bits <= IPMMU_MAX_P2M_IPA_BITS);
     tsz0 = (64 - p2m_ipa_bits) << IMTTBCR_TSZ0_SHIFT;
     ipmmu_ctx_write_root(domain, IMTTBCR, IMTTBCR_EAE | IMTTBCR_PMB |
                          IMTTBCR_SL0_LVL_1 | tsz0);
@@ -1314,23 +1315,12 @@ static __init int ipmmu_init(struct dt_device_node *node, const void *data)
         return -ENODEV;
     }
     else
-    {
         /*
-         * As 4-level translation table is not supported in IPMMU, we need
-         * to check IPA size used for P2M table beforehand to be sure it is
-         * 3-level and the IPMMU will be able to use it.
-         *
-         * TODO: First initialize the IOMMU and gather the requirements and
-         * then initialize the P2M. In the P2M code, take into the account
-         * the IOMMU requirements and restrict "pa_range" if necessary.
+         * Set max Stage-2 input size supported by the IPMMU. We expect
+         * the P2M code will take into the account the IOMMU requirements and
+         * restrict "pa_range" if necessary.
          */
-        if ( IPMMU_MAX_P2M_IPA_BITS < p2m_ipa_bits )
-        {
-            printk_once(XENLOG_ERR "ipmmu: P2M IPA size is not supported (P2M=%u IPMMU=%u)!\n",
-                        p2m_ipa_bits, IPMMU_MAX_P2M_IPA_BITS);
-            return -ENODEV;
-        }
-    }
+        p2m_ipa_bits = IPMMU_MAX_P2M_IPA_BITS;
 
     ret = ipmmu_probe(node);
     if ( ret )
diff --git a/xen/drivers/passthrough/arm/smmu.c b/xen/drivers/passthrough/arm/smmu.c
index 8ae986a..9b3867d 100644
--- a/xen/drivers/passthrough/arm/smmu.c
+++ b/xen/drivers/passthrough/arm/smmu.c
@@ -1110,6 +1110,7 @@ static void arm_smmu_init_context_bank(struct arm_smmu_domain *smmu_domain)
 			reg = TTBCR_TG0_64K;
 
 		if (!stage1) {
+			ASSERT(p2m_ipa_bits <= smmu->s2_input_size);
 			reg |= (64 - smmu->s2_input_size) << TTBCR_T0SZ_SHIFT;
 
 			switch (smmu->s2_output_size) {
@@ -2198,13 +2199,12 @@ static int arm_smmu_device_cfg_probe(struct arm_smmu_device *smmu)
 	size = arm_smmu_id_size_to_bits((id >> ID2_IAS_SHIFT) & ID2_IAS_MASK);
 	smmu->s1_output_size = min_t(unsigned long, PHYS_MASK_SHIFT, size);
 
-	/* Xen: Stage-2 input size has to match p2m_ipa_bits.  */
-	if (size < p2m_ipa_bits) {
-		dev_err(smmu->dev,
-			"P2M IPA size not supported (P2M=%u SMMU=%lu)!\n",
-			p2m_ipa_bits, size);
-		return -ENODEV;
-	}
+	/*
+	 * Xen: Set max Stage-2 input size supported by the SMMU. We expect
+	 * the P2M code will take into the account the IOMMU requirements and
+	 * restrict "pa_range" if necessary.
+	 */
+	p2m_ipa_bits = size;
 	smmu->s2_input_size = p2m_ipa_bits;
 #if 0
 	/* Stage-2 input size limited due to pgd allocation (PTRS_PER_PGD) */
