add record_stream sanitizer-pytorch-AtomGit

已合并

add record_stream sanitizer #35838

bellatan创建于 5月16日从bellatan/pytorch： v2.12.0_record_stream_sanitizer 合入到 Ascend/pytorch： v2.12.0

add record_stream sanitizer #35838

已合并

bellatan创建于 5月16日

从bellatan/pytorch v2.12.0_record_stream_sanitizer合入到Ascend/pytorch v2.12.0

共 10 个文件变更+1173-31

Mtest/npu/test_sanitizer.py+64-7

	@@ -1,19 +1,31 @@
1	import os	1	import os
2	-import atexit
3	from unittest.mock import MagicMock	2	from unittest.mock import MagicMock
4	from unittest.mock import patch	3	from unittest.mock import patch
5		4
6	-import torch.cuda._sanitizer as csan	5	+import torch
		6	+import torch.utils.cpp_extension
		7	+
7	import torch_npu	8	import torch_npu
8	-import torch_npu.utils._npu_trace as npu_trace
9	-import torch_npu.npu._stream_check as stream_check
10	-import torch_npu.npu._kernel_check as kernel_check
11	-from torch_npu.utils.utils import _print_warn_log
12	import torch_npu.npu._sanitizer as sanitizer	9	import torch_npu.npu._sanitizer as sanitizer
13	from torch_npu.testing.testcase import TestCase, run_tests	10	from torch_npu.testing.testcase import TestCase, run_tests
14		11
15		12
		13	+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
		14	+PYTORCH_INSTALL_PATH = os.path.dirname(os.path.realpath(torch.__file__))
		15	+PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))
		16	+
		17	+
16	class TestSanitizer(TestCase):	18	class TestSanitizer(TestCase):
		19	+ def tearDown(self):
		20	+ if sanitizer.npu_sanitizer.dispatch is not None:
		21	+ try:
		22	+ sanitizer.npu_sanitizer.dispatch.__exit__(None, None, None)
		23	+ except Exception:
		24	+ pass
		25	+ sanitizer.npu_sanitizer.dispatch = None
		26	+ sanitizer.npu_sanitizer.event_handler = None
		27	+ sanitizer.npu_sanitizer.enabled = False
		28	+
17	def test_del_method_with_dispatch(self):	29	def test_del_method_with_dispatch(self):
18	mock_dispatch = MagicMock()	30	mock_dispatch = MagicMock()
19	sanitizer.npu_sanitizer.dispatch = mock_dispatch	31	sanitizer.npu_sanitizer.dispatch = mock_dispatch
	@@ -38,7 +50,9 @@ class TestSanitizer(TestCase):
38	patch('torch_npu.utils._npu_trace.register_callback_for_npu_stream_creation'), \	50	patch('torch_npu.utils._npu_trace.register_callback_for_npu_stream_creation'), \
39	patch('torch_npu.utils._npu_trace.register_callback_for_npu_device_synchronization'), \	51	patch('torch_npu.utils._npu_trace.register_callback_for_npu_device_synchronization'), \
40	patch('torch_npu.utils._npu_trace.register_callback_for_npu_stream_synchronization'), \	52	patch('torch_npu.utils._npu_trace.register_callback_for_npu_stream_synchronization'), \
41	- patch('torch_npu.utils._npu_trace.register_callback_for_npu_event_synchronization'):	53	+ patch('torch_npu.utils._npu_trace.register_callback_for_npu_event_synchronization'), \
		54	+ patch("torch_npu.utils._npu_trace.register_callback_for_npu_record_stream"), \
		55	+ patch("torch_npu.utils._npu_trace.register_callback_for_npu_erase_stream"):
42	mock_dispatch_instance = mock_dispach.return_value	56	mock_dispatch_instance = mock_dispach.return_value
43	mock_dispatch_instance.__enter__.return_value = None	57	mock_dispatch_instance.__enter__.return_value = None
44	sanitizer.npu_sanitizer.enable()	58	sanitizer.npu_sanitizer.enable()
	@@ -46,5 +60,48 @@ class TestSanitizer(TestCase):
46	self.assertEqual(sanitizer.npu_sanitizer.mode, sanitizer.SanitizerMode.STREAM)	60	self.assertEqual(sanitizer.npu_sanitizer.mode, sanitizer.SanitizerMode.STREAM)
47		61
48		62
		63	+class TestSanitizerRegressionGuards(TestCase):
		64	+ """Regression guards for sanitizer integration issues."""
		65	+
		66	+ def test_pluggable_allocator_traces_alloc_and_free(self):
		67	+ """Custom allocator must emit alloc/free traces, not only record_stream traces."""
		68	+ source_path = os.path.join(
		69	+ REPO_ROOT, "torch_npu", "csrc", "npu", "NPUPluggableAllocator.cpp"
		70	+ )
		71	+ with open(source_path, encoding="utf-8") as source_file:
		72	+ source = source_file.read()
		73	+
		74	+ self.assertIn(
		75	+ "traceNpuMemoryAllocation",
		76	+ source,
		77	+ "NPUPluggableAllocator allocations must be visible to NPU Sanitizer",
		78	+ )
		79	+ self.assertIn(
		80	+ "traceNpuMemoryDeallocation",
		81	+ source,
		82	+ "NPUPluggableAllocator frees must trigger deferred record_stream checks",
		83	+ )
		84	+
		85	+ def test_caching_allocator_traces_erase_stream_paths(self):
		86	+ """NPUCachingAllocator eraseStream paths should be visible to NPU Sanitizer."""
		87	+ source_path = os.path.join(
		88	+ REPO_ROOT, "torch_npu", "csrc", "core", "npu", "NPUCachingAllocator.cpp"
		89	+ )
		90	+ with open(source_path, encoding="utf-8") as source_file:
		91	+ source = source_file.read()
		92	+
		93	+ self.assertIn(
		94	+ "traceNpuEraseStream",
		95	+ source,
		96	+ "NPUCachingAllocator eraseStream / eraseStreamWithBlockPtr should be visible "
		97	+ "to NPU Sanitizer so recorded_streams can be cleared.",
		98	+ )
		99	+ self.assertIn(
		100	+ "eraseStreamWithBlockPtr",
		101	+ source,
		102	+ "MULTI_STREAM_MEMORY_REUSE=3 optimized block-ptr erase path should remain covered.",
		103	+ )
		104	+
		105	+
49	if __name__ == "__main__":	106	if __name__ == "__main__":
50	run_tests()	107	run_tests()

Atest/npu/test_sanitizer_record_stream.py+391-0

