nt-lite/code/kernel/ke_smp.c

// SMP bring-up and CPU topology management. This file owns secondary-processor admission and topology-derived scheduling metadata.

// The bootloader is intentionally not involved in secondary-processor ownership.
// The kernel asks the HAL for platform processor identities, prepares the x86
// bootstrap state, asks the HAL to start each processor, and waits for it to
// enter KiApKernelEntry() and park in kernel-owned state.
#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include <string.h>
#include <intrin.h>
#include <ntlite/status.h>
#include <kernel/ke_hal.h>
#include <kernel/ke_smp.h>
#include <kernel/ke_cpu.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_idt.h>
#include <kernel/ke_sched.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_event.h>
#include <kernel/ob_object.h>
#include <kernel/mm_phys.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_virt.h>

#define AP_TRAMPOLINE_MAX_PHYS       0x000A0000ULL
#define AP_TRAMPOLINE_PREFERRED_MIN  0x00070000ULL
#define AP_TRAMPOLINE_FALLBACK_MIN   0x00010000ULL
#define AP_TRAMPOLINE_BOOT_OFF       0x0800u
#define AP_STACK_PAGES               8u
#define AP_START_TIMEOUT_US          1000000ull
#define AP_SCHED_TIMEOUT_US          1000000ull

#define AP_CR4_PAE                   0x00000020u
#define AP_CR4_OSFXSR                0x00000200u
#define AP_CR4_OSXMMEXCPT            0x00000400u
#define AP_CR4_REQUIRED              (AP_CR4_PAE | AP_CR4_OSFXSR | AP_CR4_OSXMMEXCPT)
#define AP_CR0_PE                    0x00000001u
#define AP_CR0_MP                    0x00000002u
#define AP_CR0_EM                    0x00000004u
#define AP_CR0_TS                    0x00000008u
#define AP_CR0_NE                    0x00000020u
#define AP_CR0_PG                    0x80000000u

#define AP_BOOT_MAGIC                0x3436544F4F425041ULL  // 'APBOOT64'
#define AP_STATE_EMPTY               0u
#define AP_STATE_READY               1u
#define AP_STATE_LONG_MODE           2u
#define AP_STATE_KERNEL_ENTRY        3u
#define AP_STATE_KERNEL_PARKED       4u
#define AP_STATE_FAILED              0xFFFFFFFFu

// Returns the HAL monotonic counter in microseconds for pre-scheduler hardware deadlines.
static uint64_t KiSmpQueryTicks(void) {
    const HAL_API* hal = KeGetHal();
    return hal && hal->Timer.QueryTicks ? hal->Timer.QueryTicks() : 0;
}

// Returns a saturated deadline from the current HAL monotonic counter.
static uint64_t KiSmpDeadlineAfter(uint64_t timeoutUs) {
    uint64_t now = KiSmpQueryTicks();
    return now > UINT64_MAX - timeoutUs ? UINT64_MAX : now + timeoutUs;
}

#define AP_STAGE_NONE                0x00000000u
#define AP_STAGE_LONG64              0x4C4D3634u  // 'LM64'
#define AP_STAGE_ENTRY               0x454E5452u  // 'ENTR'
#define AP_STAGE_GS                  0x47534241u  // 'GSBA'
#define AP_STAGE_GDT                 0x47445421u  // 'GDT!'
#define AP_STAGE_IDT                 0x49445421u  // 'IDT!'
#define AP_STAGE_PARK                0x5041524Bu  // 'PARK'

#pragma pack(push, 1)
typedef struct AP_BOOT_BLOCK {
    uint64_t Magic;
    uint64_t KernelCr3;
    uint64_t LongModeEntry;
    uint64_t Kcpu;
    uint64_t StackTop;
    volatile uint64_t Heartbeat;
    volatile uint32_t State;
    volatile uint32_t Stage;
    volatile uint32_t ApIndex;
} AP_BOOT_BLOCK;
#pragma pack(pop)

_Static_assert(offsetof(AP_BOOT_BLOCK, KernelCr3) == 8, "AP_BOOT_BLOCK.KernelCr3 offset");
_Static_assert(offsetof(AP_BOOT_BLOCK, StackTop) == 32, "AP_BOOT_BLOCK.StackTop offset");
_Static_assert(offsetof(AP_BOOT_BLOCK, State) == 48, "AP_BOOT_BLOCK.State offset");
_Static_assert(offsetof(AP_BOOT_BLOCK, Stage) == 52, "AP_BOOT_BLOCK.Stage offset");

static PHYSICAL_ADDRESS gApTrampolinePhys;
static PHYSICAL_ADDRESS gReservedApTrampolinePhys;
static AP_BOOT_BLOCK* gApBootBlock;

// Returns whether investigation-only SMP diagnostics are enabled.
static bool KiSmpDiagnosticsEnabled(void) {
    return (KeGetDebugFlags() & KE_DEBUG_SELFTEST) != 0;
}

// Validates pool state around AP admission so SMP races fail near their trigger.
static void KiSmpPoolCheckpoint(const char* phase, uint32_t cpuIndex) {
    if (!KiSmpDiagnosticsEnabled())
        return;

    DbgPrint("SMP: pool checkpoint %s AP%u\n", phase ? phase : "unknown", cpuIndex);
    (void)MmValidateKernelPool(phase ? phase : "smp");
}

// Implements the emit8 operation.
static void KiEmit8(uint8_t** p, uint8_t v) { *(*p)++ = v; }
// Implements the emit16 operation.
static void KiEmit16(uint8_t** p, uint16_t v) { memcpy(*p, &v, sizeof(v)); *p += sizeof(v); }
// Implements the emit32 operation.
static void KiEmit32(uint8_t** p, uint32_t v) { memcpy(*p, &v, sizeof(v)); *p += sizeof(v); }
// Implements the emit64 operation.
static void KiEmit64(uint8_t** p, uint64_t v) { memcpy(*p, &v, sizeof(v)); *p += sizeof(v); }

// Implements the patch16 operation.
static void KiPatch16(uint8_t* base, size_t off, uint16_t v) { memcpy(base + off, &v, sizeof(v)); }
// Implements the patch32 operation.
static void KiPatch32(uint8_t* base, size_t off, uint32_t v) { memcpy(base + off, &v, sizeof(v)); }
// Implements the patch64 operation.
static void KiPatch64(uint8_t* base, size_t off, uint64_t v) { memcpy(base + off, &v, sizeof(v)); }

