Enabling thread pool to be numa-aware (#13778)

The PR enables ort thread pool to be numa-aware, so that threads could
be evenly created and distributed among numa nodes.
In addition, to facilitate performance tuning, the PR opens a new API
allowing customers to attach threads to certain logical processors.
Please check the API
[definition](https://github.com/microsoft/onnxruntime/pull/13778/files#diff-5845a5c76fb64abdc8f0cffe21b37f8da1712674eb3abc4cd87190891be1bd48)
for details.

Co-authored-by: Randy Shuai <rashuai@microsoft.com>

cmake/onnxruntime_common.cmake

@@ -115,7 +115,7 @@ target_include_directories(onnxruntime_common
         ${OPTIONAL_LITE_INCLUDE_DIR})
 
 
-target_link_libraries(onnxruntime_common PUBLIC safeint_interface ${GSL_TARGET})
+target_link_libraries(onnxruntime_common PUBLIC safeint_interface ${GSL_TARGET} ${ABSEIL_LIBS})
 
 add_dependencies(onnxruntime_common ${onnxruntime_EXTERNAL_DEPENDENCIES})
 

include/onnxruntime/core/common/logging/logging.h

@@ -136,6 +136,11 @@ class LoggingManager final {
   */
   static const Logger& DefaultLogger();
 
+  /**
+    Return a boolean indicating if the default logger has been initialized
+  */
+  static bool HasDefaultLogger() { return nullptr != s_default_logger_; }
+
   /**
      Change the minimum severity level for log messages to be output by the default logger.
      @param severity The severity.

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -3614,6 +3614,27 @@ struct OrtApi {
    */
   ORT_API2_STATUS(UpdateEnvWithCustomLogLevel, _In_ OrtEnv* ort_env, OrtLoggingLevel log_severity_level);
 
+ /*  \brief Set affinities for intra op threads
+  *
+  * Affinity string follows format:
+  * logical_processor_id,logical_processor_id;logical_processor_id,logical_processor_id
+  * Semicolon isolates configurations among threads, while comma split processors where ith thread expected to attach to.
+  * e.g. 1,2,3;4,5
+  * specifies affinities for two threads, with the 1st thread attach to the 1st, 2nd, and 3rd processor, and 2nd thread to the 4th and 5th.
+  * To ease the configuration, an "interval" is also allowed:
+  * e.g. 1-8;8-16;17-24
+  * orders that the 1st thread runs on first eight processors, 2nd thread runs on next eight processors, and so forth.
+  * Note:
+  * 1. Once set, the number of thread affinities must equal to intra_op_num_threads - 1,
+  *    ort does not set affinity on the main thread which is started and managed by the calling app;
+  * 2. For windows, ort will infer the group id from a logical processor id, for example, assuming there are two groups with each has 64 logical processors,
+  *    an id of 64 will be inferred as the last processor of the 1st group, while 65 will be interpreted as the 1st processor of the second group.
+  *    Hence 64-65 is an invalid configuration, because a windows thread cannot be attached to processors across group boundary.
+  * 
+  *  \since Version 1.14
+  */
+  ORT_API2_STATUS(SetGlobalIntraOpThreadAffinity, _Inout_ OrtThreadingOptions* tp_options, const char* affinity_string);
+
 #ifdef __cplusplus
   OrtApi(const OrtApi&)=delete; // Prevent users from accidentally copying the API structure, it should always be passed as a pointer
 #endif

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -158,3 +158,20 @@ static const char* const kOrtSessionOptionsConfigForceSpinningStop = "session.fo
 // "0": in some cases warnings will be logged but processing will continue. The default.
 // May be useful to expose bugs in models.
 static const char* const kOrtSessionOptionsConfigStrictShapeTypeInference = "session.strict_shape_type_inference";
+
+// This Option allows setting affinities for intra op threads.
+// Affinity string follows format:
+// logical_processor_id,logical_processor_id;logical_processor_id,logical_processor_id
+// Semicolon isolates configurations among threads, while comma split processors where ith thread expected to attach to.
+// e.g.1,2,3;4,5
+// specifies affinities for two threads, with the 1st thread attach to the 1st, 2nd, and 3rd processor, and 2nd thread to the 4th and 5th.
+// To ease the configuration, an "interval" is also allowed:
+// e.g. 1-8;8-16;17-24
+// orders that the 1st thread runs on first eight processors, 2nd thread runs on next eight processors, and so forth.
+// Note:
+// 1. Once set, the number of thread affinities must equal to intra_op_num_threads - 1, since ort does not set affinity on the main thread which
+//    is started and managed by the calling app;
+// 2. For windows, ort will infer the group id from a logical processor id, for example, assuming there are two groups with each has 64 logical processors,
+//    an id of 64 will be inferred as the last processor of the 1st group, while 65 will be interpreted as the 1st processor of the second group.
+//    Hence 64-65 is an invalid configuration, because a windows thread cannot be attached to processors across group boundary.
+static const char* const kOrtSessionOptionsConfigIntraOpThreadAffinities = "session.intra_op_thread_affinities";

onnxruntime/core/common/threadpool.cc

@@ -383,10 +383,10 @@ ThreadPool::ThreadPool(Env* env,
   if (degree_of_parallelism >= 2) {
     int threads_to_create = degree_of_parallelism - 1;
 
-    if (!thread_options_.affinity.empty()) {
+    if (!thread_options_.affinities.empty()) {
       // Remove first affinity element as designated for the caller thread
-      thread_options_.affinity.erase(thread_options_.affinity.begin());
-      assert(thread_options_.affinity.size() >= size_t(threads_to_create));
+      thread_options_.affinities.erase(thread_options_.affinities.begin());
+      assert(thread_options_.affinities.size() >= size_t(threads_to_create));
     }
 
     extended_eigen_threadpool_ =

onnxruntime/core/framework/config_options.cc

@@ -35,8 +35,10 @@ Status ConfigOptions::AddConfigEntry(const char* config_key, const char* config_
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Config key is empty or longer than maximum length 128");
 
   std::string val(config_value);
-  if (val.length() > 1024)
-    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Config value is longer than maximum length 1024");
+  if (val.length() > onnxruntime::kMaxStrLen)
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Config value is longer than maximum length: ",
+                           onnxruntime::kMaxStrLen);
 
   auto iter = configurations.find(config_key);
   if (iter != configurations.cend()) {

onnxruntime/core/platform/env.h

@@ -82,7 +82,7 @@ struct ThreadOptions {
   // that are contained in a given physical core with the same index as the thread. ORT does not set any affinity
   // to the thread that is considered main (the thread that initiates the creation of the TP).
   // The process that owns the thread may consider setting its affinity.
-  std::vector<LogicalProcessors> affinity;
+  std::vector<LogicalProcessors> affinities;
 
   // Set or unset denormal as zero.
   bool set_denormal_as_zero = false;
@@ -139,8 +139,7 @@ class Env {
   /// <returns>Number of physical cores</returns>
   virtual int GetNumPhysicalCpuCores() const = 0;
 
-  // This function currently doesn't support systems with more than 64 logical processors on Windows
-  virtual std::vector<LogicalProcessors> GetThreadAffinityMasks() const = 0;
+  virtual std::vector<LogicalProcessors> GetDefaultThreadAffinities() const = 0;
 
   /// \brief Returns the number of micro-seconds since the Unix epoch.
   virtual uint64_t NowMicros() const {

onnxruntime/core/platform/posix/env.cc

@@ -175,8 +175,8 @@ class PosixThread : public EnvThread {
     custom_join_thread_fn = thread_options.custom_join_thread_fn;
 
     auto param_ptr = std::make_unique<Param>(name_prefix, index, start_address, param);
-    if (narrow<size_t>(index) < thread_options.affinity.size()) {
-      param_ptr->affinity = thread_options.affinity[index];
+    if (narrow<size_t>(index) < thread_options.affinities.size()) {
+      param_ptr->affinity = thread_options.affinities[index];
     }
 
     if (custom_create_thread_fn) {
@@ -233,12 +233,21 @@ class PosixThread : public EnvThread {
       if (p->affinity.has_value() && !p->affinity->empty()) {
         cpu_set_t cpuset;
         CPU_ZERO(&cpuset);
-        for(auto id : *p->affinity) {
-          CPU_SET(id, &cpuset);
+        for (auto id : *p->affinity) {
+          if (id > -1 && id < CPU_SETSIZE) {
+            CPU_SET(id, &cpuset);
+          } else {
+            // Logical processor id starts from 0 internally, but in ort API, it starts from 1,
+            // that's why id need to increase by 1 when logging.
+            LOGS_DEFAULT(ERROR) << "cpu " << id + 1 << " does not exist, skipping it for affinity setting";
+          }
         }
-        // pthread_setaffinity_np() does not set errno, it returns it.
         auto ret = pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
-        if (ret != 0) {
+        if (0 == ret) {
+          LOGS_DEFAULT(VERBOSE) << "pthread_setaffinity_np succeed for thread: " << syscall(SYS_gettid)
+                                << ", index: " << p->index
+                                << ", mask: " << *p->affinity;
+        } else {
           auto [err_no, err_msg] = GetSystemError(ret);
           LOGS_DEFAULT(ERROR) << "pthread_setaffinity_np failed for thread: " << syscall(SYS_gettid)
                               << ", index: " << p->index
@@ -290,7 +299,7 @@ class PosixEnv : public Env {
     return DefaultNumCores();
   }
 
-  std::vector<LogicalProcessors> GetThreadAffinityMasks() const override {
+  std::vector<LogicalProcessors> GetDefaultThreadAffinities() const override {
 
     std::vector<LogicalProcessors> ret;
 #ifdef ORT_USE_CPUINFO

onnxruntime/core/platform/windows/env.h

@@ -0,0 +1,139 @@
+/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// Portions Copyright (c) Microsoft Corporation
+
+#include "core/platform/env.h"
+#include "core/platform/windows/telemetry.h"
+#include "core/common/inlined_containers.h"
+#include <Windows.h>
+
+namespace onnxruntime {
+
+/*
+Logical processor information:
+1. its belonging group;
+2. its id within that belonging group.
+{-1,-1} stands for an invalid processor info
+*/
+struct ProcessorInfo {
+  int group_id = -1;
+  int local_processor_id = -1;
+};
+
+/*
+GlobalProcessorInfoMap is a map between global processor id
+and pair<groupid, local processor id>, the latter is required
+to be present during affinity setup.
+*/
+using GlobalProcessorInfoMap = InlinedHashMap<int, ProcessorInfo>;
+
+class WindowsEnv : public Env {
+ public:
+#if defined(_MSC_VER) && !defined(__clang__)
+#pragma warning(push)
+#pragma warning(disable : 26409)
+#endif
+  EnvThread* CreateThread(_In_opt_z_ const ORTCHAR_T* name_prefix, int index,
+                          unsigned (*start_address)(int id, Eigen::ThreadPoolInterface* param),
+                          Eigen::ThreadPoolInterface* param, const ThreadOptions& thread_options);
+#if defined(_MSC_VER) && !defined(__clang__)
+#pragma warning(pop)
+#endif
+  void WindowsEnv::SleepForMicroseconds(int64_t micros) const override;
+  static int DefaultNumCores();
+  int GetNumPhysicalCpuCores() const override;
+  std::vector<LogicalProcessors> GetDefaultThreadAffinities() const override;
+  static WindowsEnv& Instance();
+  PIDType GetSelfPid() const override;
+  Status GetFileLength(_In_z_ const ORTCHAR_T* file_path, size_t& length) const override;
+  common::Status GetFileLength(int fd, /*out*/ size_t& file_size) const override;
+  Status ReadFileIntoBuffer(_In_z_ const ORTCHAR_T* const file_path, const FileOffsetType offset, const size_t length,
+                            const gsl::span<char> buffer) const override;
+  Status MapFileIntoMemory(_In_z_ const ORTCHAR_T* file_path,
+                           FileOffsetType offset,
+                           size_t length,
+                           MappedMemoryPtr& mapped_memory) const override;
+  bool FolderExists(const std::wstring& path) const override;
+  bool FolderExists(const std::string& path) const override;
+  common::Status CreateFolder(const std::wstring& path) const override;
+  common::Status CreateFolder(const std::string& path) const override;
+  common::Status DeleteFolder(const PathString& path) const override;
+  common::Status FileOpenRd(const std::wstring& path, /*out*/ int& fd) const override;
+  common::Status FileOpenWr(const std::wstring& path, /*out*/ int& fd) const override;
+  common::Status FileOpenRd(const std::string& path, /*out*/ int& fd) const override;
+  common::Status FileOpenWr(const std::string& path, /*out*/ int& fd) const override;
+  common::Status FileClose(int fd) const override;
+  common::Status GetCanonicalPath(const PathString& path, PathString& canonical_path) const override;
+  std::string GetRuntimePath() const override;
+  Status LoadDynamicLibrary(const std::string& library_filename, bool /*global_symbols*/, void** handle) const override;
+  Status UnloadDynamicLibrary(void* handle) const override;
+  Status GetSymbolFromLibrary(void* handle, const std::string& symbol_name, void** symbol) const override;
+  std::string FormatLibraryFileName(const std::string& name, const std::string& version) const override;
+  const Telemetry& GetTelemetryProvider() const override;
+  std::string GetEnvironmentVar(const std::string& var_name) const override;
+  ProcessorInfo GetProcessorAffinityMask(int global_processor_id) const;
+
+ protected:
+  /*
+  * "cores_" host all physical cores dicoverred in a windows system.
+  * Every LogicalProcessors represent a core of logical processors.
+  * Specifically, LogicalProcessors is a vector of global processor id.
+  * E.g.
+  * Assume we have a system of 4 cores, each has 2 logical processors.
+  * Then "cores_" will be like:
+  * {
+  * {0,1} // core 1
+  * {2,3} // core 2
+  * {4,5} // core 3
+  * {6,7} // core 4
+  * }
+  * Further, assume we have a system of two groups, each has 4 cores like above,
+  * then "cores_" will be like:
+  * {
+  * {0,1}   // core 1, group 1
+  * {2,3}   // core 2, group 1
+  * {4,5}   // core 3, group 1
+  * {6,7}   // core 4, group 1
+  * {8,9}   // core 5, group 2
+  * {10,11} // core 6, group 2
+  * {12,13} // core 7, group 2
+  * {14,15} // core 8, group 2
+  * }
+  */
+  std::vector<LogicalProcessors> cores_;
+  /*
+  * "global_processor_info_map_" is a map of:
+  * global_processor_id <--> (group_id, local_processor_id)
+  * "global_processor_id" means "index" of a logical processor in the system,
+  * "local_processor_id" refers to "index" of a logical processor in its belonging group,
+  * E.g.
+  * Assume the system have two groups,
+  * each group has two cores,
+  * and each core has two logical processors.
+  * then the global processor ids will be like:
+  * 0,1,2,3,4,5,6,7
+  * the local processor ids will be like:
+  * 0,1,2,3,0,1,2,3
+  */
+  GlobalProcessorInfoMap global_processor_info_map_;
+  WindowsEnv();
+
+ private:
+  void InitializeCpuInfo();
+  typedef VOID(WINAPI* FnGetSystemTimePreciseAsFileTime)(LPFILETIME);
+  WindowsTelemetry telemetry_provider_;
+};
+
+}  // namespace onnxruntime

onnxruntime/core/session/inference_session.cc

@@ -285,9 +285,6 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options,
         to.set_denormal_as_zero = set_denormal_as_zero;
         // If the thread pool can use all the processors, then
         // we set affinity of each thread to each processor.
-        to.auto_set_affinity = to.thread_pool_size == 0 &&
-                               session_options_.execution_mode == ExecutionMode::ORT_SEQUENTIAL &&
-                               to.affinity_vec_len == 0;
         to.allow_spinning = allow_intra_op_spinning;
         to.dynamic_block_base_ = std::stoi(session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigDynamicBlockBase, "0"));
         LOGS(*session_logger_, INFO) << "Dynamic block base set to " << to.dynamic_block_base_;
@@ -296,10 +293,17 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options,
         to.custom_create_thread_fn = session_options_.custom_create_thread_fn;
         to.custom_thread_creation_options = session_options.custom_thread_creation_options;
         to.custom_join_thread_fn = session_options_.custom_join_thread_fn;
+        if (session_options_.config_options.TryGetConfigEntry(kOrtSessionOptionsConfigIntraOpThreadAffinities, to.affinity_str)) {
+          ORT_ENFORCE(!to.affinity_str.empty(), "Affinity string must not be empty");
+        }
+        to.auto_set_affinity = to.thread_pool_size == 0 &&
+                               session_options_.execution_mode == ExecutionMode::ORT_SEQUENTIAL &&
+                               to.affinity_str.empty();
 
         if (to.custom_create_thread_fn) {
           ORT_ENFORCE(to.custom_join_thread_fn, "custom join thread function not set for intra op thread pool");
         }
+
         thread_pool_ =
             concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP);
       }
@@ -309,10 +313,7 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options,
         bool allow_inter_op_spinning =
             session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigAllowInterOpSpinning, "1") == "1";
         OrtThreadPoolParams to = session_options_.inter_op_param;
-        // If the thread pool can use all the processors, then
-        // we set thread affinity.
-        to.auto_set_affinity =
-            to.thread_pool_size == 0 && session_options_.execution_mode == ExecutionMode::ORT_SEQUENTIAL;
+        to.auto_set_affinity = to.thread_pool_size == 0 && session_options_.execution_mode == ExecutionMode::ORT_SEQUENTIAL;
         std::basic_stringstream<ORTCHAR_T> ss;
         if (to.name) {
           ss << to.name << ORT_TSTR("-");

onnxruntime/core/session/onnxruntime_c_api.cc

@@ -2614,6 +2614,7 @@ static constexpr OrtApi ort_api_1_to_12 = {
     // Start of Version 14 API in progress, safe to modify/rename/rearrange until we ship
     &OrtApis::MemoryInfoGetDeviceType,
     &OrtApis::UpdateEnvWithCustomLogLevel,
+    &OrtApis::SetGlobalIntraOpThreadAffinity,
 };
 
 

onnxruntime/core/session/ort_apis.h

@@ -404,4 +404,6 @@ ORT_API(void, ReleaseCANNProviderOptions, _Frees_ptr_opt_ OrtCANNProviderOptions
 ORT_API(void, MemoryInfoGetDeviceType, _In_ const OrtMemoryInfo* ptr, _Out_ OrtMemoryInfoDeviceType* out);
 
 ORT_API_STATUS_IMPL(UpdateEnvWithCustomLogLevel, _In_ OrtEnv* ort_env, OrtLoggingLevel log_severity_level);
+
+ORT_API_STATUS_IMPL(SetGlobalIntraOpThreadAffinity, _Inout_ OrtThreadingOptions* tp_options, const char* affinity_string);
 }  // namespace OrtApis

onnxruntime/core/util/thread_utils.cc

@@ -10,40 +10,100 @@
 #endif
 #include <thread>
 #include "core/session/ort_apis.h"
+#include "core/common/string_utils.h"
+#include "core/common/logging/logging.h"
 
 namespace onnxruntime {
 namespace concurrency {
+
+// Extract affinity from affinity string.
+// Processor id from affinity string starts from 1,
+// but internally, processor id starts from 0, so here we minus the id by 1
+std::vector<LogicalProcessors> ReadThreadAffinityConfig(const std::string& affinity_str) {
+  ORT_TRY {
+    std::vector<LogicalProcessors> logical_processors_vector;
+    auto affinities = utils::SplitString(affinity_str, ";");
+
+    for (const auto& affinity : affinities) {
+      LogicalProcessors logical_processors;
+      auto processor_interval = utils::SplitString(affinity, "-");
+
+      if (processor_interval.size() == 2) {
+        ORT_ENFORCE(std::all_of(processor_interval[0].begin(), processor_interval[0].end(), ::isdigit) &&
+                        std::all_of(processor_interval[1].begin(), processor_interval[1].end(), ::isdigit),
+                    std::string{"Processor id must consist of only digits: "} + std::string{affinity});
+
+        auto processor_from = std::stoi(std::string{processor_interval[0]});
+        auto processor_to = std::stoi(std::string{processor_interval[1]});
+
+        ORT_ENFORCE(processor_from > 0 && processor_to > 0,
+                    std::string{"Processor id must start from 1: "} + std::string{affinity});
+        ORT_ENFORCE(processor_from <= processor_to,
+                    std::string{"Invalid processor interval: "} + std::string{affinity});
+
+        logical_processors.resize(static_cast<size_t>(1ULL + processor_to - processor_from));
+        std::iota(logical_processors.begin(), logical_processors.end(), processor_from - 1);
+
+      } else {
+        for (const auto& processor_str : utils::SplitString(affinity, ",")) {
+          ORT_ENFORCE(std::all_of(processor_str.begin(), processor_str.end(), ::isdigit),
+                      std::string{"Processor id must consist of only digits: "} + std::string{processor_str});
+
+          auto processor_id = std::stoi(std::string{processor_str});
+          ORT_ENFORCE(processor_id > 0, std::string{"Processor id must start from 1: "} + std::string{processor_str});
+          logical_processors.push_back(processor_id - 1);
+        }
+      }
+      logical_processors_vector.push_back(std::move(logical_processors));
+    }
+    return logical_processors_vector;
+  }
+  ORT_CATCH(const std::invalid_argument&) {
+    LOGS_DEFAULT(ERROR) << "Found invalid processor id in affinity string: "
+                        << affinity_str << ", skip affinity setting";
+  }
+  ORT_CATCH(const std::out_of_range&) {
+    LOGS_DEFAULT(ERROR) << "Found out-of-range processor id in affinity string: "
+                        << affinity_str << ", skip affinity setting";
+  }
+  ORT_THROW("Failed to read affinities from affinity string");
+}
+
 static std::unique_ptr<ThreadPool>
 CreateThreadPoolHelper(Env* env, OrtThreadPoolParams options) {
-  if (options.thread_pool_size == 1)
-    return nullptr;
   ThreadOptions to;
-
-  if (options.affinity_vec_len != 0) {
-    // Currently, the affinities are passed in as bit masks and they need to be converted to integers.
-    // We when create a public API, bit-masks must be done away with because of the following reasons:
-    // 1) integers have a limited number of bits
-    // 2) bit-masks of integers can only represent numbers 0 -63, but on VMs the actual logical processor numbering
-    //    may not start with zero for a given core and may be way beyond 63.
-    // 3) Customers would be forced to concoct bit-masks which is far less convenient than simply an array of processor integers. 
-    to.affinity.reserve(options.affinity_vec_len);
-    std::transform(options.affinity_vec, options.affinity_vec + options.affinity_vec_len, std::back_inserter(to.affinity),
-                   [](size_t affinity) {
-                     return LogicalProcessors{static_cast<int>(affinity)};
-                   });
-  }
-
   if (options.thread_pool_size <= 0) {  // default
-    auto cpu_list = Env::Default().GetThreadAffinityMasks();
-    if (cpu_list.empty() || cpu_list.size() == 1)
+    auto default_affinities = Env::Default().GetDefaultThreadAffinities();
+    if (default_affinities.size() <= 1) {
       return nullptr;
-    options.thread_pool_size = static_cast<int>(cpu_list.size());
-    if (options.auto_set_affinity)
-      to.affinity = cpu_list;
+    }
+    options.thread_pool_size = static_cast<int>(default_affinities.size());
+    if (options.auto_set_affinity) {
+      to.affinities = std::move(default_affinities);
+    }
+  }
+  if (options.thread_pool_size <= 1) {
+    return nullptr;
+  }
+  // override affinity setting if specified from customer
+  if (!options.affinity_str.empty()) {
+    to.affinities = ReadThreadAffinityConfig(options.affinity_str);
+    // Limiting the number of affinities to be of thread_pool_size - 1,
+    // for the fact that the main thread is a special "member" of the threadpool,
+    // which onnxruntime has no control.
+    auto actual_num_affinities = to.affinities.size();
+    ORT_ENFORCE(actual_num_affinities == static_cast<size_t>(options.thread_pool_size) - 1,
+                (std::string{"Number of affinities does not equal to thread_pool_size minus one, affinities: "} +
+                 std::to_string(actual_num_affinities) +
+                 std::string{", thread_pool_size: "} +
+                 std::to_string(options.thread_pool_size))
+                    .c_str());
+    // prepend with an empty affinity as placeholder for the main thread,
+    // it will be dropped later during threadpool creation.
+    to.affinities.insert(to.affinities.begin(), LogicalProcessors{});
   }
 
   to.set_denormal_as_zero = options.set_denormal_as_zero;
-
   // set custom thread management members
   to.custom_create_thread_fn = options.custom_create_thread_fn;
   to.custom_thread_creation_options = options.custom_thread_creation_options;
@@ -144,4 +204,22 @@ ORT_API_STATUS_IMPL(SetGlobalCustomJoinThreadFn, _Inout_ OrtThreadingOptions* tp
   return nullptr;
 }
 
+ORT_API_STATUS_IMPL(SetGlobalIntraOpThreadAffinity, _Inout_ OrtThreadingOptions* tp_options, const char* affinity_string) {
+  if (!tp_options) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Received null OrtThreadingOptions");
+  }
+  if (!affinity_string) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, "Received null affinity_string");
+  }
+  auto len = strnlen(affinity_string, onnxruntime::kMaxStrLen + 1);
+  if (0 == len || len > onnxruntime::kMaxStrLen) {
+    return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT,
+                                 (std::string{"Size of affinity string must be between 1 and "} +
+                                  std::to_string(onnxruntime::kMaxStrLen))
+                                     .c_str());
+  }
+  tp_options->intra_op_thread_pool_params.affinity_str = affinity_string;
+  return nullptr;
+}
+
 }  // namespace OrtApis

onnxruntime/core/util/thread_utils.h

@@ -12,21 +12,31 @@ struct OrtThreadPoolParams {
   //1: Don't create thread pool
   //n: Create a thread pool with n threads.
   int thread_pool_size = 0;
+
   //If it is true and thread_pool_size = 0, populate the thread affinity information in ThreadOptions.
   //Otherwise if the thread_options has affinity information, we'll use it and set it.
   //In the other case, don't set affinity
   bool auto_set_affinity = false;
+
   //If it is true, the thread pool will spin a while after the queue became empty.
   bool allow_spinning = true;
+
   //It it is non-negative, thread pool will split a task by a decreasing block size
   //of remaining_of_total_iterations / (num_of_threads * dynamic_block_base_)
   int dynamic_block_base_ = 0;
 
   unsigned int stack_size = 0;
-  //Index is thread id, value is processor ID
-  //If the vector is empty, no explict affinity binding
-  size_t* affinity_vec = nullptr;
-  size_t affinity_vec_len = 0;
+
+  // A utf-8 string of affinity settings, format be like:
+  // <1st_thread_affinity_config>;<2nd_thread_affinity_config>;<3rd_thread_affinity_config>...
+  // ith_thread_affinity_config could be:
+  // 1,2,3
+  // meaing ith thread attach to logical processor 1,2,3
+  // or
+  // 1-8
+  // meaning ith thread will be attached to first 8 logical processors
+  std::string affinity_str;
+
   const ORTCHAR_T* name = nullptr;
 
   // Set or unset denormal as zero

onnxruntime/test/framework/math_test.cc

@@ -28,7 +28,7 @@ namespace onnxruntime {
 //parameter is thread pool size
 class MathGemmTest : public testing::TestWithParam<int> {
  protected:
-  constexpr static OrtThreadPoolParams CreateThreadPoolOptions(int size) {
+  static OrtThreadPoolParams CreateThreadPoolOptions(int size) {
     OrtThreadPoolParams option;
     option.thread_pool_size = size;
     return option;

onnxruntime/test/global_thread_pools/test_main.cc

@@ -26,6 +26,11 @@ std::unique_ptr<Ort::Env> ort_env;
     return -1;                             \
   }
 
+#define ORT_RETURN_IF_NULL_STATUS(arg) \
+  if (!arg) {                              \
+    return -1;                             \
+  }
+
 namespace TestGlobalCustomThreadHooks {
 
 std::vector<std::thread> threads;
@@ -57,6 +62,14 @@ using namespace TestGlobalCustomThreadHooks;
 int main(int argc, char** argv) {
   int status = 0;
   const int thread_pool_size = std::thread::hardware_concurrency();
+
+  //compose affinity string
+  std::stringstream affinity_stream;
+  //skip the 1st logical processor
+  for (int i = 2; i <= thread_pool_size; ++i) {
+    affinity_stream << (i == 2 ? "" : ";") << i;
+  }
+
   ORT_TRY {
     ::testing::InitGoogleTest(&argc, argv);
     const OrtApi* g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION);
@@ -72,6 +85,18 @@ int main(int argc, char** argv) {
     st_ptr.reset(g_ort->SetGlobalIntraOpNumThreads(tp_options, thread_pool_size));
     ORT_RETURN_IF_NON_NULL_STATUS(st_ptr);
 
+    // test with an empty affinity string, error status expected
+    st_ptr.reset(g_ort->SetGlobalIntraOpThreadAffinity(tp_options, ""));
+    ORT_RETURN_IF_NULL_STATUS(st_ptr);
+
+    // test with an oversized affinity string, error status expected
+    std::string long_affinity_str(onnxruntime::kMaxStrLen + 1, '0');
+    st_ptr.reset(g_ort->SetGlobalIntraOpThreadAffinity(tp_options, long_affinity_str.c_str()));
+    ORT_RETURN_IF_NULL_STATUS(st_ptr);
+
+    st_ptr.reset(g_ort->SetGlobalIntraOpThreadAffinity(tp_options, affinity_stream.str().c_str()));
+    ORT_RETURN_IF_NON_NULL_STATUS(st_ptr);
+
     st_ptr.reset(g_ort->SetGlobalCustomCreateThreadFn(tp_options, CreateThreadCustomized));
     ORT_RETURN_IF_NON_NULL_STATUS(st_ptr);
 

onnxruntime/test/perftest/command_args_parser.cc

@@ -94,6 +94,11 @@ namespace perftest {
       "\t    [SNPE only] [buffer_type]: options: 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. default: ITENSOR'. \n"
       "\t [Usage]: -e <provider_name> -i '<key1>|<value1> <key2>|<value2>' \n\n"
       "\t [Example] [For SNPE EP] -e snpe -i \"runtime|CPU priority|low\" \n\n"
+      "\t-T [Set intra op thread affinities]: Specify intra op thread affinity string\n"
+      "\t [Example]: -T 1,2;3,4;5,6 or -T 1-2;3-4;5-6' \n"
+      "\t\t Use semicolon to separate configuration between threads.\n"
+      "\t\t E.g. 1,2;3,4;5,6 specifies affinities for three threads, the first thread will be attached to the first and second logical processor.\n"
+      "\t\t The number of affinities must be equal to intra_op_num_threads - 1\n\n"
       "\t-h: help\n");
 }
 #ifdef _WIN32
@@ -123,7 +128,7 @@ static bool ParseDimensionOverride(std::basic_string<ORTCHAR_T>& dim_identifier,
 
 /*static*/ bool CommandLineParser::ParseArguments(PerformanceTestConfig& test_config, int argc, ORTCHAR_T* argv[]) {
   int ch;
-  while ((ch = getopt(argc, argv, ORT_TSTR("b:m:e:r:t:p:x:y:c:d:o:u:i:f:F:S:AMPIvhsqz"))) != -1) {
+  while ((ch = getopt(argc, argv, ORT_TSTR("b:m:e:r:t:p:x:y:c:d:o:u:i:f:F:S:T:AMPIvhsqz"))) != -1) {
     switch (ch) {
       case 'f': {
         std::basic_string<ORTCHAR_T> dim_name;
@@ -287,6 +292,9 @@ static bool ParseDimensionOverride(std::basic_string<ORTCHAR_T>& dim_identifier,
       case 'i':
         test_config.run_config.ep_runtime_config_string = optarg;
         break;
+      case 'T':
+        test_config.run_config.intra_op_thread_affinities = ToUTF8String(optarg);
+        break;
       case '?':
       case 'h':
       default:

onnxruntime/test/perftest/ort_test_session.cc

@@ -524,6 +524,11 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)");
     session_options.SetIntraOpNumThreads(performance_test_config.run_config.intra_op_num_threads);
   }
 
+  if (!performance_test_config.run_config.intra_op_thread_affinities.empty()) {
+    fprintf(stdout, "Setting intra op thread affinity as %s\n", performance_test_config.run_config.intra_op_thread_affinities.c_str());
+    session_options.AddConfigEntry("session.intra_op_thread_affinities", performance_test_config.run_config.intra_op_thread_affinities.c_str());
+  }
+
   if (performance_test_config.run_config.execution_mode == ExecutionMode::ORT_PARALLEL && performance_test_config.run_config.inter_op_num_threads > 0) {
     fprintf(stdout, "Setting inter_op_num_threads to %d\n", performance_test_config.run_config.inter_op_num_threads);
     session_options.SetInterOpNumThreads(performance_test_config.run_config.inter_op_num_threads);

onnxruntime/test/perftest/test_configuration.h

@@ -58,6 +58,7 @@ struct RunConfig {
   std::basic_string<ORTCHAR_T> ep_runtime_config_string;
   std::map<std::basic_string<ORTCHAR_T>, int64_t> free_dim_name_overrides;
   std::map<std::basic_string<ORTCHAR_T>, int64_t> free_dim_denotation_overrides;
+  std::string intra_op_thread_affinities;
 };
 
 struct PerformanceTestConfig {

onnxruntime/test/platform/threadpool_test.cc

@@ -4,6 +4,11 @@
 #include "core/platform/threadpool.h"
 #include "core/platform/EigenNonBlockingThreadPool.h"
 #include "core/platform/ort_mutex.h"
+#include "core/util/thread_utils.h"
+#ifdef _WIN32
+#include "test/platform/windows/env.h"
+#include <Windows.h>
+#endif
 
 #include "gtest/gtest.h"
 #include <algorithm>
@@ -539,4 +544,128 @@ TEST(ThreadPoolTest, TestStackSize) {
 #endif
 #endif
 
+#ifndef ORT_NO_EXCEPTIONS
+TEST(ThreadPoolTest, TestAffinityStringMisshaped) {
+  OrtThreadPoolParams tp_params;
+  tp_params.thread_pool_size = 3;
+  const char* wrong_formats[] = {
+      ",",     //1st and 2nd processor id are empty strings
+      "1,",    //2nd processor id is an empty string
+      ";",     //affinity settings for both threads are empty
+      ";1",    //missing the affinity setting for the 1st thread
+      "a",     //invalid char, must be digit
+      "a;b",   //invalid char, must be digit
+      "1;a",   //invalid char, must be digit
+      "0;1",   //processor string must start from 1
+      "-;2",   //invalid char, must be digit
+      "--",    //invalid char, must be digit
+      "2-1;3", //invalid interval, "from" must be equal to or smaller than "to"
+      "5;3a"   //invalid processor id containing non-digit as suffix
+  };
+  for (const auto* wrong_format : wrong_formats) {
+    tp_params.affinity_str = wrong_format;
+    ASSERT_THROW(concurrency::CreateThreadPool(&onnxruntime::Env::Default(),
+                                               tp_params,
+                                               concurrency::ThreadPoolType::INTRA_OP),
+                 std::exception);
+  }
+  const char* less_than_expected_vec[] = {"1", "1,2", "1-2"};
+  for (const auto* less_than_expected : less_than_expected_vec) {
+    tp_params.affinity_str = less_than_expected;
+    ASSERT_THROW(concurrency::CreateThreadPool(&onnxruntime::Env::Default(),
+                                               tp_params,
+                                               concurrency::ThreadPoolType::INTRA_OP),
+                 std::exception);
+  }
+  const char* more_than_expected_vec[] = {"1;2;3", "1-2;2-2;3-4", "1;2;3;4;5"};
+  for (const auto* more_than_expected : more_than_expected_vec) {
+    tp_params.affinity_str = more_than_expected;
+    ASSERT_THROW(concurrency::CreateThreadPool(&onnxruntime::Env::Default(),
+                                               tp_params,
+                                               concurrency::ThreadPoolType::INTRA_OP),
+                 std::exception);
+  }
+}
+#endif
+
+TEST(ThreadPoolTest, TestAffinityStringWellShaped) {
+  OrtThreadPoolParams tp_params;
+  auto default_tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(),
+                                                  tp_params,
+                                                  concurrency::ThreadPoolType::INTRA_OP);
+  if (concurrency::ThreadPool::DegreeOfParallelism(default_tp.get()) < 3) {
+    return;
+  }
+  tp_params.thread_pool_size = 3;
+  const char* good_formats[] = {"1;1",
+                                "2;2",
+                                "1-1;2-2",
+                                "1-2;1-2"};
+  for (const auto* good_format : good_formats) {
+    tp_params.affinity_str = good_format;
+    auto non_default_tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(),
+                                                        tp_params,
+                                                        concurrency::ThreadPoolType::INTRA_OP);
+    ASSERT_TRUE(concurrency::ThreadPool::DegreeOfParallelism(non_default_tp.get()) == 3);
+  }
+}
+
+#ifdef _WIN32
+TEST(ThreadPoolTest, TestDefaultAffinity) {
+  test::CpuGroup cpu_group = {{0, 1},
+                              {2, 3},
+                              {4, 5},
+                              {6, 7}};
+  // 2 logical processors per core, single group
+  test::CpuInfo cpu_info_single = {cpu_group};
+  test::WindowsEnvTester win_env;
+  win_env.SetCpuInfo(cpu_info_single);
+  auto default_affinities = win_env.GetDefaultThreadAffinities();
+  ASSERT_TRUE(default_affinities.size() == 4);
+  for (int i = 0; i < 4; ++i) {
+    ASSERT_TRUE(default_affinities[i].size() == 2);
+    for (int j = 0; j < 2; ++j) {
+      ASSERT_TRUE(default_affinities[i][j] == i * 2 + j);
+    }
+  }
+  // 2 logical processors per core, two groups
+  test::CpuInfo cpu_info_double = {cpu_group, cpu_group};
+  win_env.SetCpuInfo(cpu_info_double);
+  default_affinities = win_env.GetDefaultThreadAffinities();
+  ASSERT_TRUE(default_affinities.size() == 8);
+  for (int i = 0; i < 8; ++i) {
+    ASSERT_TRUE(default_affinities[i].size() == 2);
+    for (int j = 0; j < 2; ++j) {
+      ASSERT_TRUE(default_affinities[i][j] == i * 2 + j);
+    }
+  }
+  // 4 logical processors per core, single group
+  cpu_group = {{0, 1, 2, 3},
+               {4, 5, 6, 7},
+               {8, 9, 10, 11},
+               {12, 13, 14, 15}};
+  cpu_info_single = {cpu_group};
+  win_env.SetCpuInfo(cpu_info_single);
+  default_affinities = win_env.GetDefaultThreadAffinities();
+  ASSERT_TRUE(default_affinities.size() == 4);
+  for (int i = 0; i < 4; ++i) {
+    ASSERT_TRUE(default_affinities[i].size() == 4);
+    for (int j = 0; j < 4; ++j) {
+      ASSERT_TRUE(default_affinities[i][j] == i * 4 + j);
+    }
+  }
+  // 4 logical processors per core, two groups
+  cpu_info_double = {cpu_group, cpu_group};
+  win_env.SetCpuInfo(cpu_info_double);
+  default_affinities = win_env.GetDefaultThreadAffinities();
+  ASSERT_TRUE(default_affinities.size() == 8);
+  for (int i = 0; i < 8; ++i) {
+    ASSERT_TRUE(default_affinities[i].size() == 4);
+    for (int j = 0; j < 4; ++j) {
+      ASSERT_TRUE(default_affinities[i][j] == i * 4 + j);
+    }
+  }
+}
+#endif
+
 }  // namespace onnxruntime

onnxruntime/test/platform/windows/env.cc

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "test/platform/windows/env.h"
+#include <vector>
+
+namespace onnxruntime {
+namespace test {
+
+/*
+Mock cpu layout for windows
+"global_processor_id" means "index" of a logical processor in the system,
+"local_processor_id" refers to "index" of a logical processor in its belonging group,
+the "local" here means group-local.
+Unlike the definition of affinity API, the id here starts from 0, not 1.
+E.g.
+assume cpu_info is like:
+{
+{{0,1},{2,3}} // group 1 of two cores
+{{0,1},{2,3}} // group 2 of two cores
+}
+then we will have eight logical processors with global id be like:
+0,1,2,3,4,5,6,7
+and local id be like:
+0,1,2,3,0,1,2,3
+*/
+bool WindowsEnvTester::SetCpuInfo(const CpuInfo& cpu_info) {
+  if (cpu_info.empty()) {
+    return false;
+  }
+  cores_.clear();
+  global_processor_info_map_.clear();
+  int global_processor_id = 0;
+  for (int group_id = 0; group_id < static_cast<int>(cpu_info.size()); ++group_id) {
+    int local_processor_id = 0;
+    for (int core_id = 0; core_id < static_cast<int>(cpu_info[group_id].size()); ++core_id) {
+      onnxruntime::LogicalProcessors logical_processors;
+      for (int i = 0; i < static_cast<int>(cpu_info[group_id][core_id].size()); ++i) {
+        logical_processors.push_back(global_processor_id);
+        global_processor_info_map_[global_processor_id] = {group_id, local_processor_id};
+        local_processor_id++;
+        global_processor_id++;
+      }
+      cores_.push_back(std::move(logical_processors));
+    }
+  }
+  return true;
+}
+
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/platform/windows/env.h

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/platform/windows/env.h"
+
+namespace onnxruntime {
+namespace test {
+
+using CpuLogicProcessorId = int;                   // an id of a logical processor starting from 0
+using CpuCore = std::vector<CpuLogicProcessorId>;  // a core of multiple logical processors
+using CpuGroup = std::vector<CpuCore>;             // core group
+using CpuInfo = std::vector<CpuGroup>;             // groups
+
+class WindowsEnvTester : public WindowsEnv {
+ public:
+  WindowsEnvTester() = default;
+  ~WindowsEnvTester() = default;
+  ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(WindowsEnvTester);
+  bool SetCpuInfo(const CpuInfo& cpu_info);
+};
+
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/shared_lib/test_model_loading.cc

@@ -55,6 +55,20 @@ TEST(CApiTest, model_from_array) {
   create_session(so);
 #endif
 }
+
+TEST(CApiTest, session_options_empty_affinity_string) {
+  Ort::SessionOptions options;
+  options.AddConfigEntry(kOrtSessionOptionsConfigIntraOpThreadAffinities, "");
+  constexpr auto model_path = ORT_TSTR("testdata/matmul_1.onnx");
+
+  try {
+    Ort::Session session(*ort_env.get(), model_path, options);
+    ASSERT_TRUE(false) << "Creation of session should have thrown exception";
+  } catch (const std::exception& ex) {
+    ASSERT_THAT(ex.what(), testing::HasSubstr("Affinity string must not be empty"));
+  }
+}
+
 #endif
 
 #ifdef DISABLE_EXTERNAL_INITIALIZERS

onnxruntime/test/shared_lib/test_session_options.cc

@@ -2,8 +2,9 @@
 // Licensed under the MIT License.
 
 #include "core/session/onnxruntime_cxx_api.h"
+#include "core/session/onnxruntime_session_options_config_keys.h"
 #include "core/optimizer/graph_transformer_level.h"
-#include <gtest/gtest.h>
+#include "gmock/gmock.h"
 
 using namespace onnxruntime;
 
@@ -12,3 +13,18 @@ TEST(CApiTest, session_options_graph_optimization_level) {
   Ort::SessionOptions options;
   options.SetGraphOptimizationLevel(ORT_ENABLE_EXTENDED);
 }
+
+#if !defined(ORT_MINIMAL_BUILD) && !defined(ORT_NO_EXCEPTIONS)
+
+TEST(CApiTest, session_options_oversized_affinity_string) {
+  Ort::SessionOptions options;
+  std::string long_affinity_str(onnxruntime::kMaxStrLen + 1, '0');
+  try {
+    options.AddConfigEntry(kOrtSessionOptionsConfigIntraOpThreadAffinities, long_affinity_str.c_str());
+    ASSERT_TRUE(false) << "Creation of config should have thrown exception";
+  } catch (const std::exception& ex) {
+    ASSERT_THAT(ex.what(), testing::HasSubstr("Config value is longer than maximum length: "));
+  }
+}
+
+#endif

winml/adapter/winml_adapter_c_api.h

@@ -66,10 +66,6 @@ struct OrtThreadPoolOptions {
   int dynamic_block_base_ = 0;
 
   unsigned int stack_size = 0;
-  //Index is thread id, value is processor ID
-  //If the vector is empty, no explict affinity binding
-  size_t* affinity_vec = nullptr;
-  size_t affinity_vec_len = 0;
   const ORTCHAR_T* name = nullptr;
 
   // Set or unset denormal as zero

winml/adapter/winml_adapter_model.cpp

@@ -816,8 +816,6 @@ ORT_API_STATUS_IMPL(winmla::CreateThreadPool,
   params.allow_spinning = options->allow_spinning;
   params.dynamic_block_base_ = options->dynamic_block_base_;
   params.stack_size = options->stack_size;
-  params.affinity_vec = options->affinity_vec;
-  params.affinity_vec_len = options->affinity_vec_len;
   params.name = options->name;
   params.set_denormal_as_zero = options->set_denormal_as_zero;