@@ -0,0 +1,391 @@
	1	+# Owner(s): ["module: npu"]
	2	+"""
	3	+Tests for NPU Sanitizer record_stream detection.
	4	+
	5	+record_stream detection is separate from data race detection:
	6	+ - Data race: detected at kernel launch time (raises CUDASanitizerErrors)
	7	+ - Missing record_stream: detected at deallocation or via flush_record_stream_warnings()
	8	+
	9	+Per PyTorch docs (torch.Tensor.record_stream), record_stream is NOT needed when
	10	+creation_stream has synced with usage_stream before tensor deallocation:
	11	+ - creation_stream.wait_stream(usage_stream)
	12	+ - creation_stream.wait_event(event recorded on usage_stream)
	13	+ - torch_npu.npu.synchronize() (device-level sync covers all directions)
	14	+
	15	+Conversely, usage_stream.wait_stream(creation_stream) only resolves data races
	16	+but does NOT guarantee memory safety — record_stream is still needed.
	17	+
	18	+Test matrix:
	19	+┌───────────────────────────────────┬──────────────┬───────────────────────┐
	20	+│ Scenario │ Data race? │ record_stream needed? │
	21	+├───────────────────────────────────┼──────────────┼───────────────────────┤
	22	+│ No sync at all │ Yes │ Yes │
	23	+│ usage.wait(creation) only │ No │ Yes │
	24	+│ creation.wait(usage) only │ Possible │ No │
	25	+│ Both directions synced │ No │ No │
	26	+│ device synchronize │ No │ No │
	27	+│ record_stream called │ Unresolved │ No (recorded) │
	28	+│ Same stream │ No │ No │
	29	+└───────────────────────────────────┴──────────────┴───────────────────────┘
	30	+"""
	31	+
	32	+import gc
	33	+import os
	34	+
	35	+import torch
	36	+import torch.cuda._sanitizer as csan
	37	+import torch.distributed as dist
	38	+
	39	+import torch_npu
	40	+from torch_npu.testing.testcase import TestCase, run_tests
	41	+
	42	+
	43	+def setup_sanitizer():
	44	+ """Enable sanitizer with record_stream checking."""
	45	+ os.environ['TORCH_NPU_SANITIZER'] = '1'
	46	+ import torch_npu.npu._sanitizer as sanitizer
	47	+ if not sanitizer.npu_sanitizer.enabled:
	48	+ sanitizer.npu_sanitizer.enable()
	49	+
	50	+
	51	+def reset_sanitizer():
	52	+ """Reset sanitizer state between tests."""
	53	+ import torch_npu.npu._sanitizer as sanitizer
	54	+ if sanitizer.npu_sanitizer.dispatch is not None:
	55	+ try:
	56	+ sanitizer.npu_sanitizer.dispatch.__exit__(None, None, None)
	57	+ except Exception:
	58	+ pass
	59	+ sanitizer.npu_sanitizer.dispatch = None
	60	+ sanitizer.npu_sanitizer.event_handler = None
	61	+ sanitizer.npu_sanitizer.enabled = False
	62	+
	63	+
	64	+def get_event_handler():
	65	+ import torch_npu.npu._sanitizer as sanitizer
	66	+ return sanitizer.npu_sanitizer.event_handler
	67	+
	68	+
	69	+class SanitizerRecordStreamTestBase(TestCase):
	70	+ def setUp(self):
	71	+ reset_sanitizer()
	72	+ setup_sanitizer()
	73	+
	74	+ def tearDown(self):
	75	+ reset_sanitizer()
	76	+ if dist.is_available() and dist.is_initialized():
	77	+ dist.destroy_process_group()
	78	+
	79	+ @staticmethod
	80	+ def _get_tracked_tensor_info(tensor):
	81	+ storage_ptr = tensor.untyped_storage().data_ptr()
	82	+ info = get_event_handler()._npu_tensors.get(storage_ptr)
	83	+ if info is None:
	84	+ raise AssertionError(
	85	+ f"Tensor storage {storage_ptr} is not tracked by NPU sanitizer."
	86	+ )
	87	+ return storage_ptr, info
	88	+
	89	+
	90	+class TestDataRaceDetection(SanitizerRecordStreamTestBase):
	91	+ def test_write_read_race(self):
	92	+ """Unsynchronized cross-stream read after write should raise data race."""
	93	+ x = torch.randn(100, device="npu")
	94	+ stream = torch_npu.npu.Stream()
	95	+
	96	+ with self.assertRaises(csan.CUDASanitizerErrors):
	97	+ with torch_npu.npu.stream(stream):
	98	+ _ = x + 1
	99	+
	100	+ def test_record_stream_does_not_fix_data_race(self):
	101	+ """record_stream should not hide a real cross-stream data race."""
	102	+ x = torch.randn(100, device="npu")
	103	+ stream = torch_npu.npu.Stream()
	104	+ x.record_stream(stream)
	105	+
	106	+ with self.assertRaises(csan.CUDASanitizerErrors):
	107	+ with torch_npu.npu.stream(stream):
	108	+ _ = x + 1
	109	+
	110	+
	111	+class TestMissingRecordStream(SanitizerRecordStreamTestBase):
	112	+ def test_cross_stream_no_record_stream(self):
	113	+ """Cross-stream use without record_stream should report missing record_stream."""
	114	+ x = torch.randn(100, device="npu")
	115	+ stream = torch_npu.npu.Stream()
	116	+ default_stream = torch_npu.npu.default_stream()
	117	+
	118	+ stream.wait_stream(default_stream)
	119	+ with torch_npu.npu.stream(stream):
	120	+ _ = x + 1
	121	+
	122	+ warnings = get_event_handler().flush_record_stream_warnings()
	123	+ self.assertGreater(len(warnings), 0)
	124	+
	125	+ def test_multiple_streams_need_record_stream(self):
	126	+ """Each non-creation stream needs its own record_stream coverage."""
	127	+ x = torch.randn(100, device="npu")
	128	+ stream1 = torch_npu.npu.Stream()
	129	+ stream2 = torch_npu.npu.Stream()
	130	+ default_stream = torch_npu.npu.default_stream()
	131	+
	132	+ stream1.wait_stream(default_stream)
	133	+ with torch_npu.npu.stream(stream1):
	134	+ _ = x + 1
	135	+
	136	+ stream2.wait_stream(default_stream)
	137	+ with torch_npu.npu.stream(stream2):
	138	+ _ = x + 2
	139	+
	140	+ warnings = get_event_handler().flush_record_stream_warnings()
	141	+ self.assertGreaterEqual(len(warnings), 2)
	142	+
	143	+
	144	+class TestRecordStreamNotNeeded(SanitizerRecordStreamTestBase):
	145	+ def test_record_stream_suppresses_warning(self):
	146	+ """record_stream should suppress missing-record_stream warning for that stream."""
	147	+ x = torch.randn(100, device="npu")
	148	+ stream = torch_npu.npu.Stream()
	149	+ default_stream = torch_npu.npu.default_stream()
	150	+
	151	+ x.record_stream(stream)
	152	+ stream.wait_stream(default_stream)
	153	+ with torch_npu.npu.stream(stream):
	154	+ _ = x + 1
	155	+
	156	+ warnings = get_event_handler().flush_record_stream_warnings()
	157	+ self.assertEqual(len(warnings), 0)
	158	+
	159	+ def test_creation_waits_usage_via_wait_stream(self):
	160	+ """creation_stream.wait_stream(usage_stream) should make record_stream unnecessary."""
	161	+ x = torch.randn(100, device="npu")
	162	+ stream = torch_npu.npu.Stream()
	163	+ default_stream = torch_npu.npu.default_stream()
	164	+
	165	+ stream.wait_stream(default_stream)
	166	+ with torch_npu.npu.stream(stream):
	167	+ _ = x + 1
	168	+
	169	+ default_stream.wait_stream(stream)
	170	+ warnings = get_event_handler().flush_record_stream_warnings()
	171	+ self.assertEqual(len(warnings), 0)
	172	+
	173	+ def test_creation_waits_usage_via_event(self):
	174	+ """creation_stream.wait_event(event_on_usage) should make record_stream unnecessary."""
	175	+ x = torch.randn(100, device="npu")
	176	+ stream = torch_npu.npu.Stream()
	177	+ default_stream = torch_npu.npu.default_stream()
	178	+
	179	+ stream.wait_stream(default_stream)
	180	+ with torch_npu.npu.stream(stream):
	181	+ _ = x + 1
	182	+ event = torch_npu.npu.Event()
	183	+ event.record(stream)
	184	+
	185	+ default_stream.wait_event(event)
	186	+ warnings = get_event_handler().flush_record_stream_warnings()
	187	+ self.assertEqual(len(warnings), 0)
	188	+
	189	+ def test_device_sync(self):
	190	+ """Device synchronize should cover prior cross-stream uses."""
	191	+ x = torch.randn(100, device="npu")
	192	+ stream = torch_npu.npu.Stream()
	193	+ default_stream = torch_npu.npu.default_stream()
	194	+
	195	+ stream.wait_stream(default_stream)
	196	+ with torch_npu.npu.stream(stream):
	197	+ _ = x + 1
	198	+
	199	+ torch_npu.npu.synchronize()
	200	+
	201	+ warnings = get_event_handler().flush_record_stream_warnings()
	202	+ self.assertEqual(len(warnings), 0)
	203	+
	204	+
	205	+class TestRecordStreamSequenceBoundaries(SanitizerRecordStreamTestBase):
	206	+ def test_creation_waits_usage_only_covers_prior_uses(self):
	207	+ """A mid-sequence reverse wait should not cover later cross-stream uses."""
	208	+ x = torch.randn(100, device="npu")
	209	+ stream = torch_npu.npu.Stream()
	210	+ default_stream = torch_npu.npu.default_stream()
	211	+
	212	+ stream.wait_stream(default_stream)
	213	+ with torch_npu.npu.stream(stream):
	214	+ _ = x + 1
	215	+
	216	+ default_stream.wait_stream(stream)
	217	+ with torch_npu.npu.stream(stream):
	218	+ _ = x + 2
	219	+
	220	+ warnings = get_event_handler().flush_record_stream_warnings()
	221	+ self.assertGreater(len(warnings), 0)
	222	+
	223	+ def test_creation_wait_event_only_covers_prior_uses(self):
	224	+ """A mid-sequence event wait should not cover later cross-stream uses."""
	225	+ x = torch.randn(100, device="npu")
	226	+ stream = torch_npu.npu.Stream()
	227	+ default_stream = torch_npu.npu.default_stream()
	228	+
	229	+ stream.wait_stream(default_stream)
	230	+ with torch_npu.npu.stream(stream):
	231	+ _ = x + 1
	232	+ event = torch_npu.npu.Event()
	233	+ event.record(stream)
	234	+
	235	+ default_stream.wait_event(event)
	236	+ with torch_npu.npu.stream(stream):
	237	+ _ = x + 2
	238	+ warnings = get_event_handler().flush_record_stream_warnings()
	239	+ self.assertGreater(len(warnings), 0)
	240	+
	241	+ def test_record_stream_partial_coverage(self):
	242	+ """record_stream for one stream should not cover another stream."""
	243	+ x = torch.randn(100, device="npu")
	244	+ stream1 = torch_npu.npu.Stream()
	245	+ stream2 = torch_npu.npu.Stream()
	246	+ default_stream = torch_npu.npu.default_stream()
	247	+
	248	+ x.record_stream(stream1)
	249	+
	250	+ stream1.wait_stream(default_stream)
	251	+ with torch_npu.npu.stream(stream1):
	252	+ _ = x + 1
	253	+
	254	+ stream2.wait_stream(default_stream)
	255	+ with torch_npu.npu.stream(stream2):
	256	+ _ = x + 2
	257	+
	258	+ warnings = get_event_handler().flush_record_stream_warnings()
	259	+ stream1_warnings = [
	260	+ w for w in warnings if w.usage_stream == int(stream1.npu_stream)
	261	+ ]
	262	+ stream2_warnings = [
	263	+ w for w in warnings if w.usage_stream == int(stream2.npu_stream)
	264	+ ]
	265	+ self.assertEqual(len(stream1_warnings), 0)
	266	+ self.assertGreater(len(stream2_warnings), 0)
	267	+
	268	+
	269	+class TestViewAndSlice(SanitizerRecordStreamTestBase):
	270	+ def test_view_cross_stream_no_record_stream_warns(self):
	271	+ """Cross-stream use of a view should be tracked at storage level."""
	272	+ x = torch.randn(100, device="npu")
	273	+ view = x[10:50]
	274	+ stream = torch_npu.npu.Stream()
	275	+ default_stream = torch_npu.npu.default_stream()
	276	+
	277	+ stream.wait_stream(default_stream)
	278	+ with torch_npu.npu.stream(stream):
	279	+ _ = view + 1
	280	+
	281	+ warnings = get_event_handler().flush_record_stream_warnings()
	282	+ self.assertGreater(len(warnings), 0)
	283	+
	284	+ def test_full_record_stream_covers_view_use(self):
	285	+ """record_stream on base tensor should cover view usage."""
	286	+ x = torch.randn(100, device="npu")
	287	+ stream = torch_npu.npu.Stream()
	288	+ default_stream = torch_npu.npu.default_stream()
	289	+
	290	+ x.record_stream(stream)
	291	+ view = x[10:50]
	292	+
	293	+ stream.wait_stream(default_stream)
	294	+ with torch_npu.npu.stream(stream):
	295	+ _ = view + 1
	296	+
	297	+ warnings = get_event_handler().flush_record_stream_warnings()
	298	+ self.assertEqual(len(warnings), 0)
	299	+
	300	+ def test_view_record_stream_covers_full_use(self):
	301	+ """record_stream on a view should cover base tensor usage."""
	302	+ x = torch.randn(100, device="npu")
	303	+ view = x[10:50]
	304	+ stream = torch_npu.npu.Stream()
	305	+ default_stream = torch_npu.npu.default_stream()
	306	+
	307	+ view.record_stream(stream)
	308	+ stream.wait_stream(default_stream)
	309	+ with torch_npu.npu.stream(stream):
	310	+ _ = x + 1
	311	+
	312	+ warnings = get_event_handler().flush_record_stream_warnings()
	313	+ self.assertEqual(len(warnings), 0)
	314	+
	315	+
	316	+class TestMemoryReuse(SanitizerRecordStreamTestBase):
	317	+ def test_no_stale_recorded_streams_after_realloc(self):
	318	+ """A new allocation should not inherit old recorded-stream state."""
	319	+ stream = torch_npu.npu.Stream()
	320	+ default_stream = torch_npu.npu.default_stream()
	321	+
	322	+ x = torch.randn(100, device="npu")
	323	+ x.record_stream(stream)
	324	+ del x
	325	+
	326	+ y = torch.randn(100, device="npu")
	327	+ stream.wait_stream(default_stream)
	328	+
	329	+ with torch_npu.npu.stream(stream):
	330	+ _ = y + 1
	331	+
	332	+ warnings = get_event_handler().flush_record_stream_warnings()
	333	+ self.assertGreater(len(warnings), 0)
	334	+
	335	+ def test_realloc_with_explicit_record_stream(self):
	336	+ """A reallocated tensor with explicit record_stream should not warn."""
	337	+ stream = torch_npu.npu.Stream()
	338	+ default_stream = torch_npu.npu.default_stream()
	339	+
	340	+ x = torch.randn(100, device="npu")
	341	+ del x
	342	+
	343	+ y = torch.randn(100, device="npu")
	344	+ y.record_stream(stream)
	345	+
	346	+ stream.wait_stream(default_stream)
	347	+ with torch_npu.npu.stream(stream):
	348	+ _ = y + 1
	349	+
	350	+ warnings = get_event_handler().flush_record_stream_warnings()
	351	+ self.assertEqual(len(warnings), 0)
	352	+
	353	+
	354	+class TestSanitizerDisabled(TestCase):
	355	+ """Behavior when sanitizer is disabled."""
	356	+ def setUp(self):
	357	+ reset_sanitizer()
	358	+ os.environ.pop("TORCH_NPU_SANITIZER", None)
	359	+
	360	+ def test_no_errors_when_disabled(self):
	361	+ """Disabled sanitizer should not report cross-stream issues."""
	362	+ x = torch.randn(100, device="npu")
	363	+ stream = torch_npu.npu.Stream()
	364	+ error_raised = False
	365	+ try:
	366	+ with torch_npu.npu.stream(stream):
	367	+ _ = x + 1
	368	+ torch_npu.npu.synchronize()
	369	+ except Exception:
	370	+ error_raised = True
	371	+
	372	+ self.assertFalse(error_raised)
	373	+
	374	+
	375	+class TestFlushBehavior(SanitizerRecordStreamTestBase):
	376	+ def test_dealloc_records_into_error_log(self):
	377	+ """Deallocation-time missing-record_stream warnings should be retained."""
	378	+ stream = torch_npu.npu.Stream()
	379	+ default_stream = torch_npu.npu.default_stream()
	380	+
	381	+ x = torch.randn(100, device="npu")
	382	+ stream.wait_stream(default_stream)
	383	+ with torch_npu.npu.stream(stream):
	384	+ _ = x + 1
	385	+ del x
	386	+ gc.collect()
	387	+ self.assertGreater(len(get_event_handler().record_stream_errors), 0)
	388	+
	389	+
	390	+if __name__ == "__main__":
	391	+ run_tests()