// Implements the gdt Descriptor operation.
static uint64_t KiGdtDescriptor(uint32_t base, uint32_t limit, uint8_t access, uint8_t flags) {
    return ((uint64_t)(limit & 0xFFFF)) |
           ((uint64_t)(base & 0xFFFFFF) << 16) |
           ((uint64_t)access << 40) |
           ((uint64_t)((limit >> 16) & 0xF) << 48) |
           ((uint64_t)(flags & 0xF) << 52) |
           ((uint64_t)((base >> 24) & 0xFF) << 56);
}

// Updates application Processor Trampoline.
void KeSetApplicationProcessorTrampoline(uint64_t trampolinePhys) {
    if (trampolinePhys >= PAGE_SIZE &&
        trampolinePhys < AP_TRAMPOLINE_MAX_PHYS &&
        !(trampolinePhys & (PAGE_SIZE - 1))) {
        gReservedApTrampolinePhys = (PHYSICAL_ADDRESS)trampolinePhys;
    } else {
        gReservedApTrampolinePhys = 0;
    }
}

// Builds ap Trampoline.
static KSTATUS KiBuildApTrampoline(void) {
    if (gApTrampolinePhys)
        return NTLSTATUS_SUCCESS;

    PHYSICAL_ADDRESS frame = gReservedApTrampolinePhys;
    if (frame) {
        if (frame < AP_TRAMPOLINE_FALLBACK_MIN ||
            frame >= AP_TRAMPOLINE_MAX_PHYS ||
            (frame & (PAGE_SIZE - 1))) {
            if (KiSmpDiagnosticsEnabled()) {
                DbgPrint("SMP: rejecting bootloader AP trampoline frame 0x%016llX\n",
                         (uint64_t)frame);
            }
            frame = 0;
        } else {
            if (KiSmpDiagnosticsEnabled()) {
                DbgPrint("SMP: using bootloader-reserved AP trampoline frame 0x%016llX\n",
                         (uint64_t)frame);
            }
        }
    }

    if (!frame) {
        if (KiSmpDiagnosticsEnabled()) {
            DbgPrint("SMP: trying PMM AP trampoline range 0x%016llX-0x%016llX\n",
                     (uint64_t)AP_TRAMPOLINE_PREFERRED_MIN,
                     (uint64_t)AP_TRAMPOLINE_MAX_PHYS);
        }
        frame = MmPfnAllocateInRange(AP_TRAMPOLINE_PREFERRED_MIN,
                                       AP_TRAMPOLINE_MAX_PHYS,
                                       MmPfnUseKernel);
    }
    if (!frame) {
        if (KiSmpDiagnosticsEnabled()) {
            DbgPrint("SMP: trying PMM AP trampoline range 0x%016llX-0x%016llX\n",
                     (uint64_t)AP_TRAMPOLINE_FALLBACK_MIN,
                     (uint64_t)AP_TRAMPOLINE_PREFERRED_MIN);
        }
        frame = MmPfnAllocateInRange(AP_TRAMPOLINE_FALLBACK_MIN,
                                       AP_TRAMPOLINE_PREFERRED_MIN,
                                       MmPfnUseKernel);
    }
    if (!frame) {
        DbgPrint("SMP: no free PMM frame available below 0x%016llX for AP trampoline\n",
                 (uint64_t)AP_TRAMPOLINE_MAX_PHYS);
        return NTLSTATUS_NO_MEMORY;
    }

    if (KiSmpDiagnosticsEnabled())
        DbgPrint("SMP: AP trampoline frame 0x%016llX selected\n", (uint64_t)frame);

    uint8_t* page = (uint8_t*)(PHYS_MAP_BASE + frame);
    memset(page, 0, PAGE_SIZE);

    uint32_t physBase = (uint32_t)frame;
    uint8_t* p = page;

    // --- 16-bit real-mode SIPI entry ---.
    KiEmit8(&p, 0xFA);                    // cli
    KiEmit8(&p, 0x8C); KiEmit8(&p, 0xC8); // mov ax, cs
    KiEmit8(&p, 0x8E); KiEmit8(&p, 0xD8); // mov ds, ax
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x01); KiEmit8(&p, 0x16);
    size_t gdtr16Patch = (size_t)(p - page); KiEmit16(&p, 0); // lgdt [gdtr]
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x20); KiEmit8(&p, 0xC0);  // mov eax, cr0
    KiEmit8(&p, 0x66); KiEmit8(&p, 0x83); KiEmit8(&p, 0xC8); KiEmit8(&p, 0x01); // or eax, 1
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x22); KiEmit8(&p, 0xC0);  // mov cr0, eax
    KiEmit8(&p, 0x66); KiEmit8(&p, 0xEA);                     // jmp far 32-bit offset
    size_t prot32Patch = (size_t)(p - page); KiEmit32(&p, 0);
    KiEmit16(&p, 0x0008);                                      // selector 32-bit code

    while (((uintptr_t)(p - page)) & 0x0F) KiEmit8(&p, 0x90);
    uint32_t prot32Off = (uint32_t)(p - page);

    // --- 32-bit protected-mode transition to long mode ---.
    KiEmit8(&p, 0x66); KiEmit8(&p, 0xB8); KiEmit16(&p, 0x0010); // mov ax, 0x10
    KiEmit8(&p, 0x8E); KiEmit8(&p, 0xD8); // mov ds, ax
    KiEmit8(&p, 0x8E); KiEmit8(&p, 0xC0); // mov es, ax
    KiEmit8(&p, 0x8E); KiEmit8(&p, 0xD0); // mov ss, ax
    KiEmit8(&p, 0x8E); KiEmit8(&p, 0xE0); // mov fs, ax
    KiEmit8(&p, 0x8E); KiEmit8(&p, 0xE8); // mov gs, ax
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x20); KiEmit8(&p, 0xE0); // mov eax, cr4
    KiEmit8(&p, 0x0D); KiEmit32(&p, AP_CR4_REQUIRED);          // or eax, PAE|OSFXSR|OSXMMEXCPT
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x22); KiEmit8(&p, 0xE0); // mov cr4, eax
    KiEmit8(&p, 0xB9); KiEmit32(&p, 0xC0000080u);            // mov ecx, IA32_EFER
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x32);                    // rdmsr
    KiEmit8(&p, 0x0D); KiEmit32(&p, 0x00000900u);            // or eax, LME|NXE
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x30);                    // wrmsr
    KiEmit8(&p, 0xA1); size_t cr3Patch = (size_t)(p - page); KiEmit32(&p, 0); // mov eax,[cr3]
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x22); KiEmit8(&p, 0xD8); // mov cr3, eax
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x20); KiEmit8(&p, 0xC0); // mov eax, cr0
    KiEmit8(&p, 0x25); KiEmit32(&p, ~(AP_CR0_EM | AP_CR0_TS)); // and eax, ~(EM|TS)
    KiEmit8(&p, 0x0D); KiEmit32(&p, AP_CR0_PG | AP_CR0_PE | AP_CR0_MP | AP_CR0_NE); // or eax, PG|PE|MP|NE
    KiEmit8(&p, 0x0F); KiEmit8(&p, 0x22); KiEmit8(&p, 0xC0); // mov cr0, eax
    KiEmit8(&p, 0xEA); size_t long64Patch = (size_t)(p - page); KiEmit32(&p, 0);
    KiEmit16(&p, 0x0018);                                    // selector 64-bit code

    while (((uintptr_t)(p - page)) & 0x0F) KiEmit8(&p, 0x90);
    uint32_t long64Off = (uint32_t)(p - page);

    // --- 64-bit entry. Jump into high-half KiApKernelEntry(KCPU*). ---.
    KiEmit8(&p, 0x48); KiEmit8(&p, 0xBB); size_t bootBlockPatch = (size_t)(p - page); KiEmit64(&p, 0); // mov rbx, bootblock
    KiEmit8(&p, 0xC7); KiEmit8(&p, 0x83); KiEmit32(&p, (uint32_t)offsetof(AP_BOOT_BLOCK, Stage)); KiEmit32(&p, AP_STAGE_LONG64);
    KiEmit8(&p, 0xC7); KiEmit8(&p, 0x83); KiEmit32(&p, (uint32_t)offsetof(AP_BOOT_BLOCK, State)); KiEmit32(&p, AP_STATE_LONG_MODE);
    KiEmit8(&p, 0x48); KiEmit8(&p, 0x8B); KiEmit8(&p, 0xA3); KiEmit32(&p, (uint32_t)offsetof(AP_BOOT_BLOCK, StackTop)); // mov rsp,[rbx+StackTop]
    KiEmit8(&p, 0x48); KiEmit8(&p, 0x83); KiEmit8(&p, 0xE4); KiEmit8(&p, 0xF0); // and rsp,-16
    KiEmit8(&p, 0x48); KiEmit8(&p, 0x83); KiEmit8(&p, 0xEC); KiEmit8(&p, 0x28); // sub rsp,28h
    KiEmit8(&p, 0x48); KiEmit8(&p, 0x8B); KiEmit8(&p, 0x8B); KiEmit32(&p, (uint32_t)offsetof(AP_BOOT_BLOCK, Kcpu)); // mov rcx,[rbx+Kcpu]
    KiEmit8(&p, 0x48); KiEmit8(&p, 0x8B); KiEmit8(&p, 0x83); KiEmit32(&p, (uint32_t)offsetof(AP_BOOT_BLOCK, LongModeEntry)); // mov rax,[rbx+LongModeEntry]
    KiEmit8(&p, 0xFF); KiEmit8(&p, 0xE0); // jmp rax

    while (((uintptr_t)(p - page)) & 0x07) KiEmit8(&p, 0);
    uint32_t gdtOff = (uint32_t)(p - page);
    uint64_t* gdt = (uint64_t*)p;
    gdt[0] = 0;
    gdt[1] = KiGdtDescriptor(0, 0xFFFFF, 0x9A, 0xC); // 32-bit code
    gdt[2] = KiGdtDescriptor(0, 0xFFFFF, 0x92, 0xC); // 32-bit data
    gdt[3] = KiGdtDescriptor(0, 0xFFFFF, 0x9A, 0xA); // 64-bit code
    gdt[4] = KiGdtDescriptor(0, 0xFFFFF, 0x92, 0xC); // 64-bit data
    p += 5 * sizeof(uint64_t);

    uint32_t gdtrOff = (uint32_t)(p - page);
    KiEmit16(&p, (uint16_t)(5 * sizeof(uint64_t) - 1));
    KiEmit32(&p, physBase + gdtOff);

    if ((uintptr_t)(p - page) >= AP_TRAMPOLINE_BOOT_OFF)
        return NTLSTATUS_DEVICE_ERROR;

    AP_BOOT_BLOCK* boot = (AP_BOOT_BLOCK*)(page + AP_TRAMPOLINE_BOOT_OFF);
    memset(boot, 0, sizeof(*boot));
    boot->Magic = AP_BOOT_MAGIC;

    KiPatch16(page, gdtr16Patch, (uint16_t)gdtrOff);
    KiPatch32(page, prot32Patch, physBase + prot32Off);
    KiPatch32(page, cr3Patch, physBase + AP_TRAMPOLINE_BOOT_OFF + (uint32_t)offsetof(AP_BOOT_BLOCK, KernelCr3));
    KiPatch32(page, long64Patch, physBase + long64Off);
    KiPatch64(page, bootBlockPatch, (uint64_t)(physBase + AP_TRAMPOLINE_BOOT_OFF));

    gApTrampolinePhys = frame;
    gApBootBlock = boot;
    return NTLSTATUS_SUCCESS;
}

