Fix NUMA support on Windows (#15049)

* osNumaNodes now returns the right number of nodes
* thread affinity is now correctly set

TODO: no noticeable performance improvement.
does windows already distribute threads in a NUMA-aware fashion?

Test Plan:
* validate
* local tests on a NUMA machine

Reviewers: bgamari, erikd, simonmar

Reviewed By: bgamari, simonmar

Subscribers: thomie, carter

Differential Revision: https://phabricator.haskell.org/D4607

2 files changed, 17 insertions, 9 deletions

diff --git a/rts/win32/OSMem.c b/rts/win32/OSMem.c
index 534cd15fa6..d05151ce61 100644
--- a/rts/win32/OSMem.c
+++ b/rts/win32/OSMem.c
@@ -510,9 +510,18 @@ uint32_t osNumaNodes(void)
     static ULONG numNumaNodes = 0;
 
     /* Cache the amount of NUMA nodes. */
-    if (!numNumaNodes && !GetNumaHighestNodeNumber(&numNumaNodes))
+    if (!numNumaNodes)
     {
-        numNumaNodes = 1;
+        if (GetNumaHighestNodeNumber(&numNumaNodes))
+        {
+            // GetNumaHighestNodeNumber returns the highest node number
+            // i.e: 0 for a non-NUMA system, and >0 for a NUMA system, so add a 1.
+            numNumaNodes += 1;
+        }
+        else
+        {
+            numNumaNodes = 1;
+        }
     }
 
     return numNumaNodes;
@@ -520,12 +529,12 @@ uint32_t osNumaNodes(void)
 
 uint64_t osNumaMask(void)
 {
-    uint64_t numaMask;
-    if (!GetNumaNodeProcessorMask(0, &numaMask))
-    {
-        return 1;
+    // the concept of a numa node mask (c.f. numa_get_mems_allowed on POSIX)
+    // doesn't exist on Windows. Thus, all nodes are allowed.
+    if (osNumaNodes() > sizeof(StgWord)*8) {
+        barf("osNumaMask: too many NUMA nodes (%d)", osNumaNodes());
     }
-    return numaMask;
+    return (1 << osNumaNodes()) - 1;
 }
 
 void osBindMBlocksToNode(
diff --git a/rts/win32/OSThreads.c b/rts/win32/OSThreads.c
index b1a98cecbd..cc673532d7 100644
--- a/rts/win32/OSThreads.c
+++ b/rts/win32/OSThreads.c
@@ -579,8 +579,7 @@ void setThreadNode (uint32_t node)
     if (osNumaAvailable())
     {
         StgWord mask = 0;
-        mask |= 1 << node;
-        if (!SetThreadAffinityMask(GetCurrentThread(), mask))
+        if (!GetNumaNodeProcessorMask(node, &mask) && !SetThreadAffinityMask(GetCurrentThread(), mask))
         {
             sysErrorBelch(
                 "setThreadNode: Error setting affinity of thread to NUMA node `%u': %lu.",