Make VM init code more robust

Now it can handle the case that the physical load address overlaps the
desired virtual address.

arch/riscv/kernel/head.S

@@ -23,36 +23,53 @@ ENTRY(_start)
 	/* By SBI convention, hart ID is in a0 */
 	bnez a0, .Lsecondary_start
 
-	/* Clear the .bss segment */
-	la a0, __bss_start
-	li a1, 0
-	la a2, __bss_stop
-	sub a2, a2, a0
-	call memset
-
-	/* Initialize stack pointer */
+	/* Initialize page tables and relocate to virtual addresses */
 	la sp, init_thread_union + THREAD_SIZE
+	call setup_vm
+	call relocate
 
-	/* Initialize virtual memory */
+	/* Start the kernel */
+	la sp, init_thread_union + THREAD_SIZE
+	la tp, init_task
+	tail start_kernel
+
+relocate:
+	/* Relocate return address */
+	li a1, PAGE_OFFSET
+	la a0, _start
+	sub a1, a1, a0
+	add ra, ra, a1
+
+	/* Point stvec to virtual address of intruction after sptbr write */
 	la a0, 1f
-	tail setup_vm
+	add a0, a0, a1
+	csrw stvec, a0
 
+	/* Compute sptbr for kernel page tables, but don't load it yet */
+	la a2, swapper_pg_dir
+	srl a2, a2, PAGE_SHIFT
+	li a1, SPTBR_MODE
+	or a2, a2, a1
+
+	/* Load trampoline page directory, which will cause us to trap to
+	   stvec if VA != PA, or simply fall through if VA == PA */
+	la a0, trampoline_pg_dir
+	srl a0, a0, PAGE_SHIFT
+	or a0, a0, a1
+	sfence.vma
+	csrw sptbr, a0
 1:
-	/* Temporary trap vector */
+	/* Set trap vector to spin forever to help debug */
 	la a0, .Lsecondary_park
 	csrw stvec, a0
 
 	/* Load the global pointer */
-	auipc gp, %pcrel_hi(_gp)
-	addi gp, gp, %pcrel_lo(1b)
-
-	/* Re-initialize stack pointer */
-	la sp, init_thread_union + THREAD_SIZE
+	la gp, _gp
 
-	/* Initialize current task_struct pointer */
-	la tp, init_task
+	/* Switch to kernel page tables */
+	csrw sptbr, a2
 
-	tail start_kernel
+	ret
 
 .Lsecondary_start:
 #ifdef CONFIG_SMP
@@ -68,6 +85,9 @@ ENTRY(_start)
 	beqz sp, .Lwait_for_cpu_up
 	fence
 
+	/* Enable virtual memory and relocate to virtual address */
+	call relocate
+
 	/* Initialize task_struct pointer */
 	li tp, -THREAD_SIZE
 	add tp, tp, sp

arch/riscv/kernel/setup.c

@@ -85,13 +85,17 @@ early_param("mem", early_mem);
 pgd_t swapper_pg_dir[PTRS_PER_PGD] __page_aligned_bss;
 pmd_t swapper_pmd[PTRS_PER_PMD] __page_aligned_bss;
 
-asmlinkage void __init setup_vm(uintptr_t next_fn)
+pgd_t trampoline_pg_dir[PTRS_PER_PGD] __initdata __aligned(PAGE_SIZE);
+pmd_t trampoline_pmd[PTRS_PER_PGD] __initdata __aligned(PAGE_SIZE);
+
+asmlinkage void __init setup_vm(void)
 {
 	extern char _start;
 	uintptr_t i;
 	uintptr_t pa = (uintptr_t) &_start;
 	uintptr_t size = roundup((uintptr_t) _end - pa, PMD_SIZE);
 	uintptr_t pmd_pfn = PFN_DOWN((uintptr_t)swapper_pmd);
+	uintptr_t prot = pgprot_val(PAGE_KERNEL) | _PAGE_EXEC;
 	pgd_t pmd = __pgd((pmd_pfn << _PAGE_PFN_SHIFT) | _PAGE_TABLE);
 
 	/* Sanity check alignment and size */
@@ -103,15 +107,12 @@ asmlinkage void __init setup_vm(uintptr_t next_fn)
 	pfn_base = PFN_DOWN(pa);
 	mem_size = size;
 
-	swapper_pg_dir[((uintptr_t)PAGE_OFFSET >> PGDIR_SHIFT) % PTRS_PER_PGD] = pmd;
-	for (i = 0; i < size / PMD_SIZE; i++, pa += PMD_SIZE) {
-		uintptr_t prot = pgprot_val(PAGE_KERNEL) | _PAGE_EXEC;
-		swapper_pmd[i] = __pmd((PFN_DOWN(pa) << _PAGE_PFN_SHIFT) | prot);
-	}
+	trampoline_pg_dir[(PAGE_OFFSET >> PGDIR_SHIFT) % PTRS_PER_PGD] = pmd;
+	trampoline_pmd[0] = __pmd((PFN_DOWN(pa) << _PAGE_PFN_SHIFT) | prot);
 
-	csr_write(stvec, next_fn + va_pa_offset);
-	local_flush_tlb_all();
-	csr_write(sptbr, PFN_DOWN((uintptr_t)swapper_pg_dir) | SPTBR_MODE);
+	swapper_pg_dir[(PAGE_OFFSET >> PGDIR_SHIFT) % PTRS_PER_PGD] = pmd;
+	for (i = 0; i < size / PMD_SIZE; i++, pa += PMD_SIZE)
+		swapper_pmd[i] = __pmd((PFN_DOWN(pa) << _PAGE_PFN_SHIFT) | prot);
 }
 
 static void __init setup_bootmem(void)