_Noreturn void KiApKernelEntry(KCPU* cpu) {
    (void)KeDisableInterrupts();

    if (!cpu)
        KeStopProcessor();

    // Secondary processors enter through the x86 bootstrap path, so do not assume
    // firmware CR4 feature bits are present. The scheduler and context-switch paths
    // use FXSAVE/FXRSTOR on every CPU; without OSFXSR, those instructions
    // raise #UD as soon as an AP enters its idle thread.
    uintptr_t cr4 = __readcr4();
    if ((cr4 & AP_CR4_REQUIRED) != AP_CR4_REQUIRED)
        __writecr4(cr4 | AP_CR4_REQUIRED);

    KeInitializeMemoryTypes();

    if (gApBootBlock) {
        gApBootBlock->Stage = AP_STAGE_ENTRY;
        gApBootBlock->State = AP_STATE_KERNEL_ENTRY;
        gApBootBlock->Heartbeat++;
    }

    cpu->Self = cpu;
    KeLoadKernelGsBase(cpu);
    KeSetUserGsBase(0);
    if (gApBootBlock) gApBootBlock->Stage = AP_STAGE_GS;

    KeInitializeCurrentProcessorInterruptController();

    KeInitializeGdtForCpu(cpu);
    KeLoadGdtForCurrentCpu(cpu);
    if (gApBootBlock) gApBootBlock->Stage = AP_STAGE_GDT;

    KeLoadIdtForCurrentCpu();
    if (gApBootBlock) gApBootBlock->Stage = AP_STAGE_IDT;

    _InterlockedExchange(&cpu->State, KeCpuStarted);
    cpu->CurrentIrql = PASSIVE_LEVEL;
    cpu->InterruptDepth = 0;
    cpu->ReschedulePending = 0;

    _InterlockedExchange(&cpu->State, KeCpuParked);
    if (gApBootBlock) {
        gApBootBlock->Stage = AP_STAGE_PARK;
        gApBootBlock->State = AP_STATE_KERNEL_PARKED;
    }

    while (true) {
        cpu->SchedulerTicks++;

        if (_InterlockedCompareExchange(&cpu->State, 0, 0) == KeCpuAdmissionRequested) {
            THREAD* idle = cpu->IdleThread;
            if (idle) {
                cpu->CurrentThread = idle;
                idle->State = ThreadRunning;
                idle->RunningCpu = cpu->CpuIndex;
                idle->Queued = 0;
                // Do not arm this processor's periodic timer while it is still in the raw
                // park loop. Timer IRQs are scheduler-owned events and require
                // a current idle thread before they may be delivered.
                KeConfigureCurrentProcessorTimer();

                _InterlockedExchange(&cpu->State, KeCpuSchedulerActive);
                KeEnableInterrupts();
                KeStartInitialThread(idle);
            }
        }

        // The boot processor admits parked secondary processors by publishing AdmissionRequested and
        // sending a wake IPI. Avoid the classic lost wakeup where the state
        // changes between the poll and HLT: check once with interrupts
        // disabled, then use STI;HLT as a single sleep transition.
        uintptr_t interruptFlags = KeDisableInterrupts();
        if (_InterlockedCompareExchange(&cpu->State, 0, 0) == KeCpuAdmissionRequested) {
            KeRestoreInterrupts(interruptFlags);
            continue;
        }
        KeWaitForInterrupt();
    }
}