Mtest/test_npu_sanitizer.py+79-10

	@@ -3,9 +3,13 @@
3	import sys	3	import sys
4		4
5	import torch	5	import torch
6	-import torch.cuda._sanitizer as csan
7	from torch.testing._internal.common_utils import NoTest, run_tests, TEST_PRIVATEUSE1, TestCase	6	from torch.testing._internal.common_utils import NoTest, run_tests, TEST_PRIVATEUSE1, TestCase
8	-from torch_npu.npu._stream_check import apply_sanitizer_patch	7	+
		8	+from torch_npu.npu._stream_check import (
		9	+ NPUArgumentHandler,
		10	+ NPURecordStreamHandler,
		11	+ NPUTensorInfo,
		12	+)
9		13
10		14
11	if not TEST_PRIVATEUSE1:	15	if not TEST_PRIVATEUSE1:
	@@ -19,7 +23,7 @@ class TestArgumentHandler(TestCase):
19	a = torch.ones(5, 3, device="npu")	23	a = torch.ones(5, 3, device="npu")
20	b = torch.randn(5, 3, device="npu")	24	b = torch.randn(5, 3, device="npu")
21		25
22	- argument_handler = csan.ArgumentHandler()	26	+ argument_handler = NPUArgumentHandler()
23	argument_handler.parse_inputs(add_func._schema, (a, b), {}, is_factory=False)	27	argument_handler.parse_inputs(add_func._schema, (a, b), {}, is_factory=False)
24	c = torch.add(a, b)	28	c = torch.add(a, b)
25	argument_handler.parse_outputs(add_func._schema, c, is_factory=False)	29	argument_handler.parse_outputs(add_func._schema, c, is_factory=False)
	@@ -33,7 +37,7 @@ class TestArgumentHandler(TestCase):
33	b = torch.zeros(2, 1, 5, device="npu")	37	b = torch.zeros(2, 1, 5, device="npu")
34	c = torch.rand(2, 7, 5, device="npu")	38	c = torch.rand(2, 7, 5, device="npu")
35		39
36	- argument_handler = csan.ArgumentHandler()	40	+ argument_handler = NPUArgumentHandler()
37	argument_handler.parse_inputs(	41	argument_handler.parse_inputs(
38	cat_func._schema, ([a, b, c], 1), {}, is_factory=False	42	cat_func._schema, ([a, b, c], 1), {}, is_factory=False
39	)	43	)
	@@ -49,7 +53,7 @@ class TestArgumentHandler(TestCase):
49	split_func = torch.ops.aten.split.Tensor	53	split_func = torch.ops.aten.split.Tensor
50	a = torch.arange(10, device="npu").reshape(5, 2)	54	a = torch.arange(10, device="npu").reshape(5, 2)
51		55
52	- argument_handler = csan.ArgumentHandler()	56	+ argument_handler = NPUArgumentHandler()
53	argument_handler.parse_inputs(split_func._schema, (a, 2), {}, is_factory=False)	57	argument_handler.parse_inputs(split_func._schema, (a, 2), {}, is_factory=False)
54	out = torch.split(a, 2)	58	out = torch.split(a, 2)
55	argument_handler.parse_outputs(split_func._schema, out, is_factory=False)	59	argument_handler.parse_outputs(split_func._schema, out, is_factory=False)
	@@ -63,7 +67,7 @@ class TestArgumentHandler(TestCase):
63	add_inplace_func = torch.ops.aten.add_.Tensor	67	add_inplace_func = torch.ops.aten.add_.Tensor
64	a = torch.rand(4, 2, device="npu")	68	a = torch.rand(4, 2, device="npu")
65		69
66	- argument_handler = csan.ArgumentHandler()	70	+ argument_handler = NPUArgumentHandler()
67	argument_handler.parse_inputs(	71	argument_handler.parse_inputs(
68	add_inplace_func._schema, (a, 5), {}, is_factory=False	72	add_inplace_func._schema, (a, 5), {}, is_factory=False
69	)	73	)
	@@ -78,7 +82,7 @@ class TestArgumentHandler(TestCase):
78	a = torch.arange(8, device="npu")	82	a = torch.arange(8, device="npu")
79	b = torch.empty(8, device="npu")	83	b = torch.empty(8, device="npu")
80		84
81	- argument_handler = csan.ArgumentHandler()	85	+ argument_handler = NPUArgumentHandler()
82	argument_handler.parse_inputs(	86	argument_handler.parse_inputs(
83	mul_out_func._schema, (a, 3), {"out": b}, is_factory=False	87	mul_out_func._schema, (a, 3), {"out": b}, is_factory=False
84	)	88	)
	@@ -92,7 +96,7 @@ class TestArgumentHandler(TestCase):
