[PATCH 2/3] LoongArch: Code cleanup in function pcpu_populate_pte

[Bibo Mao <maobibo@loongson.cn>]

There are some code cleanups in function pcpu_populate_pte:
1. Replace memblock_alloc with memblock_alloc_raw for pud and pmd since
they will be reinitialized with pud_init and pmd_init.

2. Add memory allocation failure handling

3. Replace pgd_populate with p4d_populate, it will be useful if there
four-level page tables.

Signed-off-by: Bibo Mao <maobibo@loongson.cn>
---
 arch/loongarch/kernel/numa.c | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/arch/loongarch/kernel/numa.c b/arch/loongarch/kernel/numa.c
index 708665895b47..778e1c20bfb0 100644
--- a/arch/loongarch/kernel/numa.c
+++ b/arch/loongarch/kernel/numa.c
@@ -73,33 +73,40 @@ void __init pcpu_populate_pte(unsigned long addr)
 	pmd_t *pmd;
 
 	if (p4d_none(*p4d)) {
-		pud_t *new;
-
-		new = memblock_alloc(PAGE_SIZE, PAGE_SIZE);
-		pgd_populate(&init_mm, pgd, new);
+		pud = memblock_alloc_raw(PAGE_SIZE, PAGE_SIZE);
+		if (!pud)
+			goto err_alloc;
+		p4d_populate(&init_mm, p4d, pud);
 #ifndef __PAGETABLE_PUD_FOLDED
-		pud_init(new);
+		pud_init(pud);
 #endif
 	}
 
 	pud = pud_offset(p4d, addr);
 	if (pud_none(*pud)) {
-		pmd_t *new;
-
-		new = memblock_alloc(PAGE_SIZE, PAGE_SIZE);
-		pud_populate(&init_mm, pud, new);
+		pmd = memblock_alloc_raw(PAGE_SIZE, PAGE_SIZE);
+		if (!pmd)
+			goto err_alloc;
+		pud_populate(&init_mm, pud, pmd);
 #ifndef __PAGETABLE_PMD_FOLDED
-		pmd_init(new);
+		pmd_init(pmd);
 #endif
 	}
 
 	pmd = pmd_offset(pud, addr);
 	if (!pmd_present(*pmd)) {
-		pte_t *new;
+		pte_t *pte;
 
-		new = memblock_alloc(PAGE_SIZE, PAGE_SIZE);
-		pmd_populate_kernel(&init_mm, pmd, new);
+		pte = memblock_alloc(PAGE_SIZE, PAGE_SIZE);
+		if (!pte)
+			goto err_alloc;
+		pmd_populate_kernel(&init_mm, pmd, pte);
 	}
+
+	return;
+
+err_alloc:
+	panic("%s: Failed to allocate memory\n", __func__);
 }
 
 void __init setup_per_cpu_areas(void)
-- 
2.27.0