// Prepares cpu Record.
static void KiPrepareCpuRecord(uint32_t cpuIndex, const HAL_PROCESSOR_INFO* info) {
    KCPU* cpu = &gCpus[cpuIndex];
    memset(cpu, 0, sizeof(*cpu));
    cpu->Self = cpu;
    cpu->CpuIndex = cpuIndex;
    cpu->HardwareId = info->HardwareId;
    cpu->FirmwareId = info->FirmwareId;
    cpu->BootProcessorIndex = cpuIndex;
    cpu->State = info->Enabled ? KeCpuPresent : KeCpuAbsent;
    cpu->IsBootProcessor = info->IsBootProcessor ? 1 : 0;
    cpu->PackageId = 0;
    cpu->CoreId = cpuIndex;
    cpu->ThreadId = 0;
    cpu->PhysicalCoreIndex = cpuIndex;
    cpu->CoreSiblingMask = (cpuIndex < 64) ? (1ULL << cpuIndex) : 0ULL;
    cpu->PackageMask = cpu->CoreSiblingMask;
    cpu->IsPrimaryThread = 1;
    cpu->SmtThreadCount = 1;
    cpu->TopologyFlags = 0;
    cpu->CurrentIrql = PASSIVE_LEVEL;
    cpu->InterruptDepth = 0;
    cpu->ReschedulePending = 0;
}

// Implements the ceil Log2 U32 operation.
static uint32_t KiCeilLog2U32(uint32_t value) {
    if (value <= 1u)
        return 0u;
    uint32_t shift = 0;
    uint32_t n = value - 1u;
    while (n) {
        shift++;
        n >>= 1;
    }
    return shift;
}

// Implements the mask Width operation.
static uint32_t KiMaskWidth(uint32_t shift) {
    if (shift >= 32u)
        return 0xFFFFFFFFu;
    if (shift == 0u)
        return 0u;
    return (1u << shift) - 1u;
}

typedef struct KI_CPUID_TOPOLOGY_PARAMS {
    uint32_t SmtShift;
    uint32_t CoreShift;
    uint32_t ThreadsPerCore;
    uint32_t LogicalPerPackage;
    const char* Source;
} KI_CPUID_TOPOLOGY_PARAMS;

// Returns cpuid Topology Leaf.
static void KiReadCpuidTopologyLeaf(uint32_t leaf, KI_CPUID_TOPOLOGY_PARAMS* out) {
    if (!out)
        return;

    for (uint32_t sub = 0; sub < 32u; sub++) {
        int r[4] = {0, 0, 0, 0};
        __cpuidex(r, (int)leaf, (int)sub);

        uint32_t levelShift = (uint32_t)r[0] & 0x1Fu;
        uint32_t logicals   = (uint32_t)r[1] & 0xFFFFu;
        uint32_t levelType  = ((uint32_t)r[2] >> 8) & 0xFFu;

        if (levelType == 0u || logicals == 0u)
            break;

        if (levelType == 1u) { // SMT level
            out->SmtShift = levelShift;
            out->ThreadsPerCore = logicals ? logicals : 1u;
        } else if (levelType == 2u) { // Core/package level
            out->CoreShift = levelShift;
            out->LogicalPerPackage = logicals ? logicals : 1u;
            break;
        }
    }
}

// Implements the detect Cpuid Topology Params operation.
static KI_CPUID_TOPOLOGY_PARAMS KiDetectCpuidTopologyParams(void) {
    KI_CPUID_TOPOLOGY_PARAMS topo;
    memset(&topo, 0, sizeof(topo));
    topo.ThreadsPerCore = 1;
    topo.LogicalPerPackage = 1;
    topo.Source = "flat";

    int maxRegs[4] = {0, 0, 0, 0};
    __cpuid(maxRegs, 0);
    uint32_t maxLeaf = (uint32_t)maxRegs[0];

    if (maxLeaf >= 0x1Fu) {
        KiReadCpuidTopologyLeaf(0x1Fu, &topo);
        if (topo.CoreShift >= topo.SmtShift && topo.CoreShift != 0u) {
            topo.Source = "cpuid.1f";
            return topo;
        }
        memset(&topo, 0, sizeof(topo));
        topo.ThreadsPerCore = 1;
        topo.LogicalPerPackage = 1;
        topo.Source = "flat";
    }

    if (maxLeaf >= 0x0Bu) {
        KiReadCpuidTopologyLeaf(0x0Bu, &topo);
        if (topo.CoreShift >= topo.SmtShift && topo.CoreShift != 0u) {
            topo.Source = "cpuid.0b";
            return topo;
        }
        memset(&topo, 0, sizeof(topo));
        topo.ThreadsPerCore = 1;
        topo.LogicalPerPackage = 1;
        topo.Source = "flat";
    }

    int r1[4] = {0, 0, 0, 0};
    __cpuid(r1, 1);
    uint32_t logicalPerPackage = ((uint32_t)r1[1] >> 16) & 0xFFu;
    if (logicalPerPackage == 0u)
        logicalPerPackage = 1u;

    uint32_t coresPerPackage = 1u;
    if (maxLeaf >= 4u) {
        int r4[4] = {0, 0, 0, 0};
        __cpuidex(r4, 4, 0);
        coresPerPackage = (((uint32_t)r4[0] >> 26) & 0x3Fu) + 1u;
        if (coresPerPackage == 0u)
            coresPerPackage = 1u;
    }

    uint32_t threadsPerCore = logicalPerPackage / coresPerPackage;
    if (threadsPerCore == 0u)
        threadsPerCore = 1u;

    topo.ThreadsPerCore = threadsPerCore;
    topo.LogicalPerPackage = logicalPerPackage;
    topo.SmtShift = KiCeilLog2U32(threadsPerCore);
    topo.CoreShift = KiCeilLog2U32(logicalPerPackage);
    if (topo.CoreShift < topo.SmtShift)
        topo.CoreShift = topo.SmtShift;
    topo.Source = "cpuid.1/4";
    return topo;
}

