From 36a077cd9e1957453b50dde442ef9a18d407af73 Mon Sep 17 00:00:00 2001
From: gtxzsxxk <2524395907@qq.com>
Date: Sat, 11 May 2024 22:14:01 +0800
Subject: [PATCH] port: refactor cpu yielding to keep low cpu usage when free
 and keep low latency

---
 include/port/os_yield_cpu.h |  8 +++++-
 src/machine.c               |  7 ++++-
 src/port/os_yield_cpu.c     | 57 +++++++++++++++++++++++++++++++++++--
 3 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/include/port/os_yield_cpu.h b/include/port/os_yield_cpu.h
index 2a30995..7f52a56 100644
--- a/include/port/os_yield_cpu.h
+++ b/include/port/os_yield_cpu.h
@@ -5,7 +5,13 @@
 #ifndef TEMU_OS_YIELD_CPU_H
 #define TEMU_OS_YIELD_CPU_H
 
-#define PORT_OS_YIELD_CPU_EVERY_MS  1
+#define PORT_OS_YIELD_CPU_EVERY_MS  50
+#define PORT_OS_YIELD_CPU_INTC_LEN  16
+
+void
+port_os_yield_cpu_init(uint64_t (*time_conv)(uint64_t), uint64_t *zicnt_timer_ticks, uint64_t *zicnt_timer_compare);
+
+void port_os_yield_cpu_add_interrupt(void (*intc)(void));
 
 void port_os_yield_cpu(void);
 
diff --git a/src/machine.c b/src/machine.c
index 0f11452..28c0337 100644
--- a/src/machine.c
+++ b/src/machine.c
@@ -11,6 +11,7 @@
 #include "zicsr.h"
 #include "port/console.h"
 #include "perf.h"
+#include "port/os_yield_cpu.h"
 
 //#define RISCV_ISA_TESTS
 
@@ -50,8 +51,12 @@ _Noreturn void machine_start(uint32_t start, int printreg) {
 static void machine_pre_boot(uint32_t start) {
     program_counter = start;
 
-    port_os_console_init();
+    zicnt_init();
     uart8250_init();
+
+    port_os_console_init();
+    port_os_yield_cpu_add_interrupt(zicnt_time_tick);
+    port_os_yield_cpu_add_interrupt(uart8250_tick);
 }
 
 static void machine_tick(void) {
diff --git a/src/port/os_yield_cpu.c b/src/port/os_yield_cpu.c
index 7fe2196..11d0b46 100644
--- a/src/port/os_yield_cpu.c
+++ b/src/port/os_yield_cpu.c
@@ -2,16 +2,69 @@
 // Created by hanyuan on 2024/5/11.
 //
 
+#include <stdint.h>
 #include <unistd.h>
+#include "trap.h"
 #include "port/system_timer.h"
 #include "port/os_yield_cpu.h"
 
+#ifndef NULL
+#define NULL (void*)0
+#endif
+
+static uint64_t *emulation_timer_tick = NULL;
+
+static uint64_t *emulation_timer_compare = NULL;
+
+static uint64_t (*emulation_timer_conv)(uint64_t) = NULL;
+
+static void (*interrupt_checks[PORT_OS_YIELD_CPU_INTC_LEN])(void) = {0};
+
+static uint8_t intc_counter = 0;
+
+void
+port_os_yield_cpu_init(uint64_t (*time_conv)(uint64_t), uint64_t *zicnt_timer_ticks, uint64_t *zicnt_timer_compare) {
+    emulation_timer_conv = time_conv;
+    emulation_timer_tick = zicnt_timer_ticks;
+    emulation_timer_compare = zicnt_timer_compare;
+}
+
+void port_os_yield_cpu_add_interrupt(void (*intc)(void)) {
+    interrupt_checks[intc_counter++] = intc;
+}
+
 void port_os_yield_cpu(void) {
 #ifndef BARE_METAL_PLATFORM
     static port_clock_t last_yield_tick = 0;
-    port_clock_t current_tick = port_system_timer_get_ticks();
+    port_clock_t before_sleep, after_sleep, sleep_counter = 0, current_tick = port_system_timer_get_ticks();
+
     if (current_tick - last_yield_tick > PORT_OS_YIELD_CPU_EVERY_MS * (PORT_CLOCKS_PER_SEC / 1000)) {
-        usleep(10000);
+        int64_t ticks_should_elapse = (int64_t) (*emulation_timer_compare - *emulation_timer_tick);
+
+        if (ticks_should_elapse > 0) {
+            uint64_t sleep_us = emulation_timer_conv(ticks_should_elapse);
+
+            if (sleep_us > 50) {
+                sleep_us /= 35;
+
+                while (sleep_counter < sleep_us) {
+                    before_sleep = port_system_timer_get_ticks() * (1000000 / PORT_CLOCKS_PER_SEC);
+                    usleep(50);
+
+                    /* After the cpu is resumed we should check if we have any interrupt immediately */
+                    for (uint8_t i = 0; i < intc_counter; i++) {
+                        interrupt_checks[i]();
+                    }
+                    if (trap_is_pending()) {
+                        return;
+                    }
+
+                    after_sleep = port_system_timer_get_ticks() * (1000000 / PORT_CLOCKS_PER_SEC);
+                    sleep_counter += after_sleep - before_sleep;
+                }
+            }
+        }
+
         last_yield_tick = current_tick;
     }
 #endif