wehub-resource-sync
269 lines
8.5 KiB
C++
269 lines
8.5 KiB
C++
/**
|
|
* Copyright (c) 2016 by Contributors
|
|
* @file dgl/runtime/device_api.h
|
|
* @brief Abstract device memory management API
|
|
*/
|
|
#ifndef DGL_RUNTIME_DEVICE_API_H_
|
|
#define DGL_RUNTIME_DEVICE_API_H_
|
|
|
|
#include <string>
|
|
|
|
#include "c_runtime_api.h"
|
|
#include "packed_func.h"
|
|
|
|
namespace dgl {
|
|
namespace runtime {
|
|
/**
|
|
* @brief the query type into GetAttr
|
|
*/
|
|
enum DeviceAttrKind : int {
|
|
kExist = 0,
|
|
kMaxThreadsPerBlock = 1,
|
|
kWarpSize = 2,
|
|
kMaxSharedMemoryPerBlock = 3,
|
|
kComputeVersion = 4,
|
|
kDeviceName = 5,
|
|
kMaxClockRate = 6,
|
|
kMultiProcessorCount = 7,
|
|
kMaxThreadDimensions = 8
|
|
};
|
|
|
|
/** @brief Number of bytes each allocation must align to */
|
|
constexpr int kAllocAlignment = 64;
|
|
|
|
/** @brief Number of bytes each allocation must align to in temporary allocation
|
|
*/
|
|
constexpr int kTempAllocaAlignment = 64;
|
|
|
|
/** @brief Maximum size that can be allocated on stack */
|
|
constexpr int kMaxStackAlloca = 1024;
|
|
|
|
/**
|
|
* @brief DGL Runtime Device API, abstracts the device
|
|
* specific interface for memory management.
|
|
*/
|
|
class DeviceAPI {
|
|
public:
|
|
/** @brief virtual destructor */
|
|
virtual ~DeviceAPI() {}
|
|
/**
|
|
* @brief Check whether the device is available.
|
|
*/
|
|
virtual bool IsAvailable() { return true; }
|
|
|
|
/**
|
|
* @brief Set the environment device id to ctx
|
|
* @param ctx The context to be set.
|
|
*/
|
|
virtual void SetDevice(DGLContext ctx) = 0;
|
|
|
|
/**
|
|
* @brief Get attribute of specified device.
|
|
* @param ctx The device context
|
|
* @param kind The result kind
|
|
* @param rv The return value.
|
|
* @sa DeviceAttrKind
|
|
*/
|
|
virtual void GetAttr(
|
|
DGLContext ctx, DeviceAttrKind kind, DGLRetValue* rv) = 0;
|
|
|
|
/**
|
|
* @brief Allocate a data space on device.
|
|
* @param ctx The device context to perform operation.
|
|
* @param nbytes The number of bytes in memory.
|
|
* @param alignment The alignment of the memory.
|
|
* @param type_hint The type of elements. Only needed by certain backends such
|
|
* as OpenGL, as nbytes & alignment are sufficient for most backends.
|
|
* @return The allocated device pointer.
|
|
*/
|
|
virtual void* AllocDataSpace(
|
|
DGLContext ctx, size_t nbytes, size_t alignment,
|
|
DGLDataType type_hint) = 0;
|
|
|
|
/**
|
|
* @brief Free a data space on device.
|
|
* @param ctx The device context to perform operation.
|
|
* @param ptr The data space.
|
|
*/
|
|
virtual void FreeDataSpace(DGLContext ctx, void* ptr) = 0;
|
|
|
|
/**
|
|
* @brief copy data from one place to another
|
|
* @param from The source array.
|
|
* @param from_offset The byte offeset in the from.
|
|
* @param to The target array.
|
|
* @param to_offset The byte offset in the to.
|
|
* @param num_bytes The size of the memory in bytes.
|
|
* @param ctx_from The source context.
|
|
* @param ctx_to The target context.
|
|
* @param type_hint The type of elements, only needed by certain backends,
|
|
* can be useful for cross device endian converison.
|
|
*/
|
|
virtual void CopyDataFromTo(
|
|
const void* from, size_t from_offset, void* to, size_t to_offset,
|
|
size_t num_bytes, DGLContext ctx_from, DGLContext ctx_to,
|
|
DGLDataType type_hint) = 0;
|
|
|
|
/**
|
|
* @brief copy data between device and CPU while recording the event.
|
|
* @param from The source array.
|
|
* @param from_offset The byte offeset in the from.
|
|
* @param to The target array.
|
|
* @param to_offset The byte offset in the to.
|
|
* @param num_bytes The size of the memory in bytes.
|
|
* @param ctx_from The source context.
|
|
* @param ctx_to The target context.
|
|
* @param type_hint The type of elements, only needed by certain backends,
|
|
* can be useful for cross device endian converison.
|
|
* @param pytorch_ctx The context pointer from PyTorch's CachingHostAllocator.
|
|
* @note This function only works when PyTorch CachingHostAllocator is
|
|
* available.
|
|
*/
|
|
virtual void RecordedCopyDataFromTo(
|
|
void* from, size_t from_offset, void* to, size_t to_offset,
|
|
size_t num_bytes, DGLContext ctx_from, DGLContext ctx_to,
|
|
DGLDataType type_hint, void* pytorch_ctx) = 0;
|
|
|
|
/**
|
|
* @brief Create a new stream of execution.
|
|
*
|
|
* @param ctx The context of allocation.
|
|
*/
|
|
DGL_DLL virtual DGLStreamHandle CreateStream(DGLContext ctx);
|
|
|
|
/**
|
|
* @brief Free a stream of execution
|
|
*
|
|
* @param ctx The context of the stream
|
|
* @param stream The pointer to be freed.
|
|
*/
|
|
DGL_DLL virtual void FreeStream(DGLContext ctx, DGLStreamHandle stream);
|
|
|
|
/**
|
|
* @brief Synchronize the stream
|
|
* @param ctx The context to perform operation.
|
|
* @param stream The stream to be sync.
|
|
*/
|
|
virtual void StreamSync(DGLContext ctx, DGLStreamHandle stream) = 0;
|
|
|
|
/**
|
|
* @brief Set the stream
|
|
* @param ctx The context to set stream.
|
|
* @param stream The stream to be set.
|
|
*/
|
|
virtual void SetStream(DGLContext ctx, DGLStreamHandle stream) {}
|
|
|
|
/**
|
|
* @brief Get the stream
|
|
*/
|
|
virtual DGLStreamHandle GetStream() const { return nullptr; }
|
|
|
|
/**
|
|
* @brief Synchronize 2 streams of execution.
|
|
*
|
|
* An event is created in event_src stream that the second then
|
|
* stream waits on. Neither event_src or event_dst need to be of
|
|
* the same device ID as the context, but they must be of the same
|
|
* device type.
|
|
*
|
|
* @param ctx The context of the streams.
|
|
* @param event_src The source stream to synchronize.
|
|
* @param event_dst The destination stream to synchronize.
|
|
*/
|
|
DGL_DLL virtual void SyncStreamFromTo(
|
|
DGLContext ctx, DGLStreamHandle event_src, DGLStreamHandle event_dst);
|
|
|
|
/**
|
|
* @brief Pin host memory using cudaHostRegister().
|
|
*
|
|
* @param ptr The host memory pointer to be pinned.
|
|
* @param nbytes The size to be pinned.
|
|
* @return false when pinning an empty tensor. true otherwise.
|
|
*/
|
|
DGL_DLL virtual bool PinData(void* ptr, size_t nbytes);
|
|
|
|
/**
|
|
* @brief Unpin host memory using cudaHostUnregister().
|
|
*
|
|
* @param ptr The host memory pointer to be unpinned.
|
|
*/
|
|
DGL_DLL virtual void UnpinData(void* ptr);
|
|
|
|
/**
|
|
* @brief Allocate the pinned memory using PyTorch CachingHostAllocator.
|
|
*
|
|
* @param nbytes The size to be pinned.
|
|
* @param ctx Pointer to the context pointer from PyTorch's
|
|
* CachingHostAllocator.
|
|
* @param deleter Pointer to the deleter function from PyTorch's
|
|
* CachingHostAllocator.
|
|
*/
|
|
DGL_DLL virtual void* AllocPinnedDataSpace(
|
|
size_t nbytes, void** ctx, void** deleter);
|
|
|
|
/**
|
|
* @brief 'Deallocate' the pinned memory from PyTorch CachingHostAllocator.
|
|
* @note It avoids unnecessary cudaFreeHost calls and puts the memory
|
|
* block into CachingHostAllocator's free list.
|
|
* @param deleter Pointer to the deleter function from PyTorch's
|
|
* CachingHostAllocator.
|
|
*/
|
|
DGL_DLL virtual void FreePinnedDataSpace(void** deleter);
|
|
|
|
/**
|
|
* @brief Check whether the memory is in pinned memory.
|
|
*/
|
|
DGL_DLL virtual bool IsPinned(const void* ptr) { return false; }
|
|
|
|
/**
|
|
* @brief Allocate temporal workspace for backend execution.
|
|
*
|
|
* \note We have the following assumption about backend temporal
|
|
* workspace allocation, and backend will optimize for such assumption:
|
|
*
|
|
* - Only a few allocation will happen, and space will be released after use.
|
|
* - The release order is usually in reverse order of allocate (stack style).
|
|
* - Repeative pattern of same allocations over different runs.
|
|
* - Workspace should not overlap between different threads(i.e. be
|
|
* threadlocal)
|
|
*
|
|
* @param ctx The context of allocation.
|
|
* @param nbytes The size to be allocated.
|
|
* @param type_hint The type of elements. Only needed by certain backends such
|
|
* as OpenGL, as nbytes is sufficient for most backends.
|
|
*/
|
|
DGL_DLL virtual void* AllocWorkspace(
|
|
DGLContext ctx, size_t nbytes, DGLDataType type_hint = {});
|
|
|
|
/**
|
|
* @brief Free temporal workspace in backend execution.
|
|
*
|
|
* @param ctx The context of allocation.
|
|
* @param ptr The pointer to be freed.
|
|
*/
|
|
DGL_DLL virtual void FreeWorkspace(DGLContext ctx, void* ptr);
|
|
|
|
/**
|
|
* @brief Get device API based on context.
|
|
* @param ctx The context
|
|
* @param allow_missing Whether allow missing
|
|
* @return The corresponding device API.
|
|
*/
|
|
DGL_DLL static DeviceAPI* Get(DGLContext ctx, bool allow_missing = false);
|
|
|
|
/**
|
|
* @brief Get device API based on device type.
|
|
* @param dev_type The device type
|
|
* @param allow_missing Whether allow missing
|
|
* @return The corresponding device API.
|
|
*/
|
|
DGL_DLL static DeviceAPI* Get(
|
|
DGLDeviceType dev_type, bool allow_missing = false);
|
|
};
|
|
|
|
/** @brief The device type bigger than this is RPC device */
|
|
constexpr int kRPCSessMask = 128;
|
|
} // namespace runtime
|
|
} // namespace dgl
|
|
#endif // DGL_RUNTIME_DEVICE_API_H_
|