// Builds processor Topology.
void KeBuildProcessorTopology(void) {
    KI_CPUID_TOPOLOGY_PARAMS topo = KiDetectCpuidTopologyParams();

    uint32_t smtMask = KiMaskWidth(topo.SmtShift);
    uint32_t coreBits = (topo.CoreShift > topo.SmtShift) ?
                        (topo.CoreShift - topo.SmtShift) : 0u;
    uint32_t coreMask = KiMaskWidth(coreBits);

    for (uint32_t i = 0; i < gCpuCount && i < KE_MAX_CPUS; i++) {
        KCPU* cpu = &gCpus[i];
        if (cpu->State == KeCpuAbsent)
            continue;

        uint32_t hardwareId = cpu->HardwareId;
        cpu->ThreadId = smtMask ? (hardwareId & smtMask) : 0u;
        cpu->CoreId = coreBits ? ((hardwareId >> topo.SmtShift) & coreMask) :
                                  (hardwareId >> topo.SmtShift);
        cpu->PackageId = (topo.CoreShift < 32u) ? (hardwareId >> topo.CoreShift) : 0u;
        cpu->PhysicalCoreIndex = KE_INVALID_CPU;
        cpu->CoreSiblingMask = 0;
        cpu->PackageMask = 0;
        cpu->IsPrimaryThread = 0;
        cpu->SmtThreadCount = 1;
        cpu->TopologyFlags = (topo.ThreadsPerCore > 1u) ? 1u : 0u;
    }

    uint32_t nextCoreIndex = 0;
    for (uint32_t i = 0; i < gCpuCount && i < KE_MAX_CPUS; i++) {
        KCPU* cpu = &gCpus[i];
        if (cpu->State == KeCpuAbsent)
            continue;
        if (cpu->PhysicalCoreIndex != KE_INVALID_CPU)
            continue;

        uint32_t coreIndex = nextCoreIndex++;
        uint64_t coreMaskBits = 0;
        uint64_t packageMaskBits = 0;
        uint32_t primaryCpuIndex = i;
        uint32_t siblingCount = 0;

        for (uint32_t j = 0; j < gCpuCount && j < KE_MAX_CPUS; j++) {
            KCPU* peer = &gCpus[j];
            if (peer->State == KeCpuAbsent)
                continue;
            if (peer->PackageId == cpu->PackageId) {
                if (j < 64)
                    packageMaskBits |= (1ULL << j);
            }
            if (peer->PackageId == cpu->PackageId && peer->CoreId == cpu->CoreId) {
                peer->PhysicalCoreIndex = coreIndex;
                if (j < 64)
                    coreMaskBits |= (1ULL << j);
                if (j < primaryCpuIndex)
                    primaryCpuIndex = j;
                siblingCount++;
            }
        }

        if (siblingCount == 0)
            siblingCount = 1;

        for (uint32_t j = 0; j < gCpuCount && j < KE_MAX_CPUS; j++) {
            KCPU* peer = &gCpus[j];
            if (peer->State == KeCpuAbsent)
                continue;
            if (peer->PackageId == cpu->PackageId && peer->CoreId == cpu->CoreId) {
                peer->CoreSiblingMask = coreMaskBits;
                peer->PackageMask = packageMaskBits;
                peer->SmtThreadCount = (uint8_t)((siblingCount > 255u) ? 255u : siblingCount);
                peer->IsPrimaryThread = (j == primaryCpuIndex) ? 1u : 0u;
            }
        }
    }

    DbgPrint("SMP: SMT topology source=%s smtShift=%u coreShift=%u threads/core=%u logical/package=%u physicalCores=%u\n",
             topo.Source ? topo.Source : "?",
             topo.SmtShift,
             topo.CoreShift,
             topo.ThreadsPerCore,
             topo.LogicalPerPackage,
             nextCoreIndex);

    if (KiSmpDiagnosticsEnabled()) {
        for (uint32_t i = 0; i < gCpuCount && i < KE_MAX_CPUS; i++) {
            KCPU* cpu = &gCpus[i];
            if (cpu->State == KeCpuAbsent)
                continue;
            DbgPrint("SMP: cpu%u hardware=%u pkg=%u core=%u thread=%u primary=%u siblings=0x%llx package=0x%llx\n",
                     i,
                     cpu->HardwareId,
                     cpu->PackageId,
                     cpu->CoreId,
                     cpu->ThreadId,
                     cpu->IsPrimaryThread,
                     (unsigned long long)cpu->CoreSiblingMask,
                     (unsigned long long)cpu->PackageMask);
        }
    }
}

// Creates ap Stack.
static KSTATUS KiAllocateApStack(KCPU* cpu) {
    // AP kernel stacks do not need physically contiguous memory. Requiring
    // MmPfnAllocateContiguous(8) made SMP startup depend on a DMA-style run being
    // available, which is fragile on real machines and unrelated to AP startup.
    size_t stackBytes = (size_t)AP_STACK_PAGES * PAGE_SIZE;
    void* stack = MmAllocatePool(stackBytes);
    if (!stack) {
        DbgPrint("SMP: AP%u stack allocation failed bytes=%llu\n",
                 cpu ? cpu->CpuIndex : 0xFFFFFFFFu, (uint64_t)stackBytes);
        return NTLSTATUS_NO_MEMORY;
    }

    cpu->KernelStack = stack;
    cpu->KernelStackTop = ((uintptr_t)stack + stackBytes) & ~(uintptr_t)15;
    memset(cpu->KernelStack, 0, stackBytes);
    memset(&cpu->Tss, 0, sizeof(cpu->Tss));
    cpu->Tss.IoMapBase = sizeof(TSS64);
    cpu->Tss.RSP0 = cpu->KernelStackTop;
    KeInitializeGdtForCpu(cpu);
    return NTLSTATUS_SUCCESS;
}

