From c803a3e2b2b74a6a21594305d7e5789ad828ce2e Mon Sep 17 00:00:00 2001
From: Mark Poliakov <mark@alnyan.me>
Date: Mon, 17 Mar 2025 17:26:40 +0200
Subject: [PATCH] Implement thread-locals for per-CPU data

---
 etc/riscv64-unknown-none.ld | 12 ++++++++++++
 src/arch/riscv64.zig        |  6 ++++++
 src/arch/riscv64/boot.zig   | 21 ++++++++++++++-------
 src/kernel.zig              |  6 ------
 src/thread.zig              | 32 ++++++++++++++++++++++++++++++++
 5 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/etc/riscv64-unknown-none.ld b/etc/riscv64-unknown-none.ld
index 8a1e884..446b96b 100644
--- a/etc/riscv64-unknown-none.ld
+++ b/etc/riscv64-unknown-none.ld
@@ -21,6 +21,12 @@ SECTIONS {
         *(.dynamic*)
     }
 
+    .tdata : ALIGN(4K) {
+        PROVIDE(__tdata_start = .);
+        *(.tdata*)
+        PROVIDE(__tdata_end = .);
+    }
+
     .rela : ALIGN(4K) {
         PROVIDE(__rela_start = .);
         *(.rela*)
@@ -42,6 +48,12 @@ SECTIONS {
         PROVIDE(__bss_end = .);
     }
 
+    .tbss : {
+        PROVIDE(__tbss_start = .);
+        *(.tbss*)
+        PROVIDE(__tbss_end = .);
+    }
+
     . = ALIGN(4K);
     PROVIDE(__kernel_end = .);
 }
diff --git a/src/arch/riscv64.zig b/src/arch/riscv64.zig
index a613e89..4d40434 100644
--- a/src/arch/riscv64.zig
+++ b/src/arch/riscv64.zig
@@ -14,6 +14,8 @@ extern fn __rv64_task_enter_kernel() callconv(.C) noreturn;
 
 pub fn arch() type {
     return struct {
+        pub threadlocal var tHartId: u32 = 0;
+
         pub const Context = extern struct {
             const STACK_SIZE: usize = 8192;
 
@@ -101,5 +103,9 @@ pub fn arch() type {
             }
             asm volatile ("":::"memory");
         }
+
+        pub inline fn setThreadPointer(tp: usize) void {
+            asm volatile ("mv tp, %[tp]"::[tp]"r"(tp):"memory");
+        }
     };
 }
diff --git a/src/arch/riscv64/boot.zig b/src/arch/riscv64/boot.zig
index 7b5ac75..bb463d1 100644
--- a/src/arch/riscv64/boot.zig
+++ b/src/arch/riscv64/boot.zig
@@ -18,7 +18,7 @@ extern const __kernel_start: u8;
 extern const __kernel_end: u8;
 
 var gDtbAddress: usize = 0;
-var gBspHartId: usize = 0;
+var gBspHartId: u32 = 0;
 
 pub export fn rv64RelocateKernel(imageBase: usize, relaStart: usize, relaEnd: usize) void {
     const elf = @import("std").elf;
@@ -27,11 +27,15 @@ pub export fn rv64RelocateKernel(imageBase: usize, relaStart: usize, relaEnd: us
     const relaCount = (relaEnd - relaStart) / @sizeOf(elf.Rela);
     const relaTable = relaTablePtr[0..relaCount];
     for (relaTable) |entry| {
-        if (entry.r_type() == 0x03) {
-            const value = @as(*isize, @ptrFromInt(imageBase + entry.r_offset));
-            value.* = @as(isize, @bitCast(imageBase)) + entry.r_addend;
-        } else {
-            arch.halt();
+        const relaType: elf.R_RISCV = @enumFromInt(entry.r_type());
+        switch (relaType) {
+            .RELATIVE => {
+                const value = @as(*isize, @ptrFromInt(imageBase + entry.r_offset));
+                value.* = @as(isize, @bitCast(imageBase)) + entry.r_addend;
+            },
+            else => {
+                arch.halt();
+            },
         }
     }
 }
@@ -94,6 +98,9 @@ fn bspUpperEntry(realAddress: usize, unused: usize) callconv(.C) noreturn {
     // Setup physical memory management
     setupMemoryFromFdt(realAddress);
 
+    kernel.thread.setupCurrentCpu();
+    arch.tHartId = gBspHartId;
+
     kernel.kernel_main();
 }
 
@@ -115,7 +122,7 @@ pub export fn rv64BspLowerEntry(realAddress: usize, bspHartId: usize, dtbAddress
     debug.log.setWriteFn(&sbi.debugPrintByte);
 
     gDtbAddress = dtbAddress;
-    gBspHartId = bspHartId;
+    gBspHartId = @truncate(bspHartId);
 
     vmm.mapEarly(realAddress);
 
diff --git a/src/kernel.zig b/src/kernel.zig
index ef82599..ba0ee4d 100644
--- a/src/kernel.zig
+++ b/src/kernel.zig
@@ -34,12 +34,6 @@ pub export fn kernel_main() callconv(.C) noreturn {
         thread.addThread(t);
     }
 
-
-    const ptr: *const u32 = @ptrFromInt(0x1111111111111114);
-    const z = ptr.*;
-    _ = z;
-    while (true) {}
-
     thread.enter();
 
     arch.halt();
diff --git a/src/thread.zig b/src/thread.zig
index c03c508..23dbfde 100644
--- a/src/thread.zig
+++ b/src/thread.zig
@@ -1,3 +1,5 @@
+const std = @import("std");
+
 const arena = @import("arena.zig");
 const arch = @import("kernel.zig").arch;
 const log = @import("debug.zig").log;
@@ -90,3 +92,33 @@ pub fn yield() void {
         next.switchFrom(curr);
     }
 }
+
+extern var __tdata_start: u8;
+extern var __tdata_end: u8;
+extern var __tbss_start: u8;
+extern var __tbss_end: u8;
+
+pub fn setupCurrentCpu() void {
+    // Assume .tbss follows .tdata
+    const tdataStart = @intFromPtr(&__tdata_start);
+    const tdataEnd = @intFromPtr(&__tdata_end);
+    const tdataSize = tdataEnd - tdataStart;
+    const tbssStart = @intFromPtr(&__tbss_start);
+    const tbssEnd = @intFromPtr(&__tbss_end);
+    const tbssSize = tbssEnd - tbssStart;
+
+    const tdataData = @as([*]u8, @ptrFromInt(tdataStart))[0..tdataSize];
+
+    const tlsSize = tdataSize + tbssSize;
+    const tlsPageCount = (tlsSize + mem.vmm.PAGE_SIZE - 1) / mem.vmm.PAGE_SIZE;
+    // Variant I: TLS block 0 follows TP after a certain displacement
+    const tlsAddress = mem.phys.alloc_pages(tlsPageCount).?.virtualize();
+    const tlsData = @as([*]u8, @ptrFromInt(tlsAddress))[0..tlsSize];
+
+    log.info("Allocated TLS @ {*}", .{ tlsData });
+
+    @memcpy(tlsData[0..tdataSize], tdataData);
+    @memset(tlsData[tdataSize..], 0);
+
+    arch.setThreadPointer(tlsAddress);
+}