From aefafc45e98fe121d8ec13fd06573b079a7b7c71 Mon Sep 17 00:00:00 2001
From: Ville Juven <ville.juven@unikie.com>
Date: Wed, 22 Jan 2025 13:13:54 +0200
Subject: [PATCH] riscv/context_switch: Set tp when a context switch occurs

This removes the need to call nxsched_self upon exception return. Later
with a bit of added logic, it is possible to use tp in kernel to store and
read the current tcb.

Note: setting tp in riscv_restorecontext is too late for this, as
this_task() points to the (next) ready-to-run task, thus tp should be
updated when the rtr list is updated (via up_update_task, which does not
exist for risc-v yet)
---
 arch/risc-v/src/common/riscv_exception_common.S | 9 ++++-----
 arch/risc-v/src/common/riscv_initialstate.c     | 6 ++++++
 arch/risc-v/src/common/riscv_internal.h         | 6 ++++++
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/arch/risc-v/src/common/riscv_exception_common.S b/arch/risc-v/src/common/riscv_exception_common.S
index a8326c5eb46..2c60e0c4e1e 100644
--- a/arch/risc-v/src/common/riscv_exception_common.S
+++ b/arch/risc-v/src/common/riscv_exception_common.S
@@ -153,6 +153,9 @@ exception_common:
   REGSTORE   s3, REG_SP(sp)
 
 #ifdef CONFIG_LIB_SYSCALL
+  csrr       tp, CSR_SCRATCH      /* Load kernel TP */
+  REGLOAD    tp, RISCV_PERCPU_TCB(tp)
+
   /* Check whether it is an exception or interrupt */
 
   blt        s2, x0, handle_irq   /* If cause < 0 it is interrupt */
@@ -177,9 +180,6 @@ exception_common:
   addi       s1, s1, 0x4          /* Must move EPC forward by +4 */
   REGSTORE   s1, REG_EPC(sp)      /* Updated EPC to user context */
 
-  csrr       tp, CSR_SCRATCH      /* Load kernel TP */
-  REGLOAD    tp, RISCV_PERCPU_TCB(tp)
-
   mv         a7, sp               /* a7 = context */
   call       x1, dispatch_syscall /* Dispatch the system call */
 
@@ -246,9 +246,8 @@ return_from_exception:
 #ifdef CONFIG_LIB_SYSCALL
   /* Store tcb to scratch register */
 
-  call       x1, nxsched_self
   csrr       s1, CSR_SCRATCH
-  REGSTORE   a0, RISCV_PERCPU_TCB(s1)
+  REGSTORE   tp, RISCV_PERCPU_TCB(s1)
 #endif
 
 #ifdef CONFIG_ARCH_KERNEL_STACK
diff --git a/arch/risc-v/src/common/riscv_initialstate.c b/arch/risc-v/src/common/riscv_initialstate.c
index 8a3527204ca..5f65b3bcf6e 100644
--- a/arch/risc-v/src/common/riscv_initialstate.c
+++ b/arch/risc-v/src/common/riscv_initialstate.c
@@ -115,6 +115,12 @@ void up_initial_state(struct tcb_s *tcb)
       /* Set idle process' initial interrupt context */
 
       riscv_set_idleintctx();
+
+#ifdef CONFIG_LIB_SYSCALL
+      /* Update current thread pointer */
+
+      __asm__ __volatile__("mv tp, %0" : : "r"(tcb));
+#endif
       return;
     }
 
diff --git a/arch/risc-v/src/common/riscv_internal.h b/arch/risc-v/src/common/riscv_internal.h
index 9e72f5bb8ba..f9208ddf5c6 100644
--- a/arch/risc-v/src/common/riscv_internal.h
+++ b/arch/risc-v/src/common/riscv_internal.h
@@ -282,6 +282,12 @@ static inline void riscv_restorecontext(struct tcb_s *tcb)
 
   riscv_restorevpu(tcb->xcp.regs, riscv_vpuregs(tcb));
 #endif
+
+#ifdef CONFIG_LIB_SYSCALL
+  /* Update current thread pointer */
+
+  __asm__ __volatile__("mv tp, %0" : : "r"(tcb));
+#endif
 }
 
 #ifdef CONFIG_ARCH_RISCV_INTXCPT_EXTENSIONS