[Support] On Windows, take the affinity mask into account

[aganea]

Summary

Now using the affinity mask when it's set, thus limiting the number of processors available to the ThreadPool.
For example:

> start /B /AFFINITY 0xF lld-link.exe ...

Would let LLD only use 4 hyper-threads.

This also fixes an issue on a 48-core Windows Server 2019, which was preventing from using both CPU sockets. In that case, ProcessorGroup::AllThreads was retrieved as 64 for each group, but ProcessorGroup::UsableThreads was set to 48 (which is the max possible for each group). The previous code in llvm/lib/Support/Windows/Threading.inc L201 was improperly assuming those values would be equal when the affinity mask was fully set. I've tested the new code on a wide variety of configurations & OSes (Windows Server 2016, 2019 and Windows 10).

[amccarth]

llvm/lib/Support/Windows/Threading.inc

200

nit: In "group boundaries", _group_ should be singular.

llvm/unittests/Support/ThreadPool.cpp

178

I think the generality introduced here may be more complexity than it's worth. It imposes a high boilerplate cost on each caller.

Both tests that currently call this function simply look at the size (population count) of the bit vector. Unless you expect to use this in more ways, it would be simpler to return a copy of the bit vector (or just its size) and let the individual tests inspect it. That would still move the assertions to the individual tests, which is valuable.

[aganea]

As suggested by Adrian.

[amccarth]

Thanks for the changes!

[aganea]

Thank you for your time! :)

llvm/include/llvm/Support/Program.h

 //===- llvm/Support/Program.h ------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the llvm::sys::Program class.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_SUPPORT_PROGRAM_H
 #define LLVM_SUPPORT_PROGRAM_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/ErrorOr.h"
 #include <chrono>
 #include <system_error>
 
 namespace llvm {
int ExecuteAndWait(
       ///< of memory can be allocated by process. If memory usage will be
       ///< higher limit, the child is killed and this call returns. If zero
       ///< - no memory limit.
       std::string *ErrMsg = nullptr, ///< If non-zero, provides a pointer to a
       ///< string instance in which error messages will be returned. If the
       ///< string is non-empty upon return an error occurred while invoking the
       ///< program.
       bool *ExecutionFailed = nullptr,
-      Optional<ProcessStatistics> *ProcStat = nullptr ///< If non-zero, provides
-      /// a pointer to a structure in which process execution statistics will be
-      /// stored.
+      Optional<ProcessStatistics> *ProcStat = nullptr, ///< If non-zero,
+      /// provides a pointer to a structure in which process execution
+      /// statistics will be stored.
+      BitVector *AffinityMask = nullptr ///< CPUs or processors the new
+                                        /// program shall run on.
   );
 
   /// Similar to ExecuteAndWait, but returns immediately.
   /// @returns The \see ProcessInfo of the newly launched process.
   /// \note On Microsoft Windows systems, users will need to either call
   /// \see Wait until the process finished execution or win32 CloseHandle() API
   /// on ProcessInfo.ProcessHandle to avoid memory leaks.
   ProcessInfo ExecuteNoWait(StringRef Program, ArrayRef<StringRef> Args,
                             Optional<ArrayRef<StringRef>> Env,
                             ArrayRef<Optional<StringRef>> Redirects = {},
                             unsigned MemoryLimit = 0,
                             std::string *ErrMsg = nullptr,
-                            bool *ExecutionFailed = nullptr);
+                            bool *ExecutionFailed = nullptr,
+                            BitVector *AffinityMask = nullptr);
 
   /// Return true if the given arguments fit within system-specific
   /// argument length limits.
   bool commandLineFitsWithinSystemLimits(StringRef Program,
                                          ArrayRef<StringRef> Args);
 
   /// Return true if the given arguments fit within system-specific
   /// argument length limits.

