more fixes

arch/x86_64/include/bits/sched.h

@@ -1,10 +1,12 @@
 #pragma once
 
+#include <zinnia/types.h>
 #include <stdint.h>
 
 struct arch_task_context {
     uint64_t rsp;
     uint8_t* fpu_region;
+    phys_t fpu_phys;
     uint16_t ds, es, fs, gs;
     uint64_t fs_base, gs_base;
 };

arch/x86_64/include/x86_64/defs.h

@@ -185,6 +185,7 @@
 #define MSR_EFER_FFXSR (1 << 14)
 #define MSR_EFER_TCE   (1 << 15)
 
+#define MSR_PAT            (0x277)
 #define MSR_EFER           (0xC0000080)
 #define MSR_STAR           (0xC0000081)
 #define MSR_LSTAR          (0xC0000082)
@@ -194,6 +195,13 @@
 #define MSR_GS_BASE        (0xC0000101)
 #define MSR_KERNEL_GS_BASE (0xC0000102)
 
+#define PAT_UC  (0x0)
+#define PAT_WC  (0x1)
+#define PAT_WT  (0x4)
+#define PAT_WP  (0x5)
+#define PAT_WB  (0x6)
+#define PAT_UCM (0x7)
+
 #define CPL_USER   (0x3)
 #define CPL_KERNEL (0x0)
 

arch/x86_64/include/x86_64/pmap.h