92	nonzero_func = torch.ops.aten.nonzero.default	96	nonzero_func = torch.ops.aten.nonzero.default
93	a = torch.ones(5, 3, 2, device="npu")	97	a = torch.ones(5, 3, 2, device="npu")
94		98
95	- argument_handler = csan.ArgumentHandler()	99	+ argument_handler = NPUArgumentHandler()
96	argument_handler.parse_inputs(	100	argument_handler.parse_inputs(
97	nonzero_func._schema, (a,), {"as_tuple": True}, is_factory=False	101	nonzero_func._schema, (a,), {"as_tuple": True}, is_factory=False
98	)	102	)
	@@ -108,7 +112,7 @@ class TestArgumentHandler(TestCase):
108	vec = torch.arange(1, 4, device="npu")	112	vec = torch.arange(1, 4, device="npu")
109	M = torch.zeros(3, 3, device="npu")	113	M = torch.zeros(3, 3, device="npu")
110		114
111	- argument_handler = csan.ArgumentHandler()	115	+ argument_handler = NPUArgumentHandler()
112	argument_handler.parse_inputs(	116	argument_handler.parse_inputs(
113	addr_func._schema, (M, vec, vec), {}, is_factory=False	117	addr_func._schema, (M, vec, vec), {}, is_factory=False
114	)	118	)
	@@ -125,7 +129,72 @@ class TestArgumentHandler(TestCase):
125	)	129	)
126	self.assertEqual({out.data_ptr()}, argument_handler.outputs)	130	self.assertEqual({out.data_ptr()}, argument_handler.outputs)
127		131
		132	+ def test_empty_like_factory_input_not_read_but_output_written(self):
		133	+ """Factory-like op: input tensor is metadata-only, but output is a real allocation."""
		134	+ empty_like_func = torch.ops.aten.empty_like.default
		135	+ a = torch.ones(5, 3, device="npu")
		136	+
		137	+ argument_handler = NPUArgumentHandler()
		138	+
		139	+ # empty_like uses a only as metadata source: shape/dtype/device/layout.
		140	+ # It should not read a's data.
		141	+ argument_handler.parse_inputs(
		142	+ empty_like_func._schema,
		143	+ (a,),
		144	+ {},
		145	+ is_factory=True,
		146	+ )
		147	+
		148	+ out = torch.empty_like(a)
		149	+
		150	+ # Even though this is a factory-like op, the output is a real newly allocated tensor
		151	+ # and should be treated as written/output.
		152	+ argument_handler.parse_outputs(
		153	+ empty_like_func._schema,
		154	+ out,
		155	+ is_factory=True,
		156	+ )
		157	+
		158	+ self.assertEqual(set(), argument_handler.dataptrs_read)
		159	+ self.assertNotIn(a.data_ptr(), argument_handler.dataptrs_read)
		160	+
		161	+ self.assertEqual({out.data_ptr()}, argument_handler.dataptrs_written)
		162	+ self.assertEqual({out.data_ptr()}, argument_handler.outputs)
		163	+
		164	+ def test_equal_reads_inputs_but_no_tensor_output_written(self):
		165	+ """Data-reading op with non-tensor output should not record tensor writes."""
		166	+ equal_func = torch.ops.aten.equal.default
		167	+ a = torch.ones(5, 3, device="npu")
		168	+ b = torch.ones(5, 3, device="npu")
		169	+
		170	+ argument_handler = NPUArgumentHandler()
		171	+ argument_handler.parse_inputs(equal_func._schema, (a, b), {}, is_factory=False)
		172	+ out = torch.equal(a, b)
		173	+ argument_handler.parse_outputs(equal_func._schema, out, is_factory=False)
		174	+
		175	+ self.assertEqual({a.data_ptr(), b.data_ptr()}, argument_handler.dataptrs_read)
		176	+ self.assertEqual(set(), argument_handler.dataptrs_written)
		177	+ self.assertEqual(set(), argument_handler.outputs)
		178	+ self.assertTrue(isinstance(out, bool))
		179	+
		180	+
		181	+class TestRecordStreamHandler(TestCase):
		182	+ def test_erase_stream_removes_recorded_stream(self):
		183	+ """Communication eraseStream should clear the matching recorded stream."""
		184	+ handler = NPURecordStreamHandler()
		185	+ handler._npu_tensors[123] = NPUTensorInfo(recorded_streams={11, 22})
		186	+
		187	+ handler._handle_erase_stream(123, 11)
		188	+
		189	+ self.assertEqual({22}, handler._npu_tensors[123].recorded_streams)
		190	+
		191	+ def test_erase_stream_unknown_tensor_is_noop(self):
		192	+ """Late eraseStream callbacks for already-freed tensors should be ignored."""
		193	+ handler = NPURecordStreamHandler()
		194	+
		195	+ handler._handle_erase_stream(123, 11)
		196	+
		197	+ self.assertEqual({}, handler._npu_tensors)
128		198
129	if __name__ == "__main__":	199	if __name__ == "__main__":
130	- apply_sanitizer_patch()
131	run_tests()	200	run_tests()

Atest/test_sanitizer_pluggable_allocator.py+142-0

@@ -0,0 +1,142 @@
	1	+# Owner(s): ["module: unknown"]
	2	+import os
	3	+import platform
	4	+import shutil
	5	+import subprocess
	6	+import unittest
	7	+
	8	+import torch
	9	+import torch.utils.cpp_extension
	10	+
	11	+import torch_npu
	12	+from torch_npu.testing.testcase import run_tests, TestCase
	13	+
	14	+
	15	+PYTORCH_INSTALL_PATH = os.path.dirname(os.path.realpath(torch.__file__))
	16	+PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.realpath(torch_npu.__file__))
	17	+IS_ARM64 = platform.machine() in ('arm64', 'aarch64')
	18	+
	19	+
	20	+def setup_sanitizer():
	21	+ os.environ["TORCH_NPU_SANITIZER"] = "1"
	22	+ import torch_npu.npu._sanitizer as sanitizer
	23	+
	24	+ if not sanitizer.npu_sanitizer.enabled:
	25	+ sanitizer.npu_sanitizer.enable()
	26	+
	27	+
	28	+def get_event_handler():
	29	+ import torch_npu.npu._sanitizer as sanitizer
	30	+
	31	+ return sanitizer.npu_sanitizer.event_handler
	32	+
	33	+
	34	+def create_build_path(build_directory):
	35	+ if os.path.exists(build_directory):
	36	+ shutil.rmtree(build_directory, ignore_errors=True)
	37	+ os.makedirs(build_directory, exist_ok=True)
	38	+
	39	+
	40	+def build_stub(base_dir):
	41	+ build_stub_cmd = [
	42	+ "sh",
	43	+ os.path.join(base_dir, "third_party/acl/libs/build_stub.sh"),
	44	+ ]
	45	+ if subprocess.call(build_stub_cmd) != 0:
	46	+ raise RuntimeError(f"Failed to build stub: {build_stub_cmd}")
	47	+
	48	+
	49	+def reset_sanitizer():
	50	+ import torch_npu.npu._sanitizer as sanitizer
	51	+
	52	+ if sanitizer.npu_sanitizer.dispatch is not None:
	53	+ try:
	54	+ sanitizer.npu_sanitizer.dispatch.__exit__(None, None, None)
	55	+ except Exception:
	56	+ pass
	57	+ sanitizer.npu_sanitizer.dispatch = None
	58	+
	59	+ sanitizer.npu_sanitizer.event_handler = None
	60	+ sanitizer.npu_sanitizer.enabled = False
	61	+
	62	+
	63	+@unittest.skipUnless(IS_ARM64, "Only working on ARM")
	64	+class TestSanitizerPluggableAllocator(TestCase):
	65	+ module = None
	66	+ build_directory = os.path.join("allocator", "build_sanitizer_pluggable")
	67	+
	68	+ @classmethod
	69	+ def setUpClass(cls):
	70	+ BASE_DIR = os.path.abspath("./../")
	71	+ build_stub(BASE_DIR)
	72	+ create_build_path(cls.build_directory)
	73	+ CANN_LIB_PATH = os.path.join(BASE_DIR, "third_party/acl/libs")
	74	+ extra_ldflags = []
	75	+ extra_ldflags.append("-lascendcl")
	76	+ extra_ldflags.append(f"-L{CANN_LIB_PATH}")
	77	+ extra_ldflags.append("-lc10")
	78	+ extra_ldflags.append(f"-L{PYTORCH_INSTALL_PATH}")
	79	+ extra_include_paths = ["cpp_extensions"]
	80	+ extra_include_paths.append(os.path.join(PYTORCH_NPU_INSTALL_PATH, "include"))
	81	+
	82	+ cls.module = torch.utils.cpp_extension.load(
	83	+ name="sanitizer_pluggable_allocator_extensions",
	84	+ sources=["cpp_extensions/pluggable_allocator_extensions.cpp"],
	85	+ extra_include_paths=extra_include_paths,
	86	+ extra_cflags=["-g"],
	87	+ extra_ldflags=extra_ldflags,
	88	+ build_directory=cls.build_directory,
	89	+ verbose=True,
	90	+ )
	91	+
	92	+ def test_pluggable_allocator_record_stream_warning_and_suppression(self):
	93	+ """Pluggable allocator should support both missing-record_stream warning and suppression."""
	94	+ os_path = os.path.join(
	95	+ self.build_directory, "sanitizer_pluggable_allocator_extensions.so"
	96	+ )
	97	+ allocator = torch_npu.npu.memory.NPUPluggableAllocator(
	98	+ os_path, "my_malloc", "my_free"
	99	+ )
	100	+ torch_npu.npu.memory.change_current_allocator(allocator)
	101	+
	102	+ # Case 1: no record_stream -> should warn, guards false negative.
	103	+ setup_sanitizer()
	104	+
	105	+ x = torch.randn(100, device="npu")
	106	+ stream1 = torch_npu.npu.Stream()
	107	+ default_stream = torch_npu.npu.default_stream()
	108	+
	109	+ stream1.wait_stream(default_stream)
	110	+ with torch_npu.npu.stream(stream1):
	111	+ _ = x + 1
	112	+
	113	+ warnings = get_event_handler().flush_record_stream_warnings()
	114	+ self.assertGreater(
	115	+ len(warnings),
	116	+ 0,
	117	+ "Missing record_stream should be detected with NPUPluggableAllocator.",
	118	+ )
	119	+
	120	+ reset_sanitizer()
	121	+
	122	+ # Case 2: record_stream -> should not warn, guards false positive.
	123	+ setup_sanitizer()
	124	+
	125	+ y = torch.randn(100, device="npu")
	126	+ stream2 = torch_npu.npu.Stream()
	127	+ default_stream = torch_npu.npu.default_stream()
	128	+
	129	+ y.record_stream(stream2)
	130	+ stream2.wait_stream(default_stream)
	131	+ with torch_npu.npu.stream(stream2):
	132	+ _ = y + 1
	133	+
	134	+ warnings = get_event_handler().flush_record_stream_warnings()
	135	+ self.assertEqual(
	136	+ len(warnings),
	137	+ 0,
	138	+ "record_stream should suppress missing-record_stream warning with NPUPluggableAllocator.",
	139	+ )
	140	+
	141	+if __name__ == "__main__":
	142	+ run_tests()

Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+24-0

	@@ -3409,6 +3409,14 @@ public:
3409	// block must not be null reaching here	3409	// block must not be null reaching here
3410	TORCH_INTERNAL_ASSERT(block != nullptr, "No allocated block can be found", PTA_ERROR(ErrCode::NOT_FOUND));	3410	TORCH_INTERNAL_ASSERT(block != nullptr, "No allocated block can be found", PTA_ERROR(ErrCode::NOT_FOUND));
3411	device_allocator[block->device]->recordStream(block, stream);	3411	device_allocator[block->device]->recordStream(block, stream);
		3412	+#ifndef BUILD_LIBTORCH
		3413	+ const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace();
		3414	+ if (C10_UNLIKELY(trigger)) {
		3415	+ trigger->traceNpuRecordStream(
		3416	+ reinterpret_cast<uintptr_t>(ptr.get()),
		3417	+ reinterpret_cast<uintptr_t>(stream.stream(false)));
		3418	+ }
		3419	+#endif
3412	}	3420	}
3413		3421
3414	void eraseStream(const c10::DataPtr &ptr, c10_npu::NPUStream stream)	3422	void eraseStream(const c10::DataPtr &ptr, c10_npu::NPUStream stream)
	@@ -3441,6 +3449,14 @@ public:
3441	}	3449	}
3442		3450
3443	device_allocator[block->device]->eraseStream(block, stream);	3451	device_allocator[block->device]->eraseStream(block, stream);
		3452	+#ifndef BUILD_LIBTORCH
		3453	+ const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
		3454	+ if (C10_UNLIKELY(trigger)) {
		3455	+ trigger->traceNpuEraseStream(
		3456	+ reinterpret_cast<uintptr_t>(ptr.get()),
		3457	+ reinterpret_cast<uintptr_t>(stream.stream(false)));
		3458	+ }
		3459	+#endif
3444	}	3460	}
3445		3461
3446	void eraseStreamWithBlockPtr(void* block_ptr, c10_npu::NPUStream stream, void* work_ptr) override	3462	void eraseStreamWithBlockPtr(void* block_ptr, c10_npu::NPUStream stream, void* work_ptr) override
	@@ -3463,6 +3479,14 @@ public:
3463	}	3479	}
3464		3480
3465	device_allocator[block->device]->eraseStream(block, stream);	3481	device_allocator[block->device]->eraseStream(block, stream);
		3482	+#ifndef BUILD_LIBTORCH
		3483	+ const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
		3484	+ if (C10_UNLIKELY(trigger)) {
		3485	+ trigger->traceNpuEraseStream(
		3486	+ reinterpret_cast<uintptr_t>(block->ptr),
		3487	+ reinterpret_cast<uintptr_t>(stream.stream(false)));
		3488	+ }
		3489	+#endif
3466	}	3490	}
3467		3491
3468	void* getBlockPtr(const c10::DataPtr& ptr) override	3492	void* getBlockPtr(const c10::DataPtr& ptr) override

Mtorch_npu/csrc/npu/NPUPluggableAllocator.cpp+37-0

	@@ -5,6 +5,9 @@
5	#include "torch_npu/csrc/npu/NPUPluggableAllocator.h"	5	#include "torch_npu/csrc/npu/NPUPluggableAllocator.h"
6	#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"	6	#include "torch_npu/csrc/core/npu/NPUCachingAllocator.h"
7	#include "torch_npu/csrc/core/npu/NPUGuard.h"	7	#include "torch_npu/csrc/core/npu/NPUGuard.h"
		8	+#ifndef BUILD_LIBTORCH
		9	+#include "torch_npu/csrc/sanitizer/NPUTrace.h"
		10	+#endif
8		11
9	namespace torch::npu::NPUPluggableAllocator {	12	namespace torch::npu::NPUPluggableAllocator {
10		13
	@@ -119,6 +122,14 @@ void* NPUPluggableAllocator::malloc(
119	const std::lock_guard<std::mutex> lock(allocator_mutex_);	122	const std::lock_guard<std::mutex> lock(allocator_mutex_);
120	allocation_metadata_.emplace(r, _AllocationMetadata(size, device, stream));	123	allocation_metadata_.emplace(r, _AllocationMetadata(size, device, stream));
121	}	124	}
		125	+#ifndef BUILD_LIBTORCH
		126	+ if (r) {
		127	+ const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace();
		128	+ if (C10_UNLIKELY(trigger)) {
		129	+ trigger->traceNpuMemoryAllocation(reinterpret_cast<uintptr_t>(r));
		130	+ }
		131	+ }
		132	+#endif
122	return r;	133	return r;
123	}	134	}
124		135
	@@ -182,6 +193,12 @@ void NPUPluggableAllocator::raw_delete(void* ptr)
182	stream = metadata.stream;	193	stream = metadata.stream;
183	allocation_metadata_.erase(ptr);	194	allocation_metadata_.erase(ptr);
184	}	195	}
		196	+#ifndef BUILD_LIBTORCH
		197	+ const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace();
		198	+ if (C10_UNLIKELY(trigger)) {
		199	+ trigger->traceNpuMemoryDeallocation(reinterpret_cast<uintptr_t>(ptr));
		200	+ }
		201	+#endif
185	free_fn_(ptr, size, device_idx, stream);	202	free_fn_(ptr, size, device_idx, stream);
186	}	203	}
187		204
	@@ -251,6 +268,16 @@ void NPUPluggableAllocator::recordStream(
251	const c10::DataPtr& ptr,	268	const c10::DataPtr& ptr,
252	streamType stream)	269	streamType stream)
253	{	270	{
		271	+#ifndef BUILD_LIBTORCH
		272	+ if (ptr.get()) {
		273	+ const c10_npu::impl::PyCallbackTrigger *trigger = c10_npu::impl::NPUTrace::getTrace();
		274	+ if (C10_UNLIKELY(trigger)) {
		275	+ trigger->traceNpuRecordStream(
		276	+ reinterpret_cast<uintptr_t>(ptr.get()),
		277	+ reinterpret_cast<uintptr_t>(stream.stream(false)));
		278	+ }
		279	+ }
		280	+#endif
254	if (record_stream_fn_) {	281	if (record_stream_fn_) {
255	record_stream_fn_(ptr.get(), stream);	282	record_stream_fn_(ptr.get(), stream);
256	}	283	}
	@@ -263,6 +290,16 @@ void NPUPluggableAllocator::eraseStream(
263	if (erase_stream_fn_) {	290	if (erase_stream_fn_) {
264	erase_stream_fn_(ptr.get(), stream);	291	erase_stream_fn_(ptr.get(), stream);
265	}	292	}
		293	+#ifndef BUILD_LIBTORCH
		294	+ if (ptr.get()) {
		295	+ const c10_npu::impl::PyCallbackTrigger* trigger = c10_npu::impl::NPUTrace::getTrace();
		296	+ if (C10_UNLIKELY(trigger)) {
		297	+ trigger->traceNpuEraseStream(
		298	+ reinterpret_cast<uintptr_t>(ptr.get()),
		299	+ reinterpret_cast<uintptr_t>(stream.stream(false)));
		300	+ }
		301	+ }
		302	+#endif
266	}	303	}
267		304
268	void NPUPluggableAllocator::eraseStreamWithBlockPtr(void* block_ptr, c10_npu::NPUStream stream, void* work_ptr)	305	void NPUPluggableAllocator::eraseStreamWithBlockPtr(void* block_ptr, c10_npu::NPUStream stream, void* work_ptr)

Mtorch_npu/csrc/sanitizer/PyCallbackTrigger.h+11-0

	@@ -125,6 +125,17 @@ struct PyCallbackTrigger {
125	CONCRETE_TRACE_NPU("traceNpuExternalEventWaitCallbacks", event, stream);	125	CONCRETE_TRACE_NPU("traceNpuExternalEventWaitCallbacks", event, stream);
126	}	126	}
127	}	127	}
		128	+ void traceNpuRecordStream(uintptr_t data_ptr, uintptr_t stream) const
		129	+ {
		130	+ if (sanitizer_mode == SanitizerMode::STREAM) {
		131	+ CONCRETE_TRACE_NPU("NPURecordStreamCallbacks", data_ptr, stream);
		132	+ }
		133	+ }
		134	+ void traceNpuEraseStream(uintptr_t data_ptr, uintptr_t stream) const {
		135	+ if (sanitizer_mode == SanitizerMode::STREAM) {
		136	+ CONCRETE_TRACE_NPU("NPUEraseStreamCallbacks", data_ptr, stream);
		137	+ }
		138	+ }