llvm/lib/Support/Program.cpp

 //===----------------------------------------------------------------------===//
 //=== WARNING: Implementation here must contain only TRULY operating system
 //===          independent code.
 //===----------------------------------------------------------------------===//
 
 static bool Execute(ProcessInfo &PI, StringRef Program,
                     ArrayRef<StringRef> Args, Optional<ArrayRef<StringRef>> Env,
                     ArrayRef<Optional<StringRef>> Redirects,
-                    unsigned MemoryLimit, std::string *ErrMsg);
+                    unsigned MemoryLimit, std::string *ErrMsg,
+                    BitVector *AffinityMask);
 
 int sys::ExecuteAndWait(StringRef Program, ArrayRef<StringRef> Args,
                         Optional<ArrayRef<StringRef>> Env,
                         ArrayRef<Optional<StringRef>> Redirects,
                         unsigned SecondsToWait, unsigned MemoryLimit,
                         std::string *ErrMsg, bool *ExecutionFailed,
-                        Optional<ProcessStatistics> *ProcStat) {
+                        Optional<ProcessStatistics> *ProcStat,
+                        BitVector *AffinityMask) {
   assert(Redirects.empty() || Redirects.size() == 3);
   ProcessInfo PI;
-  if (Execute(PI, Program, Args, Env, Redirects, MemoryLimit, ErrMsg)) {
+  if (Execute(PI, Program, Args, Env, Redirects, MemoryLimit, ErrMsg,
+              AffinityMask)) {
     if (ExecutionFailed)
       *ExecutionFailed = false;
     ProcessInfo Result =
         Wait(PI, SecondsToWait, /*WaitUntilTerminates=*/SecondsToWait == 0,
              ErrMsg, ProcStat);
     return Result.ReturnCode;
   }
 
   if (ExecutionFailed)
     *ExecutionFailed = true;
 
   return -1;
 }
 
 ProcessInfo sys::ExecuteNoWait(StringRef Program, ArrayRef<StringRef> Args,
                                Optional<ArrayRef<StringRef>> Env,
                                ArrayRef<Optional<StringRef>> Redirects,
                                unsigned MemoryLimit, std::string *ErrMsg,
-                               bool *ExecutionFailed) {
+                               bool *ExecutionFailed, BitVector *AffinityMask) {
   assert(Redirects.empty() || Redirects.size() == 3);
   ProcessInfo PI;
   if (ExecutionFailed)
     *ExecutionFailed = false;
-  if (!Execute(PI, Program, Args, Env, Redirects, MemoryLimit, ErrMsg))
+  if (!Execute(PI, Program, Args, Env, Redirects, MemoryLimit, ErrMsg,
+               AffinityMask))
     if (ExecutionFailed)
       *ExecutionFailed = true;
 
   return PI;
 }
 
 bool sys::commandLineFitsWithinSystemLimits(StringRef Program,
                                             ArrayRef<const char *> Args) {

llvm/lib/Support/Unix/Program.inc

for (StringRef S : Strings)
     Result.push_back(Saver.save(S).data());
   Result.push_back(nullptr);
   return Result;
 }
 
 static bool Execute(ProcessInfo &PI, StringRef Program,
                     ArrayRef<StringRef> Args, Optional<ArrayRef<StringRef>> Env,
                     ArrayRef<Optional<StringRef>> Redirects,
-                    unsigned MemoryLimit, std::string *ErrMsg) {
+                    unsigned MemoryLimit, std::string *ErrMsg,
+                    BitVector *AffinityMask) {
   if (!llvm::sys::fs::exists(Program)) {
     if (ErrMsg)
       *ErrMsg = std::string("Executable \"") + Program.str() +
                 std::string("\" doesn't exist!");
     return false;
   }
 
+  assert(!AffinityMask && "Starting a process with an affinity mask is "
+                          "currently not supported on Unix!");
+
   BumpPtrAllocator Allocator;
   StringSaver Saver(Allocator);
   std::vector<const char *> ArgVector, EnvVector;
   const char **Argv = nullptr;
   const char **Envp = nullptr;
   ArgVector = toNullTerminatedCStringArray(Args, Saver);
   Argv = ArgVector.data();
   if (Env) {

llvm/lib/Support/Windows/Program.inc

static HANDLE RedirectIO(Optional<StringRef> Path, int fd,
   return h;
 }
 
 }
 
 static bool Execute(ProcessInfo &PI, StringRef Program,
                     ArrayRef<StringRef> Args, Optional<ArrayRef<StringRef>> Env,
                     ArrayRef<Optional<StringRef>> Redirects,
-                    unsigned MemoryLimit, std::string *ErrMsg) {
+                    unsigned MemoryLimit, std::string *ErrMsg,
+                    BitVector *AffinityMask) {
   if (!sys::fs::can_execute(Program)) {
     if (ErrMsg)
       *ErrMsg = "program not executable";
     return false;
   }
 
   // can_execute may succeed by looking at Program + ".exe". CreateProcessW
   // will implicitly add the .exe if we provide a command line without an
static bool Execute(ProcessInfo &PI, StringRef Program,
   SmallVector<wchar_t, MAX_PATH> ProgramUtf16;
   if (std::error_code ec = sys::windows::widenPath(Program, ProgramUtf16)) {
     SetLastError(ec.value());
     MakeErrMsg(ErrMsg,
                std::string("Unable to convert application name to UTF-16"));
     return false;
   }
 
+  unsigned CreateFlags = CREATE_UNICODE_ENVIRONMENT;
+  if (AffinityMask)
+    CreateFlags |= CREATE_SUSPENDED;
+
   std::vector<wchar_t> CommandUtf16(Command.size() + 1, 0);
   std::copy(Command.begin(), Command.end(), CommandUtf16.begin());
   BOOL rc = CreateProcessW(ProgramUtf16.data(), CommandUtf16.data(), 0, 0, TRUE,
-                           CREATE_UNICODE_ENVIRONMENT,
-                           EnvBlock.empty() ? 0 : EnvBlock.data(), 0, &si, &pi);
+                           CreateFlags, EnvBlock.empty() ? 0 : EnvBlock.data(),
+                           0, &si, &pi);
   DWORD err = GetLastError();
 
   // Regardless of whether the process got created or not, we are done with
   // the handles we created for it to inherit.
   CloseHandle(si.hStdInput);
   CloseHandle(si.hStdOutput);
   CloseHandle(si.hStdError);
 
if (!success) {
       SetLastError(GetLastError());
       MakeErrMsg(ErrMsg, std::string("Unable to set memory limit"));
       TerminateProcess(pi.hProcess, 1);
       WaitForSingleObject(pi.hProcess, INFINITE);
       return false;
     }
   }
 
+  // Set the affinity mask
+  if (AffinityMask) {
+    ::SetProcessAffinityMask(pi.hProcess,
+                             (DWORD_PTR)AffinityMask->getData().front());
+    ::ResumeThread(pi.hThread);
+  }
+
   return true;
 }
 
 static bool argNeedsQuotes(StringRef Arg) {
   if (Arg.empty())
     return true;
   return StringRef::npos != Arg.find_first_of("\t \"&\'()*<>\\`^|\n");
 }

llvm/lib/Support/Windows/Threading.inc

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>();
 
-    // If there's an affinity mask set on one of the CPUs, then assume the user
-    // wants to constrain the current process to only a single CPU.
-    for (auto &G : Groups) {
-      if (G.UsableThreads != G.AllThreads) {
-        ProcessorGroup NewG{G};
+    // If there's an affinity mask set, then 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 groups 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];
+        ProcessorGroup NewG{Groups[CurrentGroupID]};
+        NewG.Affinity = ProcessAffinityMask;
+        NewG.UsableThreads = countPopulation(ProcessAffinityMask);
         Groups.clear();
         Groups.push_back(NewG);
-        break;
       }
     }
 
     return std::vector<ProcessorGroup>(Groups.begin(), Groups.end());
   };
   static auto Groups = computeGroups();
   return ArrayRef<ProcessorGroup>(Groups);
 }

