more scheduling

ntoskrnl/ke/amd64/krnlinit.c

@@ -195,10 +195,12 @@ KiSystemStartupBootStack(VOID)
     /* Set the Idle Priority to 0. This will jump into Phase 1 */
     KeSetPriorityThread(Thread, 0);
 
+#if 1
     /* If there's no thread scheduled, put this CPU in the Idle summary */
     KiAcquirePrcbLock(Prcb);
-    if (!Prcb->NextThread) KiIdleSummary |= (ULONG_PTR)1 << Prcb->Number;
+    if (!Prcb->NextThread) InterlockedBitTestAndSetAffinity(&KiIdleSummary, Prcb->Number);
     KiReleasePrcbLock(Prcb);
+#endif
 
     /* Raise back to HIGH_LEVEL and clear the PRCB for the loader block */
     KfRaiseIrql(HIGH_LEVEL);

ntoskrnl/ke/amd64/stubs.c

@@ -101,6 +101,9 @@ KiIdleLoop(VOID)
     PKPRCB Prcb = KeGetCurrentPrcb();
     PKTHREAD OldThread, NewThread;
 
+    /* Set idle summary bit */
+    InterlockedBitTestAndSetAffinity(&KiIdleSummary, Prcb->Number);
+
     /* Now loop forever */
     while (TRUE)
     {
@@ -144,10 +147,15 @@ KiIdleLoop(VOID)
             KfRaiseIrql(SYNCH_LEVEL);
 #endif
 
+            /* Clear idle summary bit */
+            InterlockedBitTestAndResetAffinity(&KiIdleSummary, Prcb->Number);
+
             /* Switch away from the idle thread */
             KiSwapContext(APC_LEVEL, OldThread);
 
-#ifdef CONFIG_SMP
+            /* Set idle summary bit */
+            InterlockedBitTestAndSetAffinity(&KiIdleSummary, Prcb->Number);
+
             /* Go back to DISPATCH_LEVEL */
             KeLowerIrql(DISPATCH_LEVEL);
 #endif

ntoskrnl/ke/thrdschd.c

@@ -79,15 +79,21 @@ ULONG
 KiSelectBestProcessor(
     _In_ PKTHREAD Thread)
 {
-    KAFFINITY Affinity;
+    KAFFINITY Affinity, PreferredSet;
     ULONG Processor;
 
     /* Get the affinity */
     Affinity = Thread->Affinity;
 
-    /* For now just return the first set bit */
-    Processor = RtlFindLeastSignificantBit(Affinity);
-    ASSERT(Processor < MAXIMUM_PROCESSORS);
+    PreferredSet = Affinity & KiIdleSummary;
+    if (PreferredSet == 0)
+    {
+        PreferredSet = Affinity;
+    }
+
+    /* Return the first set bit */
+    NT_VERIFY(BitScanForwardAffinity(&Processor, PreferredSet) != FALSE);
+    ASSERT(Processor < KeNumberProcessors);
 
     return Processor;
 }
@@ -254,16 +260,24 @@ KiDeferredReadyThread(IN PKTHREAD Thread)
     Prcb = KiProcessorBlock[Processor];
     KiAcquirePrcbLock(Prcb);
 
-    /* Check if we have an idle summary */
-    if (KiIdleSummary)
+    /* Check if the processor is idle */
+    if (KiIdleSummary & Prcb->SetMember)
     {
         /* Clear it and set this thread as the next one */
-        KiIdleSummary = 0;
+        InterlockedBitTestAndResetAffinity(&KiIdleSummary, Prcb->Number);
         Thread->State = Standby;
         Prcb->NextThread = Thread;
 
         /* Unlock the PRCB and return */
         KiReleasePrcbLock(Prcb);
+
+        /* Check if we're running on another CPU */
+        if (KeGetCurrentProcessorNumber() != Processor)
+        {
+            /* We are, send an IPI */
+            KiIpiSend(AFFINITY_MASK(Processor), IPI_DPC);
+        }
+
         return;
     }