【feat】NPUAffinityController support multiple rangs of affinity cpu core-pytorch-AtomGit

已合并

【feat】NPUAffinityController support multiple rangs of affinity cpu core #34959

zhaoyu65创建于 5月7日从zhaoyu65/pytorch-npu： v2.7.1-Affinity 合入到 Ascend/pytorch： v2.7.1

【feat】NPUAffinityController support multiple rangs of affinity cpu core #34959

已合并

zhaoyu65创建于 5月7日

从zhaoyu65/pytorch-npu v2.7.1-Affinity合入到Ascend/pytorch v2.7.1

共 8 个文件变更+3277-2742

Mtest/torch_npu_schema.json+1-1

	@@ -2274,7 +2274,7 @@
2274	"signature": "(self)"	2274	"signature": "(self)"
2275	},	2275	},
2276	"torch_npu.utils.set_thread_affinity": {	2276	"torch_npu.utils.set_thread_affinity": {
2277	- "signature": "(core_range: List[int] = None)"	2277	+ "signature": "(core_range: list[int] \| list[list[int]] \| None = None)"
2278	},	2278	},
2279	"torch_npu.utils.reset_thread_affinity": {	2279	"torch_npu.utils.reset_thread_affinity": {
2280	"signature": "()"	2280	"signature": "()"

Mtest/utils/test_affinity.py+18-13

	@@ -1,10 +1,10 @@
		1	+# Owner(s): ["module: tests"]
1	import torch_npu	2	import torch_npu
2	-from torch_npu.testing.testcase import TestCase, run_tests	3	+from torch_npu.testing.testcase import run_tests, TestCase
3	-from torch_npu.utils.affinity import _set_thread_affinity, _reset_thread_affinity	4	+from torch_npu.utils.affinity import _reset_thread_affinity, _set_thread_affinity
4		5
5		6
6	class TestAffinity(TestCase):	7	class TestAffinity(TestCase):
7	-
8	def test_reset_thread_affinity(self):	8	def test_reset_thread_affinity(self):
9	original_func = torch_npu._C._npu_reset_thread_affinity	9	original_func = torch_npu._C._npu_reset_thread_affinity
10	call_count = 0	10	call_count = 0
	@@ -20,36 +20,40 @@ class TestAffinity(TestCase):
20	finally:	20	finally:
21	torch_npu._C._npu_reset_thread_affinity = original_func	21	torch_npu._C._npu_reset_thread_affinity = original_func
22		22
23	-
24	def test_set_thread_affinity_invalid_length(self):	23	def test_set_thread_affinity_invalid_length(self):
25	with self.assertRaises(ValueError) as context:	24	with self.assertRaises(ValueError) as context:
26	_set_thread_affinity([1, 2, 3])	25	_set_thread_affinity([1, 2, 3])
27	- self.assertIn("The length of input list of set_thread_affinity should be 2", str(context.exception))	26	+ self.assertIn("Invalid core range", str(context.exception))
28		27
29	with self.assertRaises(ValueError) as context:	28	with self.assertRaises(ValueError) as context:
30	_set_thread_affinity([])	29	_set_thread_affinity([])
31	- self.assertIn("The length of input list of set_thread_affinity should be 2", str(context.exception))	30	+ self.assertIn("Invalid core range", str(context.exception))
32		31
33	def test_set_thread_affinity_negative_values(self):	32	def test_set_thread_affinity_negative_values(self):
34	with self.assertRaises(ValueError) as context:	33	with self.assertRaises(ValueError) as context:
35	_set_thread_affinity([-1, 5])	34	_set_thread_affinity([-1, 5])
36	- self.assertIn("Core range should be nonnegative", str(context.exception))	35	+ self.assertIn("Invalid core range", str(context.exception))
37		36
38	with self.assertRaises(ValueError) as context:	37	with self.assertRaises(ValueError) as context:
39	_set_thread_affinity([2, -3])	38	_set_thread_affinity([2, -3])
40	- self.assertIn("Core range should be nonnegative", str(context.exception))	39	+ self.assertIn("Invalid core range", str(context.exception))
41		40
42	def test_set_thread_affinity_valid_range(self):	41	def test_set_thread_affinity_valid_range(self):
43	original_func = torch_npu._C._npu_set_thread_affinity	42	original_func = torch_npu._C._npu_set_thread_affinity
44	call_args = []	43	call_args = []
45		44
46	- def mock_npu_set_thread_affinity(start, end):	45	+ def mock_npu_set_thread_affinity(cores):
47	- call_args.append((start, end))	46	+ nonlocal call_args
		47	+ call_args = cores
48		48
49	torch_npu._C._npu_set_thread_affinity = mock_npu_set_thread_affinity	49	torch_npu._C._npu_set_thread_affinity = mock_npu_set_thread_affinity
50	try:	50	try:
51	_set_thread_affinity([2, 5])	51	_set_thread_affinity([2, 5])
52	- self.assertEqual(call_args, [(2, 5)])	52	+ self.assertEqual(call_args, [2, 3, 4, 5])
		53	+ _set_thread_affinity([[2, 5], [7, 9]])
		54	+ self.assertEqual(call_args, [2, 3, 4, 5, 7, 8, 9])
		55	+ _set_thread_affinity([[2, 7], [4, 9]])
		56	+ self.assertEqual(call_args, [2, 3, 4, 5, 6, 7, 8, 9])
53	finally:	57	finally:
54	torch_npu._C._npu_set_thread_affinity = original_func	58	torch_npu._C._npu_set_thread_affinity = original_func
55		59
	@@ -67,5 +71,6 @@ class TestAffinity(TestCase):
67	finally:	71	finally:
68	torch_npu._C._npu_set_thread_affinity = original_func	72	torch_npu._C._npu_set_thread_affinity = original_func
69		73
70	-if __name__ == '__main__':	74	+
71	- run_tests()	75	+if __name__ == "__main__":
		76	+ run_tests()

Mtorch_npu/csrc/core/npu/GetAffinityCPUInfo.cpp+134-112

	@@ -1,131 +1,153 @@
		1	+#include <c10/core/Device.h>
1	#include <climits>	2	#include <climits>
		3	+#include <iterator>
2	#include <unordered_map>	4	#include <unordered_map>
3	-#include "torch_npu/csrc/core/npu/interface/DcmiInterface.h"
4	-#include "torch_npu/csrc/core/npu/NPUException.h"
5	-#include "torch_npu/csrc/core/npu/NPUAffinityController.h"
6		5
		6	+#include <torch_npu/csrc/core/npu/NPUAffinityController.h>
		7	+#include <torch_npu/csrc/core/npu/NPUException.h>
		8	+#include <torch_npu/csrc/core/npu/interface/DcmiInterface.h>
		9	+
		10	+namespace {
7	constexpr int NPU_OK = 0;	11	constexpr int NPU_OK = 0;
		12	+using c10_npu::CoreId;
		13	+using c10_npu::CoreIdList;
		14	+using c10_npu::isAllDigits;
8		15
9	-static int DcmiInit()	16	+std::unordered_map<int, CoreIdList> CardIdAffinityCPU;
10	-{	17	+
		18	+void DcmiInit() {
		19	+ static bool initialized = false;
		20	+ if (!initialized) {
11	int ret = c10_npu::dcmi::DcmiInit();	21	int ret = c10_npu::dcmi::DcmiInit();
12	- if (ret != NPU_OK) {	22	+ TORCH_CHECK(
13	- TORCH_CHECK(false, "Failed to init dcmi. ", PTA_ERROR(ErrCode::INTERNAL));	23	+ ret == NPU_OK,
14	- }	24	+ "Failed to init dcmi. Error code: ",
15	- return ret;	25	+ ret,
		26	+ PTA_ERROR(ErrCode::ACL));
		27	+ initialized = true;
		28	+ }
16	}	29	}
17		30
18	-std::string GetAffinityCPUBaseInfo(int card_id)	31	+std::string GetAffinityCPUBaseInfo(int card_id) {
19	-{	32	+ int device_id = 0;
20	- int ret = DcmiInit();	33	+ int device_id_max = 0;
21	- int device_id = 0;	34	+ int mcu_id = 0;
22	- int device_id_max = 0;	35	+ int cpu_id = 0;
23	- int mcu_id = 0;	36	+ int ret = c10_npu::dcmi::DcmiGetDeviceIdInCard(
24	- int cpu_id = 0;	37	+ card_id, &device_id_max, &mcu_id, &cpu_id);
25	- ret = c10_npu::dcmi::DcmiGetDeviceIdInCard(card_id, &device_id_max, &mcu_id, &cpu_id);	38	+ if (ret != NPU_OK) {
26	- if (ret != NPU_OK) {	39	+ TORCH_NPU_WARN_ONCE(
27	- TORCH_NPU_WARN_ONCE("dcmi_get_device_id_in_card is not supported. "	40	+ "dcmi get device id in card is not supported. "
28	- "The npu_affine configuration of CPU_AFFINITY_CONF will be disabled.");	41	+ "The npu_affine configuration of CPU_AFFINITY_CONF will be disabled.");
29	- return "";
30	- }
31	- device_id = std::max(0, device_id_max - 1);
32	- char affinity_cpu[TOPO_INFO_MAX_LENTH] = {0};
33	- int length = 0;
34	- ret = c10_npu::dcmi::DcmiGetAffinityCpuInfoByDeviceId(card_id, device_id, affinity_cpu, &length);
35	- if (ret == NPU_OK) {
36	- return affinity_cpu;
37	- }
38	- TORCH_NPU_WARN_ONCE("dcmi_get_affinity_cpu_info_by_device_id is not supported. "
39	- "The npu_affine configuration of CPU_AFFINITY_CONF will be disabled.");
40	return "";	42	return "";
		43	+ }
		44	+ device_id = std::max(0, device_id_max - 1);
		45	+ char affinity_cpu[TOPO_INFO_MAX_LENTH] = {0};
		46	+ int length = 0;
		47	+ ret = c10_npu::dcmi::DcmiGetAffinityCpuInfoByDeviceId(
		48	+ card_id, device_id, affinity_cpu, &length);
		49	+ if (ret == NPU_OK) {
		50	+ return affinity_cpu;
		51	+ }
		52	+ TORCH_NPU_WARN_ONCE(
		53	+ "dcmi get affinity cpu info by device id is not supported. "
		54	+ "The npu_affine configuration of CPU_AFFINITY_CONF will be disabled.");
		55	+ return "";
41	}	56	}
42		57
43	-std::unordered_map<int, c10_npu::CoreIdRange> CardIdAffinityCPU;	58	+CoreIdList parseAffinityCores(const std::string cpuString) {
44	-	59	+ CoreIdList cores;
45	-c10_npu::CoreIdRange parseAffinityCPU(const std::string cpuString)	60	+ TORCH_CHECK(
46	-{	61	+ !cpuString.empty(),
47	- size_t pos = cpuString.find("-");	62	+ "Affinity cpu string is empty",
48	- if (pos != std::string::npos) {	63	+ PTA_ERROR(ErrCode::VALUE));
49	- std::string start = cpuString.substr(0, pos);	64	+ std::stringstream ss_value(cpuString);
50	- std::string end = cpuString.substr(pos + 1);	65	+ std::string range;
51	-	66	+ while (std::getline(ss_value, range, ',')) {
52	- char* start_endptr = nullptr;	67	+ size_t dashPos = range.find('-');
53	- char* end_endptr = nullptr;	68	+ if (dashPos != std::string::npos) {
54	-	69	+ std::string startStr = range.substr(0, dashPos);
55	- errno = 0;	70	+ std::string endStr = range.substr(dashPos + 1);
56	- long startNum = strtol(start.c_str(), &start_endptr, 10);	71	+ if (isAllDigits(startStr) && isAllDigits(endStr)) {
57	- long endNum = strtol(end.c_str(), &end_endptr, 10);	72	+ CoreId start = static_cast<CoreId>(std::stoi(startStr));
58	- if (start_endptr != start.c_str() && *start_endptr == '\0' &&	73	+ CoreId end = static_cast<CoreId>(std::stoi(endStr));
59	- end_endptr != end.c_str() && *end_endptr == '\0' &&	74	+ for (CoreId id = start; id <= end; ++id) {
60	- startNum >= 0 && endNum >= 0 && startNum <= INT_MAX && endNum <= INT_MAX &&	75	+ cores.insert(id);
61	- errno != ERANGE &&
62	- startNum < endNum) {
63	- return c10_npu::CoreIdRange{static_cast<int>(startNum), static_cast<int>(endNum)};
64	}	76	}
		77	+ }
65	}	78	}
66	- TORCH_CHECK(false, "affinity cpu " + cpuString + " is error ", PTA_ERROR(ErrCode::VALUE));	79	+ }
		80	+ return cores;
67	}	81	}
68		82
69	-void GetExclusiveAffinityCPU()	83	+void GetExclusiveAffinityCPU() {
70	-{	84	+ static bool initialized = false;
71	- int ret = DcmiInit();	85	+ if (initialized) {
72	- int device_count = 0;	86	+ return;
73	- int card_id_list[16];	87	+ }
74	- int list_len = 16;	88	+ initialized = true;
75	- ret = c10_npu::dcmi::DcmiGetCardNumList(&device_count, card_id_list, list_len);	89	+ DcmiInit();
76	- std::unordered_map<std::string, int> SameAffinityCpuNum;
77	- std::map<int, std::string> CardIdAffinityCpuDefault;
78	- for (int i = 0; i < device_count; i++) {
79	- std::string affinity_cpu = GetAffinityCPUBaseInfo(i);
80	- if (affinity_cpu.find(",") != std::string::npos) {
81	- ASCEND_LOGW("torch_npu not support affinity cpu format:%s when set npu_affine:1.", affinity_cpu.c_str());
82	- CardIdAffinityCPU.clear();
83	- return;
84	- }
85	- if (affinity_cpu.empty()) {
86	- return;
87	- }
88	- CardIdAffinityCpuDefault[i] = affinity_cpu;
89	- auto it = SameAffinityCpuNum.find(affinity_cpu);
90	- if (it != SameAffinityCpuNum.end()) {
91	- SameAffinityCpuNum[affinity_cpu] = it->second + 1;
92	- } else {
93	- SameAffinityCpuNum[affinity_cpu] = 1;
94	- }
95	- }
96	- std::unordered_map<std::string, int> offsetMap;
97	- for (const auto& it : CardIdAffinityCpuDefault) {
98	- int card_id = it.first;
99	- std::string affinity_cpu = it.second;
100	- int same_num = 1;
101	- auto find_same_affinity_cpu = SameAffinityCpuNum.find(affinity_cpu);
102	- if (find_same_affinity_cpu != SameAffinityCpuNum.end()) {
103	- same_num = find_same_affinity_cpu->second;
104	- }
105	- int offset = 0;
106	- auto find_offset = offsetMap.find(affinity_cpu);
107	- if (find_offset != offsetMap.end()) {
108	- offset = find_offset->second;
109	- }
110	- c10_npu::CoreIdRange cpu_range = parseAffinityCPU(affinity_cpu);
111	- unsigned int length = (cpu_range.end - cpu_range.start + 1) / static_cast<unsigned int>(same_num);
112	- c10_npu::CoreIdRange exclusiveAffinityCpu = {
113	- cpu_range.start + static_cast<unsigned int>(offset) * length,
114	- (cpu_range.start + length - 1) + static_cast<unsigned int>(offset) * length};
115	- offsetMap[affinity_cpu] = offset + 1;
116	- CardIdAffinityCPU[card_id] = exclusiveAffinityCpu;
117	- }
118	-}
119		90
120	-c10_npu::CoreIdRange GetAssignAffinityCPU(int card_id)	91	+ int device_count = 0;
121	-{	92	+ int card_id_list[16];
122	- GetExclusiveAffinityCPU();	93	+ int list_len = 16;
123	- if (CardIdAffinityCPU.empty()) {	94	+ TORCH_CHECK(
124	- return {0, 0};	95	+ c10_npu::dcmi::DcmiGetCardNumList(
		96	+ &device_count, card_id_list, list_len) == NPU_OK,
		97	+ "Dcmi get card num failed.");
		98	+ std::unordered_map<std::string, std::vector<int>> rangeAndDevs;
		99	+ for (int i = 0; i < device_count; i++) {
		100	+ std::string affinity_range = GetAffinityCPUBaseInfo(i);
		101	+ if (affinity_range.empty()) {
		102	+ TORCH_NPU_WARN_ONCE(
		103	+ "Get affinity cpu info by device id is not supported, "
		104	+ "The npu_affine configuration of CPU_AFFINITY_CONF will be disabled.");
		105	+ return;
125	}	106	}
126	- auto it = CardIdAffinityCPU.find(card_id);	107	+ rangeAndDevs[affinity_range].push_back(i);
127	- if (it != CardIdAffinityCPU.end()) {	108	+ ASCEND_LOGD(
128	- return it->second;	109	+ "Device_id: %d, affinity_range: %s.", i, affinity_range.c_str());
		110	+ }
		111	+ for (auto& [affinity_range, dev_list] : rangeAndDevs) {
		112	+ CoreIdList cores = parseAffinityCores(affinity_range);
		113	+ int per_dev_cores = cores.size() / dev_list.size();
		114	+ if (per_dev_cores == 0) {
		115	+ ASCEND_LOGD("Insufficient cores for device allocation.");
		116	+ for (int i = 0; i < dev_list.size(); ++i) {
		117	+ CardIdAffinityCPU[dev_list[i]].insert(cores.begin(), cores.end());
		118	+ }
		119	+ continue;
129	}	120	}
130	- TORCH_CHECK(false, "card_id ", std::to_string(card_id), " is invalid.", PTA_ERROR(ErrCode::VALUE));	121	+ for (int i = 0; i < dev_list.size(); ++i) {
		122	+ auto& dev_i = CardIdAffinityCPU[dev_list[i]];
		123	+ auto it = cores.begin();
		124	+ std::advance(it, i * per_dev_cores);
		125	+ auto end_it = cores.begin();
		126	+ if (i == dev_list.size() - 1) {
		127	+ end_it = cores.end();
		128	+ } else {
		129	+ std::advance(end_it, (i + 1) * per_dev_cores);
		130	+ }
		131	+ dev_i.insert(it, end_it);
		132	+ }
		133	+ }
131	}	134	}
		135	+} // namespace
		136	+
		137	+namespace c10_npu {
		138	+CoreIdList GetAffinityCores(int card_id) {
		139	+ GetExclusiveAffinityCPU();
		140	+ if (CardIdAffinityCPU.empty()) {
		141	+ return CoreIdList{};
		142	+ }
		143	+ auto it = CardIdAffinityCPU.find(card_id);
		144	+ if (it != CardIdAffinityCPU.end()) {
		145	+ return it->second;
		146	+ }
		147	+ TORCH_CHECK(
		148	+ false,
		149	+ "Can't get affinity cores for card_id ",
		150	+ std::to_string(card_id),
		151	+ PTA_ERROR(ErrCode::VALUE));
		152	+}
		153	+} // namespace c10_npu

Mtorch_npu/csrc/core/npu/GetAffinityCPUInfo.h+7-6

	@@ -1,10 +1,11 @@
1	#ifndef THNP_GETAFFINITY_INC	1	#ifndef THNP_GETAFFINITY_INC
2	#define THNP_GETAFFINITY_INC	2	#define THNP_GETAFFINITY_INC
3	-#include <string>	3	+#include <set>
4		4
5	-std::string GetAffinityCPUBaseInfo(int card_id);	5	+namespace c10_npu {
6	-c10_npu::CoreIdRange parseAffinityCPU(const std::string cpuString);	6	+using CoreId = unsigned int;
7	-void GetExclusiveAffinityCPU();	7	+using CoreIdList = std::set<CoreId>;
8	-c10_npu::CoreIdRange GetAssignAffinityCPU(int card_id);
9		8
10	-#endif	9	+CoreIdList GetAffinityCores(int card_id);
		10	+} // namespace c10_npu
		11	+#endif

Mtorch_npu/csrc/core/npu/NPUAffinityController.cpp+500-388

	@@ -1,445 +1,557 @@
1	-#include "torch_npu/csrc/core/npu/NPUAffinityController.h"	1	+#include <torch_npu/csrc/core/npu/NPUAffinityController.h>
2	-#include "torch_npu/csrc/core/npu/NPUFunctions.h"
3	-#include "torch_npu/csrc/core/npu/GetAffinityCPUInfo.h"
4	-#include "torch_npu/csrc/core/npu/NpuVariables.h"
5		2
6	#include <pthread.h>	3	#include <pthread.h>
7	-#include <unistd.h>
8	-#include <sys/syscall.h>
9	#include <sys/prctl.h>	4	#include <sys/prctl.h>
		5	+#include <sys/syscall.h>
		6	+#include <unistd.h>
		7	+#include <mutex>
		8	+#include <regex>
		9	+#include <set>
10	#include <string>	10	#include <string>
11	#include <unordered_map>	11	#include <unordered_map>
12	-#include <mutex>	12	+
		13	+#include <torch_npu/csrc/core/npu/GetAffinityCPUInfo.h>
		14	+#include <torch_npu/csrc/core/npu/NPUFunctions.h>
		15	+#include <torch_npu/csrc/core/npu/NpuVariables.h>
13		16
14	namespace c10_npu {	17	namespace c10_npu {
15		18
16	-static thread_local ThreadType local_thread = ThreadType::MAIN_THREAD;	19	+namespace {
17		20
18	-static pthread_t main_thread;	21	+thread_local ThreadType local_thread = ThreadType::MAIN_THREAD;
19	-static bool start_main_thread_bind = false;
20	-static std::mutex core_map_mutex;
21	-static bool lazy_bind = true;
22	-static bool force_bind = false;
23		22
		23	+pthread_t main_thread;
		24	+bool start_main_thread_bind = false;
		25	+std::mutex core_map_mutex;
		26	+bool lazy_bind = true;
		27	+bool force_bind = false;
24		28
25	-using ThreadCoreMap = std::unordered_map<ThreadType, CoreIdRange>;	29	+using ThreadCoreMap = std::unordered_map<ThreadType, CoreIdList>;
26		30
27	-static uint32_t cpu_affinity_mode;	31	+uint32_t cpu_affinity_mode;
28	-static std::vector<CoreIdRange> device_ranges;	32	+std::vector<CoreIdList> devices_aff_cores;
29	-static std::unordered_map<c10::DeviceIndex, ThreadCoreMap> device_thread_core_maps;	33	+std::unordered_map<c10::DeviceIndex, ThreadCoreMap> device_thread_core_maps;
30		34
31	-const std::initializer_list<ThreadType> threadTypeList = {	35	+const std::initializer_list<ThreadType> threadTypeList =
32	- MAIN_THREAD, ACL_THREAD, RELEASE_THREAD, WATCHDOG_THREAD, OTHER_THREAD};	36	+ {MAIN_THREAD, ACL_THREAD, RELEASE_THREAD, WATCHDOG_THREAD, OTHER_THREAD};
33		37
34	const std::unordered_map<ThreadType, std::string> threadTypeToNameMap = {	38	const std::unordered_map<ThreadType, std::string> threadTypeToNameMap = {
35	- {MAIN_THREAD, "main_thread"},	39	+ {MAIN_THREAD, "main_thread"},
36	- {ACL_THREAD, "acl_thread"},	40	+ {ACL_THREAD, "acl_thread"},
37	- {RELEASE_THREAD, "release_thread"},	41	+ {RELEASE_THREAD, "release_thread"},
38	- {WATCHDOG_THREAD, "hccl_watchdog_t"},	42	+ {WATCHDOG_THREAD, "hccl_watchdog_t"},
39	- {OTHER_THREAD, "other_thread"}};	43	+ {OTHER_THREAD, "other_thread"}};
40		44
41	-CoreIdRange getCPUDefaultRange(c10::DeviceIndex device_id)	45	+std::string formatCoreRange(const CoreIdList& cores) {
42	-{	46	+ if (cores.empty()) {
43	- static int core_nums = sysconf(_SC_NPROCESSORS_ONLN);	47	+ return "";
44	- int device_nums = device_count_ensure_non_zero();	48	+ }
45	- int block_size = (core_nums > 0 && device_nums > 0) ? core_nums / device_nums : 0;
46	- return CoreIdRange{static_cast<CoreId>(device_id * block_size),
47	- static_cast<CoreId>((device_id + 1) * block_size - 1)};
48	-}
49		49
50	-inline bool isAllDigits(const std::string &str)	50	+ std::ostringstream oss;
51	-{	51	+ auto it = cores.begin();
52	- if (str.empty()) {	52	+ while (it != cores.end()) {
53	- return false;	53	+ CoreId start = *it;
54	- }	54	+ CoreId end = start;
55	- return std::all_of(str.begin(), str.end(), [](unsigned char c) {	55	+ auto next_it = std::next(it);
56	- return std::isdigit(c);	56	+ while (next_it != cores.end() && *next_it == end + 1) {
57	- });	57	+ ++end;
58	-}	58	+ ++next_it;
59	-
60	-void parseCPUAffinityConf(uint32_t &mode, std::vector<CoreIdRange> &ranges)
61	-{
62	- // init
63	- int device_nums = device_count_ensure_non_zero();
64	- ranges.clear();
65	- ranges.resize(device_nums);
66	- for (int i = 0; i < device_nums; ++i) {
67	- ranges[i] = getCPUDefaultRange(i);
68	- }
69	- mode = 0;
70	-
71	- const char *input = c10_npu::option::OptionsManager::GetCpuAffinityConf();
72	- if (input == nullptr \|\| strlen(input) == 0) {
73	- return;
74	}	59	}
75		60
76	- std::string inputStr(input);	61	+ if (start == end) {
77	- std::istringstream stream(inputStr);	62	+ oss << start;
78	- std::string option;
79	-
80	- std::regex pattern("npu_affine:(\\d)");
81	- std::smatch match;
82	- if (std::regex_search(inputStr, match, pattern)) {
83	- int isAffinity = std::stoi(match[1].str());
84	- if (isAffinity != 0) {
85	- for (int i = 0; i < device_nums; i++) {
86	- CoreIdRange getRange = GetAssignAffinityCPU(i);
87	- if (getRange.start == 0 && getRange.end == 0) {
88	- break;
89	- }
90	- ranges[i] = getRange;
91	- }
92	- }
93	- }
94	-
95	- std::regex pattern_for_lazy_bind("lazy_bind:(\\d)");
96	- std::smatch match_for_lazy_bind;
97	- if (std::regex_search(inputStr, match_for_lazy_bind, pattern_for_lazy_bind)) {
98	- int lazy_bind_val = std::stoi(match_for_lazy_bind[1].str());
99	- if (lazy_bind_val == 0) {
100	- lazy_bind = false;
101	- }
102	- }
103	-
104	- std::regex pattern_for_force("force:(\\d)");
105	- std::smatch match_for_force;
106	- if (std::regex_search(inputStr, match_for_force, pattern_for_force)) {
107	- int force_val = std::stoi(match_for_force[1].str());
108	- if (force_val != 0 && force_val != 1) {
109	- ASCEND_LOGE("force value must be 0 or 1, got: %d", force_val);
110	- } else {
111	- force_bind = (force_val != 0);
112	- }
113	} else {	63	} else {
114	- std::regex pattern_for_force_check("force:([^,]+)");	64	+ oss << start << "-" << end;
115	- std::smatch match_for_force_check;
116	- if (std::regex_search(inputStr, match_for_force_check, pattern_for_force_check)) {
117	- ASCEND_LOGE("force value must be 0 or 1, got: %s", match_for_force_check[1].str().c_str());
118	- }
119	}	65	}
120		66
121	- // Handle cases where only `mode` is provided, or `mode:` without value	67	+ it = next_it;
122	- if (isAllDigits(inputStr)) {	68	+ if (it != cores.end()) {
123	- mode = static_cast<uint32_t>(std::stoi(inputStr));	69	+ oss << ",";
124	- return; // Return directly, `mode` has already been processed
125	- }
126	-
127	- // Parse each option
128	- while (std::getline(stream, option, ',')) {
129	- // Split `option` based on colon
130	- size_t colonPos = option.find(':');
131	- if (colonPos != std::string::npos) {
132	- std::string key = option.substr(0, colonPos);
133	- std::string value = option.substr(colonPos + 1);
134	-
135	- // Process `mode`
136	- if (key == "mode") {
137	- if (isAllDigits(value)) {
138	- mode = static_cast<uint32_t>(std::stoi(value));
139	- } else {
140	- ASCEND_LOGW("mode is %s, should be all digits", value.c_str());
141	- }
142	- } else if (key.rfind("npu", 0) == 0) {
143	- // Handle NPU core binding range
144	- // The key is like 'npu:0', so skip first 3 chars.
145	- if (isAllDigits(key.substr(3))) {
146	- int device_id = std::stoi(key.substr(3)); // Parse NPU device ID
147	- if (device_id < device_nums) {
148	- size_t dashPos = value.find('-');
149	- if (dashPos != std::string::npos) {
150	- std::string startStr = value.substr(0, dashPos);
151	- std::string endStr = value.substr(dashPos + 1);
152	- if (isAllDigits(startStr) && isAllDigits(endStr)) {
153	- CoreId start = static_cast<CoreId>(std::stoi(startStr));
154	- CoreId end = static_cast<CoreId>(std::stoi(endStr));
155	- ranges[device_id] = {start, end};
156	- } else {
157	- ASCEND_LOGW("core range is %s-%s, should be all digits", startStr.c_str(), endStr.c_str());
158	- }
159	- } else {
160	- if (isAllDigits(value)) {
161	- CoreId singleCore = static_cast<CoreId>(std::stoi(value));
162	- ranges[device_id] = {singleCore, singleCore};
163	- } else {
164	- ASCEND_LOGW("core range is string : %s, should be all digits", value.c_str());
165	- }
166	- }
167	- }
168	- }
169	- }
170	- } else if (isAllDigits(option)) {
171	- // If no colon and the value is a number, use it directly as `mode`
172	- mode = static_cast<uint32_t>(std::stoi(option));
173	- }
174	}	70	}
		71	+ }
		72	+ return oss.str();
175	}	73	}
176		74
177	-void printCoreRanges(const uint32_t mode, const std::vector<CoreIdRange> &ranges)	75	+CoreIdList getNPUDefaultCores(c10::DeviceIndex device_id) {
178	-{	76	+ static int core_nums = sysconf(_SC_NPROCESSORS_ONLN);
179	- std::ostringstream oss;	77	+ int device_nums = device_count_ensure_non_zero();
180	- oss << "Mode: " << mode << ". Core range for each device ID: ";	78	+ int block_size =
181	-	79	+ (core_nums > 0 && device_nums > 0) ? core_nums / device_nums : 0;
182	- for (size_t i = 0; i < ranges.size(); ++i) {	80	+ CoreIdList cores;
183	- oss << "Device " << i << ": [" << ranges[i].start << ", " << ranges[i].end << "]";	81	+ for (int i = 0; i < block_size; ++i) {
184	- if (i != ranges.size() - 1) {	82	+ cores.insert(static_cast<CoreId>(device_id * block_size + i));
185	- oss << "; ";	83	+ }
186	- } else {	84	+ return cores;
187	- oss << ".";
188	- }
189	- }
190	-
191	- ASCEND_LOGD("Read CPU affinity config: %s", oss.str().c_str());
192	}	85	}
193		86
194	-bool getThreadAffinityInfo()	87	+// Parse npu_affine setting from CPU_AFFINITY_CONF
195	-{	88	+void parseNpuAffineMode(
196	- parseCPUAffinityConf(cpu_affinity_mode, device_ranges);	89	+ const std::string& inputStr,
197	- printCoreRanges(cpu_affinity_mode, device_ranges);	90	+ int device_nums,
		91	+ std::vector<CoreIdList>& devices_aff_cores) {
		92	+ static const std::regex pattern("npu_affine:(\\d)");
		93	+ std::smatch match;
		94	+ if (std::regex_search(inputStr, match, pattern) &&
		95	+ std::stoi(match[1].str()) != 0) {
		96	+ ASCEND_LOGD(
		97	+ "Get npu_affine mode: %s, device_nums: %d, set affinity cores for each device.",
		98	+ match[1].str().c_str(),
		99	+ device_nums);
		100	+ for (int i = 0; i < device_nums; i++) {
		101	+ CoreIdList cores = GetAffinityCores(i);
		102	+ if (cores.size() == 0) {
		103	+ ASCEND_LOGW("Device-%d has no affinity cores.", i);
		104	+ continue;
		105	+ }
		106	+ devices_aff_cores[i] = cores;
		107	+ ASCEND_LOGD(
		108	+ "Device-%d affinity cores [%s].",
		109	+ i,
		110	+ formatCoreRange(devices_aff_cores[i]).c_str());
		111	+ }
		112	+ }
		113	+}
198		114
199	- if (cpu_affinity_mode == 0) {	115	+// Parse lazy_bind setting from CPU_AFFINITY_CONF
200	- return false;	116	+void parseLazyBindMode(const std::string& inputStr) {
201	- }	117	+ std::regex pattern_for_lazy_bind("lazy_bind:(\\d)");
		118	+ std::smatch match_for_lazy_bind;
		119	+ if (std::regex_search(inputStr, match_for_lazy_bind, pattern_for_lazy_bind)) {
		120	+ lazy_bind = std::stoi(match_for_lazy_bind[1].str()) == 0 ? false : true;
		121	+ }
		122	+}
202		123
203	- if (force_bind) {	124	+// Parse force setting from CPU_AFFINITY_CONF
204	- ASCEND_LOGI("CPU affinity force mode enabled, skipping affinity conflict detection, applying CPU_AFFINITY_CONF binding.");	125	+void parseForceMode(const std::string& inputStr) {
205	- return true;	126	+ std::regex pattern_for_force("force:(\\d)");
		127	+ std::smatch match_for_force;
		128	+ if (std::regex_search(inputStr, match_for_force, pattern_for_force)) {
		129	+ int force_val = std::stoi(match_for_force[1].str());
		130	+ if (force_val != 0 && force_val != 1) {
		131	+ ASCEND_LOGE("force value must be 0 or 1, got: %d", force_val);
		132	+ } else {
		133	+ force_bind = (force_val != 0);
206	}	134	}
		135	+ } else {
		136	+ std::regex pattern_for_force_check("force:([^,]+)");
		137	+ std::smatch match_for_force_check;
		138	+ if (std::regex_search(
		139	+ inputStr, match_for_force_check, pattern_for_force_check)) {
		140	+ ASCEND_LOGE(
		141	+ "force value must be 0 or 1, got: %s",
		142	+ match_for_force_check[1].str().c_str());
		143	+ }
		144	+ }
		145	+}
207		146
208	- cpu_set_t mask;	147	+// Parse mode from CPU_AFFINITY_CONF when only digits or mode:xxx is provided
209	- pthread_getaffinity_np(pthread_self(), sizeof(mask), &mask);	148	+bool parseModeOnly(const std::string& inputStr, uint32_t& mode) {
210	-	149	+ // Handle cases where only `mode` is provided, or `mode:` without value
211	- std::ostringstream affinity_oss;	150	+ if (isAllDigits(inputStr)) {
212	- affinity_oss << "[";	151	+ mode = static_cast<uint32_t>(std::stoi(inputStr));
213	- int cpu_count = sysconf(_SC_NPROCESSORS_ONLN);
214	- int range_start = -1;
215	- bool first_range = true;
216	- for (int k = 0; k < cpu_count; ++k) {
217	- if (CPU_ISSET(k, &mask)) {
218	- if (range_start == -1) {
219	- range_start = k;
220	- }
221	- } else {
222	- if (range_start != -1) {
223	- if (!first_range) affinity_oss << ", ";
224	- if (range_start == k - 1) {
225	- affinity_oss << range_start;
226	- } else {
227	- affinity_oss << range_start << "-" << (k - 1);
228	- }
229	- range_start = -1;
230	- first_range = false;
231	- }
232	- }
233	- }
234	- if (range_start != -1) {
235	- if (!first_range) affinity_oss << ", ";
236	- if (range_start == cpu_count - 1) {
237	- affinity_oss << range_start;
238	- } else {
239	- affinity_oss << range_start << "-" << (cpu_count - 1);
240	- }
241	- }
242	- affinity_oss << "]";
243	- ASCEND_LOGI("Current thread CPU affinity mask: %s", affinity_oss.str().c_str());
244	-
245	- for (auto &range : device_ranges) {
246	- for (unsigned int i = range.start; i <= range.end; i++) {
247	- if (!CPU_ISSET(i, &mask)) {
248	- ASCEND_LOGW("Thread affinity conflict detected! Expected core %u (in config range [%u, %u]) is NOT in current thread affinity mask. %s",
249	- i, range.start, range.end, affinity_oss.str().c_str());
250	- ASCEND_LOGW("Thread affinity is already set. Use force:1 to skip this check and force bind.");
251	- return false;
252	- }
253	- }
254	- }
255	return true;	152	return true;
256	-}	153	+ }
257		154
258	-inline bool needToSetThreadAffinity()	155	+ std::istringstream stream(inputStr);
259	-{	156	+ std::string option;
260	- static bool need_to_set_affinity = getThreadAffinityInfo();	157	+ while (std::getline(stream, option, ',')) {
261	- return need_to_set_affinity;	158	+ size_t colonPos = option.find(':');
262	-}	159	+ if (colonPos == std::string::npos && isAllDigits(option)) {
263	-	160	+ mode = static_cast<uint32_t>(std::stoi(option));
264	-void SetThreadType(ThreadType type)	161	+ return false;
265	-{
266	- // Called at the start of the thread's execution to avoid frequent triggering of this function.
267	- local_thread = type;
268	- if (type == ThreadType::OTHER_THREAD \|\| type == ThreadType::MAIN_THREAD) {
269	- return;
270	}	162	}
271	- if (prctl(PR_SET_NAME, threadTypeToNameMap.at(type).c_str()) != 0) {	163	+ std::string key = option.substr(0, colonPos);
272	- ASCEND_LOGW("Set thread name to %s failed!", threadTypeToNameMap.at(type).c_str());	164	+ if (key == "mode") {
		165	+ std::string value = option.substr(colonPos + 1);
		166	+ if (isAllDigits(value)) {
		167	+ mode = static_cast<uint32_t>(std::stoi(value));
		168	+ } else {
		169	+ ASCEND_LOGW("mode is %s, should be all digits", value.c_str());
		170	+ }
		171	+ return false;
273	}	172	}
		173	+ }
		174	+ return false;
274	}	175	}
275		176
276	-std::string getAffinityMapAsString(c10::DeviceIndex device_id, const ThreadCoreMap &threadCoreMap)	177	+// Parse device-specific core range from CPU_AFFINITY_CONF (e.g., npu0:0-1)
277	-{	178	+void parseDeviceCoreRange(
278	- std::ostringstream oss;	179	+ const std::string& inputStr,
279	- for (auto thread_type : threadTypeList) {	180	+ int device_nums,
280	- oss << threadTypeToNameMap.at(thread_type) << ": ["	181	+ std::vector<CoreIdList>& devices_aff_cores) {
281	- << threadCoreMap.at(thread_type).start << ", "	182	+ std::istringstream stream(inputStr);
282	- << threadCoreMap.at(thread_type).end << "]";	183	+ std::string option;
283	- if (thread_type != OTHER_THREAD) {	184	+ std::set<int> user_def_devices;
284	- oss << "; ";	185	+ while (std::getline(stream, option, ',')) {
285	- } else {	186	+ size_t colonPos = option.find(':');
286	- oss << ".";	187	+ if (colonPos == std::string::npos) {
		188	+ continue;
		189	+ }
		190	+ std::string key = option.substr(0, colonPos);
		191	+ std::string value = option.substr(colonPos + 1);
		192	+
		193	+ std::regex npuPattern("^npu[0-9]{1,2}$");
		194	+ if (!std::regex_match(key, npuPattern)) {
		195	+ ASCEND_LOGW("Invalid device name: %s", key.c_str());
		196	+ continue;
		197	+ }
		198	+ int device_id = std::stoi(key.substr(3)); // Skip first 3 chars ("npu").
		199	+ if (device_id >= device_nums \|\| device_id < 0) {
		200	+ ASCEND_LOGW(
		201	+ "device_id in CPU_AFFINITY_CONF is %d, should be in range [0, %d)",
		202	+ device_id,
		203	+ device_nums);
		204	+ continue;
		205	+ }
		206	+ if (user_def_devices.count(device_id) == 0) {
		207	+ user_def_devices.insert(device_id);
		208	+ devices_aff_cores[device_id].clear();
		209	+ }
		210	+ if (isAllDigits(value)) {
		211	+ CoreId singleCore = static_cast<CoreId>(std::stoi(value));
		212	+ devices_aff_cores[device_id].insert(singleCore);
		213	+ continue;
		214	+ }
		215	+ size_t dashPos = value.find('-');
		216	+ if (dashPos != std::string::npos) {
		217	+ std::string startStr = value.substr(0, dashPos);
		218	+ std::string endStr = value.substr(dashPos + 1);
		219	+ if (isAllDigits(startStr) && isAllDigits(endStr)) {
		220	+ CoreId start = static_cast<CoreId>(std::stoi(startStr));
		221	+ CoreId end = static_cast<CoreId>(std::stoi(endStr));
		222	+ for (CoreId core = start; core <= end; ++core) {
		223	+ devices_aff_cores[device_id].insert(core);
287	}	224	}
		225	+ } else {
		226	+ ASCEND_LOGW(
		227	+ "core range is %s-%s, should be all digits",
		228	+ startStr.c_str(),
		229	+ endStr.c_str());
		230	+ }
288	}	231	}
289	- return oss.str();	232	+ }
290	}	233	}
291		234
292	-ThreadCoreMap getCpuAffinityMap(c10::DeviceIndex device_id, const std::vector<CoreIdRange> &deviceRanges)	235	+void parseCPUAffinityConf(
293	-{	236	+ uint32_t& mode,
294	- ThreadCoreMap threadCoreMap;	237	+ std::vector<CoreIdList>& devices_aff_cores) {
295	- CoreIdRange range = deviceRanges[device_id];	238	+ int device_nums = device_count_ensure_non_zero();
296	- unsigned int core_nums = range.end - range.start + 1;	239	+ devices_aff_cores.clear();
297	- if (core_nums < threadTypeList.size()) {	240	+ devices_aff_cores.resize(device_nums);
298	- ASCEND_LOGW("Device %d available core numbers (%d) are insufficient for all %zu thread types and will bind available cores to all threads.",	241	+ ASCEND_LOGD("Get device nums: %d by aclrtGetDeviceCount.", device_nums);
299	- device_id, core_nums, threadTypeList.size());	242	+ for (int i = 0; i < device_nums; ++i) {
300	- for (auto thread_type : threadTypeList) {	243	+ devices_aff_cores[i] = getNPUDefaultCores(i);
301	- threadCoreMap[thread_type] = range;	244	+ }
302	- }	245	+ mode = 0;
303	- return threadCoreMap;
304	- }
305		246
306	- CoreId now = range.start;	247	+ const char* input = c10_npu::option::OptionsManager::GetCpuAffinityConf();
307	- for (auto thread_type : threadTypeList) {	248	+ if (input == nullptr \|\| strlen(input) == 0) {
308	- if (thread_type != ThreadType::OTHER_THREAD) {	249	+ return; // CPU_AFFINITY_CONF is not set, use default cores
309	- threadCoreMap[thread_type] = CoreIdRange{now, now};	250	+ }
310	- } else {	251	+ ASCEND_LOGD("Get env var CPU_AFFINITY_CONF: %s", input);
311	- threadCoreMap[ThreadType::OTHER_THREAD] = CoreIdRange{now, range.end};
312	- }
313	- now++;
314	- }
315		252
316	- ASCEND_LOGD("Device %d thread affinity map: %s", device_id, getAffinityMapAsString(device_id, threadCoreMap).c_str());	253	+ const std::string inputStr(input);
317	- return threadCoreMap;	254	+
		255	+ // Parse mode if only digits or mode:xxx is provided
		256	+ if (parseModeOnly(inputStr, mode)) {
		257	+ ASCEND_LOGD("Only mode is provided, mode: %d", mode);
		258	+ return;
		259	+ }
		260	+
		261	+ parseNpuAffineMode(inputStr, device_nums, devices_aff_cores);
		262	+ parseLazyBindMode(inputStr);
		263	+ parseForceMode(inputStr);
		264	+ // Parse device-specific core ranges defined by user
		265	+ parseDeviceCoreRange(inputStr, device_nums, devices_aff_cores);
318	}	266	}
319		267
320	-bool setThreadAffinityImpl(pthread_t thread, CoreIdRange core_range)	268	+void printCoreRanges(
321	-{	269	+ const uint32_t mode,
322	- cpu_set_t mask;	270	+ const std::vector<CoreIdList>& devices_aff_cores) {
323	- CPU_ZERO(&mask);	271	+ std::ostringstream oss;
324	- for (auto i = core_range.start; i <= core_range.end; i++) {	272	+ oss << "Mode: " << mode << ". Core range for each device ID: ";
325	- CPU_SET(i, &mask);	273	+
		274	+ for (size_t i = 0; i < devices_aff_cores.size(); ++i) {
		275	+ oss << "Device " << i << ": [" << formatCoreRange(devices_aff_cores[i])
		276	+ << "]";
		277	+ std::string end_str = (i == devices_aff_cores.size() - 1) ? "." : "; ";
		278	+ oss << end_str;
		279	+ }
		280	+ ASCEND_LOGD("Read CPU affinity config: %s", oss.str().c_str());
		281	+}
		282	+
		283	+std::string formatCPUSetMask(const cpu_set_t& mask) {
		284	+ CoreIdList cores;
		285	+ int cpu_count = sysconf(_SC_NPROCESSORS_ONLN);
		286	+ for (int i = 0; i < cpu_count; ++i) {
		287	+ if (CPU_ISSET(i, &mask)) {
		288	+ cores.insert(i);
326	}	289	}
327	- if (!pthread_setaffinity_np(thread, sizeof(mask), &mask)) {	290	+ }
328	- return true;	291	+ return formatCoreRange(cores);
329	- } else {	292	+}
		293	+
		294	+bool checkThreadAffinityConflict(
		295	+ const std::vector<CoreIdList>& devices_aff_cores) {
		296	+ cpu_set_t mask;
		297	+ pthread_getaffinity_np(pthread_self(), sizeof(mask), &mask);
		298	+ std::string affinity_mask_str = formatCPUSetMask(mask);
		299	+ ASCEND_LOGI(
		300	+ "Current thread CPU affinity mask: %s", affinity_mask_str.c_str());
		301	+
		302	+ for (auto& cores : devices_aff_cores) {
		303	+ for (auto& core : cores) {
		304	+ if (!CPU_ISSET(core, &mask)) {
		305	+ ASCEND_LOGW(
		306	+ "Thread affinity conflict detected! Expected core %u (in config range [%s]) is NOT in current thread affinity mask. %s",
		307	+ core,
		308	+ formatCoreRange(cores).c_str(),
		309	+ affinity_mask_str.c_str());
		310	+ ASCEND_LOGW(
		311	+ "Thread affinity is already set. Use force:1 to skip this check and force bind.");
330	return false;	312	return false;
		313	+ }
331	}	314	}
		315	+ }
		316	+ return true;
332	}	317	}
333		318
334	-CoreIdRange getCoreRange(c10::DeviceIndex device_id, ThreadType type)	319	+bool getThreadAffinityInfo() {
335	-{	320	+ parseCPUAffinityConf(cpu_affinity_mode, devices_aff_cores);
336	- CoreIdRange core_range;	321	+ printCoreRanges(cpu_affinity_mode, devices_aff_cores);
337	- if (cpu_affinity_mode == 0 \|\| cpu_affinity_mode == 1) {	322	+
338	- core_range = device_ranges[device_id];	323	+ for (int i = 0; i < devices_aff_cores.size(); ++i) {
339	- } else {	324	+ if (devices_aff_cores[i].size() > 0) {
340	- std::lock_guard<std::mutex> lock(core_map_mutex);	325	+ ASCEND_LOGD(
341	- if (device_thread_core_maps.find(device_id) == device_thread_core_maps.end()) {	326	+ "Device %d get cores %s.",
342	- device_thread_core_maps.emplace(device_id, getCpuAffinityMap(device_id, device_ranges));	327	+ i,
343	- }	328	+ formatCoreRange(devices_aff_cores[i]).c_str());
344	- core_range = device_thread_core_maps.at(device_id).at(type);
345	}	329	}
346	- return core_range;	330	+ }
347	-}
348		331
349	-void SetThreadAffinity(c10::DeviceIndex device_id)	332	+ if (cpu_affinity_mode == 0) {
350	-{
351	- if (!needToSetThreadAffinity() \|\| local_thread == ThreadType::USER_THREAD) {
352	- return;
353	- }
354	-
355	- CoreIdRange core_range = getCoreRange(device_id, local_thread);
356	- if (setThreadAffinityImpl(pthread_self(), core_range)) {
357	- ASCEND_LOGD("Device %d set %s affinity to %d-%d success.",
358	- device_id, threadTypeToNameMap.at(local_thread).c_str(), core_range.start, core_range.end);
359	- } else {
360	- ASCEND_LOGE("Device %d set %s affinity to %d-%d failed.",
361	- device_id, threadTypeToNameMap.at(local_thread).c_str(), core_range.start, core_range.end);
362	- }
363	-}
364	-
365	-void SetThreadAffinity(ThreadType type)
366	-{
367	- if (!needToSetThreadAffinity()) {
368	- return;
369	- }
370	- int device_index;
371	- NPU_CHECK_ERROR_WITHOUT_UCE(GetDevice(&device_index));
372	- c10::DeviceIndex device = static_cast<c10::DeviceIndex>(device_index);
373	- local_thread = type;
374	- if (local_thread == ThreadType::MAIN_THREAD) {
375	- start_main_thread_bind = true;
376	- }
377	- SetThreadAffinity(device);
378	-}
379	-
380	-void SetThreadAffinity(int core_start, int core_end)
381	-{
382	- if (!needToSetThreadAffinity()) {
383	- return;
384	- }
385	-
386	- static int core_nums = sysconf(_SC_NPROCESSORS_ONLN);
387	- CoreIdRange core_range;
388	- core_range.start = static_cast<CoreId>(std::min(core_start, core_nums));
389	- core_range.end = static_cast<CoreId>(std::min(core_end, core_nums));
390	- local_thread = ThreadType::USER_THREAD;
391	-
392	- if (setThreadAffinityImpl(pthread_self(), core_range)) {
393	- ASCEND_LOGD("Set thread affinity to user-defined range %d-%d success.", core_range.start, core_range.end);
394	- } else {
395	- ASCEND_LOGE("Set thread affinity to user-defined range %d-%d failed.", core_range.start, core_range.end);
396	- }
397	-}
398	-
399	-void SetMainThread()
400	-{
401	- main_thread = pthread_self();
402	-}
403	-
404	-bool NeedMainThreadBind()
405	-{
406	- return start_main_thread_bind && (local_thread == ThreadType::MAIN_THREAD);
407	-}
408	-
409	-bool SetThreadAffinityInInitialize()
410	-{
411	- if (needToSetThreadAffinity() && !lazy_bind) {
412	- return true;
413	- }
414	return false;	333	return false;
		334	+ }
		335	+
		336	+ if (force_bind) {
		337	+ ASCEND_LOGI(
		338	+ "CPU affinity force mode enabled, skipping affinity conflict detection, applying CPU_AFFINITY_CONF binding.");
		339	+ return true;
		340	+ }
		341	+
		342	+ return checkThreadAffinityConflict(devices_aff_cores);
415	}	343	}
416		344
417	-void StartMainThreadBind(c10::DeviceIndex device_id)	345	+std::string getAffinityMapAsString(
418	-{	346	+ c10::DeviceIndex device_id,
419	- if (!needToSetThreadAffinity() \|\| local_thread == ThreadType::USER_THREAD) {	347	+ const ThreadCoreMap& threadCoreMap) {
420	- return;	348	+ std::ostringstream oss;
421	- }	349	+ for (auto thread_type : threadTypeList) {
422	-	350	+ oss << threadTypeToNameMap.at(thread_type) << ": ["
423	- static thread_local bool seted = false;	351	+ << formatCoreRange(threadCoreMap.at(thread_type)) << "]";
424	- if (!seted) {	352	+ std::string end_str =
425	- seted = true;	353	+ (thread_type == ThreadType::OTHER_THREAD) ? "." : "; ";
426	- if (syscall(SYS_gettid) != getpid()) {	354	+ oss << end_str;
427	- start_main_thread_bind = true;	355	+ }
428	-	356	+ return oss.str();
429	- SetThreadAffinity(device_id);
430	-
431	- CoreIdRange core_range = getCoreRange(device_id, ThreadType::MAIN_THREAD);
432	- if (setThreadAffinityImpl(main_thread, core_range)) {
433	- ASCEND_LOGD("Device %d set %s affinity to %d-%d success.",
434	- device_id, threadTypeToNameMap.at(ThreadType::MAIN_THREAD).c_str(),
435	- core_range.start, core_range.end);
436	- } else {
437	- ASCEND_LOGE("Device %d set %s affinity to %d-%d failed.",
438	- device_id, threadTypeToNameMap.at(ThreadType::MAIN_THREAD).c_str(),
439	- core_range.start, core_range.end);
440	- }
441	- }
442	- }
443	}	357	}
444		358
445	-} // namespace c10_npu	359	+ThreadCoreMap getCpuAffinityMap(
		360	+ c10::DeviceIndex device_id,
		361	+ const std::vector<CoreIdList>& devices_aff_cores) {
		362	+ ThreadCoreMap threadCoreMap;
		363	+ CoreIdList cores = devices_aff_cores[device_id];
		364	+ if (cores.size() < threadTypeList.size()) {
		365	+ ASCEND_LOGW(
		366	+ "Device %d available core numbers (%zu) are insufficient for all %zu thread types and will bind available cores to all threads.",
		367	+ device_id,
		368	+ cores.size(),
		369	+ threadTypeList.size());
		370	+ for (auto thread_type : threadTypeList) {
		371	+ threadCoreMap[thread_type] = cores;
		372	+ }
		373	+ return threadCoreMap;
		374	+ }
		375	+ for (auto thread_type : threadTypeList) {
		376	+ if (thread_type != ThreadType::OTHER_THREAD) {
		377	+ CoreId first_core = *cores.begin();
		378	+ threadCoreMap[thread_type].insert(first_core);
		379	+ cores.erase(first_core);
		380	+ } else {
		381	+ threadCoreMap[ThreadType::OTHER_THREAD] = cores;
		382	+ }
		383	+ }
		384	+
		385	+ ASCEND_LOGD(
		386	+ "Device %d thread affinity map: %s",
		387	+ device_id,
		388	+ getAffinityMapAsString(device_id, threadCoreMap).c_str());
		389	+ return threadCoreMap;
		390	+}
		391	+
		392	+CoreIdList getCoreList(c10::DeviceIndex device_id, ThreadType type) {
		393	+ CoreIdList core_list;
		394	+ if (cpu_affinity_mode == 0 \|\| cpu_affinity_mode == 1) {
		395	+ core_list = devices_aff_cores[device_id];
		396	+ } else {
		397	+ std::lock_guard<std::mutex> lock(core_map_mutex);
		398	+ if (device_thread_core_maps.find(device_id) ==
		399	+ device_thread_core_maps.end()) {
		400	+ device_thread_core_maps.emplace(
		401	+ device_id, getCpuAffinityMap(device_id, devices_aff_cores));
		402	+ }
		403	+ core_list = device_thread_core_maps.at(device_id).at(type);
		404	+ }
		405	+ return core_list;
		406	+}
		407	+
		408	+bool setThreadAffinityImpl(pthread_t thread, CoreIdList core_list) {
		409	+ cpu_set_t mask;
		410	+ CPU_ZERO(&mask);
		411	+ for (auto core : core_list) {
		412	+ CPU_SET(core, &mask);
		413	+ }
		414	+ return pthread_setaffinity_np(thread, sizeof(mask), &mask) == 0;
		415	+}
		416	+
		417	+inline bool needToSetThreadAffinity() {
		418	+ static bool need_to_set_affinity = getThreadAffinityInfo();
		419	+ return need_to_set_affinity;
		420	+}
		421	+
		422	+} // namespace
		423	+
		424	+void SetThreadType(ThreadType type) {
		425	+ // Called at the start of the thread's execution to avoid frequent triggering
		426	+ // of this function.
		427	+ local_thread = type;
		428	+ if (type == ThreadType::OTHER_THREAD \|\| type == ThreadType::MAIN_THREAD) {
		429	+ return;
		430	+ }
		431	+ if (prctl(PR_SET_NAME, threadTypeToNameMap.at(type).c_str()) != 0) {
		432	+ ASCEND_LOGW(
		433	+ "Set thread name to %s failed!", threadTypeToNameMap.at(type).c_str());
		434	+ }
		435	+ ASCEND_LOGD(
		436	+ "Set thread name to %s success.", threadTypeToNameMap.at(type).c_str());
		437	+}
		438	+
		439	+void SetThreadAffinity(c10::DeviceIndex device_id) {
		440	+ if (!needToSetThreadAffinity() \|\| local_thread == ThreadType::USER_THREAD) {
		441	+ return;
		442	+ }
		443	+
		444	+ CoreIdList core_list = getCoreList(device_id, local_thread);
		445	+ std::string range_str = formatCoreRange(core_list);
		446	+ if (setThreadAffinityImpl(pthread_self(), core_list)) {
		447	+ ASCEND_LOGD(
		448	+ "Device %d set %s affinity to %s success.",
		449	+ device_id,
		450	+ threadTypeToNameMap.at(local_thread).c_str(),
		451	+ range_str.c_str());
		452	+ } else {
		453	+ ASCEND_LOGE(
		454	+ "Device %d set %s affinity to %s failed.",
		455	+ device_id,
		456	+ threadTypeToNameMap.at(local_thread).c_str(),
		457	+ range_str.c_str());
		458	+ }
		459	+}
		460	+
		461	+void SetThreadAffinity(ThreadType type) {
		462	+ if (!needToSetThreadAffinity()) {
		463	+ return;
		464	+ }
		465	+ int device_index;
		466	+ NPU_CHECK_ERROR_WITHOUT_UCE(GetDevice(&device_index));
		467	+ c10::DeviceIndex device = static_cast<c10::DeviceIndex>(device_index);
		468	+ local_thread = type;
		469	+ if (local_thread == ThreadType::MAIN_THREAD) {
		470	+ start_main_thread_bind = true;
		471	+ }
		472	+ SetThreadAffinity(device);
		473	+}
		474	+
		475	+void SetThreadAffinity(const CoreIdList core_ids) {
		476	+ if (!needToSetThreadAffinity()) {
		477	+ return;
		478	+ }
		479	+ CoreIdList processed_core_ids = core_ids;
		480	+ static int core_nums = sysconf(_SC_NPROCESSORS_ONLN);
		481	+ for (auto it = processed_core_ids.begin(); it != processed_core_ids.end();) {
		482	+ if (static_cast<int>(*it) >= core_nums) {
		483	+ ASCEND_LOGW(
		484	+ "core id %d >= core_nums %d, it will be ignored when setting thread affinity.",
		485	+ *it,
		486	+ core_nums);
		487	+ it = processed_core_ids.erase(it);
		488	+ } else {
		489	+ ++it;
		490	+ }
		491	+ }
		492	+ local_thread = ThreadType::USER_THREAD;
		493	+ if (setThreadAffinityImpl(pthread_self(), processed_core_ids)) {
		494	+ ASCEND_LOGD(
		495	+ "Set thread affinity to user-defined range %s success.",
		496	+ formatCoreRange(processed_core_ids).c_str());
		497	+ } else {
		498	+ ASCEND_LOGE(
		499	+ "Set thread affinity to user-defined range %s failed.",
		500	+ formatCoreRange(processed_core_ids).c_str());
		501	+ }
		502	+}
		503	+
		504	+void SetThreadAffinity(int core_start, int core_end) {
		505	+ CoreIdList core_list;
		506	+ for (int i = core_start; i <= core_end; ++i) {
		507	+ core_list.insert(static_cast<CoreId>(i));
		508	+ }
		509	+ SetThreadAffinity(core_list);
		510	+}
		511	+
		512	+void SetMainThread() {
		513	+ main_thread = pthread_self();
		514	+}
		515	+
		516	+bool NeedMainThreadBind() {
		517	+ return start_main_thread_bind && (local_thread == ThreadType::MAIN_THREAD);
		518	+}
		519	+
		520	+bool SetThreadAffinityInInitialize() {
		521	+ if (needToSetThreadAffinity() && !lazy_bind) {
		522	+ return true;
		523	+ }
		524	+ return false;
		525	+}
		526	+
		527	+void StartMainThreadBind(c10::DeviceIndex device_id) {
		528	+ if (!needToSetThreadAffinity() \|\| local_thread == ThreadType::USER_THREAD) {
		529	+ return;
		530	+ }
		531	+
		532	+ static thread_local bool seted = false;
		533	+ if (seted) {
		534	+ return;
		535	+ }
		536	+ seted = true;
		537	+ if (syscall(SYS_gettid) != getpid()) {
		538	+ start_main_thread_bind = true;
		539	+ SetThreadAffinity(device_id);
		540	+ CoreIdList core_list = getCoreList(device_id, ThreadType::MAIN_THREAD);
		541	+ if (setThreadAffinityImpl(main_thread, core_list)) {
		542	+ ASCEND_LOGD(
		543	+ "Device %d set %s affinity to %s success.",
		544	+ device_id,
		545	+ threadTypeToNameMap.at(ThreadType::MAIN_THREAD).c_str(),
		546	+ formatCoreRange(core_list).c_str());
		547	+ } else {
		548	+ ASCEND_LOGE(
		549	+ "Device %d set %s affinity to %s failed.",
		550	+ device_id,
		551	+ threadTypeToNameMap.at(ThreadType::MAIN_THREAD).c_str(),
		552	+ formatCoreRange(core_list).c_str());
		553	+ }
		554	+ }
		555	+}
		556	+
		557	+} // namespace c10_npu

Mtorch_npu/csrc/core/npu/NPUAffinityController.h+21-14

	@@ -1,27 +1,26 @@
1	#pragma once	1	#pragma once
2	#include <c10/core/Device.h>	2	#include <c10/core/Device.h>
		3	+#include <torch_npu/csrc/core/npu/GetAffinityCPUInfo.h>
		4	+#include <set>
3		5
4	namespace c10_npu {	6	namespace c10_npu {
5		7
6	-using CoreId = unsigned int;
7	-struct CoreIdRange {
8	- CoreId start;
9	- CoreId end;
10	-};
11	-
12	enum ThreadType {	8	enum ThreadType {
13	- MAIN_THREAD = 0, // 1st performance hotspot, responsible for operator dispatching.	9	+ MAIN_THREAD =
14	- ACL_THREAD = 1, // 2rd performance hotspot in PTA, responsible for handling the task queue.	10	+ 0, // 1st performance hotspot, responsible for operator dispatching.
15	- RELEASE_THREAD = 2, // Thread responsible for resource release.	11	+ ACL_THREAD = 1, // 2nd performance hotspot in PTA, responsible for handling
16	- WATCHDOG_THREAD = 3, // Thread responsible for HCCL communication monitoring.	12	+ // the task queue.
17	- OTHER_THREAD = 4, // Mostly refers to threads in PyTorch's motorized sleep thread pool, which	13	+ RELEASE_THREAD = 2, // Thread responsible for resource release.
18	- // are not considered in PTA.	14	+ WATCHDOG_THREAD = 3, // Thread responsible for HCCL communication monitoring.
19	- USER_THREAD = 5, // Thread responsible for user.	15	+ OTHER_THREAD = 4, // Mostly refers to threads in PyTorch's motorized sleep
		16	+ // thread pool, which are not considered in PTA.
		17	+ USER_THREAD = 5, // Thread responsible for user.
20	};	18	};
21		19
22	void SetThreadType(ThreadType type);	20	void SetThreadType(ThreadType type);
23	void SetThreadAffinity(c10::DeviceIndex device);	21	void SetThreadAffinity(c10::DeviceIndex device);
24	void SetThreadAffinity(ThreadType type);	22	void SetThreadAffinity(ThreadType type);
		23	+void SetThreadAffinity(const CoreIdList core_ids);
25	void SetThreadAffinity(int core_start, int core_end);	24	void SetThreadAffinity(int core_start, int core_end);
26		25
27	void SetMainThread();	26	void SetMainThread();
	@@ -29,4 +28,12 @@ bool NeedMainThreadBind();
29	bool SetThreadAffinityInInitialize();	28	bool SetThreadAffinityInInitialize();
30	void StartMainThreadBind(c10::DeviceIndex device_id);	29	void StartMainThreadBind(c10::DeviceIndex device_id);
31		30
32	-} // namespace c10_npu	31	+inline bool isAllDigits(const std::string& str) {
		32	+ if (str.empty()) {
		33	+ return false;
		34	+ }
		35	+ return std::all_of(
		36	+ str.begin(), str.end(), [](unsigned char c) { return std::isdigit(c); });
		37	+}
		38	+
		39	+} // namespace c10_npu