llvm/unittests/Support/ThreadPool.cpp

 //========- unittests/Support/ThreadPools.cpp - ThreadPools.h tests --========//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/ThreadPool.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #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"
 
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 // Fixture for the unittests, allowing to *temporarily* disable the unittests
void setMainThreadReady() {
       std::unique_lock<std::mutex> LockGuard(WaitMainThreadMutex);
       MainThreadReady = true;
     }
     WaitMainThread.notify_all();
   }
 
   void SetUp() override { MainThreadReady = false; }
 
-  void RunOnAllSockets(ThreadPoolStrategy S);
+  void RunOnAllSockets(
+      ThreadPoolStrategy S,
+      llvm::function_ref<void(llvm::DenseSet<llvm::BitVector> &)> Fn);
 
   std::condition_variable WaitMainThread;
   std::mutex WaitMainThreadMutex;
   bool MainThreadReady = false;
 };
 
 #define CHECK_UNSUPPORTED() \
   do { \
std::atomic_int checked_in{0};
     ASSERT_EQ(0, checked_in);
     setMainThreadReady();
   }
   ASSERT_EQ(5, checked_in);
 }
 
 #if LLVM_ENABLE_THREADS == 1
 
-void ThreadPoolTest::RunOnAllSockets(ThreadPoolStrategy S) {
+void ThreadPoolTest::RunOnAllSockets(
+    ThreadPoolStrategy S,
+    llvm::function_ref<void(llvm::DenseSet<llvm::BitVector> &)> Fn) {
   // FIXME: Skip these tests on non-Windows because multi-socket system were not
   // tested on Unix yet, and llvm::get_thread_affinity_mask() isn't implemented
   // for Unix.
   Triple Host(Triple::normalize(sys::getProcessTriple()));
   if (!Host.isOSWindows())
     return;
 
   llvm::DenseSet<llvm::BitVector> ThreadsUsed;
std::mutex Lock;
     ASSERT_EQ(true, ThreadsUsed.empty());
     {
       std::unique_lock<std::mutex> Guard(AllThreadsLock);
       AllThreads.wait(Guard,
                       [&]() { return Active == S.compute_thread_count(); });
     }
     setMainThreadReady();
   }
-  ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size());
+  Fn(ThreadsUsed);
 }
 
 TEST_F(ThreadPoolTest, AllThreads_UseAllRessources) {
   CHECK_UNSUPPORTED();
-  RunOnAllSockets({});
+  RunOnAllSockets({}, [](llvm::DenseSet<llvm::BitVector> &ThreadsUsed) {
+    ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size());
+  });
 }
 
 TEST_F(ThreadPoolTest, AllThreads_OneThreadPerCore) {
   CHECK_UNSUPPORTED();
-  RunOnAllSockets(llvm::heavyweight_hardware_concurrency());
+  RunOnAllSockets(llvm::heavyweight_hardware_concurrency(),
+                  [](llvm::DenseSet<llvm::BitVector> &ThreadsUsed) {
+                    ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size());
+                  });
 }
 
