d22c4cc353
MemPool is a separate pool of memory handled by the caching allocator. This PR adds the option let the caching allocator try to use this pool as a last resort instead of OOMing by associating a use_on_oom bool with each MemPool.
Usage:
Users can optionally specify a ``use_on_oom`` bool (which is False by default) during MemPool creation. If true, then the CUDACachingAllocator will be able to use memory in this pool as a last resort instead of OOMing.
```
pool = torch.cuda.MemPool(allocator, use_on_oom=True)
with torch.cuda.use_mem_pool(pool):
a = torch.randn(40 * 1024 * 1024, dtype=torch.uint8, device="cuda")
del a
# at the memory limit, this will succeed by using pool's memory in order to avoid the oom
b = torch.randn(40 * 1024 * 1024, dtype=torch.uint8, device="cuda")
```
Testing:
```
python test/test_cuda.py -k test_mempool_limited_memory_with_allocator
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/151487
Approved by: https://github.com/eqy, https://github.com/syed-ahmed, https://github.com/ngimel
c10/cuda is a core library with CUDA functionality. It is distinguished from c10 in that it links against the CUDA library, but like c10 it doesn't contain any kernels, and consists solely of core functionality that is generally useful when writing CUDA code; for example, C++ wrappers for the CUDA C API.
Important notes for developers. If you want to add files or functionality to this folder, TAKE NOTE. The code in this folder is very special, because on our AMD GPU build, we transpile it into c10/hip to provide a ROCm environment. Thus, if you write:
// c10/cuda/CUDAFoo.h
namespace c10 { namespace cuda {
void my_func();
}}
this will get transpiled into:
// c10/hip/HIPFoo.h
namespace c10 { namespace hip {
void my_func();
}}
Thus, if you add new functionality to c10, you must also update C10_MAPPINGS
torch/utils/hipify/cuda_to_hip_mappings.py to transpile
occurrences of cuda::my_func to hip::my_func. (At the moment,
we do NOT have a catch all cuda:: to hip:: namespace conversion,
as not all cuda namespaces are converted to hip::, even though
c10's are.)
Transpilation inside this folder is controlled by CAFFE2_SPECIFIC_MAPPINGS
(oddly enough.) C10_MAPPINGS apply to ALL source files.
If you add a new directory to this folder, you MUST update both c10/cuda/CMakeLists.txt and c10/hip/CMakeLists.txt