dmlc--dgl
Hongzhi (Steve), Chen
bcd3768426
* replace * blabla * balbla * blabla Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
74 行
1.6 KiB
C++
74 行
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_
|