[Support] On Windows 11, fix an affinity mask issue on large core count machines

[aganea]

Summary

Recent Windows 11 and Windows Server 2022 changed the way they assign 'processor groups' to a starting PE. Before Windows 11 and Windows Server2022, only one processor group was assigned by default, then the program was responsible for dispatching its own threads on more 'processor groups'. That is what D71775 was doing, allowing LLVM programs use all threads on many cores machines.

After Windows 11 and Windows Server 2022, the OS takes care of that. This has an adverse effect reported in PR56618 which is that using ::GetProcessAffinityMask() API in some edge cases seems buggy now. That API was used to detect if an affinity mask was set, and adjust accordingly the available threads for a ThreadPool.

From the looks of it this seems to be a TOCTOU bug, the OS assigns a default affinity mask on the process group where the PE is started, however later the PE's main() thread runs on a different-sized process group. On Windows, the max size for an affinity mask is 64. In our case, when running on a n2d-highcpu-224 GCE instance, we're seeing 4 processor groups, 2 of size 64 and 2 others of size 48, which makes a total 224 vCPUs. The Windows OS randomly assigns a starting process mask of either (2^64)-1 or (2^48)-1 bits. In some edge cases, the main thread calling ::GetProcessAffinityMask() randomly runs on a different process group, thus making hard for a program to determine if a custom affinity mask was set or not. This wasn't happening before Windows 11, since only a single 'processor group' was used on PE startup, even on machines with asymmetric processor groups.

With this patch, on one hand, on Windows 11 & Windows Server 2022 we disable manual dispatching of threads on processor groups, and instead let the Windows OS do that. On the other hand, a workaround was added to mitigate the ::GetProcessAffinityMask() issue described above (see Threading.inc, L226).

Fixes PR56618.

[tschuett]

llvm/lib/Support/Windows/Process.inc

480–484

You can use std::option<RTL_OSVERSIONINFOEXW>.

502

This looks like cute way to model nullopt. Would returning std::option<lvm::VersionTuple> be an improvement?

[aganea]

As suggested by @tschuett.

llvm/lib/Support/Windows/Process.inc

502

I made the changes so you can see how it looks from the caller perspective. I think std::option internally in Process.inc makes sense. However externally, returning std::option from llvm::GetWindowsOSVersion() leaks the fact that the Win32 APIs can fail, which clients don't really care about, and requires a bit more client-side code. Semantically clients only need to compare < or > against a version number, and 0 is fine. Please let me know your thoughts.

[tschuett]

It is completely up to you.

But it is kind of dangerous to return 0.

[tschuett]

Thinking about this again. Let's say you are running on victious Windows 5000 machine, fail to detect the version, and return 0. All version comparisons are bogus. I would prefer to return an optional to communicate that this API is fallible.

[rnk]

I'm preparing for leave and don't have much time to review, so I'm not likely to get to this on a helpful timeline.

[aaron.ballman]

llvm/lib/Support/Windows/Process.inc

485–504

RtlGetVersion came with Windows 2000, which is older than the oldest version of Windows we support. I think we can safely assert that this interface is loaded.

491–493