// Implements the ap Stage Name operation.
static const char* KiApStageName(uint32_t stage) {
    switch (stage) {
    case AP_STAGE_NONE:   return "none";
    case AP_STAGE_LONG64: return "long64";
    case AP_STAGE_ENTRY:  return "entry";
    case AP_STAGE_GS:     return "gs";
    case AP_STAGE_GDT:    return "gdt";
    case AP_STAGE_IDT:    return "idt";
    case AP_STAGE_PARK:   return "park";
    default:              return "unknown";
    }
}

// Fills boot Block For Cpu.
static void KiFillBootBlockForCpu(KCPU* cpu, uint32_t apIndex) {
    memset(gApBootBlock, 0, sizeof(*gApBootBlock));
    gApBootBlock->Magic = AP_BOOT_MAGIC;
    gApBootBlock->KernelCr3 = (uint64_t)__readcr3();
    gApBootBlock->LongModeEntry = (uint64_t)(uintptr_t)KiApKernelEntry;
    gApBootBlock->Kcpu = (uint64_t)(uintptr_t)cpu;
    gApBootBlock->StackTop = cpu->KernelStackTop;
    gApBootBlock->State = AP_STATE_READY;
    gApBootBlock->Stage = AP_STAGE_NONE;
    gApBootBlock->ApIndex = apIndex;
}

// Starts application Processors.
KSTATUS KeStartApplicationProcessors(void) {
    HAL_PROCESSOR_TOPOLOGY topo;
    KSTATUS st = KeQueryProcessorTopology(&topo);
    if (!KSUCCESS(st)) return st;

    if (topo.Count == 0 || topo.Count > KE_MAX_CPUS)
        return NTLSTATUS_DEVICE_ERROR;

    uint32_t bootHardwareId = topo.BootProcessorHardwareId;
    KCPU* bsp = KeGetCurrentCpu();
    bsp->HardwareId = bootHardwareId;
    bsp->IsBootProcessor = 1;
    bsp->State = KeCpuStarted;

    DbgPrint("SMP: platform processors boot-id=%u firmware-count=%u\n",
             bootHardwareId, topo.Count);

    uint32_t count = 1;
    for (uint32_t i = 0; i < topo.Count && count < KE_MAX_CPUS; i++) {
        HAL_PROCESSOR_INFO* info = &topo.Processors[i];
        if (KiSmpDiagnosticsEnabled()) {
            DbgPrint("SMP: topology[%u] firmware=%u hardware=%u flags=0x%x enabled=%u boot=%u startable=%u\n",
                     i, info->FirmwareId, info->HardwareId, info->Flags,
                     info->Enabled, info->IsBootProcessor, info->Startable);
        }

        if (!info->Enabled) continue;
        if (info->HardwareId == bootHardwareId) {
            bsp->FirmwareId = info->FirmwareId;
            continue;
        }
        if (!info->Startable) {
            DbgPrint("SMP: skipping processor firmware=%u hardware=%u; HAL cannot start it\n",
                     info->FirmwareId, info->HardwareId);
            continue;
        }
        KiPrepareCpuRecord(count, info);
        count++;
    }
    gCpuCount = count;

    DbgPrint("SMP: discovered %u usable processor(s), boot hardware ID=%u\n", gCpuCount, bootHardwareId);
    if (gCpuCount <= 1) {
        DbgPrint("SMP: no application processors present\n");
        return NTLSTATUS_SUCCESS;
    }

    if ((__readcr3() >> 32) != 0) {
        DbgPrint("SMP: disabled; bootstrap CR3 is above 4GiB (cr3=0x%016llX), AP trampoline cannot load it yet\n",
                 (uint64_t)__readcr3());
        gCpuCount = 1;
        return NTLSTATUS_SUCCESS;
    }

    if (KiSmpDiagnosticsEnabled()) {
        DbgPrint("SMP: bootloader AP trampoline phys=0x%016llX\n",
                 (uint64_t)gReservedApTrampolinePhys);
    }
    st = KiBuildApTrampoline();
    if (!KSUCCESS(st)) {
        DbgPrint("SMP: disabled; AP trampoline allocation/build failed status=0x%x\n", st);
        gCpuCount = 1;
        return NTLSTATUS_SUCCESS;
    }

    uint32_t claimed = 0;

    for (uint32_t i = 1; i < gCpuCount; i++) {
        KCPU* cpu = &gCpus[i];
        if (cpu->State == KeCpuAbsent) continue;

        st = KiAllocateApStack(cpu);
        if (!KSUCCESS(st)) {
            DbgPrint("SMP: AP%u hardware=%u stack allocation failed status=0x%x; disabling AP and continuing\n",
                     i, cpu->HardwareId, st);
            _InterlockedExchange(&cpu->State, KeCpuAbsent);
            continue;
        }

        KiFillBootBlockForCpu(cpu, i);
        if (KiSmpDiagnosticsEnabled())
            DbgPrint("SMP: starting AP%u hardware=%u entry=0x%016llX\n",
                     i, cpu->HardwareId, (uint64_t)gApTrampolinePhys);

        st = KeStartProcessor(cpu->HardwareId, gApTrampolinePhys);
        if (!KSUCCESS(st)) {
            DbgPrint("SMP: AP%u hardware=%u startup IPI failed status=0x%x; disabling AP and continuing\n",
                     i, cpu->HardwareId, st);
            _InterlockedExchange(&cpu->State, KeCpuAbsent);
            continue;
        }

        uint64_t lastHeartbeat = gApBootBlock->Heartbeat;
        uint64_t startDeadline = KiSmpDeadlineAfter(AP_START_TIMEOUT_US);
        int parked = 0;
        for (;;) {
            if (gApBootBlock->State == AP_STATE_KERNEL_PARKED) {
                parked = 1;
                break;
            }
            lastHeartbeat = gApBootBlock->Heartbeat;
            if (KiSmpQueryTicks() >= startDeadline)
                break;
            _mm_pause();
        }

        if (!parked) {
            DbgPrint("SMP: AP%u failed to reach kernel park state hardware=%u state=%u stage=0x%x(%s) heartbeat=%llu last=%llu; disabling AP and continuing\n",
                     i, cpu->HardwareId, gApBootBlock->State, gApBootBlock->Stage,
                     KiApStageName(gApBootBlock->Stage), gApBootBlock->Heartbeat, lastHeartbeat);
            _InterlockedExchange(&cpu->State, KeCpuAbsent);
            continue;
        }

        if (_InterlockedCompareExchange(&cpu->State, 0, 0) != KeCpuParked) {
            KePanic(KE_PANIC_CODE_INIT,
                    "SMP: AP%u boot block parked without CPU state transition state=%ld\n",
                    i, cpu->State);
        }
        claimed++;
        if (KiSmpDiagnosticsEnabled())
            DbgPrint("SMP: AP%u reached kernel park\n", i);
    }

    if (claimed == 0 && gCpuCount > 1)
        DbgPrint("SMP: no APs reached kernel-owned state; continuing boot-processor-only\n");

    KeBuildProcessorTopology();
    DbgPrint("SMP: kernel claimed %u APs\n", claimed);
    KiSmpPoolCheckpoint("smp after AP park", 0);
    return NTLSTATUS_SUCCESS;
}

