pytorch/torch/cuda
Syed Tousif Ahmed 1637a40796 Adds snapshot API for MemPools to get pool memory segments (#133601)
Canonically, the snapshot API returns the entire memory state of the CUDACachingAllocator (using `get_all_blocks`). There is no API that can only return the memory state of a given pool.

In this PR, we extend the functionality of snapshot API such that it can only return the memory addresses of an active pool. When snapshot API is called under a MemPoolContext, we only return the blocks that correspond to the pool id of the active pool.

Part of https://github.com/pytorch/pytorch/issues/124807.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/133601
Approved by: https://github.com/ezyang
2024-10-29 01:01:47 +00:00
..
amp
__init__.py Remove unused Python variables in torch/[b-z]* (#136963) 2024-10-19 16:45:22 +00:00
_gpu_trace.py
_memory_viz.py Remove unused Python variables in torch/[b-z]* (#136963) 2024-10-19 16:45:22 +00:00
_sanitizer.py Add proper handling for view and factory function for csan (#138236) 2024-10-18 14:04:18 +00:00
_utils.py
comm.py
error.py
gds.py [Reland] Add wrappers for synchronous GPUDirect Storage APIs (#133489) 2024-08-15 17:11:52 +00:00
graphs.py Remove unused Python variables in torch/[b-z]* (#136963) 2024-10-19 16:45:22 +00:00
jiterator.py
memory.py Adds snapshot API for MemPools to get pool memory segments (#133601) 2024-10-29 01:01:47 +00:00
nccl.py
nvtx.py
profiler.py
random.py
sparse.py
streams.py Use torch.Stream&torch.Event for Dynamo capature (#134850) 2024-10-02 14:15:33 +00:00
tunable.py [ROCm] Tunableop record untuned (#128813) 2024-10-09 21:59:03 +00:00