From 6503f98d55e49adb00767fc17c751ec7c3f74143 Mon Sep 17 00:00:00 2001
From: Luca Colagrande <luca.colagrande3@gmail.com>
Date: Sat, 23 Dec 2023 18:32:20 +0100
Subject: [PATCH] sw: Align with reduction-based exit routine

---
 target/sim/run.py                               |  3 +--
 target/sim/sw/device/runtime/src/occamy_start.c | 14 ++++++++++++--
 target/sim/sw/host/apps/offload/src/offload.c   |  7 ++-----
 target/sim/sw/host/runtime/host.c               |  8 +++++++-
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/target/sim/run.py b/target/sim/run.py
index fd263be3a..45681c7ce 100755
--- a/target/sim/run.py
+++ b/target/sim/run.py
@@ -20,10 +20,9 @@
 
 def main():
     args = parser('vsim', SIMULATORS.keys()).parse_args()
-    simulations = get_simulations(args.testlist, SIMULATORS[args.simulator])
+    simulations = get_simulations(args.testlist, SIMULATORS[args.simulator], run_dir=args.run_dir)
     return run_simulations(simulations,
                            n_procs=args.n_procs,
-                           run_dir=Path(args.run_dir),
                            dry_run=args.dry_run,
                            early_exit=args.early_exit)
 
diff --git a/target/sim/sw/device/runtime/src/occamy_start.c b/target/sim/sw/device/runtime/src/occamy_start.c
index 0dec4735b..e16e9812e 100644
--- a/target/sim/sw/device/runtime/src/occamy_start.c
+++ b/target/sim/sw/device/runtime/src/occamy_start.c
@@ -10,12 +10,22 @@
 #define SNRT_CRT0_PRE_BARRIER
 #define SNRT_INVOKE_MAIN
 #define SNRT_CRT0_POST_BARRIER
-#define SNRT_CRT0_CALLBACK7
+#define SNRT_CRT0_EXIT
+#define SNRT_CRT0_ALTERNATE_EXIT
 
 static inline void snrt_crt0_callback3() {
     _snrt_cluster_hw_barrier = cluster_hw_barrier_addr(snrt_cluster_idx());
 }
 
-static inline void snrt_crt0_callback7() { return_to_cva6(SYNC_CLUSTERS); }
+static inline uint32_t* snrt_exit_code_destination() {
+    return soc_ctrl_scratch_ptr(3);
+}
+
+static inline void snrt_exit_default(int exit_code);
+
+static inline void snrt_exit(int exit_code) {
+    snrt_exit_default(exit_code);
+    if (snrt_global_core_idx() == 0) set_host_sw_interrupt();
+}
 
 #include "start.c"
diff --git a/target/sim/sw/host/apps/offload/src/offload.c b/target/sim/sw/host/apps/offload/src/offload.c
index dab7df9b1..4717d1b83 100644
--- a/target/sim/sw/host/apps/offload/src/offload.c
+++ b/target/sim/sw/host/apps/offload/src/offload.c
@@ -18,9 +18,6 @@ int main() {
     // Start Snitches
     wakeup_snitches_cl();
 
-    // Wait for job done
-    wait_snitches_done();
-
-    // Exit routine
-    mcycle();
+    // Wait for job done and return Snitch exit code
+    return wait_snitches_done();
 }
diff --git a/target/sim/sw/host/runtime/host.c b/target/sim/sw/host/runtime/host.c
index 78bc4c830..454cf4a3d 100644
--- a/target/sim/sw/host/runtime/host.c
+++ b/target/sim/sw/host/runtime/host.c
@@ -258,9 +258,15 @@ void wakeup_master_snitches() {
 /**
  * @brief Waits until snitches are done executing
  */
-static inline void wait_snitches_done() {
+static inline int wait_snitches_done() {
     wait_sw_interrupt();
     clear_host_sw_interrupt();
+    int retval = *soc_ctrl_scratch_ptr(3);
+    // LSB signals completion
+    if (retval & 1)
+        return retval >> 1;
+    else
+        return -1;
 }
 
 static inline volatile uint32_t* get_shared_lock() {