// Implements the admit Application Processors To Scheduler operation.
KSTATUS KeAdmitApplicationProcessorsToScheduler(void) {
    if (gCpuCount <= 1) {
        DbgPrint("SMP: no AP scheduler admission needed\n");
        return NTLSTATUS_SUCCESS;
    }

    uint32_t admitted = 0;
    for (uint32_t i = 1; i < gCpuCount; i++) {
        KCPU* cpu = &gCpus[i];
        if (cpu->State == KeCpuAbsent) continue;

        if (_InterlockedCompareExchange(&cpu->State, 0, 0) != KeCpuParked) {
            KePanic(KE_PANIC_CODE_INIT,
                    "SMP: AP%u cannot enter scheduler from state=%ld\n",
                    i, cpu->State);
        }

        KiSmpPoolCheckpoint("smp before AP idle create", i);
        KSTATUS st = KiCreateIdleThreadForCpu(cpu);
        if (!KSUCCESS(st))
            return st;
        KiSmpPoolCheckpoint("smp after AP idle create", i);

        _InterlockedExchange(&cpu->State, KeCpuAdmissionRequested);
        KiSmpPoolCheckpoint("smp before AP wake", i);

        // Parked secondary processors sleep with no periodic timer armed yet. The
        // state transition alone is not sufficient; send a fixed-vector IPI
        // only to wake the parked loop. ReschedulePending is intentionally left
        // clear, so the interrupt exit path will not try to schedule before
        // the AP has installed its idle thread.
        KeSendProcessorInterrupt(cpu->HardwareId, KE_VECTOR_RESCHEDULE_IPI);

        int active = 0;
        uint64_t schedulerDeadline = KiSmpDeadlineAfter(AP_SCHED_TIMEOUT_US);
        for (;;) {
            if (_InterlockedCompareExchange(&cpu->State, 0, 0) == KeCpuSchedulerActive) {
                active = 1;
                break;
            }
            if (KiSmpQueryTicks() >= schedulerDeadline)
                break;
            _mm_pause();
        }

        if (!active) {
            KePanic(KE_PANIC_CODE_INIT,
                    "SMP: AP%u failed scheduler admission hardware=%u state=%ld ticks=%llu\n",
                    i, cpu->HardwareId, cpu->State, cpu->SchedulerTicks);
        }

        _InterlockedExchange(&cpu->State, KeCpuOnline);
        KiSmpPoolCheckpoint("smp after AP scheduler active", i);
        admitted++;
        if (KiSmpDiagnosticsEnabled())
            DbgPrint("SMP: AP%u entered scheduler idle thread\n", i);
    }

    DbgPrint("SMP: scheduler admitted %u APs\n", admitted);
    return NTLSTATUS_SUCCESS;
}

#define SMP_AP_WORKER_ROUNDS          64u
#define SMP_AP_WORKER_TIMEOUT_US      2000000ull
#define SMP_AP_TEST_MAGIC             0x524B575041504D53ULL  // 'SMPAPWKR'

typedef struct SMP_AP_WORK_TEST SMP_AP_WORK_TEST;

typedef struct SMP_AP_WORKER_ARG {
    uint64_t Magic;
    uint32_t CpuIndex;
    uint32_t WorkerIndex;
    struct SMP_AP_WORK_TEST* Test;
    EVENT* StepEvent;
    volatile long Ready;
    volatile long WaitArmed;
    volatile long Progress;
    volatile long Done;
    volatile long Failure;
    uint32_t WrongCpu;
    uint32_t WrongRunningCpu;
} SMP_AP_WORKER_ARG;

struct SMP_AP_WORK_TEST {
    uint64_t Magic;
    uint32_t WorkerCount;
    volatile long ReadyCount;
    volatile long DoneCount;
    volatile long FailureCount;
    volatile long CpuHits[KE_MAX_CPUS];
    SMP_AP_WORKER_ARG Workers[KE_MAX_CPUS];
};

// Implements the smp Ap Controlled Worker operation.
static void KiSmpApControlledWorker(void* context) {
    SMP_AP_WORKER_ARG* arg = (SMP_AP_WORKER_ARG*)context;
    if (!arg || arg->Magic != SMP_AP_TEST_MAGIC || !arg->StepEvent) {
        KePanic(KE_PANIC_CODE_SELFTEST, "SMP AP worker received bad context 0x%016llX\n", (uint64_t)arg);
    }

    _InterlockedExchange(&arg->Ready, 1);
    _InterlockedIncrement(&arg->Test->ReadyCount);

    for (uint32_t round = 0; round < SMP_AP_WORKER_ROUNDS; round++) {
        _InterlockedIncrement(&arg->WaitArmed);
        KSTATUS waitStatus = KeWaitForEvent(arg->StepEvent, 0);
        if (!KSUCCESS(waitStatus)) {
            _InterlockedExchange(&arg->Failure, 1);
            break;
        }

        KCPU* cpu = KeGetCurrentCpu();
        THREAD* self = KeGetCurrentThread();
        if (!cpu || cpu->CpuIndex != arg->CpuIndex) {
            arg->WrongCpu = cpu ? cpu->CpuIndex : KE_INVALID_CPU;
            _InterlockedExchange(&arg->Failure, 1);
            break;
        }
        if (!self || self->RunningCpu != cpu->CpuIndex) {
            arg->WrongRunningCpu = self ? self->RunningCpu : KE_INVALID_CPU;
            _InterlockedExchange(&arg->Failure, 1);
            break;
        }
        if (self && (self->AffinityMask & (1ULL << cpu->CpuIndex)) == 0) {
            arg->WrongCpu = cpu->CpuIndex;
            _InterlockedExchange(&arg->Failure, 1);
            break;
        }

        if (cpu->CpuIndex < KE_MAX_CPUS)
            _InterlockedIncrement(&arg->Test->CpuHits[cpu->CpuIndex]);

        _InterlockedIncrement(&arg->Progress);
        if ((round & 3u) == 0)
            KeYield();
    }

    if (_InterlockedCompareExchange(&arg->Failure, 0, 0) != 0) {
        _InterlockedIncrement(&arg->Test->FailureCount);
    }

    _InterlockedExchange(&arg->Done, 1);
    _InterlockedIncrement(&arg->Test->DoneCount);
}