Mtorch_npu/csrc/npu/Module.cpp+2551-2200

Mtorch_npu/utils/affinity.py+45-8

	@@ -1,21 +1,58 @@
1	__all__ = []	1	__all__ = []
2		2
3	-from typing import List
4		3
5	import torch_npu	4	import torch_npu
6	from torch_npu.utils._error_code import ErrCode, pta_error	5	from torch_npu.utils._error_code import ErrCode, pta_error
7		6
8		7
9	-def _set_thread_affinity(core_range: List[int] = None):	8	+def _set_thread_affinity(core_range: list[int] \| list[list[int]] \| None = None):
		9	+ """Set thread CPU affinity.
		10	+
		11	+ Args:
		12	+ core_range: CPU core range to bind to, supports three forms:
		13	+ - None: reset to default affinity.
		14	+ - [start, end]: single range, binds to cores from start to end (inclusive).
		15	+ Example: [0, 3] binds to CPU cores 0, 1, 2, 3.
		16	+ - [[start1, end1], [start2, end2], ...]: multiple ranges.
		17	+ Example: [[0, 3], [8, 11]] binds to CPU cores 0, 1, 2, 3, 8, 9, 10, 11.
		18	+ """
10	if core_range is None:	19	if core_range is None:
11	torch_npu._C._npu_set_thread_affinity(-1, -1)	20	torch_npu._C._npu_set_thread_affinity(-1, -1)
12	- elif (len(core_range) == 2):	21	+ return
13	- if core_range[0] < 0 or core_range[1] < 0:	22	+
14	- raise ValueError("Core range should be nonnegative." + pta_error(ErrCode.PARAM))	23	+ # Single range: [start, end]
15	- torch_npu._C._npu_set_thread_affinity(core_range[0], core_range[1])	24	+ if (
		25	+ isinstance(core_range, list)
		26	+ and len(core_range) == 2
		27	+ and all(isinstance(x, int) for x in core_range)
		28	+ ):
		29	+ core_ranges = [core_range]
		30	+ # Multiple ranges: [[start1, end1], [start2, end2], ...]
		31	+ elif (
		32	+ isinstance(core_range, list)
		33	+ and len(core_range) > 0
		34	+ and all(isinstance(x, list) and len(x) == 2 for x in core_range)
		35	+ ):
		36	+ core_ranges = core_range
16	else:	37	else:
17	- raise ValueError("The length of input list of set_thread_affinity should be 2." + pta_error(ErrCode.PARAM))	38	+ raise ValueError(f"Invalid core range: {core_range}" + pta_error(ErrCode.PARAM))
		39	+
		40	+ for start, end in core_ranges:
		41	+ if not (
		42	+ isinstance(start, int)
		43	+ and isinstance(end, int)
		44	+ and start >= 0
		45	+ and start <= end
		46	+ ):
		47	+ raise ValueError(
		48	+ f"Invalid core range values: [{start}, {end}]"
		49	+ + pta_error(ErrCode.VALUE)
		50	+ )
		51	+
		52	+ core_list = [core for start, end in core_ranges for core in range(start, end + 1)]
		53	+ core_list = sorted(set(core_list))
		54	+ torch_npu._C._npu_set_thread_affinity(core_list)
18		55
19		56
20	def _reset_thread_affinity():	57	def _reset_thread_affinity():
21	- torch_npu._C._npu_reset_thread_affinity()	58	+ torch_npu._C._npu_reset_thread_affinity()