fix: add PyObject_ClearWeakRefs in NPU Stream/Event tp

已合并

fix: add PyObject_ClearWeakRefs in NPU Stream/Event tp_dealloc #41909

21xiu创建于 7月17日从xiu_21/pytorch： cherry-pick-mr-41678-1784250811924-auto 合入到 Ascend/pytorch： v2.12.0-26.1.0

fix: add PyObject_ClearWeakRefs in NPU Stream/Event tp_dealloc #41909

已合并

21xiu创建于 7月17日

从xiu_21/pytorch cherry-pick-mr-41678-1784250811924-auto合入到Ascend/pytorch v2.12.0-26.1.0

共 3 个文件变更+96-2

Mtest/dynamo/test_stream.py+91-0

	@@ -1,8 +1,15 @@
1	# Owner(s): ["module: dynamo"]	1	# Owner(s): ["module: dynamo"]
2	import functools	2	import functools
		3	+import gc
3	import unittest	4	import unittest
		5	+import weakref
4	import torch	6	import torch
5	import torch._dynamo.test_case	7	import torch._dynamo.test_case
		8	+from torch._dynamo.graph_bytecode_inputs import (
		9	+ CURRENT_STREAM_INDEX,
		10	+ index_to_external_object_weakref,
		11	+ reset_user_object_tracking,
		12	+)
6	import torch_npu	13	import torch_npu
7		14
8	requires_npu = functools.partial(unittest.skipIf, not torch.npu.is_available(), "requires npu")	15	requires_npu = functools.partial(unittest.skipIf, not torch.npu.is_available(), "requires npu")
	@@ -25,6 +32,90 @@ class StreamintoDynamoTests(torch._dynamo.test_case.TestCase):
25	output1 = model_1(inp)	32	output1 = model_1(inp)
26	torch.allclose(output, output1)	33	torch.allclose(output, output1)
27		34
		35	+ def _assert_weakref_callback_fires(self, factory):
		36	+ """Backend Stream/Event tp_dealloc overrides must call
		37	+ PyObject_ClearWeakRefs so that weakrefs to destroyed instances
		38	+ are properly cleared. Without it, the weakref's wr_object is
		39	+ left dangling and any later access (e.g. the dynamo external-
		40	+ object registry being cleared at interpreter finalization)
		41	+ hits a use-after-free.
		42	+
		43	+ The callback firing is the only reliable Python-level signal —
		44	+ `weakref.ref(s)() is None` can return True even with the bug
		45	+ present, because other CPython paths may null wr_object
		46	+ without invoking callbacks (Objects/weakrefobject.c).
		47	+ """
		48	+ called = []
		49	+ obj = factory()
		50	+ # ``_weakref_keepalive`` must outlive ``del obj`` so the callback
		51	+ # has a chance to fire; the binding is load-bearing.
		52	+ _weakref_keepalive = weakref.ref(obj, lambda _ref: called.append(True))
		53	+ del obj
		54	+ gc.collect()
		55	+ self.assertEqual(called, [True])
		56	+ del _weakref_keepalive
		57	+
		58	+ @requires_npu()
		59	+ def test_npu_stream_event_weakref_callback(self):
		60	+ self._assert_weakref_callback_fires(torch.npu.Stream)
		61	+ self._assert_weakref_callback_fires(torch.npu.Event)
		62	+
		63	+ @requires_npu()
		64	+ def test_dynamo_registry_no_dangling_weakref(self):
		65	+ """Natural repro of the original UAF pattern.
		66	+
		67	+ ``torch.compile(fn, backend="eager")`` with ``fn`` referencing
		68	+ ``torch.npu.current_stream()`` causes dynamo to register a
		69	+ weakref to the captured stream in
		70	+ ``index_to_external_object_weakref`` (via
		71	+ ``store_user_object_weakrefs``, which does NOT pin via
		72	+ ``keep_alive``). As soon as ``fn`` returns, the captured wrapper
		73	+ has no strong references and is freed via ``THNPStream_dealloc``.
		74	+
		75	+ At that point the patched tp_dealloc must call
		76	+ ``PyObject_ClearWeakRefs``, otherwise the registry now holds a
		77	+ weakref whose ``wr_object`` is a dangling pointer to freed
		78	+ memory. Later, when ``_PyModule_ClearDict`` tears down the
		79	+ registry at interpreter finalization, dereferencing that
		80	+ dangling pointer to clear the weakref hits a use-after-free.
		81	+ """
		82	+ # Start from a known-clean registry so the assertion below is
		83	+ # a statement about what happened in this test, not residue
		84	+ # from earlier tests in the same process.
		85	+ reset_user_object_tracking()
		86	+
		87	+ def fn(x):
		88	+ return torch.npu.current_stream()
		89	+
		90	+ x = torch.zeros(1, device="npu")
		91	+ compiled = torch.compile(fn, backend="eager")
		92	+ compiled(x)
		93	+ del compiled
		94	+ gc.collect()
		95	+
		96	+ # Tripwire: torch.compile of a function referencing
		97	+ # current_stream() must still register under
		98	+ # CURRENT_STREAM_INDEX. If this fails, the production code
		99	+ # path that originally surfaced the UAF has moved and this
		100	+ # regression test no longer covers it.
		101	+ self.assertIn(
		102	+ CURRENT_STREAM_INDEX,
		103	+ index_to_external_object_weakref,
		104	+ "torch.compile of a function referencing current_stream() must "
		105	+ "register a weakref under CURRENT_STREAM_INDEX",
		106	+ )
		107	+
		108	+ # The captured stream wrapper was freed when fn returned. The
		109	+ # patched tp_dealloc must have cleared the registry's weakref;
		110	+ # dereferencing it now must return None rather than a dangling
		111	+ # pointer into freed memory.
		112	+ self.assertIsNone(
		113	+ index_to_external_object_weakref[CURRENT_STREAM_INDEX](),
		114	+ "dynamo registry holds a weakref to a Stream wrapper that has "
		115	+ "been freed; tp_dealloc must call PyObject_ClearWeakRefs to "
		116	+ "clear it, otherwise the registry retains a dangling pointer",
		117	+ )
		118	+
