struct kmem_cache *amd_iommu_irq_cache;
 
 static void detach_device(struct device *dev);
-static int domain_enable_v2(struct protection_domain *domain, int pasids);
 
 /****************************************************************************
  *
        free_page((unsigned long)domain->gcr3_tbl);
 }
 
+/*
+ * Number of GCR3 table levels required. Level must be 4-Kbyte
+ * page and can contain up to 512 entries.
+ */
+static int get_gcr3_levels(int pasids)
+{
+       int levels;
+
+       if (pasids == -1)
+               return amd_iommu_max_glx_val;
+
+       levels = get_count_order(pasids);
+
+       return levels ? (DIV_ROUND_UP(levels, 9) - 1) : levels;
+}
+
+/* Note: This function expects iommu_domain->lock to be held prior calling the function. */
+static int setup_gcr3_table(struct protection_domain *domain, int pasids)
+{
+       int levels = get_gcr3_levels(pasids);
+
+       if (levels > amd_iommu_max_glx_val)
+               return -EINVAL;
+
+       domain->gcr3_tbl = alloc_pgtable_page(domain->nid, GFP_ATOMIC);
+       if (domain->gcr3_tbl == NULL)
+               return -ENOMEM;
+
+       domain->glx      = levels;
+       domain->flags   |= PD_IOMMUV2_MASK;
+
+       amd_iommu_domain_update(domain);
+
+       return 0;
+}
+
 static void set_dte_entry(struct amd_iommu *iommu, u16 devid,
                          struct protection_domain *domain, bool ats, bool ppr)
 {
 
        domain->domain.pgsize_bitmap = AMD_IOMMU_PGSIZES_V2;
 
-       if (domain_enable_v2(domain, 1)) {
+       if (setup_gcr3_table(domain, 1)) {
                domain_id_free(domain->id);
                return -ENOMEM;
        }
 }
 EXPORT_SYMBOL(amd_iommu_domain_direct_map);
 
-/* Note: This function expects iommu_domain->lock to be held prior calling the function. */
-static int domain_enable_v2(struct protection_domain *domain, int pasids)
-{
-       int levels;
-
-       /* Number of GCR3 table levels required */
-       for (levels = 0; (pasids - 1) & ~0x1ff; pasids >>= 9)
-               levels += 1;
-
-       if (levels > amd_iommu_max_glx_val)
-               return -EINVAL;
-
-       domain->gcr3_tbl = (void *)get_zeroed_page(GFP_ATOMIC);
-       if (domain->gcr3_tbl == NULL)
-               return -ENOMEM;
-
-       domain->glx      = levels;
-       domain->flags   |= PD_IOMMUV2_MASK;
-
-       amd_iommu_domain_update(domain);
-
-       return 0;
-}
-
 int amd_iommu_domain_enable_v2(struct iommu_domain *dom, int pasids)
 {
        struct protection_domain *pdom = to_pdomain(dom);
                goto out;
 
        if (!pdom->gcr3_tbl)
-               ret = domain_enable_v2(pdom, pasids);
+               ret = setup_gcr3_table(pdom, pasids);
 
 out:
        spin_unlock_irqrestore(&pdom->lock, flags);