已合并
lazy query events #30998
liubuyu1创建于 2月25日
lazy query events #30998
已合并
liubuyu1创建于 2月25日
5 个文件变更+288-4
@@ -46,6 +46,10 @@
46 - **单一值**:为每个内存设置相同的分段数量,例如配置为“4“。46 - **单一值**:为每个内存设置相同的分段数量,例如配置为“4“。
47 - **键值对数组**:为每个2的幂区间单独设置分段数量。例如配置为“\[256:1,512:2,1024:4,\>:8\]“时,表示为256MB以下的所有分配设置1个分段,256MB到512MB之间的分配设置2个分段,512MB到1GB之间的分配设置4个分段,以及更大的分配设置8个分段。47 - **键值对数组**:为每个2的幂区间单独设置分段数量。例如配置为“\[256:1,512:2,1024:4,\>:8\]“时,表示为256MB以下的所有分配设置1个分段,256MB到512MB之间的分配设置2个分段,512MB到1GB之间的分配设置4个分段,以及更大的分配设置8个分段。
48 48 
49+- multi\_stream\_lazy\_reclaim:<value\>,多流场景下,内存申请时延迟查询Events。
50+ 
51+ 默认值为False,即每次内存申请时都执行Events查询。当设置为True时,每次内存申请优先使用空闲内存块,当Events数量超过阈值512或者找不到可用内存块时,才触发Events查询。通过减少Events查询次数,降低CPU资源占用,提升Host侧性能。该配置仅影响Events状态查询的频率,不改变内存释放的条件,也不改变内存峰值,内存块仍需等待所有相关Events完成后才会被释放。
52+ 
49- pinned\_use\_background\_threads:<value\>,是否启用后台线程来处理events。53- pinned\_use\_background\_threads:<value\>,是否启用后台线程来处理events。
50 54 
51 默认值为False,不启用后台线程。当设置为True时,启用后台线程,在后台线程执行查询和处理events操作,减少主线程的阻塞时间。55 默认值为False,不启用后台线程。当设置为True时,启用后台线程,在后台线程执行查询和处理events操作,减少主线程的阻塞时间。
@@ -104,21 +108,27 @@ export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True,segment_size_mb:40
104 108 
105示例六:109示例六:
106 110 
111+```
112+export PYTORCH_NPU_ALLOC_CONF=multi_stream_lazy_reclaim:True
113+```
114+ 
115+示例七:
116+ 
107```117```
108export PYTORCH_NPU_ALLOC_CONF=pinned_use_background_threads:True118export PYTORCH_NPU_ALLOC_CONF=pinned_use_background_threads:True
109```119```
110 120 
111 121 
112-示例122+示例
113 123 
114```124```
115export PYTORCH_NPU_ALLOC_CONF=pin_memory_expandable_segments:True125export PYTORCH_NPU_ALLOC_CONF=pin_memory_expandable_segments:True
116```126```
117 127 
118 128 
119-示例129+示例
120 130 
121-```bash131+```
122export PYTORCH_NPU_ALLOC_CONF=pinned_mem_register:True132export PYTORCH_NPU_ALLOC_CONF=pinned_mem_register:True
123```133```
124 134 
@@ -140,6 +150,9 @@ export PYTORCH_NPU_ALLOC_CONF=pinned_mem_register:True
140- pinned_mem_register使用注意事项如下:150- pinned_mem_register使用注意事项如下:
141 - 特性要求Ascend Extension for PyTorch 26.0.0及以上版本、Ascend HDK 25.5.0及以上版本、CANN商发8.5.0及以上版本使用。151 - 特性要求Ascend Extension for PyTorch 26.0.0及以上版本、Ascend HDK 25.5.0及以上版本、CANN商发8.5.0及以上版本使用。
142 - 与pin_memory_expandable_segments特性不支持同时配置。152 - 与pin_memory_expandable_segments特性不支持同时配置。
153+- multi_stream_lazy_reclaim使用注意事项:
154+ - 特性要求在Ascend Extension for PyTorch 7.3.0以上版本上使用。
155+ - 该特性主要解决多流场景下,Host侧存在下发性能瓶颈时的系统效率问题。单流、少流场景或者非Host性能瓶颈时,该功能收益不大。
143 156 
144## 支持的型号157## 支持的型号
145 158 
@@ -0,0 +1,234 @@
1+import os
2+import time
3+import multiprocessing
4+import shutil
5+import unittest
6+import platform
7+import torch
8+import torch_npu
9+ 
10+from torch_npu.testing.testcase import TestCase, run_tests
11+ 
12+# Set multiprocessing start method to spawn because NPU cannot be re-initialized in forked subprocesses
13+try:
14+ multiprocessing.set_start_method('spawn')
15+except RuntimeError:
16+ pass # May have already been set
17+ 
18+IS_ARM64 = platform.machine() in ('arm64', 'aarch64')
19+ 
20+ 
21+def extract_aclrtQueryEventStatus_count(prof_dir):
22+ """
23+ Extract the call count of aclrtQueryEventStatus from profiler results.
24+ Uses Linux system commands (find/grep/awk) to parse api_statistic.csv.
25+ 
26+ Args:
27+ prof_dir: str, path to the profiling result directory
28+
29+ Returns:
30+ count: int, call count of aclrtQueryEventStatus, 0 if not found
31+ """
32+ import subprocess
33+ 
34+ count = 0
35+ 
36+ try:
37+ # Use find command to locate api_statistic.csv files
38+ find_result = subprocess.run(
39+ ["find", prof_dir, "-name", "api_statistic.csv", "-type", "f"],
40+ capture_output=True, text=True, timeout=30
41+ )
42+ 
43+ if find_result.returncode == 0 and find_result.stdout.strip():
44+ # Get the first CSV file found
45+ csv_file = find_result.stdout.strip().split('\n')[0]
46+ # Search for aclrtQueryEventStatus line and extract the 5th column (call count)
47+ grep_result = subprocess.run(
48+ ["grep", "aclrtQueryEventStatus", csv_file],
49+ capture_output=True, text=True, timeout=10
50+ )
51+ 
52+ if grep_result.returncode == 0 and grep_result.stdout.strip():
53+ # Use awk to extract the 5th column (call count)
54+ awk_result = subprocess.run(
55+ ["awk", "-F,", '{print $5}'],
56+ input=grep_result.stdout,
57+ capture_output=True, text=True, timeout=10
58+ )
59+ if awk_result.returncode == 0 and awk_result.stdout.strip():
60+ count = int(awk_result.stdout.strip())
61+ except (subprocess.TimeoutExpired, subprocess.CalledProcessError, ValueError):
62+ pass
63+ return count
64+ 
65+ 
66+def run_matmul_with_profiling(result_queue, enable_lazy_reclaim):
67+ """
68+ Run matmul test in a separate process and return the aclrtQueryEventStatus call count.
69+ 
70+ Args:
71+ result_queue: multiprocessing.Queue, used to return results
72+ enable_lazy_reclaim: bool, whether to enable lazy reclaim feature
73+ """
74+ # Set environment variables (must be set before importing torch_npu)
75+ if enable_lazy_reclaim:
76+ os.environ["PYTORCH_NPU_ALLOC_CONF"] = "multi_stream_lazy_reclaim:True"
77+ else:
78+ # Ensure environment variable does not exist or is False
79+ if "PYTORCH_NPU_ALLOC_CONF" in os.environ:
80+ del os.environ["PYTORCH_NPU_ALLOC_CONF"]
81+ 
82+ # Reinitialize NPU (to ensure environment variables take effect)
83+ torch.npu.init()
84+ 
85+ # Create temporary directory for profiling results (use absolute path)
86+ prof_dir = os.path.abspath("./prof_" + str(enable_lazy_reclaim))
87+ 
88+ try:
89+ stream0 = torch.npu.Stream()
90+ stream1 = torch.npu.Stream()
91+ stream2 = torch.npu.Stream()
92+ 
93+ # Configure profiler to collect Level2 data (includes all API calls)
94+ experimental_config = torch_npu.profiler._ExperimentalConfig(
95+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2)
96+ 
97+ with torch_npu.profiler.profile(
98+ activities=[torch_npu.profiler.ProfilerActivity.NPU,
99+ torch_npu.profiler.ProfilerActivity.CPU],
100+ with_stack=False, # Do not collect call stacks, reduces overhead
101+ record_shapes=False, # Do not record shapes
102+ profile_memory=False, # Do not analyze memory in detail
103+ schedule=torch_npu.profiler.schedule(
104+ wait=0, warmup=0, active=1, repeat=1, skip_first=0),
105+ experimental_config=experimental_config,
106+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(prof_dir)) as prof:
107+ 
108+ # Preallocate tensors, each 2048x2048 float32 = 16MB
109+ a = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
110+ b = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
111+ c = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
112+ d = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
113+ e = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
114+ f = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
115+ 
116+ # Record tensors on multiple streams
117+ a.record_stream(stream0)
118+ a.record_stream(stream1)
119+ a.record_stream(stream2)
120+ 
121+ # Execute matmul operations on multiple streams
122+ for _ in range(50):
123+ with torch.npu.stream(stream0):
124+ torch.matmul(a, b, out=c)
125+ with torch.npu.stream(stream1):
126+ torch.matmul(a, b, out=d)
127+ with torch.npu.stream(stream2):
128+ torch.matmul(a, b, out=e)
129+ 
130+ # Release some tensors
131+ a = None
132+ f = None
133+ 
134+ # Trigger memory allocation, may trigger process_events
135+ for _ in range(10):
136+ tmp = torch.empty((1024, 1024), dtype=torch.float32, device="npu") # 4M
137+ 
138+ # Synchronize all streams
139+ torch.npu.synchronize()
140+ 
141+ # Allocate memory again
142+ a1 = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
143+ f1 = torch.empty((2048, 2048), dtype=torch.float32, device="npu")
144+ 
145+ prof.step()
146+ 
147+ # Extract aclrtQueryEventStatus call count from profiling results
148+ count = extract_aclrtQueryEventStatus_count(prof_dir)
149+ result_queue.put(("success", count))
150+ 
151+ except Exception as e:
152+ result_queue.put(("error", str(e)))
153+ finally:
154+ # Clean up temporary directory
155+ if os.path.exists(prof_dir):
156+ shutil.rmtree(prof_dir)
157+ 
158+ 
159+@unittest.skipUnless(IS_ARM64, "Only working on ARM")
H
Hhbhu_bin3月3日