+#if defined(_WIN32) // FIXME: implement AffinityMask in Support/Unix/Program.inc
+
+// From TestMain.cpp.
+extern const char *TestMainArgv0;
+
+// Just a reachable symbol to ease resolving of the executable's path.
+static cl::opt<std::string> ThreadPoolTestStringArg1("thread-pool-string-arg1");
+
+#ifdef _MSC_VER
+#define setenv(name, var, ignore) _putenv_s(name, var)
 #endif
+
+TEST_F(ThreadPoolTest, AffinityMask) {
+  CHECK_UNSUPPORTED();
+
+  // Skip this test if less than 4 threads are available.
+  if (llvm::hardware_concurrency().compute_thread_count() < 4)
+    return;
+
+  using namespace llvm::sys;
+  if (getenv("LLVM_THREADPOOL_AFFINITYMASK")) {
+    RunOnAllSockets({}, [](llvm::DenseSet<llvm::BitVector> &ThreadsUsed) {
+      // Ensure the threads only run on CPUs 0-3.
+      for (auto &It : ThreadsUsed)
+        ASSERT_LT(It.getData().front(), 16UL);
+    });
+    return;
+  }
+  std::string Executable =
+      sys::fs::getMainExecutable(TestMainArgv0, &ThreadPoolTestStringArg1);
+  StringRef argv[] = {Executable, "--gtest_filter=ThreadPoolTest.AffinityMask"};
+
+  // Add environment variable to the environment of the child process.
+  int Res = setenv("LLVM_THREADPOOL_AFFINITYMASK", "1", false);
+  ASSERT_EQ(Res, 0);
+
+  std::string Error;
+  bool ExecutionFailed;
+  BitVector Affinity;
+  Affinity.resize(4);
+  Affinity.set(0, 4); // Use CPUs 0,1,2,3.
+  int Ret = sys::ExecuteAndWait(Executable, argv, {}, {}, 0, 0, &Error,
+                                &ExecutionFailed, nullptr, &Affinity);
+  ASSERT_EQ(0, Ret);
+}
+
+#endif // #if _WIN32
+
+#endif // #if LLVM_ENABLE_THREADS == 1