diff --git a/llvm/include/llvm/Support/Program.h b/llvm/include/llvm/Support/Program.h --- a/llvm/include/llvm/Support/Program.h +++ b/llvm/include/llvm/Support/Program.h @@ -14,6 +14,7 @@ #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" @@ -125,9 +126,11 @@ ///< string is non-empty upon return an error occurred while invoking the ///< program. bool *ExecutionFailed = nullptr, - Optional *ProcStat = nullptr ///< If non-zero, provides - /// a pointer to a structure in which process execution statistics will be - /// stored. + Optional *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. @@ -140,7 +143,8 @@ ArrayRef> 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. diff --git a/llvm/lib/Support/Program.cpp b/llvm/lib/Support/Program.cpp --- a/llvm/lib/Support/Program.cpp +++ b/llvm/lib/Support/Program.cpp @@ -26,17 +26,20 @@ static bool Execute(ProcessInfo &PI, StringRef Program, ArrayRef Args, Optional> Env, ArrayRef> Redirects, - unsigned MemoryLimit, std::string *ErrMsg); + unsigned MemoryLimit, std::string *ErrMsg, + BitVector *AffinityMask); int sys::ExecuteAndWait(StringRef Program, ArrayRef Args, Optional> Env, ArrayRef> Redirects, unsigned SecondsToWait, unsigned MemoryLimit, std::string *ErrMsg, bool *ExecutionFailed, - Optional *ProcStat) { + Optional *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 = @@ -55,12 +58,13 @@ Optional> Env, ArrayRef> 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; diff --git a/llvm/lib/Support/Unix/Program.inc b/llvm/lib/Support/Unix/Program.inc --- a/llvm/lib/Support/Unix/Program.inc +++ b/llvm/lib/Support/Unix/Program.inc @@ -174,7 +174,8 @@ static bool Execute(ProcessInfo &PI, StringRef Program, ArrayRef Args, Optional> Env, ArrayRef> 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() + @@ -182,6 +183,9 @@ return false; } + assert(!AffinityMask && "Starting a process with an affinity mask is " + "currently not supported on Unix!"); + BumpPtrAllocator Allocator; StringSaver Saver(Allocator); std::vector ArgVector, EnvVector; diff --git a/llvm/lib/Support/Windows/Program.inc b/llvm/lib/Support/Windows/Program.inc --- a/llvm/lib/Support/Windows/Program.inc +++ b/llvm/lib/Support/Windows/Program.inc @@ -171,7 +171,8 @@ static bool Execute(ProcessInfo &PI, StringRef Program, ArrayRef Args, Optional> Env, ArrayRef> 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"; @@ -277,11 +278,15 @@ return false; } + unsigned CreateFlags = CREATE_UNICODE_ENVIRONMENT; + if (AffinityMask) + CreateFlags |= CREATE_SUSPENDED; + std::vector 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 @@ -329,6 +334,13 @@ } } + // Set the affinity mask + if (AffinityMask) { + ::SetProcessAffinityMask(pi.hProcess, + (DWORD_PTR)AffinityMask->getData().front()); + ::ResumeThread(pi.hThread); + } + return true; } diff --git a/llvm/lib/Support/Windows/Threading.inc b/llvm/lib/Support/Windows/Threading.inc --- a/llvm/lib/Support/Windows/Threading.inc +++ b/llvm/lib/Support/Windows/Threading.inc @@ -195,14 +195,27 @@ if (!IterateProcInfo(RelationProcessorCore, HandleProc)) return std::vector(); - // 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, 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]; + ProcessorGroup NewG{Groups[CurrentGroupID]}; + NewG.Affinity = ProcessAffinityMask; + NewG.UsableThreads = countPopulation(ProcessAffinityMask); Groups.clear(); Groups.push_back(NewG); - break; } } diff --git a/llvm/unittests/Support/ThreadPool.cpp b/llvm/unittests/Support/ThreadPool.cpp --- a/llvm/unittests/Support/ThreadPool.cpp +++ b/llvm/unittests/Support/ThreadPool.cpp @@ -8,11 +8,13 @@ #include "llvm/Support/ThreadPool.h" -#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SetVector.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" @@ -71,7 +73,7 @@ void SetUp() override { MainThreadReady = false; } - void RunOnAllSockets(ThreadPoolStrategy S); + std::vector RunOnAllSockets(ThreadPoolStrategy S); std::condition_variable WaitMainThread; std::mutex WaitMainThreadMutex; @@ -169,15 +171,16 @@ #if LLVM_ENABLE_THREADS == 1 -void ThreadPoolTest::RunOnAllSockets(ThreadPoolStrategy S) { +std::vector +ThreadPoolTest::RunOnAllSockets(ThreadPoolStrategy S) { // 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; + return {}; - llvm::DenseSet ThreadsUsed; + llvm::SetVector ThreadsUsed; std::mutex Lock; { std::condition_variable AllThreads; @@ -198,7 +201,7 @@ ThreadsUsed.insert(Mask); }); } - ASSERT_EQ(true, ThreadsUsed.empty()); + EXPECT_EQ(true, ThreadsUsed.empty()); { std::unique_lock Guard(AllThreadsLock); AllThreads.wait(Guard, @@ -206,17 +209,67 @@ } setMainThreadReady(); } - ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size()); + return ThreadsUsed.takeVector(); } TEST_F(ThreadPoolTest, AllThreads_UseAllRessources) { CHECK_UNSUPPORTED(); - RunOnAllSockets({}); + std::vector ThreadsUsed = RunOnAllSockets({}); + ASSERT_EQ(llvm::get_cpus(), ThreadsUsed.size()); } TEST_F(ThreadPoolTest, AllThreads_OneThreadPerCore) { CHECK_UNSUPPORTED(); - RunOnAllSockets(llvm::heavyweight_hardware_concurrency()); + std::vector ThreadsUsed = + RunOnAllSockets(llvm::heavyweight_hardware_concurrency()); + 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 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")) { + std::vector ThreadsUsed = RunOnAllSockets({}); + // Ensure the threads only ran 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