Adding support for group_any in iree_thread_affinity_t. (iree-org#21089)

benvanik · web-flow · commit fd69ba343507 · 2025-06-12T23:54:07.000Z
This allows code launching threads to specify that a thread should be
assigned to a processor associated with a specific group (Windows
GROUP_AFFINITY, Linux NUMA node ID, etc) instead of a specific processor
within the group. This is useful for threads created in various parts of
the codebase that don't globally coordinate (service workers/etc), while
things like the task system that explicitly layout large collections of
threads will continue to explicitly assign them.
diff --git a/runtime/src/iree/base/internal/threading.c b/runtime/src/iree/base/internal/threading.c
@@ -41,6 +41,13 @@ void iree_thread_affinity_set_any(iree_thread_affinity_t* out_thread_affinity) {
   memset(out_thread_affinity, 0x00, sizeof(*out_thread_affinity));
 }
 
+void iree_thread_affinity_set_group_any(
+    uint32_t group, iree_thread_affinity_t* out_thread_affinity) {
+  memset(out_thread_affinity, 0x00, sizeof(*out_thread_affinity));
+  out_thread_affinity->group_any = 1;
+  out_thread_affinity->group = group;
+}
+
 //==============================================================================
 // iree_thread_override_list_t
 //==============================================================================
diff --git a/runtime/src/iree/base/internal/threading.h b/runtime/src/iree/base/internal/threading.h
@@ -87,26 +87,45 @@ typedef enum iree_thread_priority_class_e {
 //       id: GROUP_AFFINITY::Mask bit/PROCESSOR_NUMBER::Number.
 //      smt: whether to set both the base ID and the subsequent ID in Mask.
 typedef struct iree_thread_affinity_t {
-  // When 0 the affinity is undefined and the system may place the thread
-  // anywhere and migrate it as much as it likes. In practice it may do that
-  // even when specified.
-  uint32_t specified : 1;
-  // When 1 and the specified processor is part of an SMT set all logical cores
-  // in the set should be reserved for the thread to avoid contention.
-  uint32_t smt : 1;
+  // When 1 the processor ID will be ignored and the platform will choose any
+  // processor associated with the specified group (NUMA node ID).
+  uint32_t group_any : 1;
   // Processor group the thread should be assigned to, aka NUMA node, cluster,
   // etc depending on platform. On platforms where the processor ID is unique
   // for the purposes of scheduling (e.g. Linux) this is used for related APIs
-  // like mbind/set_mempolicy.
-  uint32_t group : 7;
+  // like mbind/set_mempolicy. If group_any is set and id_assigned is not then
+  // any processor associated with the group will be used.
+  uint32_t group : 8;
+
+  uint32_t reserved : 23;
+
+  // When 0 the affinity is undefined and the system may place the thread
+  // anywhere and migrate it as much as it likes. In practice it may do that
+  // even when specified.
+  uint32_t id_assigned : 1;
   // Processor ID the thread should be scheduled on. The interpretation and
   // efficacy of this request varies per platform.
-  uint32_t id : 23;
+  uint32_t id : 30;
+  // When 1 and the specified processor ID is part of an SMT set all logical
+  // cores in the set should be reserved for the thread to avoid contention.
+  uint32_t smt : 1;
 } iree_thread_affinity_t;
 
-// Sets |thread_affinity| to match with any processor in the system.
+// Sets |out_thread_affinity| to match with any processor in the system.
 void iree_thread_affinity_set_any(iree_thread_affinity_t* out_thread_affinity);
 
+// Returns true if |thread_affinity| does not specify any particular processor.
+static inline bool iree_thread_affinity_is_unspecified(
+    iree_thread_affinity_t thread_affinity) {
+  return !thread_affinity.group_any && !thread_affinity.id_assigned;
+}
+
+// Sets |out_thread_affinity| to match all processors associated with the given
+// processor group (aka NUMA node ID). Any processor within the group may be
+// selected by the platform.
+void iree_thread_affinity_set_group_any(
+    uint32_t group, iree_thread_affinity_t* out_thread_affinity);
+
 // Thread creation parameters.
 // All are optional and the entire struct can safely be zero-initialized.
 typedef struct iree_thread_create_params_t {
diff --git a/runtime/src/iree/base/internal/threading_darwin.c b/runtime/src/iree/base/internal/threading_darwin.c
@@ -116,7 +116,7 @@ iree_status_t iree_thread_create(iree_thread_entry_t entry, void* entry_arg,
 
   // Ensure we start with the right QoS class.
   qos_class_t qos_class;
-  if (params.initial_affinity.specified && params.initial_affinity.smt) {
+  if (params.initial_affinity.id_assigned && params.initial_affinity.smt) {
     qos_class = QOS_CLASS_BACKGROUND;
   } else {
     qos_class = iree_thread_qos_class_for_priority_class(params.priority_class);
@@ -148,7 +148,7 @@ iree_status_t iree_thread_create(iree_thread_entry_t entry, void* entry_arg,
   }
 
   thread->mach_port = pthread_mach_thread_np(thread->handle);
-  if (params.initial_affinity.specified) {
+  if (!iree_thread_affinity_is_unspecified(params.initial_affinity)) {
     iree_thread_request_affinity(thread, params.initial_affinity);
   }
 
@@ -210,22 +210,29 @@ void iree_thread_override_end(iree_thread_override_t* override) {
 
 void iree_thread_request_affinity(iree_thread_t* thread,
                                   iree_thread_affinity_t affinity) {
-  if (!affinity.specified) return;
   IREE_TRACE_ZONE_BEGIN(z0);
 
-  // Use mach_task_self when the caller requesting the affinity change is the
-  // thread being changed.
-  mach_port_t thread_port =
-      thread->handle == pthread_self() ? mach_task_self() : thread->mach_port;
-
-  // See:
-  // https://gist.github.com/Coneko/4234842
-  // https://fergofrog.com/code/cbowser/xnu/osfmk/mach/thread_policy.h.html
-  // http://www.hybridkernel.com/2015/01/18/binding_threads_to_cores_osx.html
-  thread_affinity_policy_data_t policy_data = {affinity.id};
-  thread_policy_set(thread_port, THREAD_AFFINITY_POLICY,
-                    (thread_policy_t)(&policy_data),
-                    THREAD_AFFINITY_POLICY_COUNT);
+  // NOTE: group affinity is not yet supported, only ID affinity.
+  // When the ID is not assigned we should really clear the policy but that
+  // doesn't seem possible. Today we don't migrate affinities in a way where
+  // we'd ever want to do anything but assign new ones so this is ok. The kernel
+  // is allowed to totally ignore the affinity request and interpret it however
+  // it likes so this is all not critical anyway.
+  if (affinity.id_assigned) {
+    // Use mach_task_self when the caller requesting the affinity change is the
+    // thread being changed.
+    mach_port_t thread_port =
+        thread->handle == pthread_self() ? mach_task_self() : thread->mach_port;
+
+    // See:
+    // https://gist.github.com/Coneko/4234842
+    // https://fergofrog.com/code/cbowser/xnu/osfmk/mach/thread_policy.h.html
+    // http://www.hybridkernel.com/2015/01/18/binding_threads_to_cores_osx.html
+    thread_affinity_policy_data_t policy_data = {affinity.id};
+    thread_policy_set(thread_port, THREAD_AFFINITY_POLICY,
+                      (thread_policy_t)(&policy_data),
+                      THREAD_AFFINITY_POLICY_COUNT);
+  }
 
   IREE_TRACE_ZONE_END(z0);
 }
diff --git a/runtime/src/iree/base/internal/threading_pthreads.c b/runtime/src/iree/base/internal/threading_pthreads.c
@@ -179,7 +179,7 @@ iree_status_t iree_thread_create(iree_thread_entry_t entry, void* entry_arg,
   if (params.priority_class != IREE_THREAD_PRIORITY_CLASS_NORMAL) {
     iree_thread_set_priority_class(thread, params.priority_class);
   }
-  if (params.initial_affinity.specified) {
+  if (!iree_thread_affinity_is_unspecified(params.initial_affinity)) {
     iree_thread_request_affinity(thread, params.initial_affinity);
   }
 
@@ -297,17 +297,124 @@ void iree_thread_override_end(iree_thread_override_t* override) {
   IREE_TRACE_ZONE_END(z0);
 }
 
+// Sets all CPU bits in the given |out_set|.
+// The platform is allowed to place the thread on any CPU.
+static void iree_thread_make_cpu_set_all(cpu_set_t* out_set) {
+  for (uint32_t i = 0; i < CPU_SETSIZE; ++i) {
+    CPU_SET(i, out_set);
+  }
+}
+
+#if defined(IREE_PLATFORM_ANDROID) || defined(IREE_PLATFORM_LINUX)
+
+// Sets CPU bits associated with the given NUMA node ID.
+// If the platform query fails then all CPU bits are set.
+static void iree_thread_make_cpu_set_from_node_id(uint32_t node_id,
+                                                  cpu_set_t* out_set) {
+  // e.g. /sys/devices/system/node/node0/cpumap
+  char cpumap_path[256];
+  snprintf(cpumap_path, sizeof(cpumap_path),
+           "/sys/devices/system/node/node%u/cpumap", node_id);
+
+  // Open file for reading. This should succeed under hypervisors/lockdown.
+  FILE* file = fopen(cpumap_path, "r");
+  if (!file) {
+    // Permission denied or not found (not a conformant Linux kernel).
+    iree_thread_make_cpu_set_all(out_set);
+    return;
+  }
+
+  // Read the entire file to EOF and get the cpumap line.
+  // After trimming we expect |line| to be something like:
+  // 'ffffffff,ffffffff,ffffffff,00000000,00000000,00000000'
+  char line_buffer[512];
+  const size_t read_length = fread(line_buffer, 1, sizeof(line_buffer), file);
+  if (ferror(file)) {
+    // Read should never fail, but may if the CPU set grows to thousands. We'd
+    // probably want to then query the file length and allocate a heap buffer.
+    // For now all systems we can observe easily fit into our stack buffer.
+    iree_thread_make_cpu_set_all(out_set);
+    return;
+  }
+  iree_string_view_t line =
+      iree_string_view_trim(iree_make_string_view(line_buffer, read_length));
+
+  // Parse each comma-delimited segment. Segments are a base-16 encoded uint32_t
+  // value. Each segment contains 32 CPU bits and we track the current index
+  // as we walk them to get the absolute cpu_set_t index.
+  intptr_t split_index = 0;
+  iree_host_size_t cpu_index = 0;
+  do {
+    iree_string_view_t segment_str;
+    split_index = iree_string_view_split(line, ',', &segment_str, &line);
+    uint32_t segment = 0;
+    if (!iree_string_view_atoi_uint32_base(segment_str, 16, &segment)) {
+      // Failed to parse segment as an integer.
+      iree_thread_make_cpu_set_all(out_set);
+      return;
+    }
+    for (iree_host_size_t i = 0; i < 32; ++i) {
+      if (segment & (1ull << i)) {
+        CPU_SET(cpu_index + i, out_set);
+      }
+    }
+    cpu_index += 32;
+  } while (split_index != -1);
+
+  fclose(file);
+}
+
+#else
+
+// No implementation available. BSD may have some equivalent to the Linux
+// cpumap we could use.
+static void iree_thread_make_cpu_set_from_node_id(uint32_t node_id,
+                                                  cpu_set_t* out_set) {
+  iree_thread_make_cpu_set_all(out_set);
+}
+
+#endif  // IREE_PLATFORM_EMSCRIPTEN
+
+static void iree_thread_make_cpu_set_from_affinity(
+    iree_thread_affinity_t affinity, cpu_set_t* out_set) {
+  CPU_ZERO(out_set);
+
+  // Assign to any processor in the group.
+  if (affinity.group_any) {
+    iree_thread_make_cpu_set_from_node_id(affinity.group, out_set);
+    return;
+  }
+
+  // Specific processors can be set directly and optionally we also set its
+  // paired SMT processor. Note that we don't check whether SMT is enabled and
+  // assume the smt field is only assigned if it is.
+  if (affinity.id_assigned) {
+    CPU_SET(affinity.id, out_set);
+    if (affinity.smt) {
+      CPU_SET(affinity.id + 1, out_set);
+    }
+    return;
+  }
+
+  // No specific affinity specified; use any CPU.
+  iree_thread_make_cpu_set_all(out_set);
+}
+
 void iree_thread_request_affinity(iree_thread_t* thread,
                                   iree_thread_affinity_t affinity) {
-  if (!affinity.specified) return;
   IREE_TRACE_ZONE_BEGIN(z0);
+#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
+  char affinity_desc[64];
+  int affinity_desc_length =
+      snprintf(affinity_desc, IREE_ARRAYSIZE(affinity_desc),
+               "group_any=%u, group=%u, id_assigned=%u, id=%u, smt=%u",
+               affinity.group_any, affinity.group, affinity.id_assigned,
+               affinity.id, affinity.smt);
+  IREE_TRACE_ZONE_APPEND_TEXT(z0, affinity_desc, affinity_desc_length);
+#endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
 
   cpu_set_t cpu_set;
-  CPU_ZERO(&cpu_set);
-  CPU_SET(affinity.id, &cpu_set);
-  if (affinity.smt) {
-    CPU_SET(affinity.id + 1, &cpu_set);
-  }
+  iree_thread_make_cpu_set_from_affinity(affinity, &cpu_set);
 
 #if defined(IREE_PLATFORM_ANDROID)
   // `pthread_gettid_np` is only available on API 21+ and it is needed to set
diff --git a/runtime/src/iree/base/internal/threading_win32.c b/runtime/src/iree/base/internal/threading_win32.c
@@ -175,7 +175,7 @@ iree_status_t iree_thread_create(iree_thread_entry_t entry, void* entry_arg,
   if (params.priority_class != IREE_THREAD_PRIORITY_CLASS_NORMAL) {
     iree_thread_set_priority_class(thread, params.priority_class);
   }
-  if (params.initial_affinity.specified) {
+  if (!iree_thread_affinity_is_unspecified(params.initial_affinity)) {
     iree_thread_request_affinity(thread, params.initial_affinity);
   }
 
@@ -262,24 +262,41 @@ void iree_thread_override_end(iree_thread_override_t* override) {
 
 void iree_thread_request_affinity(iree_thread_t* thread,
                                   iree_thread_affinity_t affinity) {
-  if (!affinity.specified) return;
   IREE_TRACE_ZONE_BEGIN(z0);
 #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
-  char affinity_desc[32];
-  int affinity_desc_length = snprintf(
-      affinity_desc, IREE_ARRAYSIZE(affinity_desc), "group=%d, id=%d, smt=%d",
-      affinity.group, affinity.id, affinity.smt);
+  char affinity_desc[64];
+  int affinity_desc_length =
+      snprintf(affinity_desc, IREE_ARRAYSIZE(affinity_desc),
+               "group_any=%u, group=%u, id_assigned=%u, id=%u, smt=%u",
+               affinity.group_any, affinity.group, affinity.id_assigned,
+               affinity.id, affinity.smt);
   IREE_TRACE_ZONE_APPEND_TEXT(z0, affinity_desc, affinity_desc_length);
 #endif  // IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION
 
+  // TODO(benvanik): switch to the Windows 11 APIs when available (dynamically)
+  // for specifying groups with more than 64 processors. Prior to the new APIs
+  // each group was limited to 64 logical processors and that resulted in groups
+  // being sharded. We need to update our task topology code (which is the
+  // primary caller of this function) as well as others to assign the newer
+  // group IDs and this code to do the same.
+  //
+  // See:
+  // https://learn.microsoft.com/en-us/windows/win32/procthread/numa-support
+  // KeQueryNodeActiveAffinity2
+  // (probably SetThreadSelectedCpuSets?)
+
   GROUP_AFFINITY group_affinity;
   memset(&group_affinity, 0, sizeof(group_affinity));
   group_affinity.Group = affinity.group;
-  KAFFINITY affinity_mask = 1ull << affinity.id;
-  if (affinity.smt) {
-    affinity_mask |= 1ull << (affinity.id + 1);
+  if (affinity.group_any) {
+    group_affinity.Mask = (KAFFINITY)UINTPTR_MAX;
+  } else {
+    KAFFINITY affinity_mask = 1ull << affinity.id;
+    if (affinity.smt) {
+      affinity_mask |= 1ull << (affinity.id + 1);
+    }
+    group_affinity.Mask = affinity_mask;
   }
-  group_affinity.Mask = affinity_mask;
   SetThreadGroupAffinity(thread->handle, &group_affinity, NULL);
 
   // TODO(benvanik): figure out of this is a bad thing; sometimes it can result
diff --git a/runtime/src/iree/task/api.c b/runtime/src/iree/task/api.c
@@ -234,7 +234,9 @@ static void iree_task_flags_dump_task_topology(
     fprintf(stdout, "# group[%d]: '%s'\n", group->group_index, group->name);
     fprintf(stdout, "#      processor: %u\n", group->processor_index);
     fprintf(stdout, "#       affinity: ");
-    if (group->ideal_thread_affinity.specified) {
+    if (group->ideal_thread_affinity.group_any) {
+      fprintf(stdout, "group=%u (any)", group->ideal_thread_affinity.group);
+    } else if (group->ideal_thread_affinity.id_assigned) {
       fprintf(
           stdout, "group=%u, id=%u, smt=%u", group->ideal_thread_affinity.group,
           group->ideal_thread_affinity.id, group->ideal_thread_affinity.smt);
diff --git a/runtime/src/iree/task/topology_cpuinfo.c b/runtime/src/iree/task/topology_cpuinfo.c
@@ -74,7 +74,7 @@ iree_status_t iree_task_topology_initialize_from_logical_cpu_set(
     // really used on Linux today anyway.
     iree_thread_affinity_t* affinity = &group->ideal_thread_affinity;
     memset(affinity, 0, sizeof(*affinity));
-    affinity->specified = 1;
+    affinity->id_assigned = 1;
     affinity->id = cpu_ids[i];
   }
 
@@ -143,12 +143,6 @@ static void iree_task_topology_set_affinity_from_processor(
     const struct cpuinfo_processor* processor,
     iree_thread_affinity_t* out_affinity) {
   memset(out_affinity, 0, sizeof(*out_affinity));
-  out_affinity->specified = 1;
-
-  // Special bit to indicate that (if required) we want the entire core.
-  if (processor->core->processor_count > 1) {
-    out_affinity->smt = 1;
-  }
 
   // cpuinfo #ifdefs the fields we need to extract the right platform IDs.
   // We purposefully use the same exact macros they do there so that we don't
@@ -161,14 +155,21 @@ static void iree_task_topology_set_affinity_from_processor(
   // the kernel to distribute the threads so the exact bits don't matter as long
   // as they are unique per group we want isolated.
   out_affinity->group = processor->cluster->cluster_id;
+  out_affinity->id_assigned = 1;
   out_affinity->id = (uint32_t)(uintptr_t)processor;
 #elif defined(__linux__)
   out_affinity->group = processor->cluster->cluster_id;
+  out_affinity->id_assigned = 1;
   out_affinity->id = processor->linux_id;
 #else
   // WASM? Unusued today.
-  out_affinity->specified = 0;
+  out_affinity->id_assigned = 0;
 #endif  // cpuinfo-like platform field
+
+  // Special bit to indicate that (if required) we want the entire core.
+  if (processor->core->processor_count > 1) {
+    out_affinity->smt = 1;
+  }
 }
 
 // Populates |out_group| with the information from |processor|.
diff --git a/runtime/src/iree/task/topology_darwin.c b/runtime/src/iree/task/topology_darwin.c
@@ -213,8 +213,8 @@ iree_status_t iree_task_topology_initialize_from_physical_cores(
     // affinity info. Note that we pack "use efficiency cores only" into the SMT
     // bit and use that to force a QoS level that ensures only efficiency cores
     // are used when present. Probably.
-    group->ideal_thread_affinity.specified = 1;
     group->ideal_thread_affinity.group = (uint32_t)node_id;
+    group->ideal_thread_affinity.id_assigned = 1;
     group->ideal_thread_affinity.id = i;
     switch (performance_level) {
       default:
diff --git a/runtime/src/iree/task/topology_emscripten.c b/runtime/src/iree/task/topology_emscripten.c
diff --git a/runtime/src/iree/task/topology_win32.c b/runtime/src/iree/task/topology_win32.c
