riscv: Setup exception vector for nommu platform

Exception vector is missing on nommu platform and that is an issue.
This patch is tested in Sipeed Maix Bit Dev Board.

Fixes: 79b1feba5455 ("RISC-V: Setup exception vector early")
Suggested-by: Anup Patel <anup@brainfault.org>
Suggested-by: Atish Patra <atishp@atishpatra.org>
Signed-off-by: Qiu Wenbo <qiuwenbo@phytium.com.cn>
Reviewed-by: Atish Patra <atish.patra@wdc.com>
Reviewed-by: Anup Patel <anup@brainfault.org>
Signed-off-by: Palmer Dabbelt <palmerdabbelt@google.com>
diff --git a/arch/riscv/kernel/head.S b/arch/riscv/kernel/head.S
index d0c5c31..0a4e81b 100644
--- a/arch/riscv/kernel/head.S
+++ b/arch/riscv/kernel/head.S
@@ -77,16 +77,10 @@
 	csrw CSR_SATP, a0
 .align 2
 1:
-	/* Set trap vector to exception handler */
-	la a0, handle_exception
+	/* Set trap vector to spin forever to help debug */
+	la a0, .Lsecondary_park
 	csrw CSR_TVEC, a0
 
-	/*
-	 * Set sup0 scratch register to 0, indicating to exception vector that
-	 * we are presently executing in kernel.
-	 */
-	csrw CSR_SCRATCH, zero
-
 	/* Reload the global pointer */
 .option push
 .option norelax
@@ -144,9 +138,23 @@
 	la a0, swapper_pg_dir
 	call relocate
 #endif
+	call setup_trap_vector
 	tail smp_callin
 #endif /* CONFIG_SMP */
 
+.align 2
+setup_trap_vector:
+	/* Set trap vector to exception handler */
+	la a0, handle_exception
+	csrw CSR_TVEC, a0
+
+	/*
+	 * Set sup0 scratch register to 0, indicating to exception vector that
+	 * we are presently executing in kernel.
+	 */
+	csrw CSR_SCRATCH, zero
+	ret
+
 .Lsecondary_park:
 	/* We lack SMP support or have too many harts, so park this hart */
 	wfi
@@ -240,6 +248,7 @@
 	call relocate
 #endif /* CONFIG_MMU */
 
+	call setup_trap_vector
 	/* Restore C environment */
 	la tp, init_task
 	sw zero, TASK_TI_CPU(tp)