这个是有啥限制么

likedislike
liubuyu1
liubuyu1
3月3日 评论:
160+class TestMultiStreamLazyReclaim(TestCase):
161+ """
162+ Test the reduction effect of multi_stream_lazy_reclaim feature on event query counts.
163+ 
164+ Principle:
165+ - eager reclaim mode: Calls process_events to query event status before every memory allocation
166+ - lazy reclaim mode: Only queries in the following cases:
167+ 1. No available memory block found (!block_found)
168+ 2. Event queue exceeds threshold kLazyQuerySize (512)
169+
170+ Test Method:
171+ Use multiprocessing to test in two separate processes:
172+ - Process 1: Enable multi_stream_lazy_reclaim
173+ - Process 2: Disable multi_stream_lazy_reclaim
174+
175+ Each process sets environment variables independently to ensure configuration takes effect.
176+ """
177+ 
178+ def test_lazy_reclaim_reduces_event_queries_counts(self):
179+ """
180+ Compare aclrtQueryEventStatus call counts between eager reclaim and lazy reclaim modes.
181+ 
182+ Validation Goal:
183+ Lazy reclaim mode should significantly reduce the number of aclrtQueryEventStatus calls.
184+ """
185+ configs = [
186+ ("eager", False),
187+ ("lazy", True)
188+ ]
189+ results = {}
190+ 
191+ for name, enable_lazy in configs:
192+ print(f"\n--- Starting {name} reclaim test ---")
193+ queue = multiprocessing.Queue()
194+ process = multiprocessing.Process(
195+ target=run_matmul_with_profiling,
196+ args=(queue, enable_lazy)
197+ )
198+ 
199+ process.start()
200+ process.join(timeout=300) #
201+ 
202+ if process.is_alive():
203+ process.terminate()
204+ process.join()
205+ self.fail(f"{name} reclaim process timed out and was terminated.")
206+ 
207+ status, result = queue.get()
208+ self.assertEqual(status, "success", f"{name} reclaim process failed: {result}")
209+ print(f"---mode {name}------count:{result}")
210+ results[name] = result
211+ 
212+ #
213+ time.sleep(2)
214+ 
215+ eager_counts = results["eager"]
216+ lazy_counts = results["lazy"]
217+
218+ # Output comparison results
219+ print(f"\n========== Event Query Count Comparison ==========")
220+ print(f"Eager reclaim (multi_stream_lazy_reclaim:False): {eager_counts}")
221+ print(f"Lazy reclaim (multi_stream_lazy_reclaim:True): {lazy_counts}")
222+ 
223+ # Core validation: aclrtQueryEventStatus call count in lazy mode must be less than eager mode
224+ # This is direct evidence that multi_stream_lazy_reclaim feature is working
225+ self.assertLessEqual(
226+ lazy_counts,
227+ eager_counts,
228+ f"Lazy reclaim mode should reduce event queries. "
229+ f"Eager: {eager_counts}, Lazy: {lazy_counts}. "
230+ f"If lazy >= eager, the optimization may not be working."
231+ )
232+ 
233+if __name__ == '__main__':
234+ run_tests()
@@ -263,6 +263,19 @@ namespace c10_npu {
263 return i;263 return i;
264 }264 }
265 265 
266+ size_t CachingAllocatorConfig::parseMultiStreamLazyReclaim(const std::vector<std::string> &config, size_t i)
267+ {
268+ consumeToken(config, ++i, ':');
269+ if (++i < config.size()) {
270+ TORCH_CHECK(i < config.size() && (config[i] == "True" || config[i] == "False"),
271+ "Expected a single True/False argument for multi_stream_lazy_reclaim", PTA_ERROR(ErrCode::PARAM));
272+ m_multi_stream_lazy_reclaim = (config[i] == "True");
273+ } else {
274+ TORCH_CHECK(false, "Error, expecting multi_stream_lazy_reclaim value", PTA_ERROR(ErrCode::PARAM));
275+ }
276+ return i;
277+ }
278+ 
266 size_t CachingAllocatorConfig::roundup_power2_divisions(size_t size)279 size_t CachingAllocatorConfig::roundup_power2_divisions(size_t size)
267 {280 {
268 if (size == 0 || instance().m_roundup_power2_divisions.empty()) {281 if (size == 0 || instance().m_roundup_power2_divisions.empty()) {
@@ -323,6 +336,8 @@ namespace c10_npu {
323 i = parseSegmentSizeMb(config, i);336 i = parseSegmentSizeMb(config, i);
324 } else if (config[i] == "roundup_power2_divisions") {337 } else if (config[i] == "roundup_power2_divisions") {
325 i = parseRoundUpPower2Divisions(config, i);338 i = parseRoundUpPower2Divisions(config, i);
339+ } else if (config[i] == "multi_stream_lazy_reclaim") {
340+ i = parseMultiStreamLazyReclaim(config, i);
326 } else {341 } else {
327 TORCH_CHECK(false, "Unrecognized CachingAllocator option: ", config[i], OPS_ERROR(ErrCode::PARAM));342 TORCH_CHECK(false, "Unrecognized CachingAllocator option: ", config[i], OPS_ERROR(ErrCode::PARAM));
328 }343 }
@@ -50,6 +50,10 @@ namespace c10_npu {
50 {50 {
51 return instance().m_page_size_1g;51 return instance().m_page_size_1g;
52 }52 }
53+ static bool multi_stream_lazy_reclaim()
54+ {
55+ return instance().m_multi_stream_lazy_reclaim;
56+ }
53 57 
54 static size_t segment_size_mb()58 static size_t segment_size_mb()
55 {59 {
@@ -88,6 +92,8 @@ namespace c10_npu {
88 92 
89 bool m_page_size_1g = false; // 新增1G页配置标志93 bool m_page_size_1g = false; // 新增1G页配置标志
90 94 
95+ bool m_multi_stream_lazy_reclaim = false;
96+ 
91 size_t m_segment_size_mb;97 size_t m_segment_size_mb;
92 98 
93 std::vector<size_t> m_roundup_power2_divisions;99 std::vector<size_t> m_roundup_power2_divisions;
@@ -124,6 +130,8 @@ namespace c10_npu {
124 size_t parseSegmentSizeMb(const std::vector<std::string> &config, size_t i);130 size_t parseSegmentSizeMb(const std::vector<std::string> &config, size_t i);
125 131 
126 size_t parseRoundUpPower2Divisions(const std::vector<std::string> &config, size_t i);132 size_t parseRoundUpPower2Divisions(const std::vector<std::string> &config, size_t i);
133+ 
134+ size_t parseMultiStreamLazyReclaim(const std::vector<std::string> &config, size_t i);
127 };135 };
128 } // namespace NPUCachingAllocator136 } // namespace NPUCachingAllocator
129} // namespace c10_npu137} // namespace c10_npu
@@ -102,6 +102,7 @@ const std::string kMinCannVersion = "8.1.RC1"; // minimum cann version wh
102const std::string kMinDriverVersion = "25.0.RC1"; // minimum driver version which supports 1g mem 25.0.RC1102const std::string kMinDriverVersion = "25.0.RC1"; // minimum driver version which supports 1g mem 25.0.RC1
103const std::string kCannModule = "CANN"; // cann module name103const std::string kCannModule = "CANN"; // cann module name
104constexpr int kPrecision = 4; // precision of the memory usage information104constexpr int kPrecision = 4; // precision of the memory usage information
105+constexpr size_t kLazyQuerySize = 512; // lazy query event size
105 106 
106static char SHAREABLE_HANDLE_VERSION = 1;107static char SHAREABLE_HANDLE_VERSION = 1;
107enum ShareableHandleType : char {108enum ShareableHandleType : char {
@@ -1171,7 +1172,7 @@ public:
1171 NPU_CHECK_ERROR(c10_npu::GetDevice(&device));1172 NPU_CHECK_ERROR(c10_npu::GetDevice(&device));
1172 }1173 }
1173 1174 
1174- if (C10_LIKELY(captures_underway.empty())) {1175+ if (!CachingAllocatorConfig::multi_stream_lazy_reclaim() && C10_LIKELY(captures_underway.empty())) {
1175 // Processes end-of-life events for outstanding allocations used on1176 // Processes end-of-life events for outstanding allocations used on
1176 // multiple streams (checks if their NPU-side uses are complete and1177 // multiple streams (checks if their NPU-side uses are complete and
1177 // recycles their memory if so)1178 // recycles their memory if so)
@@ -1200,6 +1201,19 @@ public:
1200 get_free_block(params) ||1201 get_free_block(params) ||
1201 // Trigger callbacks and retry search1202 // Trigger callbacks and retry search
1202 (trigger_free_memory_callbacks(params) && get_free_block(params));1203 (trigger_free_memory_callbacks(params) && get_free_block(params));
1204+ if (CachingAllocatorConfig::multi_stream_lazy_reclaim() && C10_LIKELY(captures_underway.empty())) {
1205+ // Lazy process events and free memory
1206+ size_t sum = 0;
1207+ for (auto it = npu_events.begin(); it != npu_events.end(); ++it) {
1208+ sum += it->second.size();
1209+ }
1210+ if (!block_found || sum > kLazyQuerySize) {
1211+ process_events(context);
1212+ }
1213+ if (!block_found) {
1214+ block_found = get_free_block(params);
1215+ }
1216+ }
1203 // Can't reuse an existing block; try to get a new one.1217 // Can't reuse an existing block; try to get a new one.
1204 if (!block_found) {1218 if (!block_found) {
1205 // Do garbage collection if the flag is set.1219 // Do garbage collection if the flag is set.