static void __init pmd_swap_tests(struct pgtable_debug_args *args) { }
 #endif /* CONFIG_ARCH_ENABLE_THP_MIGRATION */
 
-static void __init swap_migration_tests(void)
+static void __init swap_migration_tests(struct pgtable_debug_args *args)
 {
        struct page *page;
        swp_entry_t swp;
        if (!IS_ENABLED(CONFIG_MIGRATION))
                return;
 
-       pr_debug("Validating swap migration\n");
        /*
         * swap_migration_tests() requires a dedicated page as it needs to
         * be locked before creating a migration entry from it. Locking the
         * page that actually maps kernel text ('start_kernel') can be real
-        * problematic. Lets allocate a dedicated page explicitly for this
-        * purpose that will be freed subsequently.
+        * problematic. Lets use the allocated page explicitly for this
+        * purpose.
         */
-       page = alloc_page(GFP_KERNEL);
-       if (!page) {
-               pr_err("page allocation failed\n");
+       page = (args->pte_pfn != ULONG_MAX) ? pfn_to_page(args->pte_pfn) : NULL;
+       if (!page)
                return;
-       }
+
+       pr_debug("Validating swap migration\n");
 
        /*
         * make_migration_entry() expects given page to be
        WARN_ON(!is_migration_entry(swp));
        WARN_ON(is_writable_migration_entry(swp));
        __ClearPageLocked(page);
-       __free_page(page);
 }
 
 #ifdef CONFIG_HUGETLB_PAGE
 #endif /* CONFIG_HUGETLB_PAGE */
 
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
-static void __init pmd_thp_tests(unsigned long pfn, pgprot_t prot)
+static void __init pmd_thp_tests(struct pgtable_debug_args *args)
 {
        pmd_t pmd;
 
         * needs to return true. pmd_present() should be true whenever
         * pmd_trans_huge() returns true.
         */
-       pmd = pfn_pmd(pfn, prot);
+       pmd = pfn_pmd(args->fixed_pmd_pfn, args->page_prot);
        WARN_ON(!pmd_trans_huge(pmd_mkhuge(pmd)));
 
 #ifndef __HAVE_ARCH_PMDP_INVALIDATE
 }
 
 #ifdef CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD
-static void __init pud_thp_tests(unsigned long pfn, pgprot_t prot)
+static void __init pud_thp_tests(struct pgtable_debug_args *args)
 {
        pud_t pud;
 
                return;
 
        pr_debug("Validating PUD based THP\n");
-       pud = pfn_pud(pfn, prot);
+       pud = pfn_pud(args->fixed_pud_pfn, args->page_prot);
        WARN_ON(!pud_trans_huge(pud_mkhuge(pud)));
 
        /*
         */
 }
 #else  /* !CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD */
-static void __init pud_thp_tests(unsigned long pfn, pgprot_t prot) { }
+static void __init pud_thp_tests(struct pgtable_debug_args *args) { }
 #endif /* CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD */
 #else  /* !CONFIG_TRANSPARENT_HUGEPAGE */
-static void __init pmd_thp_tests(unsigned long pfn, pgprot_t prot) { }
-static void __init pud_thp_tests(unsigned long pfn, pgprot_t prot) { }
+static void __init pmd_thp_tests(struct pgtable_debug_args *args) { }
+static void __init pud_thp_tests(struct pgtable_debug_args *args) { }
 #endif /* CONFIG_TRANSPARENT_HUGEPAGE */
 
 static unsigned long __init get_random_vaddr(void)
        pte_swap_tests(&args);
        pmd_swap_tests(&args);
 
-       swap_migration_tests();
+       swap_migration_tests(&args);
 
-       pmd_thp_tests(pmd_aligned, prot);
-       pud_thp_tests(pud_aligned, prot);
+       pmd_thp_tests(&args);
+       pud_thp_tests(&args);
 
        hugetlb_basic_tests(&args);