dmlc--dgl
Xin Yao
960092be02
* add set_stream * add .record_stream for NDArray and HeteroGraph * refactor dgl stream Python APIs * test record_stream * add unit test for record stream * use pytorch's stream * fix lint * fix cpu build * address comments * address comments * add record stream tests for dgl.graph * record frames and update dataloder * add docstring * update frame * add backend check for record_stream * remove CUDAThreadEntry::stream * record stream for newly created formats * fix bug * fix cpp test * fix None c_void_p to c_handle
75 行
1.6 KiB
C++
75 行
1.6 KiB
C++
/*!
|
|
* Copyright (c) 2020-2022 by Contributors
|
|
* \file tensoradapter.h
|
|
* \brief Header file for functions exposed by the adapter library.
|
|
*
|
|
* Functions in this library must be exported with extern "C" so that DGL can locate
|
|
* them with dlsym(3) (or GetProcAddress on Windows).
|
|
*/
|
|
|
|
#ifndef TENSORADAPTER_H_
|
|
#define TENSORADAPTER_H_
|
|
|
|
#ifdef DGL_USE_CUDA
|
|
#include <cuda_runtime.h>
|
|
#endif // DGL_USE_CUDA
|
|
|
|
namespace tensoradapter {
|
|
|
|
extern "C" {
|
|
|
|
/*!
|
|
* \brief Allocate a piece of CPU memory via
|
|
* PyTorch's CPUAllocator
|
|
*
|
|
* \param nbytes The size to be allocated.
|
|
* \return Pointer to the allocated memory.
|
|
*/
|
|
void* CPURawAlloc(size_t nbytes);
|
|
|
|
/*!
|
|
* \brief Free the CPU memory.
|
|
*
|
|
* \param ptr Pointer to the memory to be freed.
|
|
*/
|
|
void CPURawDelete(void* ptr);
|
|
|
|
#ifdef DGL_USE_CUDA
|
|
/*!
|
|
* \brief Allocate a piece of GPU memory via
|
|
* PyTorch's THCCachingAllocator.
|
|
*
|
|
* \param nbytes The size to be allocated.
|
|
* \param stream The stream to be allocated on.
|
|
* \return Pointer to the allocated memory.
|
|
*/
|
|
void* CUDARawAlloc(size_t nbytes, cudaStream_t stream);
|
|
|
|
/*!
|
|
* \brief Free the GPU memory.
|
|
*
|
|
* \param ptr Pointer to the memory to be freed.
|
|
*/
|
|
void CUDARawDelete(void* ptr);
|
|
|
|
/*!
|
|
* \brief Get the current CUDA stream.
|
|
*/
|
|
cudaStream_t CUDACurrentStream();
|
|
|
|
/*!
|
|
* \brief Let the caching allocator know which streams are using this tensor.
|
|
*
|
|
* \param ptr Pointer of the tensor to be recorded.
|
|
* \param stream The stream that is using this tensor.
|
|
* \param device_id Device of the tensor.
|
|
*/
|
|
void RecordStream(void* ptr, cudaStream_t stream, int device_id);
|
|
#endif // DGL_USE_CUDA
|
|
|
|
}
|
|
|
|
}; // namespace tensoradapter
|
|
|
|
#endif // TENSORADAPTER_H_
|