}
 }
 
-static int copy_pte(pmd_t *dst_pmdp, pmd_t *src_pmdp, unsigned long start,
-                   unsigned long end)
+static int copy_pte(struct trans_pgd_info *info, pmd_t *dst_pmdp,
+                   pmd_t *src_pmdp, unsigned long start, unsigned long end)
 {
        pte_t *src_ptep;
        pte_t *dst_ptep;
        unsigned long addr = start;
 
-       dst_ptep = (pte_t *)get_safe_page(GFP_ATOMIC);
+       dst_ptep = trans_alloc(info);
        if (!dst_ptep)
                return -ENOMEM;
        pmd_populate_kernel(&init_mm, dst_pmdp, dst_ptep);
        return 0;
 }
 
-static int copy_pmd(pud_t *dst_pudp, pud_t *src_pudp, unsigned long start,
-                   unsigned long end)
+static int copy_pmd(struct trans_pgd_info *info, pud_t *dst_pudp,
+                   pud_t *src_pudp, unsigned long start, unsigned long end)
 {
        pmd_t *src_pmdp;
        pmd_t *dst_pmdp;
        unsigned long addr = start;
 
        if (pud_none(READ_ONCE(*dst_pudp))) {
-               dst_pmdp = (pmd_t *)get_safe_page(GFP_ATOMIC);
+               dst_pmdp = trans_alloc(info);
                if (!dst_pmdp)
                        return -ENOMEM;
                pud_populate(&init_mm, dst_pudp, dst_pmdp);
                if (pmd_none(pmd))
                        continue;
                if (pmd_table(pmd)) {
-                       if (copy_pte(dst_pmdp, src_pmdp, addr, next))
+                       if (copy_pte(info, dst_pmdp, src_pmdp, addr, next))
                                return -ENOMEM;
                } else {
                        set_pmd(dst_pmdp,
        return 0;
 }
 
-static int copy_pud(p4d_t *dst_p4dp, p4d_t *src_p4dp, unsigned long start,
+static int copy_pud(struct trans_pgd_info *info, p4d_t *dst_p4dp,
+                   p4d_t *src_p4dp, unsigned long start,
                    unsigned long end)
 {
        pud_t *dst_pudp;
        unsigned long addr = start;
 
        if (p4d_none(READ_ONCE(*dst_p4dp))) {
-               dst_pudp = (pud_t *)get_safe_page(GFP_ATOMIC);
+               dst_pudp = trans_alloc(info);
                if (!dst_pudp)
                        return -ENOMEM;
                p4d_populate(&init_mm, dst_p4dp, dst_pudp);
                if (pud_none(pud))
                        continue;
                if (pud_table(pud)) {
-                       if (copy_pmd(dst_pudp, src_pudp, addr, next))
+                       if (copy_pmd(info, dst_pudp, src_pudp, addr, next))
                                return -ENOMEM;
                } else {
                        set_pud(dst_pudp,
        return 0;
 }
 
-static int copy_p4d(pgd_t *dst_pgdp, pgd_t *src_pgdp, unsigned long start,
+static int copy_p4d(struct trans_pgd_info *info, pgd_t *dst_pgdp,
+                   pgd_t *src_pgdp, unsigned long start,
                    unsigned long end)
 {
        p4d_t *dst_p4dp;
                next = p4d_addr_end(addr, end);
                if (p4d_none(READ_ONCE(*src_p4dp)))
                        continue;
-               if (copy_pud(dst_p4dp, src_p4dp, addr, next))
+               if (copy_pud(info, dst_p4dp, src_p4dp, addr, next))
                        return -ENOMEM;
        } while (dst_p4dp++, src_p4dp++, addr = next, addr != end);
 
        return 0;
 }
 
-static int copy_page_tables(pgd_t *dst_pgdp, unsigned long start,
-                           unsigned long end)
+static int copy_page_tables(struct trans_pgd_info *info, pgd_t *dst_pgdp,
+                           unsigned long start, unsigned long end)
 {
        unsigned long next;
        unsigned long addr = start;
                next = pgd_addr_end(addr, end);
                if (pgd_none(READ_ONCE(*src_pgdp)))
                        continue;
-               if (copy_p4d(dst_pgdp, src_pgdp, addr, next))
+               if (copy_p4d(info, dst_pgdp, src_pgdp, addr, next))
                        return -ENOMEM;
        } while (dst_pgdp++, src_pgdp++, addr = next, addr != end);
 
        return 0;
 }
 
-int trans_pgd_create_copy(pgd_t **dst_pgdp, unsigned long start,
-                         unsigned long end)
+/*
+ * Create trans_pgd and copy linear map.
+ * info:       contains allocator and its argument
+ * dst_pgdp:   new page table that is created, and to which map is copied.
+ * start:      Start of the interval (inclusive).
+ * end:                End of the interval (exclusive).
+ *
+ * Returns 0 on success, and -ENOMEM on failure.
+ */
+int trans_pgd_create_copy(struct trans_pgd_info *info, pgd_t **dst_pgdp,
+                         unsigned long start, unsigned long end)
 {
        int rc;
-       pgd_t *trans_pgd = (pgd_t *)get_safe_page(GFP_ATOMIC);
+       pgd_t *trans_pgd = trans_alloc(info);
 
        if (!trans_pgd) {
                pr_err("Failed to allocate memory for temporary page tables.\n");
                return -ENOMEM;
        }
 
-       rc = copy_page_tables(trans_pgd, start, end);
+       rc = copy_page_tables(info, trans_pgd, start, end);
        if (!rc)
                *dst_pgdp = trans_pgd;