From 4474299a292e8c994b007040a520818a705d0be4 Mon Sep 17 00:00:00 2001
From: Sam Clegg <sbc@chromium.org>
Date: Mon, 15 Sep 2025 12:55:55 -0700
Subject: [PATCH] Honor CPU affinity in ThreadPool::getNumCores

This reports that number of CPUs that are actually usable by the current
process.

This means that I can do something like `taskset -c 0 emcc hello.c -O2`
and the internal call to `wasm-opt` within emcc will only use a single
core.

Ideally this would work on macOS and windows too, but I'm not even sure
how easy it is to control affinity on those OSes.
---
 .github/workflows/ci.yml             |  2 +-
 .github/workflows/create_release.yml |  2 +-
 src/support/threads.cpp              | 24 ++++++++++++++++++++++++
 test/lit/lit.cfg.py                  |  4 ++++
 test/lit/num_cores.wast              | 13 +++++++++++++
 5 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 test/lit/num_cores.wast

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f0141e663f2..43269de8ca0 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -239,7 +239,7 @@ jobs:
     - name: install packages
       run: |
         ./alpine.sh apk update
-        ./alpine.sh apk add build-base cmake git python3 py3-pip clang ninja
+        ./alpine.sh apk add build-base cmake git python3 py3-pip clang ninja util-linux
 
     - name: avoid d8 tests (jsvu is not compatible with alpine)
       run: |
diff --git a/.github/workflows/create_release.yml b/.github/workflows/create_release.yml
index 7d6bb6ed74b..ed8bb2fb6c4 100644
--- a/.github/workflows/create_release.yml
+++ b/.github/workflows/create_release.yml
@@ -135,7 +135,7 @@ jobs:
     - name: install packages
       run: |
         ./alpine.sh apk update
-        ./alpine.sh apk add build-base cmake git python3 clang ninja py3-pip
+        ./alpine.sh apk add build-base cmake git python3 py3-pip clang ninja util-linux
 
     - name: avoid d8 tests (jsvu is not compatible with alpine)
       run: |
diff --git a/src/support/threads.cpp b/src/support/threads.cpp
index 7388986389a..e1bc5d12837 100644
--- a/src/support/threads.cpp
+++ b/src/support/threads.cpp
@@ -20,12 +20,21 @@
 #include <iostream>
 #include <string>
 
+#ifdef __linux__
+#include <sched.h> // For sched_getaffinity
+#endif
+
 #include "compiler-support.h"
+#include "support/debug.h"
 #include "threads.h"
 #include "utilities.h"
 
 // debugging tools
 
+// DEBUG_TYPE is for BYN_TRACE macro.  This tracing can be enabled at runtime
+#define DEBUG_TYPE "threads"
+
+// BINARYEN_THREAD_DEBUG is a build-time setting for detailed thread tracing
 #ifdef BINARYEN_THREAD_DEBUG
 static std::mutex debug;
 #define DEBUG_THREAD(x)                                                        \
@@ -145,9 +154,24 @@ size_t ThreadPool::getNumCores() {
   return 1;
 #else
   size_t num = std::max(1U, std::thread::hardware_concurrency());
+#ifdef __linux__
+  // On linux we can do better since we can get the number of CPU that are
+  // actually usable by the current process.
+  cpu_set_t cpu_set;
+  CPU_ZERO(&cpu_set);
+  if (sched_getaffinity(0, sizeof(cpu_set), &cpu_set) == 0) {
+    num = 0;
+    for (int i = 0; i < CPU_SETSIZE; i++) {
+      if (CPU_ISSET(i, &cpu_set)) {
+        num++;
+      }
+    }
+  }
+#endif
   if (getenv("BINARYEN_CORES")) {
     num = std::stoi(getenv("BINARYEN_CORES"));
   }
+  BYN_TRACE("getNumCores: " << num << "\n");
   return num;
 #endif
 }
diff --git a/test/lit/lit.cfg.py b/test/lit/lit.cfg.py
index 7ae6c943f17..be80ea4fc98 100644
--- a/test/lit/lit.cfg.py
+++ b/test/lit/lit.cfg.py
@@ -1,4 +1,5 @@
 import os
+import sys
 import lit.formats
 
 config.name = "Binaryen lit tests"
@@ -26,6 +27,9 @@
     python = sys.executable.replace('\\', '/')
     config.substitutions.append((tool, python + ' ' + tool_file))
 
+if 'linux' in sys.platform:
+    config.available_features.add('linux')
+
 # Finds the given executable 'program' in PATH.
 # Operates like the Unix tool 'which'.
 # This is similar to script/test/shared.py, but does not use binaryen_root, and
diff --git a/test/lit/num_cores.wast b/test/lit/num_cores.wast
new file mode 100644
index 00000000000..aadbfad7f65
--- /dev/null
+++ b/test/lit/num_cores.wast
@@ -0,0 +1,13 @@
+;; REQUIRES: linux
+;; Test that getNumCores honors thread affinity on linux
+
+(module
+  (func $a)
+)
+
+;; RUN: taskset -c 0 wasm-opt -O1 --debug=threads %s 2>&1 | filecheck %s
+;; RUN: taskset -c 0,2 wasm-opt -O1 --debug=threads %s 2>&1 | filecheck %s --check-prefix=TWO
+
+;; CHECK: getNumCores: 1
+
+;; TWO: getNumCores: 2