[Reland] Add wrappers for synchronous GPUDirect Storage APIs (#133489)

Reland #130633

USE_CUFILE turned off by default in this version
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133489
Approved by: https://github.com/albanD
This commit is contained in:
Mikayla Gawarecki
2024-08-15 17:11:52 +00:00
committed by PyTorch MergeBot
parent c23dceb8f1
commit 018e48c337
19 changed files with 391 additions and 1 deletions

View File

@ -35,6 +35,7 @@
#include <torch/csrc/CudaIPCTypes.h>
#include <torch/csrc/Generator.h>
#include <torch/csrc/cuda/CUDAPluggableAllocator.h>
#include <torch/csrc/cuda/GdsFile.h>
#include <torch/csrc/cuda/THCP.h>
#include <torch/csrc/cuda/memory_snapshot.h>
#include <torch/csrc/cuda/python_comm.h>
@ -1963,6 +1964,7 @@ namespace shared {
void initCudartBindings(PyObject* module);
void initNvtxBindings(PyObject* module);
void initGdsBindings(PyObject* module);
#if defined(USE_CUDNN) || defined(USE_ROCM)
void initCudnnBindings(PyObject* module);
#endif
@ -1978,6 +1980,7 @@ void initModule(PyObject* module) {
#if defined(USE_CUDNN) || defined(USE_ROCM)
shared::initCudnnBindings(module);
#endif
shared::initGdsBindings(module);
registerCudaDeviceProperties(module);
registerCudaPluggableAllocator(module);
}