From c33459eb86497a5ff6e7044f15fae451c1c3b284 Mon Sep 17 00:00:00 2001
From: Sun Xiaoxia <xiaoxia.sun@intel.com>
Date: Thu, 18 Jul 2024 21:24:40 +0800
Subject: [PATCH] Fix the issue of cpu affinity mask changed on a NUMA machine
 with SEQ (#25455)

### Details:
- *Fix the issue of cpu affinity mask changed on a NUMA machine with
SEQ*

### Tickets:
 - *CVS-143272*
---
 .../dev/threading/cpu_streams_executor.cpp    | 50 +++++++++++++------
 1 file changed, 34 insertions(+), 16 deletions(-)

diff --git a/src/inference/src/dev/threading/cpu_streams_executor.cpp b/src/inference/src/dev/threading/cpu_streams_executor.cpp
index 6cfbeeed305a4a..2df59d27cfe0b2 100644
--- a/src/inference/src/dev/threading/cpu_streams_executor.cpp
+++ b/src/inference/src/dev/threading/cpu_streams_executor.cpp
@@ -89,22 +89,6 @@ struct CPUStreamsExecutor::Impl {
                     });
                 }
             }
-#elif OV_THREAD == OV_THREAD_SEQ
-            auto proc_type_table = get_org_proc_type_table();
-            if (get_num_numa_nodes() > 1) {
-                pin_current_thread_to_socket(_numaNodeId);
-            } else if (proc_type_table.size() == 1 && proc_type_table[0][EFFICIENT_CORE_PROC] == 0 &&
-                       _impl->_config.get_cpu_pinning()) {
-                CpuSet processMask;
-                int ncpus = 0;
-                std::tie(processMask, ncpus) = get_process_mask();
-                if (nullptr != processMask) {
-                    pin_thread_to_vacant_core(_streamId + _impl->_config.get_thread_binding_offset(),
-                                              _impl->_config.get_thread_binding_step(),
-                                              ncpus,
-                                              processMask);
-                }
-            }
 #endif
         }
         ~Stream() {
@@ -226,6 +210,9 @@ struct CPUStreamsExecutor::Impl {
         std::unique_ptr<custom::task_arena> _taskArena;
         std::unique_ptr<Observer> _observer;
         std::vector<int> _cpu_ids;
+#elif OV_THREAD == OV_THREAD_SEQ
+        CpuSet _mask = nullptr;
+        int _ncpus = 0;
 #endif
     };
     // if the thread is created by CPUStreamsExecutor, the Impl::Stream of the thread is stored by tbb Class
@@ -429,7 +416,38 @@ struct CPUStreamsExecutor::Impl {
             task();
         }
 #else
+        pin_stream_to_cpus();
         task();
+        unpin_stream_to_cpus();
+#endif
+    }
+
+    void pin_stream_to_cpus() {
+#if OV_THREAD == OV_THREAD_SEQ
+        if (_config.get_cpu_pinning()) {
+            auto stream = _streams.local();
+            auto proc_type_table = get_org_proc_type_table();
+            std::tie(stream->_mask, stream->_ncpus) = get_process_mask();
+            if (get_num_numa_nodes() > 1) {
+                pin_current_thread_to_socket(stream->_numaNodeId);
+            } else if (proc_type_table.size() == 1 && proc_type_table[0][EFFICIENT_CORE_PROC] == 0) {
+                if (nullptr != stream->_mask) {
+                    pin_thread_to_vacant_core(stream->_streamId + _config.get_thread_binding_offset(),
+                                              _config.get_thread_binding_step(),
+                                              stream->_ncpus,
+                                              stream->_mask);
+                }
+            }
+        }
+#endif
+    }
+
+    void unpin_stream_to_cpus() {
+#if OV_THREAD == OV_THREAD_SEQ
+        auto stream = _streams.local();
+        if (stream->_mask) {
+            pin_current_thread_by_mask(stream->_ncpus, stream->_mask);
+        }
 #endif
     }