28		119
29	if __name__ == "__main__":	120	if __name__ == "__main__":
30	from torch._dynamo.test_case import run_tests	121	from torch._dynamo.test_case import run_tests

Mtorch_npu/csrc/npu/Event.cpp+2-1

	@@ -118,6 +118,7 @@ static PyObject* THNPEvent_from_ipc_handle(
118	static void THNPEvent_dealloc(THNPEvent *self)	118	static void THNPEvent_dealloc(THNPEvent *self)
119	{	119	{
120	self->npu_event.~NPUEvent();	120	self->npu_event.~NPUEvent();
		121	+ PyObject_ClearWeakRefs((PyObject*)self);
121	Py_TYPE(self)->tp_free((PyObject*)self);	122	Py_TYPE(self)->tp_free((PyObject*)self);
122	}	123	}
123		124
	@@ -262,7 +263,7 @@ PyTypeObject THNPEventType = {
262	nullptr, /* tp_traverse */	263	nullptr, /* tp_traverse */
263	nullptr, /* tp_clear */	264	nullptr, /* tp_clear */
264	nullptr, /* tp_richcompare */	265	nullptr, /* tp_richcompare */
265	- 0, /* tp_weaklistoffset */	266	+ 0, /* tp_weaklistoffset (inherited from THPStreamType via tp_base) */
266	nullptr, /* tp_iter */	267	nullptr, /* tp_iter */
267	nullptr, /* tp_iternext */	268	nullptr, /* tp_iternext */
268	THNPEvent_methods, /* tp_methods */	269	THNPEvent_methods, /* tp_methods */

Mtorch_npu/csrc/npu/Stream.cpp+3-1

	@@ -76,6 +76,8 @@ static PyObject *THNPStream_pynew(
76	static void THNPStream_dealloc(THNPStream *self)	76	static void THNPStream_dealloc(THNPStream *self)
77	{	77	{
78	self->npu_stream.~NPUStream();	78	self->npu_stream.~NPUStream();
		79	+ PyObject_ClearWeakRefs((PyObject*)self);
		80	+ Py_CLEAR(self->context);
79	Py_TYPE(self)->tp_free((PyObject*)self);	81	Py_TYPE(self)->tp_free((PyObject*)self);
80	}	82	}
81		83
	@@ -192,7 +194,7 @@ PyTypeObject THNPStreamType = {
192	nullptr, /* tp_traverse */	194	nullptr, /* tp_traverse */
193	nullptr, /* tp_clear */	195	nullptr, /* tp_clear */
194	nullptr, /* tp_richcompare */	196	nullptr, /* tp_richcompare */
195	- 0, /* tp_weaklistoffset */	197	+ 0, /* tp_weaklistoffset (inherited from THPStreamType via tp_base) */
196	nullptr, /* tp_iter */	198	nullptr, /* tp_iter */
197	nullptr, /* tp_iternext */	199	nullptr, /* tp_iternext */
198	THNPStream_methods, /* tp_methods */	200	THNPStream_methods, /* tp_methods */