arm64: Fix section mismatch on alloc_init_p[mu]d()

Commit 523d6e9fae93 (arm64:mm: free the useless initial page table)
introduced a BUG_ON checking for the allocation type but it was
referring the early_alloc() function in the __init section. This patch
changes the check to slab_is_available() and also relaxes the BUG to a
WARN_ON_ONCE.

Reported-by: Will Deacon <will.deacon@arm.com>
Acked-by: Will Deacon <will.deacon@arm.com>
Signed-off-by: Catalin Marinas <catalin.marinas@arm.com>

1 files changed, 5 insertions, 4 deletions

diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index 2eeac10aa0cf..c6daaf6c6f97 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -26,6 +26,7 @@
 #include <linux/memblock.h>
 #include <linux/fs.h>
 #include <linux/io.h>
+#include <linux/slab.h>
 #include <linux/stop_machine.h>
 
 #include <asm/cputype.h>
@@ -159,8 +160,8 @@ static void alloc_init_pmd(struct mm_struct *mm, pud_t *pud,
                                 flush_tlb_all();
                                 if (pmd_table(old_pmd)) {
                                         phys_addr_t table = __pa(pte_offset_map(&old_pmd, 0));
-                                        BUG_ON(alloc != early_alloc);
-                                        memblock_free(table, PAGE_SIZE);
+                                        if (!WARN_ON_ONCE(slab_is_available()))
+                                                memblock_free(table, PAGE_SIZE);
                                 }
                         }
                 } else {
@@ -220,8 +221,8 @@ static void alloc_init_pud(struct mm_struct *mm, pgd_t *pgd,
                                 flush_tlb_all();
                                 if (pud_table(old_pud)) {
                                         phys_addr_t table = __pa(pmd_offset(&old_pud, 0));
-                                        BUG_ON(alloc != early_alloc);
-                                        memblock_free(table, PAGE_SIZE);
+                                        if (!WARN_ON_ONCE(slab_is_available()))
+                                                memblock_free(table, PAGE_SIZE);
                                 }
                         }
                 } else {