WARN_ON(pte_young(pte));
 }
 
-static void __init pte_savedwrite_tests(unsigned long pfn, pgprot_t prot)
+static void __init pte_savedwrite_tests(struct pgtable_debug_args *args)
 {
-       pte_t pte = pfn_pte(pfn, prot);
+       pte_t pte = pfn_pte(args->fixed_pte_pfn, args->page_prot_none);
 
        if (!IS_ENABLED(CONFIG_NUMA_BALANCING))
                return;
        pgtable = pgtable_trans_huge_withdraw(mm, pmdp);
 }
 
-static void __init pmd_leaf_tests(unsigned long pfn, pgprot_t prot)
+static void __init pmd_leaf_tests(struct pgtable_debug_args *args)
 {
        pmd_t pmd;
 
                return;
 
        pr_debug("Validating PMD leaf\n");
-       pmd = pfn_pmd(pfn, prot);
+       pmd = pfn_pmd(args->fixed_pmd_pfn, args->page_prot);
 
        /*
         * PMD based THP is a leaf entry.
        WARN_ON(!pmd_leaf(pmd));
 }
 
-static void __init pmd_savedwrite_tests(unsigned long pfn, pgprot_t prot)
+static void __init pmd_savedwrite_tests(struct pgtable_debug_args *args)
 {
        pmd_t pmd;
 
                return;
 
        pr_debug("Validating PMD saved write\n");
-       pmd = pfn_pmd(pfn, prot);
+       pmd = pfn_pmd(args->fixed_pmd_pfn, args->page_prot_none);
        WARN_ON(!pmd_savedwrite(pmd_mk_savedwrite(pmd_clear_savedwrite(pmd))));
        WARN_ON(pmd_savedwrite(pmd_clear_savedwrite(pmd_mk_savedwrite(pmd))));
 }
        pudp_huge_get_and_clear(mm, vaddr, pudp);
 }
 
-static void __init pud_leaf_tests(unsigned long pfn, pgprot_t prot)
+static void __init pud_leaf_tests(struct pgtable_debug_args *args)
 {
        pud_t pud;
 
                return;
 
        pr_debug("Validating PUD leaf\n");
-       pud = pfn_pud(pfn, prot);
+       pud = pfn_pud(args->fixed_pud_pfn, args->page_prot);
        /*
         * PUD based THP is a leaf entry.
         */
                                      pgprot_t prot)
 {
 }
-static void __init pud_leaf_tests(unsigned long pfn, pgprot_t prot) { }
+static void __init pud_leaf_tests(struct pgtable_debug_args *args) { }
 #endif /* CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD */
 #else  /* !CONFIG_TRANSPARENT_HUGEPAGE */
 static void __init pmd_basic_tests(struct pgtable_debug_args *args, int idx) { }
                                      pgprot_t prot)
 {
 }
-static void __init pmd_leaf_tests(unsigned long pfn, pgprot_t prot) { }
-static void __init pud_leaf_tests(unsigned long pfn, pgprot_t prot) { }
-static void __init pmd_savedwrite_tests(unsigned long pfn, pgprot_t prot) { }
+static void __init pmd_leaf_tests(struct pgtable_debug_args *args) { }
+static void __init pud_leaf_tests(struct pgtable_debug_args *args) { }
+static void __init pmd_savedwrite_tests(struct pgtable_debug_args *args) { }
 #endif /* CONFIG_TRANSPARENT_HUGEPAGE */
 
 #ifdef CONFIG_HAVE_ARCH_HUGE_VMAP
        p4d_basic_tests(&args);
        pgd_basic_tests(&args);
 
-       pmd_leaf_tests(pmd_aligned, prot);
-       pud_leaf_tests(pud_aligned, prot);
+       pmd_leaf_tests(&args);
+       pud_leaf_tests(&args);
 
-       pte_savedwrite_tests(pte_aligned, protnone);
-       pmd_savedwrite_tests(pmd_aligned, protnone);
+       pte_savedwrite_tests(&args);
+       pmd_savedwrite_tests(&args);
 
        pte_special_tests(pte_aligned, prot);
        pte_protnone_tests(pte_aligned, protnone);