mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
6b5b69a468
Fixes #153571 Summary: 1. Set annotation callback to global to include all threads 2. Only init callbacks when enable == true and callbacks are empty under mutex 3. When enable == false, check if callbacks are present and if so remove them and set handle to 0 under mutex We don't expect memory snapshots to be called from several different threads (almost always called just from main) but we make sure to add thread safety in the off case that users do want to call it from different points of entry Test Plan: Ran basic snapshot and saw that the callbacks were registered properly Reviewed By: ngimel Differential Revision: D74771491 Pull Request resolved: https://github.com/pytorch/pytorch/pull/153839 Approved by: https://github.com/ngimel, https://github.com/Skylion007 |
||
|---|---|---|
| .. | ||
| shared | ||
| comm.cpp | ||
| comm.h | ||
| CUDAPluggableAllocator.cpp | ||
| CUDAPluggableAllocator.h | ||
| device_set.h | ||
| Event.cpp | ||
| Event.h | ||
| GdsFile.cpp | ||
| GdsFile.h | ||
| Graph.cpp | ||
| memory_snapshot.cpp | ||
| memory_snapshot.h | ||
| MemPool.cpp | ||
| Module.cpp | ||
| Module.h | ||
| nccl.cpp | ||
| nccl.h | ||
| python_comm.cpp | ||
| python_comm.h | ||
| python_nccl.cpp | ||
| python_nccl.h | ||
| Stream.cpp | ||
| Stream.h | ||
| THCP.h | ||
| utils.cpp | ||
| utils.h | ||