mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
7c89ec0f7c
In this PR: - Pool id creation logic is refactored and moved to a MemPool class. `graph_pool_handle()` API now uses `torch.cuda.MemPool()` to get a unique id for a pool. Existing tests should cover this change. - MemPool holds a pointer to a CUDAAllocator as proposed in https://github.com/pytorch/pytorch/issues/124807#issuecomment-2077506997. Tests are added to show usage with CUDAPluggableAllocator. - MemPoolContext API makes a mempool active. Tests are added to show usage of this API. This API will be used in CUDACachingAllocator to route allocations to a user provided allocator. See draft here: https://github.com/pytorch/pytorch/pull/125722/ Pull Request resolved: https://github.com/pytorch/pytorch/pull/131152 Approved by: https://github.com/eqy, https://github.com/ezyang
22 lines
854 B
C++
22 lines
854 B
C++
#include <torch/csrc/python_headers.h>
|
|
|
|
#include <torch/csrc/jit/python/pybind_utils.h>
|
|
#include <torch/csrc/utils/pybind.h>
|
|
|
|
#include <c10/cuda/CUDACachingAllocator.h>
|
|
|
|
template <typename T>
|
|
using shared_ptr_class_ = py::class_<T, std::shared_ptr<T>>;
|
|
|
|
void THCPMemPool_init(PyObject* module) {
|
|
auto torch_C_m = py::handle(module).cast<py::module>();
|
|
shared_ptr_class_<::c10::cuda::MemPool>(torch_C_m, "_MemPool")
|
|
.def(py::init<c10::cuda::CUDACachingAllocator::CUDAAllocator*, bool>())
|
|
.def_property_readonly("id", &::c10::cuda::MemPool::id)
|
|
.def_property_readonly("allocator", &::c10::cuda::MemPool::allocator);
|
|
shared_ptr_class_<::c10::cuda::MemPoolContext>(torch_C_m, "_MemPoolContext")
|
|
.def(py::init<c10::cuda::MemPool*>())
|
|
.def_static(
|
|
"active_pool", &::c10::cuda::MemPoolContext::getActiveMemPool);
|
|
}
|