From e7fe5a8f8e1f4abdc76e01bbd50dd8285c5c56dc Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Sat, 22 Apr 2023 21:07:34 +0100
Subject: [PATCH] kernel: add function to query CPU clock cycle count

---
 arch/x86_64/cpu_ctrl.S   | 10 ++++++++++
 include/socks/compiler.h |  3 +++
 include/socks/cpu.h      |  4 ++++
 include/socks/types.h    |  1 +
 kernel/cpu.c             |  5 +++++
 5 files changed, 23 insertions(+)

diff --git a/arch/x86_64/cpu_ctrl.S b/arch/x86_64/cpu_ctrl.S
index c40ca13..85fdd49 100644
--- a/arch/x86_64/cpu_ctrl.S
+++ b/arch/x86_64/cpu_ctrl.S
@@ -38,3 +38,13 @@ wrmsr:
 	wrmsr
 
 	ret
+
+
+	.global get_cycles
+	.type get_cycles, @function
+
+get_cycles:
+	rdtsc
+	shl $32, %rdx
+	orq %rdx, %rax
+	ret
diff --git a/include/socks/compiler.h b/include/socks/compiler.h
index 7a21e58..18fb198 100644
--- a/include/socks/compiler.h
+++ b/include/socks/compiler.h
@@ -35,4 +35,7 @@ void write_once(volatile T *ptr, T value)
 #undef __aligned
 #define __aligned(x) __attribute__((__aligned__(x)))
 
+#undef __weak
+#define __weak __attribute__((__weak__))
+
 #endif
diff --git a/include/socks/cpu.h b/include/socks/cpu.h
index 33016a6..e676a97 100644
--- a/include/socks/cpu.h
+++ b/include/socks/cpu.h
@@ -1,7 +1,9 @@
 #ifndef SOCKS_CPU_H_
 #define SOCKS_CPU_H_
 
+#include <socks/types.h>
 #include <socks/machine/cpu.h>
+#include <stdint.h>
 #include <socks/sched.h>
 
 #ifdef __cplusplus
@@ -33,6 +35,8 @@ extern void put_cpu(struct cpu_data *cpu);
 extern void cpu_set_available(unsigned int cpu_id);
 extern void cpu_set_online(unsigned int cpu_id);
 
+extern cycles_t get_cycles(void);
+
 #define irq_enable() ml_int_enable()
 #define irq_disable() ml_int_disable()
 
diff --git a/include/socks/types.h b/include/socks/types.h
index d8a9c4b..c3ec830 100644
--- a/include/socks/types.h
+++ b/include/socks/types.h
@@ -4,5 +4,6 @@
 #include <stdint.h>
 
 typedef uintptr_t phys_addr_t;
+typedef uint64_t cycles_t;
 
 #endif
diff --git a/kernel/cpu.c b/kernel/cpu.c
index 0942ab9..e239fec 100644
--- a/kernel/cpu.c
+++ b/kernel/cpu.c
@@ -69,3 +69,8 @@ unsigned int cpu_get_highest_available(void)
 {
 	return bitmap_highest_set(cpu_available, CPU_MAX);
 }
+
+cycles_t __weak get_cycles(void)
+{
+	return 0;
+}