128	};	139	};
129		140
130	PyCallbackTrigger* getPyCallbackTrigger(const int mode);	141	PyCallbackTrigger* getPyCallbackTrigger(const int mode);

Mtorch_npu/npu/_sanitizer.py+15-6

	@@ -1,13 +1,11 @@
1	import os	1	import os
2	import atexit	2	import atexit
3		3
4	-import torch.cuda._sanitizer as csan
5	import torch_npu	4	import torch_npu
6	import torch_npu.utils._npu_trace as npu_trace	5	import torch_npu.utils._npu_trace as npu_trace
7	import torch_npu.npu._stream_check as stream_check	6	import torch_npu.npu._stream_check as stream_check
8	import torch_npu.npu._kernel_check as kernel_check	7	import torch_npu.npu._kernel_check as kernel_check
9	from torch_npu.utils.utils import _print_warn_log	8	from torch_npu.utils.utils import _print_warn_log
10	-from torch_npu.npu._stream_check import apply_sanitizer_patch
11		9
12		10
13	class SanitizerMode:	11	class SanitizerMode:
	@@ -25,6 +23,8 @@ class NPUSanitizer:
25	self.opp_debug_path = os.path.join(os.getcwd(), "opp_debug_path")	23	self.opp_debug_path = os.path.join(os.getcwd(), "opp_debug_path")
26	self.opp_debug_kernel_path = os.getenv('ASCEND_OPP_DEBUG_PATH')	24	self.opp_debug_kernel_path = os.getenv('ASCEND_OPP_DEBUG_PATH')
27	self.enabled = False	25	self.enabled = False
		26	+ # record_stream detection is enabled by TORCH_NPU_SANITIZER=1
		27	+ self.check_record_stream = os.getenv('TORCH_NPU_SANITIZER', '0') == '1'
28		28
29	def enable(self):	29	def enable(self):
30	if self.opp_debug_kernel_path:	30	if self.opp_debug_kernel_path:
	@@ -56,7 +56,8 @@ class NPUSanitizer:
56	return True	56	return True
57		57
58	def enable_stream_check(self) -> bool:	58	def enable_stream_check(self) -> bool:
59	- self.event_handler = csan.EventHandler()	59	+ # Use our extended NPURecordStreamHandler which inherits from PyTorch's EventHandler
		60	+ self.event_handler = stream_check.NPURecordStreamHandler()
60	self.dispatch = stream_check.NPUSanitizerDispatchMode(self.event_handler)	61	self.dispatch = stream_check.NPUSanitizerDispatchMode(self.event_handler)
61	self.dispatch.__enter__()	62	self.dispatch.__enter__()
62	npu_trace.register_callback_for_npu_event_creation(	63	npu_trace.register_callback_for_npu_event_creation(
	@@ -99,11 +100,20 @@ class NPUSanitizer:
99	self.event_handler._handle_event_synchronization,	100	self.event_handler._handle_event_synchronization,
100	"handle_event_synchronization"	101	"handle_event_synchronization"
101	)	102	)
		103	+ npu_trace.register_callback_for_npu_record_stream(
		104	+ self.event_handler._handle_record_stream, "handle_record_stream"
		105	+ )
		106	+ npu_trace.register_callback_for_npu_erase_stream(
		107	+ self.event_handler._handle_erase_stream, "handle_erase_stream"
		108	+ )
102	return True	109	return True
103		110
104	def __del__(self):	111	def __del__(self):
105	- if self.dispatch:	112	+ try:
106	- self.dispatch.__exit__(None, None, None)	113	+ if self.dispatch:
		114	+ self.dispatch.__exit__(None, None, None)
		115	+ except Exception:
		116	+ pass
107		117
108	def clear_debug_env(self):	118	def clear_debug_env(self):
109	if self.kernel_path_manager:	119	if self.kernel_path_manager:
	@@ -111,7 +121,6 @@ class NPUSanitizer:
111		121
112		122
113	def enable_npu_sanitizer():	123	def enable_npu_sanitizer():
114	- apply_sanitizer_patch()
115	npu_sanitizer.enable()	124	npu_sanitizer.enable()
116		125
117		126

Mtorch_npu/npu/_stream_check.py+387-8

	@@ -1,9 +1,16 @@
1	import sys	1	import sys
2	import logging	2	import logging
3	import re	3	import re
		4	+import functools
		5	+import textwrap
		6	+import traceback
		7	+import inspect
		8	+from dataclasses import dataclass, field
		9	+from typing import Dict, List, Optional, Set
4		10
5	import torch	11	import torch
6	import torch.cuda._sanitizer as csan	12	import torch.cuda._sanitizer as csan
		13	+from torch.utils import _pytree as pytree
7	from torch.utils._python_dispatch import TorchDispatchMode	14	from torch.utils._python_dispatch import TorchDispatchMode
8	import torch_npu	15	import torch_npu
9		16
	@@ -15,6 +22,358 @@ logger = logging.getLogger(__name__)
