[libcxx-commits] [libcxx] [libc++] Fixed get count threads multi-CPU system with NUMA architecture (#72267) (PR #72270)

[Herman Semenov via libcxx-commits]

https://github.com/GermanAizek updated https://github.com/llvm/llvm-project/pull/72270

>From d1fa25f37631b8b33a71fbe9eb4ea89e3a47b723 Mon Sep 17 00:00:00 2001
From: German Semenov <GermanAizek at yandex.ru>
Date: Thu, 16 Nov 2023 12:12:28 +0300
Subject: [PATCH] [libc++] Fixed get count threads for multi-cpu system with
 NUMA architecture  (#72267)

---
 libcxx/src/thread.cpp | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/libcxx/src/thread.cpp b/libcxx/src/thread.cpp
index 289c457cd5a5cbd..16ab408a459448d 100644
--- a/libcxx/src/thread.cpp
+++ b/libcxx/src/thread.cpp
@@ -81,9 +81,43 @@ thread::hardware_concurrency() noexcept
         return 0;
     return static_cast<unsigned>(result);
 #elif defined(_LIBCPP_WIN32API)
+  // This implementation supports both conventional single-cpu PC configurations
+  // and multi-cpu system on NUMA (Non-uniform_memory_access) architecture
+  DWORD length = 0;
+  const auto single_cpu_concurrency = []() noexcept -> unsigned int {
     SYSTEM_INFO info;
     GetSystemInfo(&info);
     return info.dwNumberOfProcessors;
+  };
+  if (GetLogicalProcessorInformationEx(RelationAll, nullptr, &length) != FALSE) {
+    return single_cpu_concurrency();
+  }
+  if (GetLastError() != ERROR_INSUFFICIENT_BUFFER) {
+    return single_cpu_concurrency();
+  }
+  std::unique_ptr<void, void (*)(void*)> buffer(std::malloc(length), std::free);
+  if (!buffer) {
+    return single_cpu_concurrency();
+  }
+  auto* mem = reinterpret_cast<unsigned char*>(buffer.get());
+  if (GetLogicalProcessorInformationEx(
+          RelationAll, reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>(mem), &length) == false) {
+    return single_cpu_concurrency();
+  }
+  DWORD i = 0;
+  unsigned int concurrency = 0;
+  while (i < length) {
+    const auto* proc = reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX>(mem + i);
+    if (proc->Relationship == RelationProcessorCore) {
+      for (WORD group = 0; group < proc->Processor.GroupCount; ++group) {
+        for (KAFFINITY mask = proc->Processor.GroupMask[group].Mask; mask != 0; mask >>= 1) {
+          concurrency += mask & 1;
+        }
+      }
+    }
+    i += proc->Size;
+  }
+  return concurrency;
 #else  // defined(CTL_HW) && defined(HW_NCPU)
     // TODO: grovel through /proc or check cpuid on x86 and similar
     // instructions on other architectures.