According to MSDN, this API only returns STATUS_SUCCESS (https://learn.microsoft.com/en-us/windows-hardware/drivers/ddi/wdm/nf-wdm-rtlgetversion), so I wonder if we can get away with an assert here as well, and then we don't have to worry about this interface ever resulting in no valid version information, which can propagate to the callers.

502

I think std::option internally in Process.inc makes sense.

Agreed, I think it's an improvement.

However externally, returning std::option from llvm::GetWindowsOSVersion() leaks the fact that the Win32 APIs can fail, which clients don't really care about, and requires a bit more client-side code. Semantically clients only need to compare < or > against a version number, and 0 is fine. Please let me know your thoughts.

I agree; the very first thing I did on this review was add a comment to that declaration asking for comments explaining when the result might not exist and how callers should react in that case. I'm not 100% sure we can remove the optionality from the return type, but I found some hints that suggest we can explore it.

[wjschmidt]

Our local testing with this patch shows that it solves our issue. Thanks very much! Looking forward to it landing.

[aganea]

Address comments.

[thieta]

Is there any point in adding any tests here that can be useful? Probably not since we are working with the Windows API right? I don't have much to comment on, but it seems reasonable if it solves the problem.

[aganea]

In D138747#3974797, @thieta wrote:

Is there any point in adding any tests here that can be useful? Probably not since we are working with the Windows API right? I don't have much to comment on, but it seems reasonable if it solves the problem.

The best test is the usage of the ThreadPool. However: 1. it needs a significant enough CPU load to test all sockets (as in linking clang.exe). 2. it requires a VM with WinServer 2022 and at least two or more sockets/groups. and 3. even in that case, I don't have a good way to check that all processors (as in cores) are fully utilized. The OS is now free to do what it pleases to distribute the CPU load. Unless someone has something to suggest, I don't see how to test cover this change.

[mehdi_amini]

Thanks for the detailed description, and kudos for figuring this out.
I don’t feel qualified enough about windows to review this right now, it seems fine to me from what I gathered though.
Since this is a scheduling change I think it is fine to not have unit tests for this.

Nit: you have a typo in the description “one one hand”

[aganea]

+other Windows experts that perhaps can review this @hans @zero9178 @aaron.ballman @mstorsjo @aeubanks

[aaron.ballman]

LGTM aside from a minor coding style issue; you should also add a release note about the bug fix. Thank you!

llvm/lib/Support/Windows/Threading.inc

164–176

Minor nits for coding style.

[aganea]

Thanks for reviewing @aaron.ballman! I applied the changes and updated the release notes in the commit.

llvm/docs/ReleaseNotes.rst

 Changes to the Windows Target
 -----------------------------
 
 * For MinGW, generate embedded ``-exclude-symbols:`` directives for symbols
   with hidden visibility, omitting them from automatic export of all symbols.
   This roughly makes hidden visibility work like it does for other object
   file formats.
 
+* When using multi-threaded LLVM tools (such as LLD) on a Windows host with a
+  large number of processors or CPU sockets, previously the LLVM ThreadPool
+  would span out threads to use all processors.
+  Starting with Windows Server 2022 and Windows 11, the behavior has changed,
+  the OS now spans out threads automatically to all processors. This also fixes
+  an affinity mask issue.
+  (`D138747 <https://reviews.llvm.org/D138747>`_)
+
 Changes to the X86 Backend
 --------------------------
 
 * Add support for the ``RDMSRLIST and WRMSRLIST`` instructions.
 * Add support for the ``WRMSRNS`` instruction.
 * Support ISA of ``AMX-FP16`` which contains ``tdpfp16ps`` instruction.
 * Support ISA of ``CMPCCXADD``.
 * Support ISA of ``AVX-IFMA``.

llvm/include/llvm/Support/Windows/WindowsSupport.h

 namespace llvm {
 
 /// Determines if the program is running on Windows 8 or newer. This
 /// reimplements one of the helpers in the Windows 8.1 SDK, which are intended
 /// to supercede raw calls to GetVersionEx. Old SDKs, Cygwin, and MinGW don't
 /// yet have VersionHelpers.h, so we have our own helper.
 bool RunningWindows8OrGreater();
 
+/// Determines if the program is running on Windows 11 or Windows Server 2022.
+bool RunningWindows11OrGreater();
+
 /// Returns the Windows version as Major.Minor.0.BuildNumber. Uses
 /// RtlGetVersion or GetVersionEx under the hood depending on what is available.
 /// GetVersionEx is deprecated, but this API exposes the build number which can
 /// be useful for working around certain kernel bugs.
 llvm::VersionTuple GetWindowsOSVersion();
 
 bool MakeErrMsg(std::string *ErrMsg, const std::string &prefix);
 

llvm/lib/Support/Windows/Process.inc

unsigned Process::GetRandomNumber() {
 
   // If that fails, fall back to pseudo-random numbers.
   return GetPseudoRandomNumber();
 }
 
 typedef NTSTATUS(WINAPI *RtlGetVersionPtr)(PRTL_OSVERSIONINFOW);
 #define STATUS_SUCCESS ((NTSTATUS)0x00000000L)
 
-llvm::VersionTuple llvm::GetWindowsOSVersion() {
+static RTL_OSVERSIONINFOEXW GetWindowsVer() {
+  auto getVer = []() -> RTL_OSVERSIONINFOEXW {
     HMODULE hMod = ::GetModuleHandleW(L"ntdll.dll");
-  if (hMod) {
+    assert(hMod);
+
     auto getVer = (RtlGetVersionPtr)::GetProcAddress(hMod, "RtlGetVersion");
-    if (getVer) {
-      RTL_OSVERSIONINFOEXW info{};
-      info.dwOSVersionInfoSize = sizeof(info);
-      if (getVer((PRTL_OSVERSIONINFOW)&info) == STATUS_SUCCESS) {
-        return llvm::VersionTuple(info.dwMajorVersion, info.dwMinorVersion, 0,
-                                  info.dwBuildNumber);
-      }
-    }
-  }
-  return llvm::VersionTuple(0, 0, 0, 0);
-}
+    assert(getVer);
+
+    NTSTATUS r = getVer((PRTL_OSVERSIONINFOW)&info);
+    (void)r;
+    assert(r == STATUS_SUCCESS);
+
+    return info;
+  };
+  static RTL_OSVERSIONINFOEXW info = getVer();
+  return info;
+}
+
+llvm::VersionTuple llvm::GetWindowsOSVersion() {
+  RTL_OSVERSIONINFOEXW info = GetWindowsVer();
 
 bool llvm::RunningWindows8OrGreater() {
   // Windows 8 is version 6.2, service pack 0.
   return GetWindowsOSVersion() >= llvm::VersionTuple(6, 2, 0, 0);
 }
 
+bool llvm::RunningWindows11OrGreater() {
+  RTL_OSVERSIONINFOEXW info = GetWindowsVer();
+  auto ver = llvm::VersionTuple(info.dwMajorVersion, info.dwMinorVersion, 0,
+                                info.dwBuildNumber);
+
+  // Windows Server 2022
+  if (info.wProductType == VER_NT_SERVER)
+    return ver >= llvm::VersionTuple(10, 0, 0, 20348);
+
+  // Windows 11
+  return ver >= llvm::VersionTuple(10, 0, 0, 22000);
+}
+
 [[noreturn]] void Process::ExitNoCleanup(int RetCode) {
   TerminateProcess(GetCurrentProcess(), RetCode);
   llvm_unreachable("TerminateProcess doesn't return");
 }

llvm/lib/Support/Windows/Threading.inc

for (auto *Curr = Info; Curr < End;
         continue;
       Fn(Curr);
     }
   }
   free(Info);
   return true;
 }
 
+static std::optional<std::vector<USHORT>> getActiveGroups() {
+  USHORT Count = 0;
+  if (::GetProcessGroupAffinity(GetCurrentProcess(), &Count, nullptr))
+    return std::nullopt;
+
+  if (GetLastError() != ERROR_INSUFFICIENT_BUFFER)
+    return std::nullopt;
+
+  std::vector<USHORT> Groups;
+  Groups.resize(Count);
+  if (!::GetProcessGroupAffinity(GetCurrentProcess(), &Count, Groups.data()))
+    return std::nullopt;
+
+  return Groups;
+}
+
 static ArrayRef<ProcessorGroup> getProcessorGroups() {
   auto computeGroups = []() {
     SmallVector<ProcessorGroup, 4> Groups;
 
     auto HandleGroup = [&](SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *ProcInfo) {
       GROUP_RELATIONSHIP &El = ProcInfo->Group;
       for (unsigned J = 0; J < El.ActiveGroupCount; ++J) {
         ProcessorGroup G;
auto HandleProc = [&](SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *ProcInfo) {
         NumHyperThreads = std::bitset<64>(El.GroupMask[0].Mask).count();
       unsigned I = El.GroupMask[0].Group;
       Groups[I].ThreadsPerCore = NumHyperThreads;
     };
 
     if (!IterateProcInfo(RelationProcessorCore, HandleProc))
       return std::vector<ProcessorGroup>();
 
+    auto ActiveGroups = getActiveGroups();
+    if (!ActiveGroups)
+      return std::vector<ProcessorGroup>();
+
     // If there's an affinity mask set, assume the user wants to constrain the
     // current process to only a single CPU group. On Windows, it is not
     // possible for affinity masks to cross CPU group boundaries.
     DWORD_PTR ProcessAffinityMask = 0, SystemAffinityMask = 0;
     if (::GetProcessAffinityMask(GetCurrentProcess(), &ProcessAffinityMask,
-                                 &SystemAffinityMask) &&
-        ProcessAffinityMask != SystemAffinityMask) {
-      // We don't expect more that 4 CPU groups on Windows (256 processors).
-      USHORT GroupCount = 4;
-      USHORT GroupArray[4]{};
-      if (::GetProcessGroupAffinity(GetCurrentProcess(), &GroupCount,
-                                    GroupArray)) {
-        assert(GroupCount == 1 &&
-               "On startup, a program is expected to be assigned only to "
-               "one processor group!");
-        unsigned CurrentGroupID = GroupArray[0];
+                                 &SystemAffinityMask)) {
+
+      if (ProcessAffinityMask != SystemAffinityMask) {
+        if (llvm::RunningWindows11OrGreater() && ActiveGroups->size() > 1) {
+          // The process affinity mask is spurious, due to an OS bug, ignore it.
+          return std::vector<ProcessorGroup>(Groups.begin(), Groups.end());
+        }
+
+        assert(ActiveGroups->size() == 1 &&
+               "When an affinity mask is set, the process is expected to be "
+               "assigned to a single processor group!");
+
+        unsigned CurrentGroupID = (*ActiveGroups)[0];
         ProcessorGroup NewG{Groups[CurrentGroupID]};
         NewG.Affinity = ProcessAffinityMask;
         NewG.UsableThreads = countPopulation(ProcessAffinityMask);
         Groups.clear();
         Groups.push_back(NewG);
       }
     }
-
     return std::vector<ProcessorGroup>(Groups.begin(), Groups.end());
   };
   static auto Groups = computeGroups();
   return ArrayRef<ProcessorGroup>(Groups);
 }
 
 template <typename R, typename UnaryPredicate>
 static unsigned aggregate(R &&Range, UnaryPredicate P) {
llvm::ThreadPoolStrategy::compute_cpu_socket(unsigned ThreadPoolNum) const {
 
   // Assumes the same number of hardware threads per CPU socket.
   return (ThreadPoolNum * Groups.size()) / compute_thread_count();
 }
 
 // Assign the current thread to a more appropriate CPU socket or CPU group
 void llvm::ThreadPoolStrategy::apply_thread_strategy(
     unsigned ThreadPoolNum) const {
+
+  // After Windows 11 and Windows Server 2022, let the OS do the scheduling,
+  // since a process automatically gains access to all processor groups.
+  if (llvm::RunningWindows11OrGreater())
+    return;
+
   std::optional<unsigned> Socket = compute_cpu_socket(ThreadPoolNum);
   if (!Socket)
     return;
   ArrayRef<ProcessorGroup> Groups = getProcessorGroups();
   GROUP_AFFINITY Affinity{};
   Affinity.Group = Groups[*Socket].ID;
   Affinity.Mask = Groups[*Socket].Affinity;
   SetThreadGroupAffinity(GetCurrentThread(), &Affinity, nullptr);

llvm/unittests/Support/ThreadPool.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Host.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/Threading.h"
 
+#ifdef _WIN32
+#include "llvm/Support/Windows/WindowsSupport.h"
+#endif
+
 #include <chrono>
 #include <thread>
 
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 // Fixture for the unittests, allowing to *temporarily* disable the unittests
std::mutex Lock;
     }
     setMainThreadReady();
   }
   return ThreadsUsed.takeVector();
 }
 
 TEST_F(ThreadPoolTest, AllThreads_UseAllRessources) {
   CHECK_UNSUPPORTED();
+  // After Windows 11, the OS is free to deploy the threads on any CPU socket.
+  // We cannot relibly ensure that all thread affinity mask are covered,
+  // therefore this test should not run.
+  if (llvm::RunningWindows11OrGreater())
+    return;
   std::vector<llvm::BitVector> ThreadsUsed = RunOnAllSockets({});
   ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size());
 }
 
 TEST_F(ThreadPoolTest, AllThreads_OneThreadPerCore) {
   CHECK_UNSUPPORTED();
+  // After Windows 11, the OS is free to deploy the threads on any CPU socket.
+  // We cannot relibly ensure that all thread affinity mask are covered,
+  // therefore this test should not run.
+  if (llvm::RunningWindows11OrGreater())
+    return;
   std::vector<llvm::BitVector> ThreadsUsed =
       RunOnAllSockets(llvm::heavyweight_hardware_concurrency());
   ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size());
 }
 
 // From TestMain.cpp.
 extern const char *TestMainArgv0;