// Waits for for Worker Counter.
static int KiWaitForWorkerCounter(volatile long* counter, long expected, const char* what) {
    uint64_t deadline = KiSmpDeadlineAfter(SMP_AP_WORKER_TIMEOUT_US);
    for (;;) {
        if (_InterlockedCompareExchange(counter, 0, 0) >= expected)
            return 1;
        if (KiSmpQueryTicks() >= deadline)
            break;
        KeYield();
    }
    DbgPrint("SMP scheduler selftest timeout waiting for %s expected=%ld got=%ld\n",
             what ? what : "counter", expected, _InterlockedCompareExchange(counter, 0, 0));
    return 0;
}

// Runs application Processor Scheduler Self Tests.
KSTATUS KeRunApplicationProcessorSchedulerSelfTests(void) {
    if (gCpuCount <= 1) {
        DbgPrint("SMP scheduler selftest: SKIP, no APs\n");
        return NTLSTATUS_SUCCESS;
    }

    DbgPrint("SMP scheduler selftest: begin AP-affine workers\n");

    SMP_AP_WORK_TEST* test = (SMP_AP_WORK_TEST*)MmAllocatePool(sizeof(SMP_AP_WORK_TEST));
    if (!test) return NTLSTATUS_NO_MEMORY;
    memset(test, 0, sizeof(*test));
    test->Magic = SMP_AP_TEST_MAGIC;

    uint32_t workerCount = 0;
    for (uint32_t cpuIndex = 1; cpuIndex < gCpuCount; cpuIndex++) {
        KCPU* cpu = &gCpus[cpuIndex];
        if (cpu->State != KeCpuOnline)
            continue;
        if (cpuIndex >= 64)
            continue;

        SMP_AP_WORKER_ARG* arg = &test->Workers[workerCount];
        arg->Magic = SMP_AP_TEST_MAGIC;
        arg->CpuIndex = cpuIndex;
        arg->WorkerIndex = workerCount;
        arg->Test = test;

        KSTATUS st = KeEventCreate(EventAutoReset, 0, &arg->StepEvent);
        if (!KSUCCESS(st)) {
            KePanic(KE_PANIC_CODE_SELFTEST,
                    "SMP scheduler selftest: AP%u step event create failed status=0x%x\n",
                    cpuIndex, st);
        }

        KE_THREAD_CREATE_PARAMETERS workerParameters =
            KE_THREAD_CREATE_NAMED_NORMAL(PriorityNormal, "selftest.smp");
        workerParameters.AffinityMask = 1ULL << cpuIndex;
        st = KeCreateKernelThread((uintptr_t)KiSmpApControlledWorker,
                                  (uintptr_t)arg,
                                  &workerParameters,
                                  (THREAD**)0);
        if (!KSUCCESS(st)) {
            KePanic(KE_PANIC_CODE_SELFTEST,
                    "SMP scheduler selftest: AP%u worker create failed status=0x%x\n",
                    cpuIndex, st);
        }

        workerCount++;
    }
    test->WorkerCount = workerCount;

    if (workerCount == 0) {
        MmFreePool(test);
        DbgPrint("SMP scheduler selftest: SKIP, no online AP workers\n");
        return NTLSTATUS_SUCCESS;
    }

    if (!KiWaitForWorkerCounter(&test->ReadyCount, (long)workerCount, "AP worker ready")) {
        KePanic(KE_PANIC_CODE_SELFTEST, "SMP scheduler selftest: workers failed to become ready ready=%ld workers=%u\n",
                test->ReadyCount, workerCount);
    }

    for (uint32_t round = 0; round < SMP_AP_WORKER_ROUNDS; round++) {
        for (uint32_t w = 0; w < workerCount; w++) {
            SMP_AP_WORKER_ARG* arg = &test->Workers[w];
            if (!KiWaitForWorkerCounter(&arg->WaitArmed, (long)(round + 1), "AP worker wait-armed")) {
                KePanic(KE_PANIC_CODE_SELFTEST,
                        "SMP scheduler selftest: AP%u did not arm wait for round %u armed=%ld progress=%ld\n",
                        arg->CpuIndex, round, arg->WaitArmed, arg->Progress);
            }
            KeSignalEvent(arg->StepEvent);
        }

        for (uint32_t w = 0; w < workerCount; w++) {
            SMP_AP_WORKER_ARG* arg = &test->Workers[w];
            if (!KiWaitForWorkerCounter(&arg->Progress, (long)(round + 1), "AP worker progress")) {
                KePanic(KE_PANIC_CODE_SELFTEST,
                        "SMP scheduler selftest: AP%u no progress round=%u progress=%ld failure=%ld wrongCpu=%u wrongRunningCpu=%u\n",
                        arg->CpuIndex, round, arg->Progress, arg->Failure,
                        arg->WrongCpu, arg->WrongRunningCpu);
            }
            if (_InterlockedCompareExchange(&arg->Failure, 0, 0) != 0) {
                KePanic(KE_PANIC_CODE_SELFTEST,
                        "SMP scheduler selftest: AP%u worker failed round=%u wrongCpu=%u wrongRunningCpu=%u progress=%ld\n",
                        arg->CpuIndex, round, arg->WrongCpu, arg->WrongRunningCpu, arg->Progress);
            }
        }
    }

    if (!KiWaitForWorkerCounter(&test->DoneCount, (long)workerCount, "AP worker done")) {
        KePanic(KE_PANIC_CODE_SELFTEST, "SMP scheduler selftest: workers did not finish done=%ld workers=%u\n",
                test->DoneCount, workerCount);
    }

    if (_InterlockedCompareExchange(&test->FailureCount, 0, 0) != 0) {
        KePanic(KE_PANIC_CODE_SELFTEST, "SMP scheduler selftest: failure count=%ld\n",
                test->FailureCount);
    }

    for (uint32_t w = 0; w < workerCount; w++) {
        SMP_AP_WORKER_ARG* arg = &test->Workers[w];
        long hits = (arg->CpuIndex < KE_MAX_CPUS) ? test->CpuHits[arg->CpuIndex] : 0;
        if (hits <= 0) {
            KePanic(KE_PANIC_CODE_SELFTEST,
                    "SMP scheduler selftest: AP%u recorded no CPU hits\n", arg->CpuIndex);
        }
        if (arg->StepEvent)
            ObDereferenceObject(&arg->StepEvent->Header);
    }

    DbgPrint("SMP scheduler selftest: PASS workers=%u rounds=%u\n",
             workerCount, SMP_AP_WORKER_ROUNDS);
    MmFreePool(test);
    return NTLSTATUS_SUCCESS;
}