@@ -12,7 +12,7 @@ enum : pte_t {
     ARCH_FLAG_SIZE = 1 << 7,
     ARCH_FLAG_GLOBAL = 1 << 8,
     ARCH_FLAG_AVAILABLE = 1 << 9,
-    ARCH_FLAG_ATTRIBUTE_TABLE = 1 << 10,
+    ARCH_FLAG_ATTRIBUTE_TABLE = 1 << 7,
     ARCH_FLAG_EXECUTE_DISABLE = 1LU << 63,
 
     ARCH_PTE_ADDR_MASK = 0x000F'FFFF'FFFF'F000,

arch/x86_64/percpu.c

@@ -56,6 +56,12 @@ void x86_64_cpu_setup(struct percpu* cpu) {
     // TODO: not if FRED
     idt_load();
 
+    asm_wrmsr(
+        MSR_PAT,
+        PAT_WB | PAT_WT << 8 | PAT_UCM << 16 | PAT_UC << 24 | (uint64_t)PAT_WP << 32 | (uint64_t)PAT_WC << 40 |
+            (uint64_t)PAT_UCM << 48 | (uint64_t)PAT_UC << 56
+    );
+
     // Syscall extension.
     asm_wrmsr(MSR_EFER, asm_rdmsr(MSR_EFER) | MSR_EFER_SCE);
     union ia32_star star = {0};

arch/x86_64/pmap.c

@@ -6,6 +6,10 @@ void pte_clear(pte_t* pte) {
     *pte = 0;
 }
 
+void pte_flush(uintptr_t vaddr) {
+    asm volatile("invlpg [%0]" ::"r"(vaddr) : "memory");
+}
+
 pte_t pte_build(phys_t addr, enum pte_flags flags, enum cache_mode cache) {
     pte_t result = ((pte_t)addr & ARCH_PTE_ADDR_MASK) | ARCH_FLAG_PRESENT;
 
@@ -19,6 +23,21 @@ pte_t pte_build(phys_t addr, enum pte_flags flags, enum cache_mode cache) {
             result |= ARCH_FLAG_READ_WRITE;
         if (!(flags & PTE_EXEC))
             result |= ARCH_FLAG_EXECUTE_DISABLE;
+
+        switch (cache) {
+        case CACHE_ANY:
+        case CACHE_WRITE_BACK:
+            break;
+        case CACHE_WRITE_THROUGH:
+            result |= ARCH_FLAG_WRITE_THROUGH;
+            break;
+        case CACHE_MMIO:
+            result |= ARCH_FLAG_CACHE_DISABLE | ARCH_FLAG_WRITE_THROUGH;
+            break;
+        case CACHE_WRITE_COMBINE:
+            result |= ARCH_FLAG_ATTRIBUTE_TABLE | ARCH_FLAG_WRITE_THROUGH;
+            break;
+        }
     }
     return result;
 }

arch/x86_64/sched.c

@@ -1,4 +1,5 @@
 #include <zinnia/alloc.h>
+#include <zinnia/assert.h>
 #include <zinnia/compiler.h>
 #include <zinnia/errno.h>
 #include <zinnia/percpu.h>
@@ -58,7 +59,7 @@ struct task* arch_sched_switch(struct task* from, struct task* to) {
     struct arch_task_context* from_ctx = &from->context;
     struct arch_task_context* to_ctx = &to->context;
 
-    cpu->tss.rsp0 = to->kernel_stack + KERNEL_STACK_SIZE;
+    cpu->tss.rsp0 = kernel_stack_top(&to->kernel_stack);
 
     if (from_ctx->fpu_region != nullptr) {
         cpu->fpu_save(from_ctx->fpu_region);
@@ -90,15 +91,17 @@ struct task* arch_sched_switch(struct task* from, struct task* to) {
     return perform_switch(&from_ctx->rsp, to_ctx->rsp, from);
 }
 
+// Calls task_entry with the correct parameters.
 [[__naked]]
 void task_entry_thunk() {
     asm volatile(
         "mov rdi, rbx\n"
         "mov rsi, r12\n"
+        "mov rdx, rax\n"
         "jmp %0"
         :
         : "r"(task_entry)
-        : "memory"
+        : "memory", "rdi", "rsi", "rdx"
     );
 }
 
@@ -111,8 +114,10 @@ errno_t arch_task_init(
 ) {
     struct arch_percpu* cpu = percpu_ref(arch_data);
 
+    assert((stack_start & 0xF) == 0, "Task stack is not 16-byte aligned!\n");
+
     // Prepare a dummy stack with an entry point function to return to.
-    struct task_frame* frame = (struct task_frame*)stack_start - 1;
+    struct task_frame* frame = (struct task_frame*)(stack_start - sizeof(uint64_t)) - 1;
     frame->rbx = (uint64_t)entry;
     frame->r12 = arg0;
     frame->rip = (uint64_t)task_entry_thunk;
@@ -133,6 +138,7 @@ errno_t arch_task_init(
         for (size_t i = 0; i < fpu_size; i += arch_mem_page_size())
             pmap_map(&kernel_vma.pmap, fpu_base + i, pages + i, PTE_WRITE | PTE_READ, CACHE_WRITE_BACK);
         context->fpu_region = (void*)fpu_base;
+        context->fpu_phys = pages;
 
         context->ds = asm_read_ds();
         context->es = asm_read_es();
@@ -145,6 +151,21 @@ errno_t arch_task_init(
     return 0;
 }
 
+void arch_task_fini(struct arch_task_context* context) {
+    if (context->fpu_region == nullptr)
+        return;
+
+    const size_t page_size = arch_mem_page_size();
+    const size_t fpu_size = percpu_ref(arch_data)->fpu_size;
+    const uintptr_t fpu_base = (uintptr_t)context->fpu_region;
+    for (size_t i = 0; i < fpu_size; i += page_size)
+        pmap_unmap(&kernel_vma.pmap, fpu_base + i);
+    mem_phys_free(context->fpu_phys, ALIGN_UP(fpu_size, page_size) / page_size);
+    vm_free(fpu_base, fpu_size);
+
+    context->fpu_region = nullptr;
+}
+
 [[__naked]]
 void arch_sched_preempt_disable() {
     asm volatile(

arch/x86_64/smp.c

@@ -192,7 +192,7 @@ static bool start_ap(phys_t trampoline, uint32_t temp_cr3, uint32_t lapic_id) {
         return false;
     }
 
-    log("smp: Starting CPU %zu (LAPIC %u)\n", cpu->id, lapic_id);
+    log("smp: Starting CPU %zu\n", cpu->id);
 
     memcpy(HHDM_PTR(trampoline), smp_trampoline_start, trampoline_size);