wehub-resource-sync
282 lines
9.4 KiB
C++
282 lines
9.4 KiB
C++
/**
|
|
* Copyright (c) 2020-2022 by Contributors
|
|
* @file array/tensordispatch.h
|
|
* @brief This file defines the dispatcher of tensor operators to
|
|
* framework-specific implementations.
|
|
*
|
|
* The dispatcher consists of a TensorDispatcher singleton in DGL C library and
|
|
* one separately-built shared library per supported backend.
|
|
*
|
|
* Those shared libraries contain wrappers of the framework-specific operators.
|
|
* The wrappers are defined with extern "C", meaning that the C++ compiler will
|
|
* not do name mangling for those functions so that DGL can conveniently locate
|
|
* them using dlsym(3) (or GetProcAddress in Windows).
|
|
*
|
|
* The TensorDispatcher singleton maintains a mapping from an array operator to
|
|
* the address of the corresponding symbol in the shared library. During
|
|
* initialization, the TensorDispatcher checks which backend DGL is using.
|
|
* It then locates and opens the corresponding shared library using dlopen(3)
|
|
* (or LoadLibrary in Windows), and populates the said mapping above with
|
|
* dlsym(3) (or GetProcAddress in Windows).
|
|
*
|
|
* A tensor operator in TensorDispatcher first checks whether the corresponding
|
|
* symbol address is found in the mapping. If so, it calls the function located
|
|
* at the symbol address instead, allocate/free pieces of memory on CPU/GPU. If
|
|
* not, it falls back to DeviceAPI::AllocWorkspace/FreeWorkspace.
|
|
*/
|
|
|
|
#ifndef DGL_RUNTIME_TENSORDISPATCH_H_
|
|
#define DGL_RUNTIME_TENSORDISPATCH_H_
|
|
|
|
#include <stddef.h>
|
|
#include <tensoradapter.h>
|
|
#if defined(WIN32) || defined(_WIN32)
|
|
#include <windows.h>
|
|
#endif // WIN32
|
|
#ifdef DGL_USE_CUDA
|
|
#include <cuda_runtime.h>
|
|
#endif // DGL_USE_CUDA
|
|
#include "ndarray.h"
|
|
|
|
/**
|
|
* @brief Casts a pointer \c entry to a function pointer with signature of \c
|
|
* func.
|
|
*/
|
|
#define FUNCCAST(func, entry) (*reinterpret_cast<decltype(&(func))>(entry))
|
|
|
|
namespace dgl {
|
|
namespace runtime {
|
|
|
|
/**
|
|
* @brief Dispatcher that delegates the function calls to framework-specific C++
|
|
* APIs.
|
|
*
|
|
* This class is not thread-safe.
|
|
*/
|
|
class TensorDispatcher {
|
|
public:
|
|
/** @brief Get the singleton instance. */
|
|
static TensorDispatcher* Global() {
|
|
static TensorDispatcher inst;
|
|
return &inst;
|
|
}
|
|
|
|
/** @brief Whether an adapter library is available. */
|
|
inline bool IsAvailable() { return available_; }
|
|
|
|
/** @brief Load symbols from the given tensor adapter library path. */
|
|
bool Load(const char* path_cstr);
|
|
|
|
/**
|
|
* @brief Allocate a piece of CPU memory via PyTorch's CPUAllocator.
|
|
* Used in CPUDeviceAPI::AllocWorkspace().
|
|
*
|
|
* @param nbytes The size to be allocated.
|
|
* @return Pointer to the allocated memory.
|
|
*/
|
|
inline void* CPUAllocWorkspace(size_t nbytes) {
|
|
auto entry = entrypoints_[Op::kCPURawAlloc];
|
|
return FUNCCAST(tensoradapter::CPURawAlloc, entry)(nbytes);
|
|
}
|
|
|
|
/**
|
|
* @brief Free the CPU memory.
|
|
* Used in CPUDeviceAPI::FreeWorkspace().
|
|
*
|
|
* @param ptr Pointer to the memory to be freed.
|
|
*/
|
|
inline void CPUFreeWorkspace(void* ptr) {
|
|
auto entry = entrypoints_[Op::kCPURawDelete];
|
|
FUNCCAST(tensoradapter::CPURawDelete, entry)(ptr);
|
|
}
|
|
|
|
#ifdef DGL_USE_CUDA
|
|
/**
|
|
* @brief Allocate a piece of GPU memory via
|
|
* PyTorch's THCCachingAllocator.
|
|
* Used in CUDADeviceAPI::AllocWorkspace().
|
|
*
|
|
* @note THCCachingAllocator specify the device to allocate on
|
|
* via cudaGetDevice(). Make sure to call cudaSetDevice()
|
|
* before invoking this function.
|
|
*
|
|
* @param nbytes The size to be allocated.
|
|
* @param stream The stream to be allocated on.
|
|
* @return Pointer to the allocated memory.
|
|
*/
|
|
inline void* CUDAAllocWorkspace(size_t nbytes, cudaStream_t stream) {
|
|
auto entry = entrypoints_[Op::kCUDARawAlloc];
|
|
return FUNCCAST(tensoradapter::CUDARawAlloc, entry)(nbytes, stream);
|
|
}
|
|
|
|
/**
|
|
* @brief Free the GPU memory.
|
|
* Used in CUDADeviceAPI::FreeWorkspace().
|
|
*
|
|
* @param ptr Pointer to the memory to be freed.
|
|
*/
|
|
inline void CUDAFreeWorkspace(void* ptr) {
|
|
auto entry = entrypoints_[Op::kCUDARawDelete];
|
|
FUNCCAST(tensoradapter::CUDARawDelete, entry)(ptr);
|
|
}
|
|
|
|
/**
|
|
* @brief Find the current PyTorch CUDA stream
|
|
* Used in runtime::getCurrentCUDAStream().
|
|
*
|
|
* @note PyTorch pre-allocates/sets the current CUDA stream
|
|
* on current device via cudaGetDevice(). Make sure to call cudaSetDevice()
|
|
* before invoking this function.
|
|
*
|
|
* @return cudaStream_t stream handle
|
|
*/
|
|
inline cudaStream_t CUDAGetCurrentStream() {
|
|
auto entry = entrypoints_[Op::kCUDACurrentStream];
|
|
return FUNCCAST(tensoradapter::CUDACurrentStream, entry)();
|
|
}
|
|
|
|
/**
|
|
* @brief Allocate a piece of pinned CPU memory via PyTorch
|
|
* CachingHostAllocator.
|
|
* @note Used in CUDADeviceAPI::AllocPinnedDataSpace().
|
|
* @param nbytes The size to be allocated.
|
|
* @param ctx Pointer to the PyTorch storage ctx ptr returned from the
|
|
* allocator.
|
|
* @param deleter Pointer to the delete function ptr returned from the
|
|
* allocator.
|
|
* @return Raw pointer to the allocated memory.
|
|
*/
|
|
inline void* CUDAAllocHostWorkspace(
|
|
size_t nbytes, void** ctx, void** deleter) {
|
|
auto entry = entrypoints_[Op::kCUDARawHostAlloc];
|
|
|
|
auto alloc_func = FUNCCAST(tensoradapter::CUDARawHostAlloc, entry);
|
|
return alloc_func(nbytes, ctx, deleter);
|
|
}
|
|
|
|
/**
|
|
* @brief Insert the pinned memory block (allocated via PyTorch
|
|
* CachingHostAllocator) back to the free list for future usage.(ref:
|
|
* pytorch/pytorch/blob/master/aten/src/ATen/cuda/CachingHostAllocator.cpp).
|
|
* @note Used in CUDADeviceAPI::FreePinnedDataSpace().
|
|
* @param deleter Pointer to the delete function ptr returned from the
|
|
* allocator.
|
|
*/
|
|
inline void CUDAFreeHostWorkspace(void** deleter) {
|
|
auto entry = entrypoints_[Op::kCUDARawHostDelete];
|
|
FUNCCAST(tensoradapter::CUDARawHostDelete, entry)(deleter);
|
|
}
|
|
|
|
/**
|
|
* @brief Invoke the record_event function call from PyTorch
|
|
* CachingHostAllocator.
|
|
* @note This function assoicates a CUDA stream (used by a copy kernel) to the
|
|
* pinned data. In the free path of this data, which is achieved by
|
|
* calling CUDAFreeHostWorkspace, the set of associated streams is then
|
|
* consumed to ensure proper functionlity. (ref:
|
|
* pytorch/pytorch/blob/master/aten/src/ATen/cuda/CachingHostAllocator.cpp).
|
|
* Used in CUDADeviceAPI::RecordedCopyDataFromTo().
|
|
*
|
|
* @param data Pointer of the tensor to be recorded.
|
|
* @param ctx PyTorch storage ctx ptr returned from the allocator.
|
|
* @param stream The stream that currently consumes this tensor.
|
|
* @param device_id Device of the tensor.
|
|
*/
|
|
inline void CUDARecordHostAlloc(
|
|
void* data, void* ctx, cudaStream_t stream, int device_id) {
|
|
auto entry = entrypoints_[Op::kCUDARecordHostAlloc];
|
|
auto recorded_alloc = FUNCCAST(tensoradapter::CUDARecordHostAlloc, entry);
|
|
recorded_alloc(data, ctx, stream, device_id);
|
|
}
|
|
|
|
/**
|
|
* @brief Release cached pinned memory allocations via cudaHostFree.
|
|
* @note Used in CUDADeviceAPI::PinData() before pinning any host memory by
|
|
* DGL.
|
|
*/
|
|
inline void CUDAHostAllocatorEmptyCache() {
|
|
auto entry = entrypoints_[Op::kCUDAHostAllocatorEmptyCache];
|
|
FUNCCAST(tensoradapter::CUDAHostAllocatorEmptyCache, entry)();
|
|
}
|
|
#endif // DGL_USE_CUDA
|
|
|
|
/**
|
|
* @brief Record streams that are using this tensor.
|
|
* Used in NDArray::RecordStream().
|
|
*
|
|
* @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.
|
|
*/
|
|
inline void RecordStream(void* ptr, DGLStreamHandle stream, int device_id) {
|
|
#ifdef DGL_USE_CUDA
|
|
auto entry = entrypoints_[Op::kRecordStream];
|
|
FUNCCAST(tensoradapter::RecordStream, entry)
|
|
(ptr, static_cast<cudaStream_t>(stream), device_id);
|
|
#endif
|
|
}
|
|
|
|
private:
|
|
/** @brief ctor */
|
|
TensorDispatcher() = default;
|
|
/** @brief dtor */
|
|
~TensorDispatcher();
|
|
|
|
/**
|
|
* @brief List of symbols in the adapter library.
|
|
*
|
|
* Must match the functions in tensoradapter/include/tensoradapter.h.
|
|
*/
|
|
static constexpr const char* names_[] = {
|
|
"CPURawAlloc", "CPURawDelete",
|
|
#ifdef DGL_USE_CUDA
|
|
"CUDARawAlloc", "CUDARawDelete",
|
|
"CUDACurrentStream", "RecordStream",
|
|
"CUDARawHostAlloc", "CUDARawHostDelete",
|
|
"CUDARecordHostAlloc", "CUDAHostAllocatorEmptyCache",
|
|
#endif // DGL_USE_CUDA
|
|
};
|
|
|
|
/** @brief Index of each function to the symbol list */
|
|
class Op {
|
|
public:
|
|
static constexpr int kCPURawAlloc = 0;
|
|
static constexpr int kCPURawDelete = 1;
|
|
#ifdef DGL_USE_CUDA
|
|
static constexpr int kCUDARawAlloc = 2;
|
|
static constexpr int kCUDARawDelete = 3;
|
|
static constexpr int kCUDACurrentStream = 4;
|
|
static constexpr int kRecordStream = 5;
|
|
static constexpr int kCUDARawHostAlloc = 6;
|
|
static constexpr int kCUDARawHostDelete = 7;
|
|
static constexpr int kCUDARecordHostAlloc = 8;
|
|
static constexpr int kCUDAHostAllocatorEmptyCache = 9;
|
|
#endif // DGL_USE_CUDA
|
|
};
|
|
|
|
/** @brief Number of functions */
|
|
static constexpr int num_entries_ = sizeof(names_) / sizeof(names_[0]);
|
|
|
|
/** @brief Entrypoints of each function */
|
|
void* entrypoints_[num_entries_] = {
|
|
nullptr, nullptr,
|
|
#ifdef DGL_USE_CUDA
|
|
nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
|
|
#endif // DGL_USE_CUDA
|
|
};
|
|
|
|
bool available_ = false;
|
|
#if defined(WIN32) || defined(_WIN32)
|
|
HINSTANCE handle_;
|
|
#else // !WIN32
|
|
void* handle_;
|
|
#endif // WIN32
|
|
};
|
|
|
|
}; // namespace runtime
|
|
}; // namespace dgl
|
|
|
|
#undef FUNCCAST
|
|
|
|
#endif // DGL_RUNTIME_TENSORDISPATCH_H_
|