15	FACTORY_FUNCTION_REGEX = re.compile("(new_.\|._like)")	22	FACTORY_FUNCTION_REGEX = re.compile("(new_.\|._like)")
16		23
17		24
		25	+@dataclass
		26	+class CrossStreamUsage:
		27	+ """Records when a tensor is used on a stream different from its allocation stream."""
		28	+ usage_stream: csan.StreamId
		29	+ seq_num: csan.SeqNum
		30	+ operator: str
		31	+ stack_trace: traceback.StackSummary
		32	+
		33	+
		34	+@dataclass
		35	+class NPUTensorInfo:
		36	+ """Tracks tensor allocation and cross-stream usage for record_stream detection.
		37	+
		38	+ Maintained independently from the parent EventHandler's TensorInfo, which tracks
		39	+ read/write accesses for data race detection.
		40	+ """
		41	+ allocation_stream: Optional[csan.StreamId] = None
		42	+ allocation_stack_trace: Optional[traceback.StackSummary] = None
		43	+ recorded_streams: Set[csan.StreamId] = field(default_factory=set)
		44	+ cross_stream_usages: Dict[csan.StreamId, CrossStreamUsage] = field(default_factory=dict)
		45	+
		46	+
		47	+class MissingRecordStreamError(csan.SynchronizationError):
		48	+ """Tensor was used across streams without record_stream or proper sync.
		49	+
		50	+ Detected at tensor deallocation or via flush_record_stream_warnings().
		51	+ Per PyTorch docs, record_stream is NOT required if creation_stream has been
		52	+ synchronized to wait for the usage_stream before deallocation:
		53	+ - creation_stream.wait_stream(usage_stream)
		54	+ - creation_stream.wait_event(event_on_usage_stream)
		55	+ - torch.npu.synchronize() (device-level sync covers all directions)
		56	+ """
		57	+
		58	+ def __init__(
		59	+ self,
		60	+ data_ptr: csan.DataPtr,
		61	+ allocation_stack_trace: Optional[traceback.StackSummary],
		62	+ allocation_stream: csan.StreamId,
		63	+ usage_stream: csan.StreamId,
		64	+ usage: CrossStreamUsage,
		65	+ recorded_streams: set[csan.StreamId],
		66	+ ):
		67	+ self.data_ptr = data_ptr
		68	+ self.allocation_stack_trace = allocation_stack_trace
		69	+ self.allocation_stream = allocation_stream
		70	+ self.usage_stream = usage_stream
		71	+ self.usage = usage
		72	+ self.recorded_streams = recorded_streams
		73	+
		74	+ def __repr__(self):
		75	+ return (
		76	+ f"MissingRecordStreamError(data_ptr={self.data_ptr}, "
		77	+ f"alloc_stream={self.allocation_stream}, "
		78	+ f"usage_stream={self.usage_stream}, "
		79	+ f"operator='{self.usage.operator}')"
		80	+ )
		81	+
		82	+ def __str__(self):
		83	+ result = textwrap.dedent(
		84	+ f"""\
		85	+ ============================
		86	+ NPUSanitizer: missing record_stream detected!
		87	+ Tensor (data ptr: {self.data_ptr}) allocated on stream {self.allocation_stream}
		88	+ was used on stream {self.usage_stream} without record_stream or
		89	+ creation_stream.wait_stream(usage_stream).
		90	+
		91	+ This may cause use-after-free if the caching allocator reuses memory
		92	+ on the allocation stream before the usage stream finishes.
		93	+
		94	+ Fix with ONE of:
		95	+ A) tensor.record_stream(stream) — tell allocator about the usage
		96	+ B) creation_stream.wait_stream(usage_stream) before deallocation
		97	+
		98	+ Cross-stream usage during kernel:
		99	+ {self.usage.operator}
		100	+ """
		101	+ )
		102	+ result += f"With stack trace:\n{''.join(self.usage.stack_trace.format())}\n"
		103	+ if self.recorded_streams:
		104	+ result += f"Streams recorded via record_stream: {self.recorded_streams}\n"
		105	+ else:
		106	+ result += "No streams were recorded via record_stream.\n"
		107	+ if self.allocation_stack_trace:
		108	+ result += (
		109	+ "Tensor was allocated with stack trace:\n"
		110	+ f"{''.join(self.allocation_stack_trace.format())}"
		111	+ )
		112	+ return result
		113	+
		114	+
		115	+class NPURecordStreamHandler(csan.EventHandler):
		116	+ """EventHandler with deferred record_stream detection.
		117	+
		118	+ Record_stream checks are deferred to deallocation time (or manual flush via
		119	+ flush_record_stream_warnings), because what matters for memory safety is whether
		120	+ creation_stream has synced with usage_stream BEFORE the tensor's memory is
		121	+ reused — not at the time of the cross-stream kernel launch.
		122	+
		123	+ This avoids false positives when the user plans to sync after the kernel launch
		124	+ but before tensor deallocation, which is the typical usage pattern.
		125	+ """
		126	+
		127	+ def __init__(self) -> None:
		128	+ super().__init__()
		129	+ self._npu_tensors: dict[csan.DataPtr, NPUTensorInfo] = {}
		130	+ self.record_stream_errors: list[MissingRecordStreamError] = []
		131	+
		132	+ def _handle_memory_allocation(self, data_ptr: csan.DataPtr) -> None:
		133	+ super()._handle_memory_allocation(data_ptr)
		134	+ alloc_trace = None
		135	+ try:
		136	+ alloc_trace = self.tensors_accessed.get_allocation_stack_trace(data_ptr)
		137	+ except KeyError:
		138	+ pass
		139	+ current_stream: Optional[csan.StreamId] = None
		140	+ try:
		141	+ current_stream = int(torch_npu.npu.current_stream().npu_stream)
		142	+ except RuntimeError:
		143	+ pass
		144	+ self._npu_tensors[data_ptr] = NPUTensorInfo(
		145	+ allocation_stream=current_stream,
		146	+ allocation_stack_trace=alloc_trace,
		147	+ )
		148	+
		149	+ def _handle_memory_deallocation(self, data_ptr: csan.DataPtr) -> None:
		150	+ if data_ptr in self._npu_tensors:
		151	+ for error in self._get_record_stream_errors(data_ptr):
		152	+ print(error, file=sys.stderr)
		153	+ self.record_stream_errors.append(error)
		154	+ del self._npu_tensors[data_ptr]
		155	+ super()._handle_memory_deallocation(data_ptr)
		156	+
		157	+ def _handle_kernel_launch(
		158	+ self,
		159	+ stream: csan.StreamId,
		160	+ read_only: set[csan.DataPtr],
		161	+ read_write: set[csan.DataPtr],
		162	+ outputs: set[csan.DataPtr],
		163	+ operator: str,
		164	+ tensor_aliases: dict[int, list[str]],
		165	+ storage_dataptrs_accessed: Optional[Set[csan.DataPtr]] = None
		166	+ ) -> List[csan.SynchronizationError]:
		167	+ errors = super()._handle_kernel_launch(
		168	+ stream, read_only, read_write, outputs, operator, tensor_aliases
		169	+ )
		170	+ # Use storage-level data_ptrs (matches the allocation callback and
		171	+ # the C++ recordStream trace), falling back to tensor data_ptrs
		172	+ # only when not provided by the dispatch mode.
		173	+ accessed = (
		174	+ storage_dataptrs_accessed
		175	+ if storage_dataptrs_accessed is not None
		176	+ else (read_only \| read_write)
		177	+ )
		178	+ self._record_cross_stream_usage(stream, accessed, operator)
		179	+ return errors
		180	+
		181	+ def _record_cross_stream_usage(
		182	+ self, stream: csan.StreamId, all_accessed: set[csan.DataPtr], operator: str
		183	+ ) -> None:
		184	+ """Record that tensors are being accessed on a non-allocation stream.
		185	+
		186	+ Stack trace is captured lazily (only once per kernel launch) to avoid
		187	+ redundant walks when multiple tensors are accessed in the same kernel.
		188	+ Records the current seq_num so that sync checks can verify the sync
		189	+ happened AFTER this usage, not just from stream creation inheritance.
		190	+ """
		191	+ stack_trace = None
		192	+ current_seq = self.seq_num
		193	+ for data_ptr in all_accessed:
		194	+ info = self._npu_tensors.get(data_ptr)
		195	+ if info is None or info.allocation_stream is None:
		196	+ continue
		197	+ if info.allocation_stream == stream:
		198	+ continue
		199	+ existing = info.cross_stream_usages.get(stream)
		200	+ if existing is not None:
		201	+ existing.seq_num = current_seq
		202	+ existing.operator = operator
		203	+ if stack_trace is None:
		204	+ stack_trace = traceback.StackSummary.extract(
		205	+ traceback.walk_stack(inspect.currentframe()),
		206	+ lookup_lines=False,
		207	+ )
		208	+ stack_trace.reverse()
		209	+ existing.stack_trace = stack_trace
		210	+ continue
		211	+ if stack_trace is None:
		212	+ stack_trace = traceback.StackSummary.extract(
		213	+ traceback.walk_stack(inspect.currentframe()),
		214	+ lookup_lines=False,
		215	+ )
		216	+ stack_trace.reverse()
		217	+ info.cross_stream_usages[stream] = CrossStreamUsage(
		218	+ usage_stream=stream,
		219	+ seq_num=current_seq,
		220	+ operator=operator,
		221	+ stack_trace=stack_trace,
		222	+ )
		223	+
		224	+ def _get_record_stream_errors(self, data_ptr: csan.DataPtr) -> List[MissingRecordStreamError]:
		225	+ info = self._npu_tensors.get(data_ptr)
		226	+ if info is None or info.allocation_stream is None:
		227	+ return []
		228	+ errors = []
		229	+ for usage_stream, usage in info.cross_stream_usages.items():
		230	+ if usage_stream in info.recorded_streams:
		231	+ continue
		232	+ if self._is_creation_stream_synced_to_usage(
		233	+ info.allocation_stream, usage_stream, usage.seq_num
		234	+ ):
		235	+ continue
		236	+ errors.append(MissingRecordStreamError(
		237	+ data_ptr=data_ptr,
		238	+ allocation_stack_trace=info.allocation_stack_trace,
		239	+ allocation_stream=info.allocation_stream,
		240	+ usage_stream=usage_stream,
		241	+ usage=usage,
		242	+ recorded_streams=info.recorded_streams.copy(),
		243	+ ))
		244	+ return errors
		245	+
		246	+ def flush_record_stream_warnings(self) -> List[MissingRecordStreamError]:
		247	+ """Check all tracked tensors for missing record_stream and print errors.
		248	+
		249	+ Call after all stream operations are complete (including any synchronization)
		250	+ to detect tensors used cross-stream without record_stream or
		251	+ creation-to-usage synchronization.
		252	+
		253	+ Errors are printed to stderr and appended to self.record_stream_errors.
		254	+ """
		255	+ errors = []
		256	+ for data_ptr in list(self._npu_tensors):
		257	+ new_errors = self._get_record_stream_errors(data_ptr)
		258	+ for error in new_errors:
		259	+ print(error, file=sys.stderr)
		260	+ errors.extend(new_errors)
		261	+ self.record_stream_errors.extend(errors)
		262	+ return errors
		263	+
		264	+ def _is_creation_stream_synced_to_usage(
		265	+ self,
		266	+ creation_stream: csan.StreamId,
		267	+ usage_stream: csan.StreamId,
		268	+ usage_seq_num: csan.SeqNum = 0,
		269	+ ) -> bool:
		270	+ """Check if creation stream has synced with usage stream's operations.
		271	+
		272	+ Compares against usage_seq_num to distinguish real synchronization from
		273	+ stream creation inheritance. Stream creation sets initial sync state to 0,
		274	+ but actual kernel seq_nums start at 1, so comparing >= usage_seq_num
		275	+ ensures we detect real sync operations rather than inherited initial state.
		276	+ """
		277	+ try:
		278	+ creation_state = self.syncs.current_sync_states.get(creation_stream, {})
		279	+ return creation_state.get(usage_stream, -1) >= usage_seq_num
		280	+ except (AttributeError, KeyError):
		281	+ return False
		282	+
		283	+ def _handle_record_stream(self, data_ptr: csan.DataPtr, stream: csan.StreamId) -> None:
		284	+ """Track a record_stream call for memory safety checking."""
		285	+ if data_ptr not in self._npu_tensors:
		286	+ self._npu_tensors[data_ptr] = NPUTensorInfo()
		287	+ self._npu_tensors[data_ptr].recorded_streams.add(stream)
		288	+
		289	+ def _handle_erase_stream(
		290	+ self, data_ptr: csan.DataPtr, stream: csan.StreamId
		291	+ ) -> None:
		292	+ """Track eraseStream after a communication work no longer owns a stream."""
		293	+ info = self._npu_tensors.get(data_ptr)
		294	+ if info is None:
		295	+ return
		296	+ info.recorded_streams.discard(stream)
		297	+
		298	+
		299	+class NPUArgumentHandler:
		300	+ def __init__(self):
		301	+ self.dataptrs_read: set[csan.DataPtr] = set()
		302	+ self.dataptrs_written: set[csan.DataPtr] = set()
		303	+ self.tensor_aliases: dict[int, list[str]] = {}
		304	+ self.outputs: set[csan.DataPtr] = set()
		305	+ self.storage_dataptrs_accessed: set[csan.DataPtr] = set()
		306	+
		307	+ def _handle_argument(
		308	+ self,
		309	+ value,
		310	+ is_write: bool,
		311	+ metadata_only: bool,
		312	+ name: Optional[str] = None,
		313	+ is_output: bool = False,
		314	+ ) -> None:
		315	+ if not isinstance(value, torch.Tensor) or not value.is_npu:
		316	+ return
		317	+
		318	+ # View / metadata_only tensor arguments do not represent real data access.
		319	+ # Do not record read/write/storage access for them.
		320	+ if metadata_only:
		321	+ return
		322	+
		323	+ data_ptr = value.data_ptr() if value.data_ptr() else id(value)
		324	+ if is_write:
		325	+ self.dataptrs_written.add(data_ptr)
		326	+ else:
		327	+ self.dataptrs_read.add(data_ptr)
		328	+
		329	+ self.tensor_aliases.setdefault(data_ptr, [])
		330	+ if name is not None:
		331	+ self.tensor_aliases[data_ptr].append(name)
		332	+ if is_output:
		333	+ self.outputs.add(data_ptr)
		334	+
		335	+ # Also collect the storage start for record_stream tracking.
		336	+ try:
		337	+ storage = value.untyped_storage()
		338	+ if storage is not None:
		339	+ storage_ptr = storage.data_ptr()
		340	+ if storage_ptr:
		341	+ self.storage_dataptrs_accessed.add(storage_ptr)
		342	+ except (RuntimeError, AttributeError):
		343	+ pass
		344	+
		345	+ def parse_inputs(self, schema, args, kwargs, *, is_factory: bool = False) -> None:
		346	+ from torch.cuda._sanitizer import zip_arguments
		347	+ for argument, value in zip_arguments(schema, args, kwargs):
		348	+ is_write = argument.alias_info is not None and argument.alias_info.is_write
		349	+ metadata_only = is_factory or (
		350	+ argument.alias_info is not None and not argument.alias_info.is_write
		351	+ )
		352	+ pytree.tree_map_(
		353	+ functools.partial(
		354	+ self._handle_argument,
		355	+ is_write=is_write,
		356	+ name=argument.name,
		357	+ metadata_only=metadata_only,
		358	+ ),
		359	+ value,
		360	+ )
		361	+
		362	+ def parse_outputs(self, schema, outputs, *, is_factory: bool = False) -> None:
		363	+ from torch.cuda._sanitizer import zip_arguments
		364	+ for res, value in zip(schema.returns, (outputs,)):
		365	+ metadata_only = res.alias_info is not None and not res.alias_info.is_write
		366	+ pytree.tree_map_(
		367	+ functools.partial(
		368	+ self._handle_argument,
		369	+ is_write=True,
		370	+ metadata_only=metadata_only,
		371	+ is_output=True,
		372	+ ),
		373	+ value,
		374	+ )
		375	+
		376	+
18	class NPUSanitizerDispatchMode(TorchDispatchMode):	377	class NPUSanitizerDispatchMode(TorchDispatchMode):
19		378
20	def __init__(self, event_handler: csan.EventHandler):	379	def __init__(self, event_handler: csan.EventHandler):
	@@ -33,9 +392,13 @@ class NPUSanitizerDispatchMode(TorchDispatchMode):
33	def __torch_dispatch__(self, func, types, args=(), kwargs=None):	392	def __torch_dispatch__(self, func, types, args=(), kwargs=None):
34	kwargs = {} if kwargs is None else kwargs	393	kwargs = {} if kwargs is None else kwargs
35		394
		395	+ func_name = func.__name__ if hasattr(func, '__name__') else str(func)
		396	+ if "record_stream" in func_name:
		397	+ return self._handle_record_stream_op(func, args, kwargs)
		398	+
36	is_factory = bool(FACTORY_FUNCTION_REGEX.match(func._schema.name))	399	is_factory = bool(FACTORY_FUNCTION_REGEX.match(func._schema.name))
37		400
38	- self.args_handler = csan.ArgumentHandler()	401	+ self.args_handler = NPUArgumentHandler()
39	aten_api = func.__name__.split(".")[0]	402	aten_api = func.__name__.split(".")[0]
40	self.enable_autograd(aten_api)	403	self.enable_autograd(aten_api)
41	self.parse_inputs(func._schema, args, kwargs, is_factory=is_factory)	404	self.parse_inputs(func._schema, args, kwargs, is_factory=is_factory)
	@@ -43,10 +406,17 @@ class NPUSanitizerDispatchMode(TorchDispatchMode):
43	outputs = func(args, *kwargs)	406	outputs = func(args, *kwargs)
44		407
45	self.parse_outputs(func._schema, outputs, is_factory=is_factory)	408	self.parse_outputs(func._schema, outputs, is_factory=is_factory)
		409	+ if (
		410	+ not self.args_handler.dataptrs_read
		411	+ and not self.args_handler.dataptrs_written
		412	+ and not self.args_handler.outputs
		413	+ and not self.args_handler.storage_dataptrs_accessed
		414	+ ):
		415	+ return outputs
46		416
47	npu_stream = 0	417	npu_stream = 0
48	try:	418	try:
49	- npu_stream = torch_npu.npu.current_stream().npu_stream	419	+ npu_stream = int(torch_npu.npu.current_stream().npu_stream)
50	except RuntimeError as err:	420	except RuntimeError as err:
51	logger.info(	421	logger.info(
52	"Failed to get current stream, ignore this kernel launch record. error info is: %s",	422	"Failed to get current stream, ignore this kernel launch record. error info is: %s",
	@@ -57,6 +427,18 @@ class NPUSanitizerDispatchMode(TorchDispatchMode):
57		427
58	return outputs	428	return outputs
59		429
		430	+ def _handle_record_stream_op(self, func, args, kwargs):
		431	+ """Short-circuit record_stream so it isn't treated as a regular kernel launch.
		432	+
		433	+ Tracking is done in C++ via the NPURecordStreamCallbacks trace, which is fired
		434	+ from NpuCachingAllocator::recordStream / NPUPluggableAllocator::recordStream —
		435	+ the chokepoint that all entry points (aten op, NPUGuardImpl, HCCL/LCCL, RPC,
		436	+ pluggable allocator) funnel through. Doing tracking here would only cover the
		437	+ aten-op path and would use tensor.data_ptr() (with view offset), which would
		438	+ not match the storage data_ptr that the allocation callback uses.
		439	+ """
		440	+ return func(args, *kwargs)
		441	+
60	def parse_inputs(self, schema, args, kwargs, is_factory=False):	442	def parse_inputs(self, schema, args, kwargs, is_factory=False):
61	self.args_handler.parse_inputs(schema, args, kwargs, is_factory=is_factory)	443	self.args_handler.parse_inputs(schema, args, kwargs, is_factory=is_factory)
62		444
	@@ -69,14 +451,11 @@ class NPUSanitizerDispatchMode(TorchDispatchMode):
69	self.args_handler.dataptrs_read - self.args_handler.dataptrs_written,	451	self.args_handler.dataptrs_read - self.args_handler.dataptrs_written,
70	self.args_handler.dataptrs_written,	452	self.args_handler.dataptrs_written,
71	self.args_handler.outputs,	453	self.args_handler.outputs,
72	- func._schema,	454	+ str(func._schema),
73	- self.args_handler.tensor_aliases	455	+ self.args_handler.tensor_aliases,
		456	+ storage_dataptrs_accessed=self.args_handler.storage_dataptrs_accessed,
74	)	457	)
75	if errors:	458	if errors:
76	for error in errors:	459	for error in errors:
77	print(error, file=sys.stderr)	460	print(error, file=sys.stderr)
78	raise csan.CUDASanitizerErrors(errors)	461	raise csan.CUDASanitizerErrors(errors)
79	-
80	-
81	-def apply_sanitizer_patch():
82	- torch.Tensor.is_cuda = torch.Tensor.is_npu

Mtorch_npu/utils/_npu_trace.py+23-0

	@@ -62,6 +62,12 @@ NPUStreamSynchronizationCallbacks: "CallbackRegistry" = CallbackRegistry(
62	NPUEventSynchronizationCallbacks: "CallbackRegistry" = CallbackRegistry(	62	NPUEventSynchronizationCallbacks: "CallbackRegistry" = CallbackRegistry(
63	"[stream check] NPU event synchronization"	63	"[stream check] NPU event synchronization"
64	)	64	)
		65	+NPURecordStreamCallbacks: "CallbackRegistry" = CallbackRegistry(
		66	+ "[stream check] NPU record_stream"
		67	+)
		68	+NPUEraseStreamCallbacks: "CallbackRegistry" = CallbackRegistry(
		69	+ "[stream check] NPU erase_stream"
		70	+)
65		71
66		72
67	def register_callback_for_acl_start_execution(cb: Callable[[str], None], cb_name: str) -> None:	73	def register_callback_for_acl_start_execution(cb: Callable[[str], None], cb_name: str) -> None:
	@@ -110,3 +116,20 @@ def register_callback_for_npu_stream_synchronization(cb: Callable[[int], None],
110		116
111	def register_callback_for_npu_event_synchronization(cb: Callable[[int], None], cb_name: str) -> None:	117	def register_callback_for_npu_event_synchronization(cb: Callable[[int], None], cb_name: str) -> None:
112	NPUEventSynchronizationCallbacks.add_callback(cb, cb_name)	118	NPUEventSynchronizationCallbacks.add_callback(cb, cb_name)
		119	+
		120	+
		121	+def register_callback_for_npu_record_stream(cb: Callable[[int, int], None], cb_name: str) -> None:
		122	+ """Register callback for record_stream calls.
		123	+
		124	+ Args:
		125	+ cb: Callback function taking (data_ptr, stream_id) as arguments
		126	+ cb_name: Name of the callback for debugging
		127	+ """
		128	+ NPURecordStreamCallbacks.add_callback(cb, cb_name)
		129	+
		130	+
		131	+def register_callback_for_npu_erase_stream(
		132	+ cb: Callable[[int, int], None], cb_name: str
		133	+) -> None:
		134	+ """Register callback for eraseStream calls."""
		135	+ NPUEraseStreamCallbacks.add_callback(cb, cb_name)