chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,68 @@
|
||||
/**
|
||||
* Copyright (c) 2023 by Contributors
|
||||
* @file dgl/runtime/ndarray.h
|
||||
* @brief BFloat16 CPU header
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_BFLOAT16_H_
|
||||
#define DGL_RUNTIME_BFLOAT16_H_
|
||||
|
||||
#include <cmath>
|
||||
|
||||
class BFloat16 {
|
||||
uint16_t val;
|
||||
|
||||
public:
|
||||
constexpr BFloat16() : val(0) {}
|
||||
// Disable lint "explicit" warning, since implicit usage on constructor is
|
||||
// expected.
|
||||
BFloat16(float f) { // NOLINT
|
||||
if (std::isnan(f)) {
|
||||
val = 0x7FC0;
|
||||
} else {
|
||||
union {
|
||||
uint16_t iraw16[2];
|
||||
uint32_t iraw32;
|
||||
float f32;
|
||||
};
|
||||
|
||||
f32 = f;
|
||||
const uint32_t rounding_bias = 0x00007FFF + (iraw16[1] & 0x1);
|
||||
val = static_cast<uint16_t>((iraw32 + rounding_bias) >> 16);
|
||||
}
|
||||
}
|
||||
static constexpr BFloat16 Min() {
|
||||
BFloat16 min;
|
||||
min.val = 0xFF80;
|
||||
return min;
|
||||
}
|
||||
|
||||
static constexpr BFloat16 Max() {
|
||||
BFloat16 max;
|
||||
max.val = 0x7F80;
|
||||
return max;
|
||||
}
|
||||
|
||||
BFloat16& operator-=(const float& rhs) {
|
||||
float lhs = (*this);
|
||||
(*this) = lhs - rhs;
|
||||
return *this;
|
||||
}
|
||||
|
||||
BFloat16& operator+=(const float& rhs) {
|
||||
float lhs = (*this);
|
||||
(*this) = lhs + rhs;
|
||||
return *this;
|
||||
}
|
||||
|
||||
operator float() const {
|
||||
union {
|
||||
float f;
|
||||
uint16_t raw[2];
|
||||
};
|
||||
raw[0] = 0;
|
||||
raw[1] = val;
|
||||
return f;
|
||||
}
|
||||
};
|
||||
|
||||
#endif // DGL_RUNTIME_BFLOAT16_H_
|
||||
@@ -0,0 +1,131 @@
|
||||
/**
|
||||
* Copyright (c) 2017 by Contributors
|
||||
* @file dgl/runtime/c_backend_api.h
|
||||
* @brief DGL runtime backend API.
|
||||
*
|
||||
* The functions defined in this header are intended to be
|
||||
* used by compiled dgl operators, usually user do not need to use these
|
||||
* function directly.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_C_BACKEND_API_H_
|
||||
#define DGL_RUNTIME_C_BACKEND_API_H_
|
||||
|
||||
#include "c_runtime_api.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
// Backend related functions.
|
||||
/**
|
||||
* @brief Backend function for modules to get function
|
||||
* from its environment mod_node (its imports and global function).
|
||||
* The user do should not call DGLFuncFree on func.
|
||||
*
|
||||
* @param mod_node The module handle.
|
||||
* @param func_name The name of the function.
|
||||
* @param out The result function.
|
||||
* @return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLBackendGetFuncFromEnv(
|
||||
void* mod_node, const char* func_name, DGLFunctionHandle* out);
|
||||
/**
|
||||
* @brief Backend function to register system-wide library symbol.
|
||||
*
|
||||
* @param name The name of the symbol
|
||||
* @param ptr The symbol address.
|
||||
* @return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLBackendRegisterSystemLibSymbol(const char* name, void* ptr);
|
||||
|
||||
/**
|
||||
* @brief Backend function to allocate temporal workspace.
|
||||
*
|
||||
* @note The result allocate spaced is ensured to be aligned to
|
||||
* kTempAllocaAlignment.
|
||||
*
|
||||
* @param nbytes The size of the space requested.
|
||||
* @param device_type The device type which the space will be allocated.
|
||||
* @param device_id The device id which the space will be allocated.
|
||||
* @param dtype_code_hint The type code of the array elements. Only used in
|
||||
* certain backends such as OpenGL.
|
||||
* @param dtype_bits_hint The type bits of the array elements. Only used in
|
||||
* certain backends such as OpenGL.
|
||||
* @return nullptr when error is thrown, a valid ptr if success
|
||||
*/
|
||||
DGL_DLL void* DGLBackendAllocWorkspace(
|
||||
int device_type, int device_id, uint64_t nbytes, int dtype_code_hint,
|
||||
int dtype_bits_hint);
|
||||
|
||||
/**
|
||||
* @brief Backend function to free temporal workspace.
|
||||
*
|
||||
* @param ptr The result allocated space pointer.
|
||||
* @param device_type The device type which the space will be allocated.
|
||||
* @param device_id The device id which the space will be allocated.
|
||||
* @return 0 when no error is thrown, -1 when failure happens
|
||||
*
|
||||
* @sa DGLBackendAllocWorkspace
|
||||
*/
|
||||
DGL_DLL int DGLBackendFreeWorkspace(int device_type, int device_id, void* ptr);
|
||||
|
||||
/**
|
||||
* @brief Environment for DGL parallel task.
|
||||
*/
|
||||
typedef struct {
|
||||
/**
|
||||
* @brief Auxiliary used for synchronization
|
||||
*/
|
||||
void* sync_handle;
|
||||
/** @brief total amount of task */
|
||||
int32_t num_task;
|
||||
} DGLParallelGroupEnv;
|
||||
|
||||
/**
|
||||
* @brief The callback function to execute a parallel lambda
|
||||
* @param task_id the task id of the function.
|
||||
* @param penv The parallel environment backs the execution.
|
||||
* @param cdata The supporting closure data.
|
||||
*/
|
||||
typedef int (*FDGLParallelLambda)(
|
||||
int task_id, DGLParallelGroupEnv* penv, void* cdata);
|
||||
|
||||
/**
|
||||
* @brief Backend function for running parallel jobs.
|
||||
*
|
||||
* @param flambda The parallel function to be launched.
|
||||
* @param cdata The closure data.
|
||||
* @param num_task Number of tasks to launch, can be 0, means launch
|
||||
* with all available threads.
|
||||
*
|
||||
* @return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLBackendParallelLaunch(
|
||||
FDGLParallelLambda flambda, void* cdata, int num_task);
|
||||
|
||||
/**
|
||||
* @brief BSP barrrier between parallel threads
|
||||
* @param task_id the task id of the function.
|
||||
* @param penv The parallel environment backs the execution.
|
||||
* @return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLBackendParallelBarrier(int task_id, DGLParallelGroupEnv* penv);
|
||||
|
||||
/**
|
||||
* @brief Simple static initialization fucntion.
|
||||
* Run f once and set handle to be not null.
|
||||
* This function is mainly used for test purpose.
|
||||
*
|
||||
* @param handle An global address to indicate f
|
||||
* @param f The function to be ran
|
||||
* @param cdata The closure data to pass to the function.
|
||||
* @param nbytes Number of bytes in the closure data.
|
||||
* @return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLBackendRunOnce(
|
||||
void** handle, int (*f)(void*), void* cdata, int nbytes);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // DGL_EXTERN_C
|
||||
#endif
|
||||
#endif // DGL_RUNTIME_C_BACKEND_API_H_
|
||||
@@ -0,0 +1,71 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file dgl/runtime/c_object_api.h
|
||||
*
|
||||
* @brief DGL Object C API, used to extend and prototype new CAPIs.
|
||||
*
|
||||
* @note Most API functions are registerd as PackedFunc and
|
||||
* can be grabbed via DGLFuncGetGlobal
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_C_OBJECT_API_H_
|
||||
#define DGL_RUNTIME_C_OBJECT_API_H_
|
||||
|
||||
#include "./c_runtime_api.h"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
/** @brief handle to object */
|
||||
typedef void* ObjectHandle;
|
||||
|
||||
/**
|
||||
* @brief free the object handle
|
||||
* @param handle The object handle to be freed.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLObjectFree(ObjectHandle handle);
|
||||
|
||||
/**
|
||||
* @brief Convert type key to type index.
|
||||
* @param type_key The key of the type.
|
||||
* @param out_index the corresponding type index.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLObjectTypeKey2Index(const char* type_key, int* out_index);
|
||||
|
||||
/**
|
||||
* @brief Get runtime type index of the object.
|
||||
* @param handle the object handle.
|
||||
* @param out_index the corresponding type index.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLObjectGetTypeIndex(ObjectHandle handle, int* out_index);
|
||||
|
||||
/**
|
||||
* @brief get attributes given key
|
||||
* @param handle The object handle
|
||||
* @param key The attribute name
|
||||
* @param out_value The attribute value
|
||||
* @param out_type_code The type code of the attribute.
|
||||
* @param out_success Whether get is successful.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
* @note API calls always exchanges with type bits=64, lanes=1
|
||||
*/
|
||||
DGL_DLL int DGLObjectGetAttr(
|
||||
ObjectHandle handle, const char* key, DGLValue* out_value,
|
||||
int* out_type_code, int* out_success);
|
||||
|
||||
/**
|
||||
* @brief get attributes names in the object.
|
||||
* @param handle The object handle
|
||||
* @param out_size The number of functions
|
||||
* @param out_array The array of function names.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLObjectListAttrNames(
|
||||
ObjectHandle handle, int* out_size, const char*** out_array);
|
||||
#ifdef __cplusplus
|
||||
} // DGL_EXTERN_C
|
||||
#endif
|
||||
#endif // DGL_RUNTIME_C_OBJECT_API_H_
|
||||
@@ -0,0 +1,629 @@
|
||||
/**
|
||||
* Copyright (c) 2016-2022 by Contributors
|
||||
* @file dgl/runtime/c_runtime_api.h
|
||||
* @brief DGL runtime library.
|
||||
*
|
||||
* This runtime is adapted from TVM project (commit: 2ce5277)
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_C_RUNTIME_API_H_
|
||||
#define DGL_RUNTIME_C_RUNTIME_API_H_
|
||||
|
||||
// Macros to do weak linking
|
||||
#ifdef _MSC_VER
|
||||
#define DGL_WEAK __declspec(selectany)
|
||||
#else
|
||||
#define DGL_WEAK __attribute__((weak))
|
||||
#endif
|
||||
|
||||
#ifdef __EMSCRIPTEN__
|
||||
#include <emscripten/emscripten.h>
|
||||
#define DGL_DLL EMSCRIPTEN_KEEPALIVE
|
||||
#endif
|
||||
|
||||
#ifndef DGL_DLL
|
||||
#ifdef _WIN32
|
||||
#ifdef DGL_EXPORTS
|
||||
#define DGL_DLL __declspec(dllexport)
|
||||
#else
|
||||
#define DGL_DLL __declspec(dllimport)
|
||||
#endif
|
||||
#else
|
||||
#define DGL_DLL
|
||||
#endif
|
||||
#endif
|
||||
|
||||
// DGL version
|
||||
#define DGL_VERSION "2.5"
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
|
||||
/** @brief type of array index. */
|
||||
typedef int64_t dgl_index_t;
|
||||
|
||||
/**
|
||||
* @brief The device type in DGLContext.
|
||||
*/
|
||||
#ifdef __cplusplus
|
||||
typedef enum : int32_t {
|
||||
#else
|
||||
typedef enum {
|
||||
#endif
|
||||
/** @brief CPU device */
|
||||
kDGLCPU = 1,
|
||||
/** @brief CUDA GPU device */
|
||||
kDGLCUDA = 2,
|
||||
// add more devices once supported
|
||||
} DGLDeviceType;
|
||||
|
||||
/**
|
||||
* @brief The object type code is used in DGL FFI to indicate the types of
|
||||
* objects passed between C and Python.
|
||||
*/
|
||||
typedef enum {
|
||||
kObjectInt = 0U,
|
||||
kObjectUInt = 1U,
|
||||
kObjectFloat = 2U,
|
||||
kHandle = 3U,
|
||||
kNull = 4U,
|
||||
kDGLDataType = 5U,
|
||||
kDGLContext = 6U,
|
||||
kArrayHandle = 7U,
|
||||
kObjectHandle = 8U,
|
||||
kModuleHandle = 9U,
|
||||
kFuncHandle = 10U,
|
||||
kStr = 11U,
|
||||
kBytes = 12U,
|
||||
kNDArrayContainer = 13U,
|
||||
// Extension codes for other frameworks to integrate DGL PackedFunc.
|
||||
// To make sure each framework's id do not conflict, use first and
|
||||
// last sections to mark ranges.
|
||||
// Open an issue at the repo if you need a section of code.
|
||||
kExtBegin = 15U,
|
||||
kNNVMFirst = 16U,
|
||||
kNNVMLast = 20U,
|
||||
// The following section of code is used for non-reserved types.
|
||||
kExtReserveEnd = 64U,
|
||||
kExtEnd = 128U
|
||||
} DGLObjectTypeCode;
|
||||
|
||||
/**
|
||||
* @brief The type code options DGLDataType.
|
||||
*/
|
||||
typedef enum {
|
||||
/** @brief signed integer */
|
||||
kDGLInt = 0U,
|
||||
/** @brief unsigned integer */
|
||||
kDGLUInt = 1U,
|
||||
/** @brief IEEE floating point */
|
||||
kDGLFloat = 2U,
|
||||
/** @brief bfloat16 */
|
||||
kDGLBfloat = 4U,
|
||||
// add more data types if we are going to support them
|
||||
} DGLDataTypeCode;
|
||||
|
||||
/**
|
||||
* @brief The data type the tensor can hold. The data type is assumed to follow
|
||||
* the native endian-ness. An explicit error message should be raised when
|
||||
* attempting to export an array with non-native endianness
|
||||
*
|
||||
* Examples
|
||||
* - float: type_code = 2, bits = 32, lanes=1
|
||||
* - float4(vectorized 4 float): type_code = 2, bits = 32, lanes=4
|
||||
* - int8: type_code = 0, bits = 8, lanes=1
|
||||
*/
|
||||
typedef struct {
|
||||
/**
|
||||
* @brief Type code of base types.
|
||||
* We keep it uint8_t instead of DGLDataTypeCode for minimal memory
|
||||
* footprint, but the value should be one of DGLDataTypeCode enum values.
|
||||
* */
|
||||
uint8_t code;
|
||||
/**
|
||||
* @brief Number of bits, common choices are 8, 16, 32.
|
||||
*/
|
||||
uint8_t bits;
|
||||
/** @brief Number of lanes in the type, used for vector types. */
|
||||
uint16_t lanes;
|
||||
} DGLDataType;
|
||||
|
||||
/**
|
||||
* @brief The Device information, abstract away common device types.
|
||||
*/
|
||||
typedef struct {
|
||||
/** @brief The device type used in the device. */
|
||||
DGLDeviceType device_type;
|
||||
/**
|
||||
* @brief The device index.
|
||||
* For vanilla CPU memory, pinned memory, or managed memory, this is set to 0.
|
||||
*/
|
||||
int32_t device_id;
|
||||
} DGLContext;
|
||||
|
||||
/**
|
||||
* @brief The tensor array stucture to DGL API.
|
||||
* The structure is heavily inspired by DLTensor from DLPack.
|
||||
*/
|
||||
typedef struct {
|
||||
/**
|
||||
* @brief The data pointer points to the allocated data.
|
||||
*
|
||||
* Depending on the device context, it can be a CPU pointer, or a CUDA
|
||||
* device pointer or acl_mem handle in OpenCL.
|
||||
* This pointer is always aligned to 256 bytes as in CUDA. Use the
|
||||
* `byte_offset` field to mark the beginning of the actual data (if the
|
||||
* address is not 256 byte aligned).
|
||||
*
|
||||
* Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow,
|
||||
* TVM, perhaps others) do not adhere to this 256 byte alignment requirement
|
||||
* on CPU/CUDA/ROCm, and always use `byte_offset=0`. This is likely to be
|
||||
* fixed in the future; at the moment it is recommended
|
||||
* to not rely on the data pointer being correctly aligned.
|
||||
*
|
||||
* For a DGLArray, the size of memory required to store the contents of
|
||||
* data can be calculated as follows:
|
||||
*
|
||||
* @code{.c}
|
||||
* static inline size_t GetDataSize(const DGLArray* t) {
|
||||
* size_t size = 1;
|
||||
* for (int32_t i = 0; i < t->ndim; ++i) {
|
||||
* size *= t->shape[i];
|
||||
* }
|
||||
* size *= (t->dtype.bits * t->dtype.lanes + 7) / 8;
|
||||
* return size;
|
||||
* }
|
||||
* @endcode
|
||||
*/
|
||||
void* data;
|
||||
/** @brief The device of the tensor */
|
||||
DGLContext ctx;
|
||||
/** @brief Number of dimensions */
|
||||
int32_t ndim;
|
||||
/** @brief The data type of the pointer*/
|
||||
DGLDataType dtype;
|
||||
/** @brief The shape of the tensor */
|
||||
int64_t* shape;
|
||||
/**
|
||||
* @brief strides of the tensor (in number of elements, not bytes)
|
||||
* can be NULL, indicating tensor is compact and row-majored.
|
||||
*/
|
||||
int64_t* strides;
|
||||
/** @brief The offset in bytes to the beginning pointer to data */
|
||||
uint64_t byte_offset;
|
||||
} DGLArray;
|
||||
|
||||
/** @brief the array handle */
|
||||
typedef DGLArray* DGLArrayHandle;
|
||||
|
||||
/**
|
||||
* @brief Union type of values
|
||||
* being passed through API and function calls.
|
||||
*/
|
||||
typedef union {
|
||||
int64_t v_int64;
|
||||
double v_float64;
|
||||
void* v_handle;
|
||||
const char* v_str;
|
||||
DGLDataType v_type;
|
||||
DGLContext v_ctx;
|
||||
} DGLValue;
|
||||
|
||||
/**
|
||||
* @brief Byte array type used to pass in byte array
|
||||
* When kBytes is used as data type.
|
||||
*/
|
||||
typedef struct {
|
||||
const char* data;
|
||||
size_t size;
|
||||
} DGLByteArray;
|
||||
|
||||
/** @brief Handle to DGL runtime modules. */
|
||||
typedef void* DGLModuleHandle;
|
||||
/** @brief Handle to packed function handle. */
|
||||
typedef void* DGLFunctionHandle;
|
||||
/** @brief Handle to hold return value. */
|
||||
typedef void* DGLRetValueHandle;
|
||||
/**
|
||||
* @brief The stream that is specific to device
|
||||
* can be NULL, which indicates the default one.
|
||||
*/
|
||||
typedef void* DGLStreamHandle;
|
||||
|
||||
/**
|
||||
* @brief Used for implementing C API function.
|
||||
* Set last error message before return.
|
||||
* @param msg The error message to be set.
|
||||
*/
|
||||
DGL_DLL void DGLAPISetLastError(const char* msg);
|
||||
|
||||
/**
|
||||
* @brief return str message of the last error
|
||||
* all function in this file will return 0 when success
|
||||
* and -1 when an error occured,
|
||||
* DGLGetLastError can be called to retrieve the error
|
||||
*
|
||||
* this function is threadsafe and can be called by different thread
|
||||
*
|
||||
* @return error info
|
||||
*/
|
||||
DGL_DLL const char* DGLGetLastError(void);
|
||||
/**
|
||||
* @brief Load module from file.
|
||||
* @param file_name The file name to load the module from.
|
||||
* @param format The format of the module.
|
||||
* @param out The result module
|
||||
*
|
||||
* @return 0 when success, -1 when failure happens
|
||||
* @note The resulting module do not contain import relation.
|
||||
* It can be reconstructed by DGLModImport.
|
||||
*/
|
||||
DGL_DLL int DGLModLoadFromFile(
|
||||
const char* file_name, const char* format, DGLModuleHandle* out);
|
||||
|
||||
/**
|
||||
* @brief Add dep to mod's dependency.
|
||||
* This allows functions in this module to use modules.
|
||||
*
|
||||
* @param mod The module handle.
|
||||
* @param dep The dependent module to be imported.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLModImport(DGLModuleHandle mod, DGLModuleHandle dep);
|
||||
|
||||
/**
|
||||
* @brief Get function from the module.
|
||||
* @param mod The module handle.
|
||||
* @param func_name The name of the function.
|
||||
* @param query_imports Whether to query imported modules
|
||||
* @param out The result function, can be NULL if it is not available.
|
||||
* @return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLModGetFunction(
|
||||
DGLModuleHandle mod, const char* func_name, int query_imports,
|
||||
DGLFunctionHandle* out);
|
||||
|
||||
/**
|
||||
* @brief Free front-end extension type resource.
|
||||
* @param handle The extension handle.
|
||||
* @param type_code The type of of the extension type.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLExtTypeFree(void* handle, int type_code);
|
||||
|
||||
/**
|
||||
* @brief Free the Module
|
||||
* @param mod The module to be freed.
|
||||
*
|
||||
* @note This may not free up the module's resources.
|
||||
* If there is active DGLFunctionHandle uses the module
|
||||
* Or if this module is imported by another active module.
|
||||
*
|
||||
* The all functions remains valid until DGLFuncFree is called.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLModFree(DGLModuleHandle mod);
|
||||
|
||||
/**
|
||||
* @brief Free the function when it is no longer needed.
|
||||
* @param func The function handle
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLFuncFree(DGLFunctionHandle func);
|
||||
|
||||
/**
|
||||
* @brief Call a Packed DGL Function.
|
||||
*
|
||||
* @param func node handle of the function.
|
||||
* @param arg_values The arguments
|
||||
* @param type_codes The type codes of the arguments
|
||||
* @param num_args Number of arguments.
|
||||
*
|
||||
* @param ret_val The return value.
|
||||
* @param ret_type_code the type code of return value.
|
||||
*
|
||||
* @return 0 when success, -1 when failure happens
|
||||
* @note DGL calls always exchanges with type bits=64, lanes=1
|
||||
*
|
||||
* @note API calls always exchanges with type bits=64, lanes=1
|
||||
* If API call returns container handles (e.g. FunctionHandle)
|
||||
* these handles should be managed by the front-end.
|
||||
* The front-end need to call free function (e.g. DGLFuncFree)
|
||||
* to free these handles.
|
||||
*/
|
||||
DGL_DLL int DGLFuncCall(
|
||||
DGLFunctionHandle func, DGLValue* arg_values, int* type_codes, int num_args,
|
||||
DGLValue* ret_val, int* ret_type_code);
|
||||
|
||||
/**
|
||||
* @brief Set the return value of DGLPackedCFunc.
|
||||
*
|
||||
* This function is called by DGLPackedCFunc to set the return value.
|
||||
* When this function is not called, the function returns null by default.
|
||||
*
|
||||
* @param ret The return value handle, pass by ret in DGLPackedCFunc
|
||||
* @param value The value to be returned.
|
||||
* @param type_code The type of the value to be returned.
|
||||
* @param num_ret Number of return values, for now only 1 is supported.
|
||||
*/
|
||||
DGL_DLL int DGLCFuncSetReturn(
|
||||
DGLRetValueHandle ret, DGLValue* value, int* type_code, int num_ret);
|
||||
|
||||
/**
|
||||
* @brief Inplace translate callback argument value to return value.
|
||||
* This is only needed for non-POD arguments.
|
||||
*
|
||||
* @param value The value to be translated.
|
||||
* @param code The type code to be translated.
|
||||
* @note This function will do a shallow copy when necessary.
|
||||
*
|
||||
* @return 0 when success, -1 when failure happens.
|
||||
*/
|
||||
DGL_DLL int DGLCbArgToReturn(DGLValue* value, int code);
|
||||
|
||||
/**
|
||||
* @brief C type of packed function.
|
||||
*
|
||||
* @param args The arguments
|
||||
* @param type_codes The type codes of the arguments
|
||||
* @param num_args Number of arguments.
|
||||
* @param ret The return value handle.
|
||||
* @param resource_handle The handle additional resouce handle from fron-end.
|
||||
* @return 0 if success, -1 if failure happens, set error via
|
||||
* DGLAPISetLastError.
|
||||
* @sa DGLCFuncSetReturn
|
||||
*/
|
||||
typedef int (*DGLPackedCFunc)(
|
||||
DGLValue* args, int* type_codes, int num_args, DGLRetValueHandle ret,
|
||||
void* resource_handle);
|
||||
|
||||
/**
|
||||
* @brief C callback to free the resource handle in C packed function.
|
||||
* @param resource_handle The handle additional resouce handle from fron-end.
|
||||
*/
|
||||
typedef void (*DGLPackedCFuncFinalizer)(void* resource_handle);
|
||||
|
||||
/**
|
||||
* @brief Signature for extension function declarer.
|
||||
*
|
||||
* DGL call this function to get the extension functions
|
||||
* The declarer will call register_func to register function and their name.
|
||||
*
|
||||
* @param register_func_handle The register function
|
||||
* @return 0 if success, -1 if failure happens
|
||||
*/
|
||||
typedef int (*DGLExtensionFuncDeclarer)(DGLFunctionHandle register_func_handle);
|
||||
|
||||
/**
|
||||
* @brief Wrap a DGLPackedCFunc to become a FunctionHandle.
|
||||
*
|
||||
* The resource_handle will be managed by DGL API, until the function is no
|
||||
* longer used.
|
||||
*
|
||||
* @param func The packed C function.
|
||||
* @param resource_handle The resource handle from front-end, can be NULL.
|
||||
* @param fin The finalizer on resource handle when the FunctionHandle get
|
||||
* freed, can be NULL.
|
||||
* @param out the result function handle.
|
||||
* @return 0 when success, -1 when failure happens.
|
||||
*/
|
||||
DGL_DLL int DGLFuncCreateFromCFunc(
|
||||
DGLPackedCFunc func, void* resource_handle, DGLPackedCFuncFinalizer fin,
|
||||
DGLFunctionHandle* out);
|
||||
|
||||
/**
|
||||
* @brief Register the function to runtime's global table.
|
||||
*
|
||||
* The registered function then can be pulled by the backend by the name.
|
||||
*
|
||||
* @param name The name of the function.
|
||||
* @param f The function to be registered.
|
||||
* @param override Whether allow override already registered function.
|
||||
*/
|
||||
DGL_DLL int DGLFuncRegisterGlobal(
|
||||
const char* name, DGLFunctionHandle f, int override);
|
||||
|
||||
/**
|
||||
* @brief Get a global function.
|
||||
*
|
||||
* @param name The name of the function.
|
||||
* @param out the result function pointer, NULL if it does not exist.
|
||||
*
|
||||
* @note The function handle of global function is managed by DGL runtime,
|
||||
* So DGLFuncFree is should not be called when it get deleted.
|
||||
*/
|
||||
DGL_DLL int DGLFuncGetGlobal(const char* name, DGLFunctionHandle* out);
|
||||
|
||||
/**
|
||||
* @brief List all the globally registered function name
|
||||
* @param out_size The number of functions
|
||||
* @param out_array The array of function names.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLFuncListGlobalNames(int* out_size, const char*** out_array);
|
||||
|
||||
// Array related apis for quick proptyping
|
||||
/**
|
||||
* @brief Allocate a nd-array's memory,
|
||||
* including space of shape, of given spec.
|
||||
*
|
||||
* @param shape The shape of the array, the data content will be copied to out
|
||||
* @param ndim The number of dimension of the array.
|
||||
* @param dtype_code The type code of the dtype
|
||||
* @param dtype_bits The number of bits of dtype
|
||||
* @param dtype_lanes The number of lanes in the dtype.
|
||||
* @param device_type The device type of context
|
||||
* @param device_id The device id of context.
|
||||
* @param out The output handle.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLArrayAlloc(
|
||||
const dgl_index_t* shape, int ndim, int dtype_code, int dtype_bits,
|
||||
int dtype_lanes, int device_type, int device_id, DGLArrayHandle* out);
|
||||
|
||||
/**
|
||||
* @brief Allocate a nd-array's with shared memory,
|
||||
* including space of shape, of given spec.
|
||||
*
|
||||
* @param the name of the shared memory
|
||||
* @param shape The shape of the array, the data content will be copied to out
|
||||
* @param ndim The number of dimension of the array.
|
||||
* @param dtype_code The type code of the dtype
|
||||
* @param dtype_bits The number of bits of dtype
|
||||
* @param dtype_lanes The number of lanes in the dtype.
|
||||
* @param is_create whether the shared memory is created
|
||||
* @param out The output handle.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
int DGLArrayAllocSharedMem(
|
||||
const char* mem_name, const dgl_index_t* shape, int ndim, int dtype_code,
|
||||
int dtype_bits, int dtype_lanes, bool is_create, DGLArrayHandle* out);
|
||||
|
||||
/**
|
||||
* @brief Free the DGL Array.
|
||||
* @param handle The array handle to be freed.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLArrayFree(DGLArrayHandle handle);
|
||||
|
||||
/**
|
||||
* @brief Copy array data from CPU byte array.
|
||||
* @param handle The array handle.
|
||||
* @param data the data pointer
|
||||
* @param nbytes The number of bytes to copy.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLArrayCopyFromBytes(
|
||||
DGLArrayHandle handle, void* data, size_t nbytes);
|
||||
|
||||
/**
|
||||
* @brief Copy array data to CPU byte array.
|
||||
* @param handle The array handle.
|
||||
* @param data the data pointer
|
||||
* @param nbytes The number of bytes to copy.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLArrayCopyToBytes(
|
||||
DGLArrayHandle handle, void* data, size_t nbytes);
|
||||
|
||||
/**
|
||||
* @brief Copy the array, both from and to must be valid during the copy.
|
||||
* @param from The array to be copied from.
|
||||
* @param to The target space.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLArrayCopyFromTo(DGLArrayHandle from, DGLArrayHandle to);
|
||||
|
||||
/**
|
||||
* @brief Create a new runtime stream.
|
||||
*
|
||||
* @param device_type The device type of context
|
||||
* @param device_id The device id of context
|
||||
* @param out The new stream handle
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLStreamCreate(
|
||||
int device_type, int device_id, DGLStreamHandle* out);
|
||||
|
||||
/**
|
||||
* @brief Free a created stream handle.
|
||||
*
|
||||
* @param device_type The device type of context
|
||||
* @param device_id The device id of context
|
||||
* @param stream The stream to be freed
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLStreamFree(
|
||||
int device_type, int device_id, DGLStreamHandle stream);
|
||||
|
||||
/**
|
||||
* @brief Set the runtime stream of current thread to be stream.
|
||||
* The subsequent calls to the same device_type
|
||||
* will use the setted stream handle.
|
||||
* The specific type of stream is runtime device dependent.
|
||||
*
|
||||
* @param device_type The device type of context
|
||||
* @param device_id The device id of context.
|
||||
* @param handle The stream handle.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLSetStream(
|
||||
int device_type, int device_id, DGLStreamHandle handle);
|
||||
|
||||
/**
|
||||
* @brief Get the runtime stream of current thread.
|
||||
*
|
||||
* @param device_type The device type of context
|
||||
* @param device_id The device id of context.
|
||||
* @param handle The stream handle.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLGetStream(
|
||||
int device_type, int device_id, DGLStreamHandle* handle);
|
||||
|
||||
/**
|
||||
* @brief Wait until all computations on stream completes.
|
||||
*
|
||||
* @param device_type The device type of context
|
||||
* @param device_id The device id of context.
|
||||
* @param stream The stream to be synchronized.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLSynchronize(
|
||||
int device_type, int device_id, DGLStreamHandle stream);
|
||||
|
||||
/**
|
||||
* @brief Synchronize two streams of execution.
|
||||
*
|
||||
* @param device_type The device type of context
|
||||
* @param device_id The device id of context
|
||||
* @param src The source stream to synchronize.
|
||||
* @param dst The destination stream to synchronize.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLStreamStreamSynchronize(
|
||||
int device_type, int device_id, DGLStreamHandle src, DGLStreamHandle dst);
|
||||
|
||||
/**
|
||||
* @brief Load tensor adapter.
|
||||
* @return 0 when success, -1 when failure happens.
|
||||
*/
|
||||
DGL_DLL int DGLLoadTensorAdapter(const char* path);
|
||||
|
||||
/**
|
||||
* @brief Pin host memory.
|
||||
*/
|
||||
int DGLArrayPinData(DGLArrayHandle handle, DGLContext ctx);
|
||||
|
||||
/**
|
||||
* @brief Unpin host memory.
|
||||
*/
|
||||
int DGLArrayUnpinData(DGLArrayHandle handle, DGLContext ctx);
|
||||
|
||||
/**
|
||||
* @brief Record the stream that's using this tensor.
|
||||
*/
|
||||
int DGLArrayRecordStream(DGLArrayHandle handle, DGLStreamHandle stream);
|
||||
|
||||
/**
|
||||
* @brief Bug report macro.
|
||||
*
|
||||
* This serves as a sanity check on system side to make sure the code is correct
|
||||
* by checking whether a condition always holds for complex reasons. Failing
|
||||
* the condition signifies a system bug instead of users giving invalid inputs
|
||||
* or using the functionality incorrectly.
|
||||
*
|
||||
* Hints the user to file a bug report if the condition fails.
|
||||
*/
|
||||
#define BUG_IF_FAIL(cond) \
|
||||
CHECK(cond) \
|
||||
<< "A bug has been occurred. " \
|
||||
"Please file a bug report at https://github.com/dmlc/dgl/issues. " \
|
||||
"Message: "
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // DGL_EXTERN_C
|
||||
#endif
|
||||
#endif // DGL_RUNTIME_C_RUNTIME_API_H_
|
||||
@@ -0,0 +1,32 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file runtime/config.h
|
||||
* @brief DGL runtime config
|
||||
*/
|
||||
|
||||
#ifndef DGL_RUNTIME_CONFIG_H_
|
||||
#define DGL_RUNTIME_CONFIG_H_
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
class Config {
|
||||
public:
|
||||
static Config* Global() {
|
||||
static Config config;
|
||||
return &config;
|
||||
}
|
||||
|
||||
// Enabling or disable use libxsmm for Spmm
|
||||
void EnableLibxsmm(bool);
|
||||
bool IsLibxsmmAvailable() const;
|
||||
|
||||
private:
|
||||
Config();
|
||||
bool libxsmm_;
|
||||
};
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_RUNTIME_CONFIG_H_
|
||||
@@ -0,0 +1,685 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file runtime/container.h
|
||||
* @brief Defines the container object data structures.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_CONTAINER_H_
|
||||
#define DGL_RUNTIME_CONTAINER_H_
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "object.h"
|
||||
#include "packed_func.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
/**
|
||||
* @brief value object.
|
||||
*
|
||||
* It is typically used to wrap a non-Object type to Object type.
|
||||
* Any type that is supported by DGLRetValue is supported by this.
|
||||
*/
|
||||
class ValueObject : public Object {
|
||||
public:
|
||||
/** @brief the value data */
|
||||
DGLRetValue data;
|
||||
|
||||
static constexpr const char* _type_key = "Value";
|
||||
DGL_DECLARE_OBJECT_TYPE_INFO(ValueObject, Object);
|
||||
};
|
||||
|
||||
/** @brief Construct a value object. */
|
||||
template <typename T>
|
||||
inline std::shared_ptr<ValueObject> MakeValue(T&& val) {
|
||||
auto obj = std::make_shared<ValueObject>();
|
||||
obj->data = val;
|
||||
return obj;
|
||||
}
|
||||
|
||||
/** @brief Vallue reference type */
|
||||
class Value : public ObjectRef {
|
||||
public:
|
||||
Value() {}
|
||||
explicit Value(std::shared_ptr<Object> o) : ObjectRef(o) {}
|
||||
|
||||
const ValueObject* operator->() const {
|
||||
return static_cast<const ValueObject*>(obj_.get());
|
||||
}
|
||||
|
||||
using ContainerType = ValueObject;
|
||||
};
|
||||
|
||||
/** @brief list obj content in list */
|
||||
class ListObject : public Object {
|
||||
public:
|
||||
/** @brief the data content */
|
||||
std::vector<std::shared_ptr<Object> > data;
|
||||
|
||||
void VisitAttrs(AttrVisitor* visitor) final {
|
||||
// Visitor to list have no effect.
|
||||
}
|
||||
|
||||
static constexpr const char* _type_key = "List";
|
||||
DGL_DECLARE_OBJECT_TYPE_INFO(ListObject, Object);
|
||||
};
|
||||
|
||||
/** @brief map obj content */
|
||||
class MapObject : public Object {
|
||||
public:
|
||||
void VisitAttrs(AttrVisitor* visitor) final {
|
||||
// Visitor to map have no effect.
|
||||
}
|
||||
// hash function
|
||||
struct Hash {
|
||||
size_t operator()(const std::shared_ptr<Object>& n) const {
|
||||
return std::hash<Object*>()(n.get());
|
||||
}
|
||||
};
|
||||
// comparator
|
||||
struct Equal {
|
||||
bool operator()(
|
||||
const std::shared_ptr<Object>& a,
|
||||
const std::shared_ptr<Object>& b) const {
|
||||
return a.get() == b.get();
|
||||
}
|
||||
};
|
||||
|
||||
/** @brief The corresponding conatiner type */
|
||||
using ContainerType = std::unordered_map<
|
||||
std::shared_ptr<Object>, std::shared_ptr<Object>, Hash, Equal>;
|
||||
|
||||
/** @brief the data content */
|
||||
ContainerType data;
|
||||
|
||||
static constexpr const char* _type_key = "Map";
|
||||
DGL_DECLARE_OBJECT_TYPE_INFO(MapObject, Object);
|
||||
};
|
||||
|
||||
/** @brief specialized map obj with string as key */
|
||||
class StrMapObject : public Object {
|
||||
public:
|
||||
void VisitAttrs(AttrVisitor* visitor) final {
|
||||
// Visitor to map have no effect.
|
||||
}
|
||||
/** @brief The corresponding conatiner type */
|
||||
using ContainerType =
|
||||
std::unordered_map<std::string, std::shared_ptr<Object> >;
|
||||
|
||||
/** @brief the data content */
|
||||
ContainerType data;
|
||||
|
||||
static constexpr const char* _type_key = "StrMap";
|
||||
DGL_DECLARE_OBJECT_TYPE_INFO(StrMapObject, Object);
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief iterator adapter that adapts TIter to return another type.
|
||||
* @tparam Converter a struct that contains converting function
|
||||
* @tparam TIter the content iterator type.
|
||||
*/
|
||||
template <typename Converter, typename TIter>
|
||||
class IterAdapter {
|
||||
public:
|
||||
explicit IterAdapter(TIter iter) : iter_(iter) {}
|
||||
inline IterAdapter& operator++() { // NOLINT(*)
|
||||
++iter_;
|
||||
return *this;
|
||||
}
|
||||
inline IterAdapter& operator++(int) { // NOLINT(*)
|
||||
++iter_;
|
||||
return *this;
|
||||
}
|
||||
inline IterAdapter operator+(int offset) const { // NOLINT(*)
|
||||
return IterAdapter(iter_ + offset);
|
||||
}
|
||||
inline bool operator==(IterAdapter other) const {
|
||||
return iter_ == other.iter_;
|
||||
}
|
||||
inline bool operator!=(IterAdapter other) const { return !(*this == other); }
|
||||
inline const typename Converter::ResultType operator*() const {
|
||||
return Converter::convert(*iter_);
|
||||
}
|
||||
|
||||
private:
|
||||
TIter iter_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief List container of ObjectRef.
|
||||
*
|
||||
* List implements copy on write semantics, which means list is mutable
|
||||
* but copy will happen when list is referenced in more than two places.
|
||||
*
|
||||
* That is said when using this container for runtime arguments or return
|
||||
* values, try use the constructor to create the list at once (for example
|
||||
* from an existing vector).
|
||||
*
|
||||
* operator[] only provide const access, use Set to mutate the content.
|
||||
*
|
||||
* @tparam T The content ObjectRef type.
|
||||
*
|
||||
* @note The element type must subclass \c ObjectRef. Otherwise, the
|
||||
* compiler would throw an error:
|
||||
*
|
||||
* <code>
|
||||
* error: no type named 'type' in 'struct std::enable_if<false, void>'
|
||||
* </code>
|
||||
*
|
||||
* Example:
|
||||
*
|
||||
* <code>
|
||||
* // List<int> list; // fails
|
||||
* // List<NDArray> list2; // fails
|
||||
* List<Value> list; // works
|
||||
* list.push_back(Value(MakeValue(1))); // works
|
||||
* list.push_back(Value(MakeValue(NDArray::Empty(shape, dtype, ctx)))); //
|
||||
* works
|
||||
* </code>
|
||||
*/
|
||||
template <
|
||||
typename T,
|
||||
typename =
|
||||
typename std::enable_if<std::is_base_of<ObjectRef, T>::value>::type>
|
||||
class List : public ObjectRef {
|
||||
public:
|
||||
/**
|
||||
* @brief default constructor
|
||||
*/
|
||||
List() { obj_ = std::make_shared<ListObject>(); }
|
||||
/**
|
||||
* @brief move constructor
|
||||
* @param other source
|
||||
*/
|
||||
List(List<T>&& other) { // NOLINT(*)
|
||||
obj_ = std::move(other.obj_);
|
||||
}
|
||||
/**
|
||||
* @brief copy constructor
|
||||
* @param other source
|
||||
*/
|
||||
List(const List<T>& other) : ObjectRef(other.obj_) { // NOLINT(*)
|
||||
}
|
||||
/**
|
||||
* @brief constructor from pointer
|
||||
* @param n the container pointer
|
||||
*/
|
||||
explicit List(std::shared_ptr<Object> n) : ObjectRef(n) {}
|
||||
/**
|
||||
* @brief constructor from iterator
|
||||
* @param begin begin of iterator
|
||||
* @param end end of iterator
|
||||
* @tparam IterType The type of iterator
|
||||
*/
|
||||
template <typename IterType>
|
||||
List(IterType begin, IterType end) {
|
||||
assign(begin, end);
|
||||
}
|
||||
/**
|
||||
* @brief constructor from initializer list
|
||||
* @param init The initalizer list
|
||||
*/
|
||||
List(std::initializer_list<T> init) { // NOLINT(*)
|
||||
assign(init.begin(), init.end());
|
||||
}
|
||||
/**
|
||||
* @brief constructor from vector
|
||||
* @param init The vector
|
||||
*/
|
||||
List(const std::vector<T>& init) { // NOLINT(*)
|
||||
assign(init.begin(), init.end());
|
||||
}
|
||||
/**
|
||||
* @brief Constructs a container with n elements. Each element is a copy of
|
||||
* val
|
||||
* @param n The size of the container
|
||||
* @param val The init value
|
||||
*/
|
||||
explicit List(size_t n, const T& val) {
|
||||
auto tmp_obj = std::make_shared<ListObject>();
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
tmp_obj->data.push_back(val.obj_);
|
||||
}
|
||||
obj_ = std::move(tmp_obj);
|
||||
}
|
||||
/**
|
||||
* @brief move assign operator
|
||||
* @param other The source of assignment
|
||||
* @return reference to self.
|
||||
*/
|
||||
List<T>& operator=(List<T>&& other) {
|
||||
obj_ = std::move(other.obj_);
|
||||
return *this;
|
||||
}
|
||||
/**
|
||||
* @brief copy assign operator
|
||||
* @param other The source of assignment
|
||||
* @return reference to self.
|
||||
*/
|
||||
List<T>& operator=(const List<T>& other) {
|
||||
obj_ = other.obj_;
|
||||
return *this;
|
||||
}
|
||||
/**
|
||||
* @brief reset the list to content from iterator.
|
||||
* @param begin begin of iterator
|
||||
* @param end end of iterator
|
||||
* @tparam IterType The type of iterator
|
||||
*/
|
||||
template <typename IterType>
|
||||
void assign(IterType begin, IterType end) {
|
||||
auto n = std::make_shared<ListObject>();
|
||||
for (IterType it = begin; it != end; ++it) {
|
||||
n->data.push_back((*it).obj_);
|
||||
}
|
||||
obj_ = std::move(n);
|
||||
}
|
||||
/**
|
||||
* @brief Read i-th element from list.
|
||||
* @param i The index
|
||||
* @return the i-th element.
|
||||
*/
|
||||
inline const T operator[](size_t i) const {
|
||||
return T(static_cast<const ListObject*>(obj_.get())->data[i]);
|
||||
}
|
||||
/** @return The size of the list */
|
||||
inline size_t size() const {
|
||||
if (obj_.get() == nullptr) return 0;
|
||||
return static_cast<const ListObject*>(obj_.get())->data.size();
|
||||
}
|
||||
/**
|
||||
* @brief copy on write semantics
|
||||
* Do nothing if current handle is the unique copy of the list.
|
||||
* Otherwise make a new copy of the list to ensure the current handle
|
||||
* hold a unique copy.
|
||||
*
|
||||
* @return Handle to the internal obj container(which ganrantees to be unique)
|
||||
*/
|
||||
inline ListObject* CopyOnWrite() {
|
||||
if (obj_.get() == nullptr || !obj_.unique()) {
|
||||
obj_ = std::make_shared<ListObject>(
|
||||
*static_cast<const ListObject*>(obj_.get()));
|
||||
}
|
||||
return static_cast<ListObject*>(obj_.get());
|
||||
}
|
||||
/**
|
||||
* @brief push a new item to the back of the list
|
||||
* @param item The item to be pushed.
|
||||
*/
|
||||
inline void push_back(const T& item) {
|
||||
ListObject* n = this->CopyOnWrite();
|
||||
n->data.push_back(item.obj_);
|
||||
}
|
||||
/**
|
||||
* @brief set i-th element of the list.
|
||||
* @param i The index
|
||||
* @param value The value to be setted.
|
||||
*/
|
||||
inline void Set(size_t i, const T& value) {
|
||||
ListObject* n = this->CopyOnWrite();
|
||||
n->data[i] = value.obj_;
|
||||
}
|
||||
/** @return whether list is empty */
|
||||
inline bool empty() const { return size() == 0; }
|
||||
/** @brief Copy the content to a vector */
|
||||
inline std::vector<T> ToVector() const {
|
||||
return std::vector<T>(begin(), end());
|
||||
}
|
||||
/** @brief specify container obj */
|
||||
using ContainerType = ListObject;
|
||||
|
||||
struct Ptr2ObjectRef {
|
||||
using ResultType = T;
|
||||
static inline T convert(const std::shared_ptr<Object>& n) { return T(n); }
|
||||
};
|
||||
using iterator = IterAdapter<
|
||||
Ptr2ObjectRef, std::vector<std::shared_ptr<Object> >::const_iterator>;
|
||||
|
||||
using reverse_iterator = IterAdapter<
|
||||
Ptr2ObjectRef,
|
||||
std::vector<std::shared_ptr<Object> >::const_reverse_iterator>;
|
||||
|
||||
/** @return begin iterator */
|
||||
inline iterator begin() const {
|
||||
return iterator(static_cast<const ListObject*>(obj_.get())->data.begin());
|
||||
}
|
||||
/** @return end iterator */
|
||||
inline iterator end() const {
|
||||
return iterator(static_cast<const ListObject*>(obj_.get())->data.end());
|
||||
}
|
||||
/** @return rbegin iterator */
|
||||
inline reverse_iterator rbegin() const {
|
||||
return reverse_iterator(
|
||||
static_cast<const ListObject*>(obj_.get())->data.rbegin());
|
||||
}
|
||||
/** @return rend iterator */
|
||||
inline reverse_iterator rend() const {
|
||||
return reverse_iterator(
|
||||
static_cast<const ListObject*>(obj_.get())->data.rend());
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Map container of ObjectRef->ObjectRef.
|
||||
*
|
||||
* Map implements copy on write semantics, which means map is mutable
|
||||
* but copy will happen when list is referenced in more than two places.
|
||||
*
|
||||
* That is said when using this container for runtime arguments or return
|
||||
* values, try use the constructor to create it at once (for example
|
||||
* from an existing std::map).
|
||||
*
|
||||
* operator[] only provide const acces, use Set to mutate the content.
|
||||
*
|
||||
* @tparam K The key ObjectRef type.
|
||||
* @tparam V The value ObjectRef type.
|
||||
*
|
||||
* @note The element type must subclass \c ObjectRef. Otherwise, the
|
||||
* compiler would throw an error:
|
||||
*
|
||||
* <code>
|
||||
* error: no type named 'type' in 'struct std::enable_if<false, void>'
|
||||
* </code>
|
||||
*
|
||||
* Example:
|
||||
*
|
||||
* <code>
|
||||
* // Map<std::string, int> map; // fails
|
||||
* // Map<std::string, NDArray> map2; // fails
|
||||
* Map<std::string, Value> map; // works
|
||||
* map.Set("key1", Value(MakeValue(1))); // works
|
||||
* map.Set("key2", Value(MakeValue(NDArray::Empty(shape, dtype, ctx)))); //
|
||||
* works
|
||||
* </code>
|
||||
*/
|
||||
template <
|
||||
typename K, typename V,
|
||||
typename = typename std::enable_if<
|
||||
std::is_base_of<ObjectRef, K>::value ||
|
||||
std::is_base_of<std::string, K>::value>::type,
|
||||
typename =
|
||||
typename std::enable_if<std::is_base_of<ObjectRef, V>::value>::type>
|
||||
class Map : public ObjectRef {
|
||||
public:
|
||||
/**
|
||||
* @brief default constructor
|
||||
*/
|
||||
Map() { obj_ = std::make_shared<MapObject>(); }
|
||||
/**
|
||||
* @brief move constructor
|
||||
* @param other source
|
||||
*/
|
||||
Map(Map<K, V>&& other) { // NOLINT(*)
|
||||
obj_ = std::move(other.obj_);
|
||||
}
|
||||
/**
|
||||
* @brief copy constructor
|
||||
* @param other source
|
||||
*/
|
||||
Map(const Map<K, V>& other) : ObjectRef(other.obj_) { // NOLINT(*)
|
||||
}
|
||||
/**
|
||||
* @brief constructor from pointer
|
||||
* @param n the container pointer
|
||||
*/
|
||||
explicit Map(std::shared_ptr<Object> n) : ObjectRef(n) {}
|
||||
/**
|
||||
* @brief constructor from iterator
|
||||
* @param begin begin of iterator
|
||||
* @param end end of iterator
|
||||
* @tparam IterType The type of iterator
|
||||
*/
|
||||
template <typename IterType>
|
||||
Map(IterType begin, IterType end) {
|
||||
assign(begin, end);
|
||||
}
|
||||
/**
|
||||
* @brief constructor from initializer list
|
||||
* @param init The initalizer list
|
||||
*/
|
||||
Map(std::initializer_list<std::pair<K, V> > init) { // NOLINT(*)
|
||||
assign(init.begin(), init.end());
|
||||
}
|
||||
/**
|
||||
* @brief constructor from vector
|
||||
* @param init The vector
|
||||
*/
|
||||
template <typename Hash, typename Equal>
|
||||
Map(const std::unordered_map<K, V, Hash, Equal>& init) { // NOLINT(*)
|
||||
assign(init.begin(), init.end());
|
||||
}
|
||||
/**
|
||||
* @brief move assign operator
|
||||
* @param other The source of assignment
|
||||
* @return reference to self.
|
||||
*/
|
||||
Map<K, V>& operator=(Map<K, V>&& other) {
|
||||
obj_ = std::move(other.obj_);
|
||||
return *this;
|
||||
}
|
||||
/**
|
||||
* @brief copy assign operator
|
||||
* @param other The source of assignment
|
||||
* @return reference to self.
|
||||
*/
|
||||
Map<K, V>& operator=(const Map<K, V>& other) {
|
||||
obj_ = other.obj_;
|
||||
return *this;
|
||||
}
|
||||
/**
|
||||
* @brief reset the list to content from iterator.
|
||||
* @param begin begin of iterator
|
||||
* @param end end of iterator
|
||||
* @tparam IterType The type of iterator
|
||||
*/
|
||||
template <typename IterType>
|
||||
void assign(IterType begin, IterType end) {
|
||||
auto n = std::shared_ptr<MapObject>();
|
||||
for (IterType i = begin; i != end; ++i) {
|
||||
n->data.emplace(std::make_pair(i->first.obj_, i->second.obj_));
|
||||
}
|
||||
obj_ = std::move(n);
|
||||
}
|
||||
/**
|
||||
* @brief Read element from map.
|
||||
* @param key The key
|
||||
* @return the corresonding element.
|
||||
*/
|
||||
inline const V operator[](const K& key) const {
|
||||
return V(static_cast<const MapObject*>(obj_.get())->data.at(key.obj_));
|
||||
}
|
||||
/**
|
||||
* @brief Read element from map.
|
||||
* @param key The key
|
||||
* @return the corresonding element.
|
||||
*/
|
||||
inline const V at(const K& key) const {
|
||||
return V(static_cast<const MapObject*>(obj_.get())->data.at(key.obj_));
|
||||
}
|
||||
/** @return The size of the list */
|
||||
inline size_t size() const {
|
||||
if (obj_.get() == nullptr) return 0;
|
||||
return static_cast<const MapObject*>(obj_.get())->data.size();
|
||||
}
|
||||
/** @return The size of the list */
|
||||
inline size_t count(const K& key) const {
|
||||
if (obj_.get() == nullptr) return 0;
|
||||
return static_cast<const MapObject*>(obj_.get())->data.count(key.obj_);
|
||||
}
|
||||
/**
|
||||
* @brief copy on write semantics
|
||||
* Do nothing if current handle is the unique copy of the list.
|
||||
* Otherwise make a new copy of the list to ensure the current handle
|
||||
* hold a unique copy.
|
||||
*
|
||||
* @return Handle to the internal obj container(which ganrantees to be unique)
|
||||
*/
|
||||
inline MapObject* CopyOnWrite() {
|
||||
if (obj_.get() == nullptr || !obj_.unique()) {
|
||||
obj_ = std::make_shared<MapObject>(
|
||||
*static_cast<const MapObject*>(obj_.get()));
|
||||
}
|
||||
return static_cast<MapObject*>(obj_.get());
|
||||
}
|
||||
/**
|
||||
* @brief set the Map.
|
||||
* @param key The index key.
|
||||
* @param value The value to be setted.
|
||||
*/
|
||||
inline void Set(const K& key, const V& value) {
|
||||
MapObject* n = this->CopyOnWrite();
|
||||
n->data[key.obj_] = value.obj_;
|
||||
}
|
||||
|
||||
/** @return whether list is empty */
|
||||
inline bool empty() const { return size() == 0; }
|
||||
/** @brief specify container obj */
|
||||
using ContainerType = MapObject;
|
||||
|
||||
struct Ptr2ObjectRef {
|
||||
using ResultType = std::pair<K, V>;
|
||||
static inline ResultType convert(
|
||||
const std::pair<std::shared_ptr<Object>, std::shared_ptr<Object> >& n) {
|
||||
return std::make_pair(K(n.first), V(n.second));
|
||||
}
|
||||
};
|
||||
|
||||
using iterator =
|
||||
IterAdapter<Ptr2ObjectRef, MapObject::ContainerType::const_iterator>;
|
||||
|
||||
/** @return begin iterator */
|
||||
inline iterator begin() const {
|
||||
return iterator(static_cast<const MapObject*>(obj_.get())->data.begin());
|
||||
}
|
||||
/** @return end iterator */
|
||||
inline iterator end() const {
|
||||
return iterator(static_cast<const MapObject*>(obj_.get())->data.end());
|
||||
}
|
||||
/** @return begin iterator */
|
||||
inline iterator find(const K& key) const {
|
||||
return iterator(
|
||||
static_cast<const MapObject*>(obj_.get())->data.find(key.obj_));
|
||||
}
|
||||
};
|
||||
|
||||
// specialize of string map
|
||||
template <typename V, typename T1, typename T2>
|
||||
class Map<std::string, V, T1, T2> : public ObjectRef {
|
||||
public:
|
||||
// for code reuse
|
||||
Map() { obj_ = std::make_shared<StrMapObject>(); }
|
||||
Map(Map<std::string, V>&& other) { // NOLINT(*)
|
||||
obj_ = std::move(other.obj_);
|
||||
}
|
||||
Map(const Map<std::string, V>& other) : ObjectRef(other.obj_) { // NOLINT(*)
|
||||
}
|
||||
explicit Map(std::shared_ptr<Object> n) : ObjectRef(n) {}
|
||||
template <typename IterType>
|
||||
Map(IterType begin, IterType end) {
|
||||
assign(begin, end);
|
||||
}
|
||||
Map(std::initializer_list<std::pair<std::string, V> > init) { // NOLINT(*)
|
||||
assign(init.begin(), init.end());
|
||||
}
|
||||
|
||||
template <typename Hash, typename Equal>
|
||||
Map(const std::unordered_map<std::string, V, Hash, Equal>&
|
||||
init) { // NOLINT(*)
|
||||
assign(init.begin(), init.end());
|
||||
}
|
||||
Map<std::string, V>& operator=(Map<std::string, V>&& other) {
|
||||
obj_ = std::move(other.obj_);
|
||||
return *this;
|
||||
}
|
||||
Map<std::string, V>& operator=(const Map<std::string, V>& other) {
|
||||
obj_ = other.obj_;
|
||||
return *this;
|
||||
}
|
||||
template <typename IterType>
|
||||
void assign(IterType begin, IterType end) {
|
||||
auto n = std::make_shared<StrMapObject>();
|
||||
for (IterType i = begin; i != end; ++i) {
|
||||
n->data.emplace(std::make_pair(i->first, i->second.obj_));
|
||||
}
|
||||
obj_ = std::move(n);
|
||||
}
|
||||
inline const V operator[](const std::string& key) const {
|
||||
return V(static_cast<const StrMapObject*>(obj_.get())->data.at(key));
|
||||
}
|
||||
inline const V at(const std::string& key) const {
|
||||
return V(static_cast<const StrMapObject*>(obj_.get())->data.at(key));
|
||||
}
|
||||
inline size_t size() const {
|
||||
if (obj_.get() == nullptr) return 0;
|
||||
return static_cast<const StrMapObject*>(obj_.get())->data.size();
|
||||
}
|
||||
inline size_t count(const std::string& key) const {
|
||||
if (obj_.get() == nullptr) return 0;
|
||||
return static_cast<const StrMapObject*>(obj_.get())->data.count(key);
|
||||
}
|
||||
inline StrMapObject* CopyOnWrite() {
|
||||
if (obj_.get() == nullptr || !obj_.unique()) {
|
||||
obj_ = std::make_shared<MapObject>(
|
||||
*static_cast<const MapObject*>(obj_.get()));
|
||||
}
|
||||
return static_cast<StrMapObject*>(obj_.get());
|
||||
}
|
||||
inline void Set(const std::string& key, const V& value) {
|
||||
StrMapObject* n = this->CopyOnWrite();
|
||||
n->data[key] = value.obj_;
|
||||
}
|
||||
inline bool empty() const { return size() == 0; }
|
||||
using ContainerType = StrMapObject;
|
||||
|
||||
struct Ptr2ObjectRef {
|
||||
using ResultType = std::pair<std::string, V>;
|
||||
static inline ResultType convert(
|
||||
const std::pair<std::string, std::shared_ptr<Object> >& n) {
|
||||
return std::make_pair(n.first, V(n.second));
|
||||
}
|
||||
};
|
||||
|
||||
using iterator =
|
||||
IterAdapter<Ptr2ObjectRef, StrMapObject::ContainerType::const_iterator>;
|
||||
|
||||
/** @return begin iterator */
|
||||
inline iterator begin() const {
|
||||
return iterator(static_cast<const StrMapObject*>(obj_.get())->data.begin());
|
||||
}
|
||||
/** @return end iterator */
|
||||
inline iterator end() const {
|
||||
return iterator(static_cast<const StrMapObject*>(obj_.get())->data.end());
|
||||
}
|
||||
/** @return begin iterator */
|
||||
inline iterator find(const std::string& key) const {
|
||||
return iterator(
|
||||
static_cast<const StrMapObject*>(obj_.get())->data.find(key));
|
||||
}
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Helper function to convert a List<Value> object to a vector.
|
||||
* @tparam T element type
|
||||
* @param list Input list object.
|
||||
* @return std vector
|
||||
*/
|
||||
template <typename T>
|
||||
inline std::vector<T> ListValueToVector(const List<Value>& list) {
|
||||
std::vector<T> ret;
|
||||
ret.reserve(list.size());
|
||||
for (Value val : list)
|
||||
// (BarclayII) apparently MSVC 2017 CL 19.10 had trouble parsing
|
||||
// ret.push_back(val->data)
|
||||
// So I kindly tell it how to properly parse it.
|
||||
ret.push_back(val->data.operator T());
|
||||
return ret;
|
||||
}
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_RUNTIME_CONTAINER_H_
|
||||
@@ -0,0 +1,268 @@
|
||||
/**
|
||||
* 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_
|
||||
@@ -0,0 +1,84 @@
|
||||
/**
|
||||
* Copyright (c) 2022 by Contributors
|
||||
* @file include/dgl/runtime/dlpack_convert.h
|
||||
* @brief Conversion between NDArray and DLPack.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_DLPACK_CONVERT_H_
|
||||
#define DGL_RUNTIME_DLPACK_CONVERT_H_
|
||||
|
||||
#include "c_runtime_api.h"
|
||||
#include "ndarray.h"
|
||||
|
||||
struct DLManagedTensor;
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
struct DLPackConvert {
|
||||
/**
|
||||
* @brief Create a DGL NDArray from a DLPack tensor.
|
||||
*
|
||||
* This allows us to create a NDArray using the memory
|
||||
* allocated by an external deep learning framework
|
||||
* that is DLPack compatible.
|
||||
*
|
||||
* The memory is retained until the NDArray went out of scope.
|
||||
* @param tensor The DLPack tensor to copy from.
|
||||
* @return The created NDArray view.
|
||||
*/
|
||||
static NDArray FromDLPack(DLManagedTensor* tensor);
|
||||
|
||||
/**
|
||||
* @brief Deleter for NDArray converted from DLPack.
|
||||
*
|
||||
* This is used from data which is passed from external
|
||||
* DLPack(DLManagedTensor) that are not allocated inside of DGL. This enables
|
||||
* us to create NDArray from memory allocated by other frameworks that are
|
||||
* DLPack compatible
|
||||
*/
|
||||
static void DLPackDeleter(NDArray::Container* ptr);
|
||||
|
||||
/** @brief Convert a DGL NDArray to a DLPack tensor.
|
||||
*
|
||||
* @param from The DGL NDArray.
|
||||
* @return A DLPack tensor.
|
||||
*/
|
||||
static DLManagedTensor* ToDLPack(const NDArray& from);
|
||||
};
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
/**
|
||||
* @brief Delete (free) a DLManagedTensor's data.
|
||||
* @param dltensor Pointer to the DLManagedTensor.
|
||||
*/
|
||||
DGL_DLL void DGLDLManagedTensorCallDeleter(DLManagedTensor* dltensor);
|
||||
|
||||
/**
|
||||
* @brief Produce an array from the DLManagedTensor that shares data memory
|
||||
* with the DLManagedTensor.
|
||||
* @param from The source DLManagedTensor.
|
||||
* @param out The output array handle.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLArrayFromDLPack(DLManagedTensor* from, DGLArrayHandle* out);
|
||||
|
||||
/**
|
||||
* @brief Produce a DLMangedTensor from the array that shares data memory with
|
||||
* the array.
|
||||
* @param from The source array.
|
||||
* @param out The DLManagedTensor handle.
|
||||
* @return 0 when success, -1 when failure happens
|
||||
*/
|
||||
DGL_DLL int DGLArrayToDLPack(
|
||||
DGLArrayHandle from, DLManagedTensor** out, int alignment = 0);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // DGL_EXTERN_C
|
||||
#endif
|
||||
#endif // DGL_RUNTIME_DLPACK_CONVERT_H_
|
||||
@@ -0,0 +1,175 @@
|
||||
/**
|
||||
* Copyright (c) 2017 by Contributors
|
||||
* @file dgl/runtime/module.h
|
||||
* @brief Runtime container of the functions generated by DGL,
|
||||
* This is used to support dynamically link, load and save
|
||||
* functions from different convention under unified API.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_MODULE_H_
|
||||
#define DGL_RUNTIME_MODULE_H_
|
||||
|
||||
#include <dmlc/io.h>
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "c_runtime_api.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
// The internal container of module.
|
||||
class ModuleNode;
|
||||
class PackedFunc;
|
||||
|
||||
/**
|
||||
* @brief Module container of DGL.
|
||||
*/
|
||||
class Module {
|
||||
public:
|
||||
Module() {}
|
||||
// constructor from container.
|
||||
explicit Module(std::shared_ptr<ModuleNode> n) : node_(n) {}
|
||||
/**
|
||||
* @brief Get packed function from current module by name.
|
||||
*
|
||||
* @param name The name of the function.
|
||||
* @param query_imports Whether also query dependency modules.
|
||||
* @return The result function.
|
||||
* This function will return PackedFunc(nullptr) if function do not exist.
|
||||
* @note Implemented in packed_func.cc
|
||||
*/
|
||||
inline PackedFunc GetFunction(
|
||||
const std::string& name, bool query_imports = false);
|
||||
/** @return internal container */
|
||||
inline ModuleNode* operator->();
|
||||
/** @return internal container */
|
||||
inline const ModuleNode* operator->() const;
|
||||
// The following functions requires link with runtime.
|
||||
/**
|
||||
* @brief Import another module into this module.
|
||||
* @param other The module to be imported.
|
||||
*
|
||||
* @note Cyclic dependency is not allowed among modules,
|
||||
* An error will be thrown when cyclic dependency is detected.
|
||||
*/
|
||||
DGL_DLL void Import(Module other);
|
||||
/**
|
||||
* @brief Load a module from file.
|
||||
* @param file_name The name of the host function module.
|
||||
* @param format The format of the file.
|
||||
* @note This function won't load the import relationship.
|
||||
* Re-create import relationship by calling Import.
|
||||
*/
|
||||
DGL_DLL static Module LoadFromFile(
|
||||
const std::string& file_name, const std::string& format = "");
|
||||
|
||||
private:
|
||||
std::shared_ptr<ModuleNode> node_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Base node container of module.
|
||||
* Do not create this directly, instead use Module.
|
||||
*/
|
||||
class ModuleNode {
|
||||
public:
|
||||
/** @brief virtual destructor */
|
||||
virtual ~ModuleNode() {}
|
||||
/** @return The module type key */
|
||||
virtual const char* type_key() const = 0;
|
||||
/**
|
||||
* @brief Get a PackedFunc from module.
|
||||
*
|
||||
* The PackedFunc may not be fully initialized,
|
||||
* there might still be first time running overhead when
|
||||
* executing the function on certain devices.
|
||||
* For benchmarking, use prepare to eliminate
|
||||
*
|
||||
* @param name the name of the function.
|
||||
* @param sptr_to_self The shared_ptr that points to this module node.
|
||||
*
|
||||
* @return PackedFunc(nullptr) when it is not available.
|
||||
*
|
||||
* @note The function will always remain valid.
|
||||
* If the function need resource from the module(e.g. late linking),
|
||||
* it should capture sptr_to_self.
|
||||
*/
|
||||
virtual PackedFunc GetFunction(
|
||||
const std::string& name,
|
||||
const std::shared_ptr<ModuleNode>& sptr_to_self) = 0;
|
||||
/**
|
||||
* @brief Save the module to file.
|
||||
* @param file_name The file to be saved to.
|
||||
* @param format The format of the file.
|
||||
*/
|
||||
virtual void SaveToFile(
|
||||
const std::string& file_name, const std::string& format);
|
||||
/**
|
||||
* @brief Save the module to binary stream.
|
||||
* @param stream The binary stream to save to.
|
||||
* @note It is recommended to implement this for device modules,
|
||||
* but not necessarily host modules.
|
||||
* We can use this to do AOT loading of bundled device functions.
|
||||
*/
|
||||
DGL_DLL virtual void SaveToBinary(dmlc::Stream* stream);
|
||||
/**
|
||||
* @brief Get the source code of module, when available.
|
||||
* @param format Format of the source code, can be empty by default.
|
||||
* @return Possible source code when available.
|
||||
*/
|
||||
DGL_DLL virtual std::string GetSource(const std::string& format = "");
|
||||
/**
|
||||
* @brief Get a function from current environment
|
||||
* The environment includes all the imports as well as Global functions.
|
||||
*
|
||||
* @param name name of the function.
|
||||
* @return The corresponding function.
|
||||
*/
|
||||
DGL_DLL const PackedFunc* GetFuncFromEnv(const std::string& name);
|
||||
/** @return The module it imports from */
|
||||
const std::vector<Module>& imports() const { return imports_; }
|
||||
|
||||
protected:
|
||||
friend class Module;
|
||||
/** @brief The modules this module depend on */
|
||||
std::vector<Module> imports_;
|
||||
|
||||
private:
|
||||
/** @brief Cache used by GetImport */
|
||||
std::unordered_map<std::string, std::unique_ptr<PackedFunc> > import_cache_;
|
||||
};
|
||||
|
||||
/** @brief namespace for constant symbols */
|
||||
namespace symbol {
|
||||
/** @brief Global variable to store module context. */
|
||||
constexpr const char* dgl_module_ctx = "__dgl_module_ctx";
|
||||
/** @brief Global variable to store device module blob */
|
||||
constexpr const char* dgl_dev_mblob = "__dgl_dev_mblob";
|
||||
/** @brief Number of bytes of device module blob. */
|
||||
constexpr const char* dgl_dev_mblob_nbytes = "__dgl_dev_mblob_nbytes";
|
||||
/** @brief global function to set device */
|
||||
constexpr const char* dgl_set_device = "__dgl_set_device";
|
||||
/** @brief Auxiliary counter to global barrier. */
|
||||
constexpr const char* dgl_global_barrier_state = "__dgl_global_barrier_state";
|
||||
/**
|
||||
* @brief Prepare the global barrier before kernels that uses global barrier.
|
||||
*/
|
||||
constexpr const char* dgl_prepare_global_barrier =
|
||||
"__dgl_prepare_global_barrier";
|
||||
/** @brief Placeholder for the module's entry function. */
|
||||
constexpr const char* dgl_module_main = "__dgl_main__";
|
||||
} // namespace symbol
|
||||
|
||||
// implementations of inline functions.
|
||||
inline ModuleNode* Module::operator->() { return node_.get(); }
|
||||
|
||||
inline const ModuleNode* Module::operator->() const { return node_.get(); }
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
#include "packed_func.h"
|
||||
#endif // DGL_RUNTIME_MODULE_H_
|
||||
@@ -0,0 +1,890 @@
|
||||
/**
|
||||
* Copyright (c) 2017-2022 by Contributors
|
||||
* @file dgl/runtime/ndarray.h
|
||||
* @brief Abstract device memory management API
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_NDARRAY_H_
|
||||
#define DGL_RUNTIME_NDARRAY_H_
|
||||
|
||||
#include <atomic>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "bfloat16.h"
|
||||
#include "c_runtime_api.h"
|
||||
#include "serializer.h"
|
||||
#include "shared_mem.h"
|
||||
|
||||
#ifdef DGL_USE_CUDA
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#define BF16_ENABLED (defined(CUDART_VERSION) && CUDART_VERSION >= 11000)
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#if BF16_ENABLED
|
||||
#include <cuda_bf16.h>
|
||||
#endif // BF16_ENABLED
|
||||
#endif // DGL_USE_CUDA
|
||||
|
||||
// forward declaration
|
||||
inline std::ostream& operator<<(std::ostream& os, DGLDataType t);
|
||||
|
||||
namespace dgl {
|
||||
|
||||
/**
|
||||
* @brief Type traits that converts a C type to a DGLDataType.
|
||||
*
|
||||
* Usage:
|
||||
* DGLDataTypeTraits<int>::dtype == dtype
|
||||
*/
|
||||
template <typename T>
|
||||
struct DGLDataTypeTraits {
|
||||
static constexpr DGLDataType dtype{0, 0, 0}; // dummy
|
||||
};
|
||||
#define GEN_DGLDATATYPETRAITS_FOR(T, code, bits) \
|
||||
template <> \
|
||||
struct DGLDataTypeTraits<T> { \
|
||||
static constexpr DGLDataType dtype{code, bits, 1}; \
|
||||
}
|
||||
GEN_DGLDATATYPETRAITS_FOR(int8_t, kDGLInt, 8);
|
||||
GEN_DGLDATATYPETRAITS_FOR(uint8_t, kDGLUInt, 8);
|
||||
GEN_DGLDATATYPETRAITS_FOR(int16_t, kDGLInt, 16);
|
||||
GEN_DGLDATATYPETRAITS_FOR(int32_t, kDGLInt, 32);
|
||||
GEN_DGLDATATYPETRAITS_FOR(int64_t, kDGLInt, 64);
|
||||
// XXX(BarclayII) most DL frameworks do not support unsigned int and long
|
||||
// arrays, so I'm just converting uints to signed DTypes.
|
||||
GEN_DGLDATATYPETRAITS_FOR(uint32_t, kDGLInt, 32);
|
||||
GEN_DGLDATATYPETRAITS_FOR(uint64_t, kDGLInt, 64);
|
||||
#ifdef DGL_USE_CUDA
|
||||
GEN_DGLDATATYPETRAITS_FOR(__half, kDGLFloat, 16);
|
||||
#if BF16_ENABLED
|
||||
GEN_DGLDATATYPETRAITS_FOR(__nv_bfloat16, kDGLBfloat, 16);
|
||||
#endif // BF16_ENABLED
|
||||
#endif // DGL_USE_CUDA
|
||||
GEN_DGLDATATYPETRAITS_FOR(float, kDGLFloat, 32);
|
||||
GEN_DGLDATATYPETRAITS_FOR(double, kDGLFloat, 64);
|
||||
#undef GEN_DGLDATATYPETRAITS_FOR
|
||||
|
||||
namespace runtime {
|
||||
|
||||
/**
|
||||
* @brief DLPack converter.
|
||||
*/
|
||||
struct DLPackConvert;
|
||||
|
||||
/**
|
||||
* @brief Managed NDArray.
|
||||
* The array is backed by reference counted blocks.
|
||||
*/
|
||||
class NDArray {
|
||||
public:
|
||||
// internal container type
|
||||
struct Container;
|
||||
/** @brief default constructor */
|
||||
NDArray() {}
|
||||
/**
|
||||
* @brief cosntruct a NDArray that refers to data
|
||||
* @param data The data this NDArray refers to
|
||||
*/
|
||||
explicit inline NDArray(Container* data);
|
||||
/**
|
||||
* @brief copy constructor
|
||||
* @param other The value to be copied
|
||||
*/
|
||||
inline NDArray(const NDArray& other); // NOLINT(*)
|
||||
/**
|
||||
* @brief move constructor
|
||||
* @param other The value to be moved
|
||||
*/
|
||||
NDArray(NDArray&& other) // NOLINT(*)
|
||||
: data_(other.data_) {
|
||||
other.data_ = nullptr;
|
||||
}
|
||||
/** @brief destructor */
|
||||
~NDArray() { this->reset(); }
|
||||
/**
|
||||
* @brief Swap this array with another NDArray
|
||||
* @param other The other NDArray
|
||||
*/
|
||||
void swap(NDArray& other) { // NOLINT(*)
|
||||
std::swap(data_, other.data_);
|
||||
}
|
||||
/**
|
||||
* @brief copy assignmemt
|
||||
* @param other The value to be assigned.
|
||||
* @return reference to self.
|
||||
*/
|
||||
NDArray& operator=(const NDArray& other) { // NOLINT(*)
|
||||
// copy-and-swap idiom
|
||||
NDArray(other).swap(*this); // NOLINT(*)
|
||||
return *this;
|
||||
}
|
||||
/**
|
||||
* @brief move assignmemt
|
||||
* @param other The value to be assigned.
|
||||
* @return reference to self.
|
||||
*/
|
||||
NDArray& operator=(NDArray&& other) { // NOLINT(*)
|
||||
// copy-and-swap idiom
|
||||
NDArray(std::move(other)).swap(*this); // NOLINT(*)
|
||||
return *this;
|
||||
}
|
||||
/** @return If NDArray is defined */
|
||||
bool defined() const { return data_ != nullptr; }
|
||||
/** @return If both NDArray reference the same container */
|
||||
bool same_as(const NDArray& other) const { return data_ == other.data_; }
|
||||
/** @brief reset the content of NDArray to be nullptr */
|
||||
inline void reset();
|
||||
/**
|
||||
* @return the reference counter
|
||||
* @note this number is approximate in multi-threaded setting.
|
||||
*/
|
||||
inline int use_count() const;
|
||||
/** @return Pointer to content of DGLArray */
|
||||
inline const DGLArray* operator->() const;
|
||||
/** @return True if the ndarray is contiguous. */
|
||||
bool IsContiguous() const;
|
||||
/** @return the data pointer with type. */
|
||||
template <typename T>
|
||||
inline T* Ptr() const {
|
||||
if (!defined())
|
||||
return nullptr;
|
||||
else
|
||||
return static_cast<T*>(operator->()->data);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Copy data content from/into another array.
|
||||
* @param other The source array to be copied from.
|
||||
* @note The copy runs on the dgl internal stream if it involves a GPU
|
||||
* context.
|
||||
*/
|
||||
inline void CopyFrom(DGLArray* other);
|
||||
inline void CopyFrom(const NDArray& other);
|
||||
inline void CopyTo(DGLArray* other) const;
|
||||
inline void CopyTo(const NDArray& other) const;
|
||||
|
||||
/**
|
||||
* @brief Copy the data to another context.
|
||||
* @param ctx The target context.
|
||||
* @return The array under another context.
|
||||
*/
|
||||
inline NDArray CopyTo(const DGLContext& ctx) const;
|
||||
|
||||
/**
|
||||
* @brief Return a new array with a copy of the content.
|
||||
*/
|
||||
inline NDArray Clone() const;
|
||||
|
||||
/**
|
||||
* @brief Return a copy of the current instance of NDArray in pinned
|
||||
* (page-locked) memory.
|
||||
* @note This is an out-of-place method, which utilizes PyTorch's
|
||||
* CachingHostAllocator for allocating pinned memory and copying data
|
||||
* from the current NDAarray. As a result, PyTorch is responsible for
|
||||
* managing the lifecycle of the returned NDArray, including deciding
|
||||
* when to flush the data for reuse or call cudaFreeHost. The current
|
||||
* context must be kDGLCPU, otherwise, an error will be thrown.
|
||||
*/
|
||||
inline NDArray PinMemory();
|
||||
|
||||
/**
|
||||
* @brief In-place method to pin the current array by calling PinContainer
|
||||
* on the underlying NDArray:Container.
|
||||
* @note This is an in-place method that flags the memory as page-locked by
|
||||
* utilizing cudaHostRegister at the underlying level to pin the current
|
||||
* instance of NDArray. The current context must be kDGLCPU, otherwise,
|
||||
* an error will be thrown.
|
||||
*/
|
||||
inline void PinMemory_();
|
||||
|
||||
/**
|
||||
* @brief In-place method to unpin the current array by calling UnpinContainer
|
||||
* on the underlying NDArray:Container.
|
||||
* @note This is an in-place method. Behavior depends on the current context,
|
||||
* IsPinned: will be unpinned;
|
||||
* others: directly return.
|
||||
*/
|
||||
inline void UnpinMemory_();
|
||||
|
||||
/**
|
||||
* @brief Check if the array is pinned.
|
||||
*/
|
||||
inline bool IsPinned() const;
|
||||
|
||||
/**
|
||||
* @brief Record streams that are using the underlying tensor.
|
||||
* @param stream The stream that is using the underlying tensor.
|
||||
*/
|
||||
inline void RecordStream(DGLStreamHandle stream) const;
|
||||
|
||||
/**
|
||||
* @brief Load NDArray from stream
|
||||
* @param stream The input data stream
|
||||
* @return Whether load is successful
|
||||
*/
|
||||
bool Load(dmlc::Stream* stream);
|
||||
|
||||
/**
|
||||
* @brief Save NDArray to stream
|
||||
* @param stream The output data stream
|
||||
*/
|
||||
void Save(dmlc::Stream* stream) const;
|
||||
|
||||
/**
|
||||
* @brief Create a NDArray that shares the data memory with the current one.
|
||||
* @param shape The shape of the new array.
|
||||
* @param dtype The data type of the new array.
|
||||
* @param offset The offset (in bytes) of the starting pointer.
|
||||
* @note The memory size of new array must be smaller than the current one.
|
||||
*/
|
||||
DGL_DLL NDArray
|
||||
CreateView(std::vector<int64_t> shape, DGLDataType dtype, int64_t offset = 0);
|
||||
|
||||
/**
|
||||
* @brief Create an empty NDArray.
|
||||
* @param shape The shape of the new array.
|
||||
* @param dtype The data type of the new array.
|
||||
* @param ctx The context of the array.
|
||||
* @return The created Array
|
||||
*/
|
||||
DGL_DLL static NDArray Empty(
|
||||
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
|
||||
|
||||
/**
|
||||
* @brief Create an empty NDArray in pinned memory.
|
||||
* @param shape The shape of the new array.
|
||||
* @param dtype The data type of the new array.
|
||||
* @param ctx The context of the array.
|
||||
* @return The created array.
|
||||
*/
|
||||
DGL_DLL static NDArray PinnedEmpty(
|
||||
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
|
||||
|
||||
/**
|
||||
* @brief Create an empty NDArray with shared memory.
|
||||
* @param name The name of shared memory.
|
||||
* @param shape The shape of the new array.
|
||||
* @param dtype The data type of the new array.
|
||||
* @param ctx The context of the array.
|
||||
* @param is_create whether to create shared memory.
|
||||
* @return The created Array
|
||||
*/
|
||||
DGL_DLL static NDArray EmptyShared(
|
||||
const std::string& name, std::vector<int64_t> shape, DGLDataType dtype,
|
||||
DGLContext ctx, bool is_create);
|
||||
|
||||
/**
|
||||
* @brief Get the size of the array in the number of bytes.
|
||||
*/
|
||||
size_t GetSize() const;
|
||||
|
||||
/**
|
||||
* @brief Get the number of elements in this array.
|
||||
*/
|
||||
int64_t NumElements() const;
|
||||
|
||||
/**
|
||||
* @brief Create a NDArray by copying from std::vector.
|
||||
* @tparam T Type of vector data. Determines the dtype of returned array.
|
||||
*/
|
||||
template <typename T>
|
||||
DGL_DLL static NDArray FromVector(
|
||||
const std::vector<T>& vec, DGLContext ctx = DGLContext{kDGLCPU, 0});
|
||||
|
||||
/**
|
||||
* @brief Create a NDArray from a raw pointer.
|
||||
*/
|
||||
DGL_DLL static NDArray CreateFromRaw(
|
||||
const std::vector<int64_t>& shape, DGLDataType dtype, DGLContext ctx,
|
||||
void* raw, bool auto_free);
|
||||
|
||||
/**
|
||||
* @brief Create a std::vector from a 1D NDArray.
|
||||
* @tparam T Type of vector data.
|
||||
* @note Type casting is NOT performed. The caller has to make sure that the
|
||||
* vector type matches the dtype of NDArray.
|
||||
*/
|
||||
template <typename T>
|
||||
std::vector<T> ToVector() const;
|
||||
|
||||
std::shared_ptr<SharedMemory> GetSharedMem() const;
|
||||
|
||||
/**
|
||||
* @brief Function to copy data from one array to another.
|
||||
* @param from The source array.
|
||||
* @param to The target array.
|
||||
* @param (optional) stream The stream used in copy.
|
||||
*/
|
||||
DGL_DLL static void CopyFromTo(DGLArray* from, DGLArray* to);
|
||||
DGL_DLL static void CopyFromTo(
|
||||
DGLArray* from, DGLArray* to, DGLStreamHandle stream);
|
||||
|
||||
/**
|
||||
* @brief Function to copy data between device and CPU while recording the
|
||||
* event.
|
||||
* @param from The source array.
|
||||
* @param to The target array.
|
||||
* @param pytorch_ctx The context pointer from PyTorch's CachingHostAllocator.
|
||||
* @note This function fuses data-copy and event recording to ensure
|
||||
* CachingHostAllocator works properly.
|
||||
*/
|
||||
DGL_DLL static void RecordedCopyFromTo(
|
||||
DGLArray* from, DGLArray* to, void* pytorch_ctx);
|
||||
|
||||
/**
|
||||
* @brief Function to pin the DGLArray of a Container.
|
||||
* @param ptr The container to be pinned.
|
||||
* @note Data of the given array will be pinned inplace.
|
||||
* Behavior depends on the current context,
|
||||
* kDGLCPU: will be pinned;
|
||||
* IsPinned: directly return;
|
||||
* kDGLCUDA: invalid, will throw an error.
|
||||
*/
|
||||
DGL_DLL static void PinContainer(Container* ptr);
|
||||
|
||||
/**
|
||||
* @brief Function to unpin the DGLArray of a Container.
|
||||
* @param ptr The container to be unpinned.
|
||||
* @note Data of the given array will be unpinned inplace.
|
||||
* Behavior depends on the current context,
|
||||
* IsPinned: will be unpinned;
|
||||
* others: directly return.
|
||||
*/
|
||||
DGL_DLL static void UnpinContainer(Container* ptr);
|
||||
|
||||
/**
|
||||
* @brief Function check if the DGLArray of a Container is pinned.
|
||||
* @param ptr The container to be checked.
|
||||
* @return true if pinned.
|
||||
*/
|
||||
DGL_DLL static bool IsContainerPinned(Container* ptr);
|
||||
|
||||
/**
|
||||
* @brief Record streams that are using this tensor.
|
||||
* @param ptr Pointer of the tensor to be recorded.
|
||||
* @param stream The stream that is using this tensor.
|
||||
*/
|
||||
DGL_DLL static void RecordStream(DGLArray* tensor, DGLStreamHandle stream);
|
||||
|
||||
// internal namespace
|
||||
struct Internal {
|
||||
// Default deleter for the container
|
||||
static void DefaultDeleter(NDArray::Container* ptr);
|
||||
// Local create function which allocates tensor metadata
|
||||
// but does not allocate space for the data.
|
||||
static NDArray Create(
|
||||
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
|
||||
// Implementation of API function
|
||||
static DGLArray* MoveAsDGLArray(NDArray arr);
|
||||
};
|
||||
|
||||
private:
|
||||
/** @brief Internal Data content */
|
||||
Container* data_{nullptr};
|
||||
// enable internal functions
|
||||
friend struct Internal;
|
||||
friend struct DLPackConvert;
|
||||
friend class DGLRetValue;
|
||||
friend class DGLArgsSetter;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Save a DGLArray to stream
|
||||
* @param strm The outpu stream
|
||||
* @param tensor The tensor to be saved.
|
||||
*/
|
||||
inline bool SaveDGLArray(dmlc::Stream* strm, const DGLArray* tensor);
|
||||
|
||||
/**
|
||||
* @brief Reference counted Container object used to back NDArray.
|
||||
*
|
||||
* This object is DGLArray compatible:
|
||||
* the pointer to the NDArrayContainer can be directly
|
||||
* interpreted as a DGLArray*
|
||||
*
|
||||
* @note: do not use this function directly, use NDArray.
|
||||
*/
|
||||
struct NDArray::Container {
|
||||
public:
|
||||
/** NOTE: the first part of this structure is the same as
|
||||
* DLManagedTensor, note that, however, the deleter
|
||||
* is only called when the reference counter goes to 0
|
||||
*/
|
||||
/**
|
||||
* @brief Tensor structure.
|
||||
* @note it is important that the first field is DGLArray
|
||||
* So that this data structure is DGLArray compatible.
|
||||
* The head ptr of this struct can be viewed as DGLArray*.
|
||||
*/
|
||||
DGLArray dl_tensor;
|
||||
/**
|
||||
* @brief addtional context, reserved for recycling
|
||||
* @note We can attach additional content here
|
||||
* which the current container depend on
|
||||
* (e.g. reference to original memory when creating views).
|
||||
*/
|
||||
void* manager_ctx{nullptr};
|
||||
/**
|
||||
* @brief Customized deleter
|
||||
*
|
||||
* @note The customized deleter is helpful to enable
|
||||
* different ways of memory allocator that are not
|
||||
* currently defined by the system.
|
||||
*/
|
||||
void (*deleter)(Container* self) = nullptr;
|
||||
/** @brief default constructor */
|
||||
Container() {
|
||||
dl_tensor.data = nullptr;
|
||||
dl_tensor.ndim = 0;
|
||||
dl_tensor.shape = nullptr;
|
||||
dl_tensor.strides = nullptr;
|
||||
dl_tensor.byte_offset = 0;
|
||||
}
|
||||
/** @brief pointer to shared memory */
|
||||
std::shared_ptr<SharedMemory> mem;
|
||||
/** @brief developer function, increases reference counter */
|
||||
void IncRef() { ref_counter_.fetch_add(1, std::memory_order_relaxed); }
|
||||
/** @brief developer function, decrease reference counter */
|
||||
void DecRef() {
|
||||
if (ref_counter_.fetch_sub(1, std::memory_order_release) == 1) {
|
||||
std::atomic_thread_fence(std::memory_order_acquire);
|
||||
if (this->deleter != nullptr) {
|
||||
(*this->deleter)(this);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
friend struct DLPackConvert;
|
||||
friend class NDArray;
|
||||
friend class RPCWrappedFunc;
|
||||
/**
|
||||
* @brief The shape container,
|
||||
* can be used for shape data.
|
||||
*/
|
||||
std::vector<int64_t> shape_;
|
||||
/**
|
||||
* @brief The stride container,
|
||||
* can be used for stride data.
|
||||
*/
|
||||
std::vector<int64_t> stride_;
|
||||
/** @brief The internal array object */
|
||||
std::atomic<int> ref_counter_{0};
|
||||
|
||||
/** @brief Whether underlying dl_tensor is pinned by DGL. */
|
||||
bool pinned_by_dgl_{false};
|
||||
|
||||
/** @brief Whether underlying dl_tensor is pinned by PyTorch
|
||||
* (CachingHostAllocator). */
|
||||
bool pinned_by_pytorch_{false};
|
||||
|
||||
/** @brief The PyTorch storage ctx ptr if pinned_by_pytorch_ = True. */
|
||||
void* pytorch_ctx_{nullptr};
|
||||
|
||||
/** @brief Pointer to the corresp. PyTorch deleter if pinned_by_pytorch_ =
|
||||
* True.
|
||||
*/
|
||||
void* pytorch_raw_deleter_{nullptr};
|
||||
};
|
||||
|
||||
// implementations of inline functions
|
||||
// the usages of functions are documented in place.
|
||||
inline NDArray::NDArray(Container* data) : data_(data) {
|
||||
if (data_) data_->IncRef();
|
||||
}
|
||||
|
||||
inline NDArray::NDArray(const NDArray& other) : data_(other.data_) {
|
||||
if (data_) data_->IncRef();
|
||||
}
|
||||
|
||||
inline void NDArray::reset() {
|
||||
if (data_) {
|
||||
data_->DecRef();
|
||||
data_ = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
inline void NDArray::CopyFrom(DGLArray* other) {
|
||||
CHECK(data_ != nullptr);
|
||||
CopyFromTo(other, &(data_->dl_tensor));
|
||||
}
|
||||
|
||||
inline void NDArray::CopyFrom(const NDArray& other) {
|
||||
CHECK(other.data_ != nullptr);
|
||||
// Copy between two devices
|
||||
if (data_->dl_tensor.ctx.device_type !=
|
||||
other.data_->dl_tensor.ctx.device_type) {
|
||||
CHECK(data_ != nullptr);
|
||||
auto to_ctx_type = data_->dl_tensor.ctx.device_type;
|
||||
auto cpu_data = (to_ctx_type == kDGLCPU ? data_ : other.data_);
|
||||
// Pinned by PyTorch
|
||||
if (cpu_data->pinned_by_pytorch_) {
|
||||
// To ensure correct behavior, the event must be recorded after
|
||||
// cudaMemcpyAsync as long as the memory is pinned by PyTorch.
|
||||
void* pytorch_ctx = cpu_data->pytorch_ctx_;
|
||||
RecordedCopyFromTo(
|
||||
&(other.data_->dl_tensor), &(data_->dl_tensor), pytorch_ctx);
|
||||
return;
|
||||
}
|
||||
}
|
||||
CopyFrom(&(other.data_->dl_tensor));
|
||||
}
|
||||
|
||||
inline void NDArray::CopyTo(DGLArray* other) const {
|
||||
CHECK(data_ != nullptr);
|
||||
CopyFromTo(&(data_->dl_tensor), other);
|
||||
}
|
||||
|
||||
inline void NDArray::CopyTo(const NDArray& other) const {
|
||||
CHECK(other.data_ != nullptr);
|
||||
// copy between two devices
|
||||
if (data_->dl_tensor.ctx.device_type !=
|
||||
other.data_->dl_tensor.ctx.device_type) {
|
||||
CHECK(data_ != nullptr);
|
||||
auto from_ctx_type = data_->dl_tensor.ctx.device_type;
|
||||
auto cpu_data = (from_ctx_type == kDGLCPU ? data_ : other.data_);
|
||||
// pinned by PyTorch
|
||||
if (cpu_data->pinned_by_pytorch_) {
|
||||
// To ensure correct behavior, the event must be recorded after
|
||||
// cudaMemcpyAsync as long as the memory is pinned by PyTorch.
|
||||
void* pytorch_ctx = cpu_data->pytorch_ctx_;
|
||||
RecordedCopyFromTo(
|
||||
&(data_->dl_tensor), &(other.data_->dl_tensor), pytorch_ctx);
|
||||
return;
|
||||
}
|
||||
}
|
||||
CopyTo(&(other.data_->dl_tensor));
|
||||
}
|
||||
|
||||
inline NDArray NDArray::CopyTo(const DGLContext& ctx) const {
|
||||
CHECK(data_ != nullptr);
|
||||
const DGLArray* array = operator->();
|
||||
NDArray ret = Empty(
|
||||
std::vector<int64_t>(array->shape, array->shape + array->ndim),
|
||||
array->dtype, ctx);
|
||||
this->CopyTo(ret);
|
||||
return ret;
|
||||
}
|
||||
|
||||
inline NDArray NDArray::Clone() const {
|
||||
CHECK(data_ != nullptr);
|
||||
const DGLArray* array = operator->();
|
||||
return this->CopyTo(array->ctx);
|
||||
}
|
||||
|
||||
inline NDArray NDArray::PinMemory() {
|
||||
CHECK(data_ != nullptr);
|
||||
const DGLArray* array = operator->();
|
||||
auto ctx = array->ctx;
|
||||
NDArray ret = PinnedEmpty(
|
||||
std::vector<int64_t>(array->shape, array->shape + array->ndim),
|
||||
array->dtype, ctx);
|
||||
this->CopyTo(ret);
|
||||
return ret;
|
||||
}
|
||||
|
||||
inline void NDArray::PinMemory_() {
|
||||
CHECK(data_ != nullptr);
|
||||
PinContainer(data_);
|
||||
}
|
||||
|
||||
inline void NDArray::UnpinMemory_() {
|
||||
CHECK(data_ != nullptr);
|
||||
UnpinContainer(data_);
|
||||
}
|
||||
|
||||
inline bool NDArray::IsPinned() const {
|
||||
CHECK(data_ != nullptr);
|
||||
return IsContainerPinned(data_);
|
||||
}
|
||||
|
||||
inline void NDArray::RecordStream(DGLStreamHandle stream) const {
|
||||
CHECK(data_ != nullptr);
|
||||
RecordStream(&(data_->dl_tensor), stream);
|
||||
}
|
||||
|
||||
inline int NDArray::use_count() const {
|
||||
if (data_ == nullptr) return 0;
|
||||
return data_->ref_counter_.load(std::memory_order_relaxed);
|
||||
}
|
||||
|
||||
inline const DGLArray* NDArray::operator->() const {
|
||||
return &(data_->dl_tensor);
|
||||
}
|
||||
|
||||
/** @brief Magic number for NDArray file */
|
||||
constexpr uint64_t kDGLNDArrayMagic = 0xDD5E40F096B4A13F;
|
||||
|
||||
inline bool SaveDGLArray(dmlc::Stream* strm, DGLArray* tensor) {
|
||||
uint64_t header = kDGLNDArrayMagic, reserved = 0;
|
||||
strm->Write(header);
|
||||
strm->Write(reserved);
|
||||
// Always save data as CPU context
|
||||
//
|
||||
// Parameters that get serialized should be in CPU by default.
|
||||
// So even the array's context is GPU, it will be stored as CPU array.
|
||||
// This is used to prevent case when another user loads the parameters
|
||||
// back on machine that do not have GPU or related context.
|
||||
//
|
||||
// We can always do array.CopyTo(target_ctx) to get a corresponding
|
||||
// array in the target context.
|
||||
DGLContext cpu_ctx;
|
||||
cpu_ctx.device_type = kDGLCPU;
|
||||
cpu_ctx.device_id = 0;
|
||||
strm->Write(cpu_ctx);
|
||||
strm->Write(tensor->ndim);
|
||||
strm->Write(tensor->dtype);
|
||||
int ndim = tensor->ndim;
|
||||
strm->WriteArray(tensor->shape, ndim);
|
||||
int type_bytes = tensor->dtype.bits / 8;
|
||||
int64_t num_elems = 1;
|
||||
for (int i = 0; i < ndim; ++i) {
|
||||
num_elems *= tensor->shape[i];
|
||||
}
|
||||
int64_t data_byte_size = type_bytes * num_elems;
|
||||
strm->Write(data_byte_size);
|
||||
|
||||
if (DMLC_IO_NO_ENDIAN_SWAP && tensor->ctx.device_type == kDGLCPU &&
|
||||
tensor->strides == nullptr && tensor->byte_offset == 0) {
|
||||
// quick path
|
||||
strm->Write(tensor->data, data_byte_size);
|
||||
} else {
|
||||
std::vector<uint8_t> bytes(data_byte_size);
|
||||
CHECK_EQ(
|
||||
DGLArrayCopyToBytes(tensor, dmlc::BeginPtr(bytes), data_byte_size), 0)
|
||||
<< DGLGetLastError();
|
||||
if (!DMLC_IO_NO_ENDIAN_SWAP) {
|
||||
dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
|
||||
}
|
||||
strm->Write(dmlc::BeginPtr(bytes), data_byte_size);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Convert type code to its name
|
||||
* @param type_code The type code .
|
||||
* @return The name of type code.
|
||||
*/
|
||||
inline const char* TypeCode2Str(int type_code) {
|
||||
switch (type_code) {
|
||||
case kDGLInt:
|
||||
return "int";
|
||||
case kDGLUInt:
|
||||
return "uint";
|
||||
case kDGLFloat:
|
||||
return "float";
|
||||
case kStr:
|
||||
return "str";
|
||||
case kBytes:
|
||||
return "bytes";
|
||||
case kHandle:
|
||||
return "handle";
|
||||
case kNull:
|
||||
return "NULL";
|
||||
case kObjectHandle:
|
||||
return "ObjectHandle";
|
||||
case kArrayHandle:
|
||||
return "ArrayHandle";
|
||||
case kDGLDataType:
|
||||
return "DGLDataType";
|
||||
case kDGLContext:
|
||||
return "DGLContext";
|
||||
case kFuncHandle:
|
||||
return "FunctionHandle";
|
||||
case kModuleHandle:
|
||||
return "ModuleHandle";
|
||||
case kNDArrayContainer:
|
||||
return "NDArrayContainer";
|
||||
default:
|
||||
LOG(FATAL) << "unknown type_code=" << static_cast<int>(type_code);
|
||||
return "";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Convert device type code to its name
|
||||
* @param device_type The device type code.
|
||||
* @return The name of the device.
|
||||
*/
|
||||
inline const char* DeviceTypeCode2Str(DGLDeviceType device_type) {
|
||||
switch (device_type) {
|
||||
case kDGLCPU:
|
||||
return "cpu";
|
||||
case kDGLCUDA:
|
||||
return "cuda";
|
||||
default:
|
||||
LOG(FATAL) << "Unsupported device type code="
|
||||
<< static_cast<int>(device_type);
|
||||
return "";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief convert a string to DGL type.
|
||||
* @param s The string to be converted.
|
||||
* @return The corresponding dgl type.
|
||||
*/
|
||||
inline DGLDataType String2DGLDataType(std::string s) {
|
||||
DGLDataType t;
|
||||
t.bits = 32;
|
||||
t.lanes = 1;
|
||||
const char* scan;
|
||||
if (s.substr(0, 3) == "int") {
|
||||
t.code = kDGLInt;
|
||||
scan = s.c_str() + 3;
|
||||
} else if (s.substr(0, 4) == "uint") {
|
||||
t.code = kDGLUInt;
|
||||
scan = s.c_str() + 4;
|
||||
} else if (s.substr(0, 5) == "float") {
|
||||
t.code = kDGLFloat;
|
||||
scan = s.c_str() + 5;
|
||||
} else if (s.substr(0, 6) == "handle") {
|
||||
t.code = kHandle;
|
||||
t.bits = 64; // handle uses 64 bit by default.
|
||||
scan = s.c_str() + 6;
|
||||
} else {
|
||||
scan = s.c_str();
|
||||
LOG(FATAL) << "unknown type " << s;
|
||||
}
|
||||
char* xdelim; // emulate sscanf("%ux%u", bits, lanes)
|
||||
uint8_t bits = static_cast<uint8_t>(strtoul(scan, &xdelim, 10));
|
||||
if (bits != 0) t.bits = bits;
|
||||
if (*xdelim == 'x') {
|
||||
t.lanes = static_cast<uint16_t>(strtoul(xdelim + 1, nullptr, 10));
|
||||
}
|
||||
return t;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief convert a DGL type to string.
|
||||
* @param t The type to be converted.
|
||||
* @return The corresponding dgl type in string.
|
||||
*/
|
||||
inline std::string DGLDataType2String(DGLDataType t) {
|
||||
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
|
||||
std::ostringstream os;
|
||||
os << t;
|
||||
return os.str();
|
||||
#else
|
||||
std::string repr = "";
|
||||
repr += TypeCode2Str(t.code);
|
||||
if (t.code == kHandle) return repr;
|
||||
repr += std::to_string(static_cast<int>(t.bits));
|
||||
if (t.lanes != 1) {
|
||||
repr += "x" + std::to_string(static_cast<int>(t.lanes));
|
||||
}
|
||||
return repr;
|
||||
#endif
|
||||
}
|
||||
|
||||
// macro to check type code.
|
||||
#define DGL_CHECK_TYPE_CODE(CODE, T) \
|
||||
CHECK_EQ(CODE, T) << " expected " << TypeCode2Str(T) << " but get " \
|
||||
<< TypeCode2Str(CODE)
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
namespace dmlc {
|
||||
DMLC_DECLARE_TRAITS(has_saveload, dgl::runtime::NDArray, true);
|
||||
} // namespace dmlc
|
||||
|
||||
///////////////// Operator overloading for NDArray /////////////////
|
||||
dgl::runtime::NDArray operator+(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator-(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator*(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator/(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator%(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator+(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator-(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator*(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator/(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator%(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator+(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator-(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator*(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator/(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator%(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator-(const dgl::runtime::NDArray& array);
|
||||
|
||||
dgl::runtime::NDArray operator>(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator<(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator>=(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator<=(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator==(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator!=(
|
||||
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator>(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator<(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator>=(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator<=(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator==(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator!=(const dgl::runtime::NDArray& a1, int64_t rhs);
|
||||
dgl::runtime::NDArray operator>(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator<(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator>=(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator<=(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator==(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
dgl::runtime::NDArray operator!=(int64_t lhs, const dgl::runtime::NDArray& a2);
|
||||
|
||||
std::ostream& operator<<(std::ostream& os, dgl::runtime::NDArray array);
|
||||
|
||||
///////////////// Operator overloading for DGLDataType /////////////////
|
||||
|
||||
/** @brief Check whether two data types are the same.*/
|
||||
inline bool operator==(const DGLDataType& ty1, const DGLDataType& ty2) {
|
||||
return ty1.code == ty2.code && ty1.bits == ty2.bits && ty1.lanes == ty2.lanes;
|
||||
}
|
||||
|
||||
/** @brief Check whether two data types are different.*/
|
||||
inline bool operator!=(const DGLDataType& ty1, const DGLDataType& ty2) {
|
||||
return !(ty1 == ty2);
|
||||
}
|
||||
|
||||
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
|
||||
inline std::ostream& operator<<(std::ostream& os, DGLDataType t) {
|
||||
os << dgl::runtime::TypeCode2Str(t.code);
|
||||
if (t.code == kHandle) return os;
|
||||
os << static_cast<int>(t.bits);
|
||||
if (t.lanes != 1) {
|
||||
os << 'x' << static_cast<int>(t.lanes);
|
||||
}
|
||||
return os;
|
||||
}
|
||||
#endif
|
||||
|
||||
///////////////// Operator overloading for DGLContext /////////////////
|
||||
|
||||
/** @brief Check whether two device contexts are the same.*/
|
||||
inline bool operator==(const DGLContext& ctx1, const DGLContext& ctx2) {
|
||||
return ctx1.device_type == ctx2.device_type &&
|
||||
ctx1.device_id == ctx2.device_id;
|
||||
}
|
||||
|
||||
/** @brief Check whether two device contexts are different.*/
|
||||
inline bool operator!=(const DGLContext& ctx1, const DGLContext& ctx2) {
|
||||
return !(ctx1 == ctx2);
|
||||
}
|
||||
|
||||
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
|
||||
inline std::ostream& operator<<(std::ostream& os, const DGLContext& ctx) {
|
||||
return os << dgl::runtime::DeviceTypeCode2Str(ctx.device_type) << ":"
|
||||
<< ctx.device_id;
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif // DGL_RUNTIME_NDARRAY_H_
|
||||
@@ -0,0 +1,328 @@
|
||||
/**
|
||||
* Copyright (c) 2019 by Contributors
|
||||
* @file runtime/object.h
|
||||
* @brief Defines the Object data structures.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_OBJECT_H_
|
||||
#define DGL_RUNTIME_OBJECT_H_
|
||||
|
||||
#include <dmlc/logging.h>
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
// forward declaration
|
||||
class Object;
|
||||
class ObjectRef;
|
||||
class NDArray;
|
||||
|
||||
/**
|
||||
* @brief Visitor class to each object attribute.
|
||||
* The content is going to be called for each field.
|
||||
*/
|
||||
class AttrVisitor {
|
||||
public:
|
||||
//! \cond Doxygen_Suppress
|
||||
virtual void Visit(const char* key, double* value) = 0;
|
||||
virtual void Visit(const char* key, int64_t* value) = 0;
|
||||
virtual void Visit(const char* key, uint64_t* value) = 0;
|
||||
virtual void Visit(const char* key, int* value) = 0;
|
||||
virtual void Visit(const char* key, bool* value) = 0;
|
||||
virtual void Visit(const char* key, std::string* value) = 0;
|
||||
virtual void Visit(const char* key, ObjectRef* value) = 0;
|
||||
virtual void Visit(const char* key, NDArray* value) = 0;
|
||||
template <
|
||||
typename ENum,
|
||||
typename = typename std::enable_if<std::is_enum<ENum>::value>::type>
|
||||
void Visit(const char* key, ENum* ptr) {
|
||||
static_assert(
|
||||
std::is_same<int, typename std::underlying_type<ENum>::type>::value,
|
||||
"declare enum to be enum int to use visitor");
|
||||
this->Visit(key, reinterpret_cast<int*>(ptr));
|
||||
}
|
||||
//! \endcond
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief base class of object container.
|
||||
* All object's internal is stored as std::shared_ptr<Object>
|
||||
*/
|
||||
class Object {
|
||||
public:
|
||||
/** @brief virtual destructor */
|
||||
virtual ~Object() {}
|
||||
/** @return The unique type key of the object */
|
||||
virtual const char* type_key() const = 0;
|
||||
/**
|
||||
* @brief Apply visitor to each field of the Object
|
||||
* Visitor could mutate the content of the object.
|
||||
* override if Object contains attribute fields.
|
||||
* @param visitor The visitor
|
||||
*/
|
||||
virtual void VisitAttrs(AttrVisitor* visitor) {}
|
||||
/** @return the type index of the object */
|
||||
virtual uint32_t type_index() const = 0;
|
||||
/**
|
||||
* @brief Whether this object derives from object with type_index=tid.
|
||||
* Implemented by DGL_DECLARE_OBJECT_TYPE_INFO
|
||||
*
|
||||
* @param tid The type index.
|
||||
* @return the check result.
|
||||
*/
|
||||
virtual bool _DerivedFrom(uint32_t tid) const;
|
||||
/**
|
||||
* @brief get a runtime unique type index given a type key
|
||||
* @param type_key Type key of a type.
|
||||
* @return the corresponding type index.
|
||||
*/
|
||||
static uint32_t TypeKey2Index(const char* type_key);
|
||||
/**
|
||||
* @brief get type key from type index.
|
||||
* @param index The type index
|
||||
* @return the corresponding type key.
|
||||
*/
|
||||
static const char* TypeIndex2Key(uint32_t index);
|
||||
/**
|
||||
* @return whether the type is derived from
|
||||
*/
|
||||
template <typename T>
|
||||
inline bool derived_from() const;
|
||||
/**
|
||||
* @return whether the object is of type T
|
||||
* @tparam The type to be checked.
|
||||
*/
|
||||
template <typename T>
|
||||
inline bool is_type() const;
|
||||
// object ref can see this
|
||||
friend class ObjectRef;
|
||||
static constexpr const char* _type_key = "Object";
|
||||
};
|
||||
|
||||
/** @brief base class of all reference object */
|
||||
class ObjectRef {
|
||||
public:
|
||||
/** @brief type indicate the container type */
|
||||
using ContainerType = Object;
|
||||
/**
|
||||
* @brief Comparator
|
||||
*
|
||||
* Compare with the two are referencing to the same object (compare by
|
||||
* address).
|
||||
*
|
||||
* @param other Another object ref.
|
||||
* @return the compare result.
|
||||
* @sa same_as
|
||||
*/
|
||||
inline bool operator==(const ObjectRef& other) const;
|
||||
/**
|
||||
* @brief Comparator
|
||||
*
|
||||
* Compare with the two are referencing to the same object (compare by
|
||||
* address).
|
||||
*
|
||||
* @param other Another object ref.
|
||||
* @return the compare result.
|
||||
*/
|
||||
inline bool same_as(const ObjectRef& other) const;
|
||||
/**
|
||||
* @brief Comparator
|
||||
*
|
||||
* The operator overload allows ObjectRef be used in std::map.
|
||||
*
|
||||
* @param other Another object ref.
|
||||
* @return the compare result.
|
||||
*/
|
||||
inline bool operator<(const ObjectRef& other) const;
|
||||
/**
|
||||
* @brief Comparator
|
||||
* @param other Another object ref.
|
||||
* @return the compare result.
|
||||
* @sa same_as
|
||||
*/
|
||||
inline bool operator!=(const ObjectRef& other) const;
|
||||
/** @return the hash function for ObjectRef */
|
||||
inline size_t hash() const;
|
||||
/** @return whether the expression is null */
|
||||
inline bool defined() const;
|
||||
/** @return the internal type index of Object */
|
||||
inline uint32_t type_index() const;
|
||||
/** @return the internal object pointer */
|
||||
inline const Object* get() const;
|
||||
/** @return the internal object pointer */
|
||||
inline const Object* operator->() const;
|
||||
/**
|
||||
* @brief Downcast this object to its actual type.
|
||||
* This returns nullptr if the object is not of the requested type.
|
||||
* Example usage:
|
||||
*
|
||||
* if (const Banana *banana = obj->as<Banana>()) {
|
||||
* // This is a Banana!
|
||||
* }
|
||||
* @tparam T the target type, must be subtype of Object
|
||||
*/
|
||||
template <typename T>
|
||||
inline const T* as() const;
|
||||
|
||||
/** @brief default constructor */
|
||||
ObjectRef() = default;
|
||||
explicit ObjectRef(std::shared_ptr<Object> obj) : obj_(obj) {}
|
||||
|
||||
/** @brief the internal object, do not touch */
|
||||
std::shared_ptr<Object> obj_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief helper macro to declare type information in a base object.
|
||||
*
|
||||
* This is macro should be used in abstract base class definition
|
||||
* because it does not define type_key and type_index.
|
||||
*/
|
||||
#define DGL_DECLARE_BASE_OBJECT_INFO(TypeName, Parent) \
|
||||
const bool _DerivedFrom(uint32_t tid) const override { \
|
||||
static uint32_t tidx = TypeKey2Index(TypeName::_type_key); \
|
||||
if (tidx == tid) return true; \
|
||||
return Parent::_DerivedFrom(tid); \
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief helper macro to declare type information in a terminal class
|
||||
*
|
||||
* This is macro should be used in terminal class definition.
|
||||
*
|
||||
* For example:
|
||||
*
|
||||
* // This class is an abstract class and cannot create instances
|
||||
* class SomeBaseClass : public Object {
|
||||
* public:
|
||||
* static constexpr const char* _type_key = "some_base";
|
||||
* DGL_DECLARE_BASE_OBJECT_INFO(SomeBaseClass, Object);
|
||||
* };
|
||||
*
|
||||
* // Child class that allows instantiation
|
||||
* class SomeChildClass : public SomeBaseClass {
|
||||
* public:
|
||||
* static constexpr const char* _type_key = "some_child";
|
||||
* DGL_DECLARE_OBJECT_TYPE_INFO(SomeChildClass, SomeBaseClass);
|
||||
* };
|
||||
*/
|
||||
#define DGL_DECLARE_OBJECT_TYPE_INFO(TypeName, Parent) \
|
||||
const char* type_key() const final { return TypeName::_type_key; } \
|
||||
uint32_t type_index() const final { \
|
||||
static uint32_t tidx = TypeKey2Index(TypeName::_type_key); \
|
||||
return tidx; \
|
||||
} \
|
||||
bool _DerivedFrom(uint32_t tid) const final { \
|
||||
static uint32_t tidx = TypeKey2Index(TypeName::_type_key); \
|
||||
if (tidx == tid) return true; \
|
||||
return Parent::_DerivedFrom(tid); \
|
||||
}
|
||||
|
||||
/** @brief Macro to generate common object reference class method definition */
|
||||
#define DGL_DEFINE_OBJECT_REF_METHODS(TypeName, BaseTypeName, ObjectName) \
|
||||
TypeName() {} \
|
||||
explicit TypeName(std::shared_ptr<runtime::Object> obj) \
|
||||
: BaseTypeName(obj) {} \
|
||||
const ObjectName* operator->() const { \
|
||||
return static_cast<const ObjectName*>(obj_.get()); \
|
||||
} \
|
||||
ObjectName* operator->() { return static_cast<ObjectName*>(obj_.get()); } \
|
||||
std::shared_ptr<ObjectName> sptr() const { \
|
||||
return CHECK_NOTNULL(std::dynamic_pointer_cast<ObjectName>(obj_)); \
|
||||
} \
|
||||
operator bool() const { return this->defined(); } \
|
||||
using ContainerType = ObjectName
|
||||
|
||||
/** @brief Macro to generate object reference class definition */
|
||||
#define DGL_DEFINE_OBJECT_REF(TypeName, ObjectName) \
|
||||
class TypeName : public ::dgl::runtime::ObjectRef { \
|
||||
public: \
|
||||
DGL_DEFINE_OBJECT_REF_METHODS( \
|
||||
TypeName, ::dgl::runtime::ObjectRef, ObjectName); \
|
||||
}
|
||||
|
||||
// implementations of inline functions after this
|
||||
template <typename T>
|
||||
inline bool Object::is_type() const {
|
||||
// use static field so query only happens once.
|
||||
static uint32_t type_id = Object::TypeKey2Index(T::_type_key);
|
||||
return type_id == this->type_index();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline bool Object::derived_from() const {
|
||||
// use static field so query only happens once.
|
||||
static uint32_t type_id = Object::TypeKey2Index(T::_type_key);
|
||||
return this->_DerivedFrom(type_id);
|
||||
}
|
||||
|
||||
inline const Object* ObjectRef::get() const { return obj_.get(); }
|
||||
|
||||
inline const Object* ObjectRef::operator->() const { return obj_.get(); }
|
||||
|
||||
inline bool ObjectRef::defined() const { return obj_.get() != nullptr; }
|
||||
|
||||
inline bool ObjectRef::operator==(const ObjectRef& other) const {
|
||||
return obj_.get() == other.obj_.get();
|
||||
}
|
||||
|
||||
inline bool ObjectRef::same_as(const ObjectRef& other) const {
|
||||
return obj_.get() == other.obj_.get();
|
||||
}
|
||||
|
||||
inline bool ObjectRef::operator<(const ObjectRef& other) const {
|
||||
return obj_.get() < other.obj_.get();
|
||||
}
|
||||
|
||||
inline bool ObjectRef::operator!=(const ObjectRef& other) const {
|
||||
return obj_.get() != other.obj_.get();
|
||||
}
|
||||
|
||||
inline size_t ObjectRef::hash() const {
|
||||
return std::hash<Object*>()(obj_.get());
|
||||
}
|
||||
|
||||
inline uint32_t ObjectRef::type_index() const {
|
||||
CHECK(obj_.get() != nullptr) << "null type";
|
||||
return get()->type_index();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline const T* ObjectRef::as() const {
|
||||
const Object* ptr = get();
|
||||
if (ptr && ptr->is_type<T>()) {
|
||||
return static_cast<const T*>(ptr);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
/** @brief The hash function for nodes */
|
||||
struct ObjectHash {
|
||||
size_t operator()(const ObjectRef& a) const { return a.hash(); }
|
||||
};
|
||||
|
||||
/** @brief The equal comparator for nodes */
|
||||
struct ObjectEqual {
|
||||
bool operator()(const ObjectRef& a, const ObjectRef& b) const {
|
||||
return a.get() == b.get();
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
namespace std {
|
||||
template <>
|
||||
struct hash<::dgl::runtime::ObjectRef> {
|
||||
std::size_t operator()(const ::dgl::runtime::ObjectRef& k) const {
|
||||
return k.hash();
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace std
|
||||
|
||||
#endif // DGL_RUNTIME_OBJECT_H_
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,182 @@
|
||||
/**
|
||||
* Copyright (c) 2021 by Contributors
|
||||
* @file runtime/container.h
|
||||
* @brief Defines the container object data structures.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_PARALLEL_FOR_H_
|
||||
#define DGL_RUNTIME_PARALLEL_FOR_H_
|
||||
|
||||
#include <dgl/env_variable.h>
|
||||
#include <dmlc/omp.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <atomic>
|
||||
#include <cstdlib>
|
||||
#include <exception>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace {
|
||||
int64_t divup(int64_t x, int64_t y) { return (x + y - 1) / y; }
|
||||
} // namespace
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
namespace {
|
||||
struct DefaultGrainSizeT {
|
||||
size_t grain_size;
|
||||
|
||||
DefaultGrainSizeT() : DefaultGrainSizeT(1) {}
|
||||
|
||||
explicit DefaultGrainSizeT(size_t default_grain_size) {
|
||||
auto var = dgl::kDGLParallelForGrainSize;
|
||||
|
||||
if (var) {
|
||||
grain_size = std::stoul(var);
|
||||
} else {
|
||||
grain_size = default_grain_size;
|
||||
}
|
||||
}
|
||||
|
||||
size_t operator()() { return grain_size; }
|
||||
};
|
||||
} // namespace
|
||||
|
||||
inline size_t compute_num_threads(size_t begin, size_t end, size_t grain_size) {
|
||||
#ifdef _OPENMP
|
||||
if (omp_in_parallel() || end - begin <= grain_size || end - begin == 1)
|
||||
return 1;
|
||||
|
||||
return std::min(
|
||||
static_cast<int64_t>(omp_get_max_threads()),
|
||||
divup(end - begin, grain_size));
|
||||
#else
|
||||
return 1;
|
||||
#endif
|
||||
}
|
||||
|
||||
static DefaultGrainSizeT default_grain_size;
|
||||
|
||||
/**
|
||||
* @brief OpenMP-based parallel for loop.
|
||||
*
|
||||
* It requires each thread's workload to have at least \a grain_size elements.
|
||||
* The loop body will be a function that takes in two arguments \a begin and \a
|
||||
* end, which stands for the starting (inclusive) and ending index (exclusive)
|
||||
* of the workload.
|
||||
*/
|
||||
template <typename F>
|
||||
void parallel_for(
|
||||
const size_t begin, const size_t end, const size_t grain_size, F&& f) {
|
||||
if (begin >= end) {
|
||||
return;
|
||||
}
|
||||
|
||||
#ifdef _OPENMP
|
||||
auto num_threads = compute_num_threads(begin, end, grain_size);
|
||||
// (BarclayII) the exception code is borrowed from PyTorch.
|
||||
std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
|
||||
std::exception_ptr eptr;
|
||||
|
||||
#pragma omp parallel num_threads(num_threads)
|
||||
{
|
||||
auto tid = omp_get_thread_num();
|
||||
auto chunk_size = divup((end - begin), num_threads);
|
||||
auto begin_tid = begin + tid * chunk_size;
|
||||
if (begin_tid < end) {
|
||||
auto end_tid = std::min(end, static_cast<size_t>(chunk_size + begin_tid));
|
||||
try {
|
||||
f(begin_tid, end_tid);
|
||||
} catch (...) {
|
||||
if (!err_flag.test_and_set()) eptr = std::current_exception();
|
||||
}
|
||||
}
|
||||
}
|
||||
if (eptr) std::rethrow_exception(eptr);
|
||||
#else
|
||||
f(begin, end);
|
||||
#endif
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief OpenMP-based parallel for loop with default grain size.
|
||||
*
|
||||
* parallel_for with grain size to default value, either 1 or controlled through
|
||||
* environment variable DGL_PARALLEL_FOR_GRAIN_SIZE.
|
||||
* If grain size is set to 1, the function behaves the same way as OpenMP
|
||||
* parallel for pragma with static scheduling.
|
||||
*/
|
||||
template <typename F>
|
||||
void parallel_for(const size_t begin, const size_t end, F&& f) {
|
||||
parallel_for(begin, end, default_grain_size(), std::forward<F>(f));
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief OpenMP-based two-stage parallel reduction.
|
||||
*
|
||||
* The first-stage reduction function \a f works in parallel. Each thread's
|
||||
* workload has at least \a grain_size elements. The loop body will be a
|
||||
* function that takes in the starting index (inclusive), the ending index
|
||||
* (exclusive), and the reduction identity.
|
||||
*
|
||||
* The second-stage reduction function \a sf is a binary function working in the
|
||||
* main thread. It aggregates the partially reduced result computed from each
|
||||
* thread.
|
||||
*
|
||||
* Example to compute a parallelized max reduction of an array \c a:
|
||||
*
|
||||
* parallel_reduce(
|
||||
* 0, // starting index
|
||||
* 100, // ending index
|
||||
* 1, // grain size
|
||||
* -std::numeric_limits<float>::infinity, // identity
|
||||
* [&a] (int begin, int end, float ident) { // first-stage partial
|
||||
* reducer float result = ident; for (int i = begin; i < end; ++i) result =
|
||||
* std::max(result, a[i]); return result;
|
||||
* },
|
||||
* [] (float result, float partial_result) {
|
||||
* return std::max(result, partial_result);
|
||||
* });
|
||||
*/
|
||||
template <typename DType, typename F, typename SF>
|
||||
DType parallel_reduce(
|
||||
const size_t begin, const size_t end, const size_t grain_size,
|
||||
const DType ident, const F& f, const SF& sf) {
|
||||
if (begin >= end) {
|
||||
return ident;
|
||||
}
|
||||
|
||||
int num_threads = compute_num_threads(begin, end, grain_size);
|
||||
if (num_threads == 1) {
|
||||
return f(begin, end, ident);
|
||||
}
|
||||
|
||||
std::vector<DType> results(num_threads, ident);
|
||||
std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
|
||||
std::exception_ptr eptr;
|
||||
#pragma omp parallel num_threads(num_threads)
|
||||
{
|
||||
auto tid = omp_get_thread_num();
|
||||
auto chunk_size = divup((end - begin), num_threads);
|
||||
auto begin_tid = begin + tid * chunk_size;
|
||||
if (begin_tid < end) {
|
||||
auto end_tid = std::min(end, static_cast<size_t>(chunk_size + begin_tid));
|
||||
try {
|
||||
results[tid] = f(begin_tid, end_tid, ident);
|
||||
} catch (...) {
|
||||
if (!err_flag.test_and_set()) eptr = std::current_exception();
|
||||
}
|
||||
}
|
||||
}
|
||||
if (eptr) std::rethrow_exception(eptr);
|
||||
|
||||
DType out = ident;
|
||||
for (int64_t i = 0; i < num_threads; ++i) out = sf(out, results[i]);
|
||||
return out;
|
||||
}
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_RUNTIME_PARALLEL_FOR_H_
|
||||
@@ -0,0 +1,146 @@
|
||||
/**
|
||||
* Copyright (c) 2017 by Contributors
|
||||
* @file dgl/runtime/registry.h
|
||||
* @brief This file defines the DGL global function registry.
|
||||
*
|
||||
* The registered functions will be made available to front-end
|
||||
* as well as backend users.
|
||||
*
|
||||
* The registry stores type-erased functions.
|
||||
* Each registered function is automatically exposed
|
||||
* to front-end language(e.g. python).
|
||||
*
|
||||
* Front-end can also pass callbacks as PackedFunc, or register
|
||||
* then into the same global registry in C++.
|
||||
* The goal is to mix the front-end language and the DGL back-end.
|
||||
*
|
||||
* @code
|
||||
* // register the function as MyAPIFuncName
|
||||
* DGL_REGISTER_GLOBAL(MyAPIFuncName)
|
||||
* .set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
* // my code.
|
||||
* });
|
||||
* @endcode
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_REGISTRY_H_
|
||||
#define DGL_RUNTIME_REGISTRY_H_
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "packed_func.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
/** @brief Registry for global function */
|
||||
class Registry {
|
||||
public:
|
||||
/**
|
||||
* @brief set the body of the function to be f
|
||||
* @param f The body of the function.
|
||||
*/
|
||||
DGL_DLL Registry& set_body(PackedFunc f); // NOLINT(*)
|
||||
/**
|
||||
* @brief set the body of the function to be f
|
||||
* @param f The body of the function.
|
||||
*/
|
||||
Registry& set_body(PackedFunc::FType f) { // NOLINT(*)
|
||||
return set_body(PackedFunc(f));
|
||||
}
|
||||
/**
|
||||
* @brief set the body of the function to be TypedPackedFunc.
|
||||
*
|
||||
* @code
|
||||
*
|
||||
* DGL_REGISTER_API("addone")
|
||||
* .set_body_typed<int(int)>([](int x) { return x + 1; });
|
||||
*
|
||||
* @endcode
|
||||
*
|
||||
* @param f The body of the function.
|
||||
* @tparam FType the signature of the function.
|
||||
* @tparam FLambda The type of f.
|
||||
*/
|
||||
template <typename FType, typename FLambda>
|
||||
Registry& set_body_typed(FLambda f) {
|
||||
return set_body(TypedPackedFunc<FType>(f).packed());
|
||||
}
|
||||
/**
|
||||
* @brief Register a function with given name
|
||||
* @param name The name of the function.
|
||||
* @param override Whether allow oveeride existing function.
|
||||
* @return Reference to theregistry.
|
||||
*/
|
||||
DGL_DLL static Registry& Register(
|
||||
const std::string& name, bool override = false); // NOLINT(*)
|
||||
/**
|
||||
* @brief Erase global function from registry, if exist.
|
||||
* @param name The name of the function.
|
||||
* @return Whether function exist.
|
||||
*/
|
||||
DGL_DLL static bool Remove(const std::string& name);
|
||||
/**
|
||||
* @brief Get the global function by name.
|
||||
* @param name The name of the function.
|
||||
* @return pointer to the registered function,
|
||||
* nullptr if it does not exist.
|
||||
*/
|
||||
DGL_DLL static const PackedFunc* Get(const std::string& name); // NOLINT(*)
|
||||
/**
|
||||
* @brief Get the names of currently registered global function.
|
||||
* @return The names
|
||||
*/
|
||||
DGL_DLL static std::vector<std::string> ListNames();
|
||||
|
||||
// Internal class.
|
||||
struct Manager;
|
||||
|
||||
protected:
|
||||
/** @brief name of the function */
|
||||
std::string name_;
|
||||
/** @brief internal packed function */
|
||||
PackedFunc func_;
|
||||
friend struct Manager;
|
||||
};
|
||||
|
||||
/** @brief helper macro to supress unused warning */
|
||||
#if defined(__GNUC__)
|
||||
#define DGL_ATTRIBUTE_UNUSED __attribute__((unused))
|
||||
#else
|
||||
#define DGL_ATTRIBUTE_UNUSED
|
||||
#endif
|
||||
|
||||
#define DGL_STR_CONCAT_(__x, __y) __x##__y
|
||||
#define DGL_STR_CONCAT(__x, __y) DGL_STR_CONCAT_(__x, __y)
|
||||
|
||||
#define DGL_FUNC_REG_VAR_DEF \
|
||||
static DGL_ATTRIBUTE_UNUSED ::dgl::runtime::Registry& __mk_##DGL
|
||||
|
||||
#define DGL_TYPE_REG_VAR_DEF \
|
||||
static DGL_ATTRIBUTE_UNUSED ::dgl::runtime::ExtTypeVTable* __mk_##DGLT
|
||||
|
||||
/**
|
||||
* @brief Register a function globally.
|
||||
* @code
|
||||
* DGL_REGISTER_GLOBAL("MyPrint")
|
||||
* .set_body([](DGLArgs args, DGLRetValue* rv) {
|
||||
* });
|
||||
* @endcode
|
||||
*/
|
||||
#define DGL_REGISTER_GLOBAL(OpName) \
|
||||
DGL_STR_CONCAT(DGL_FUNC_REG_VAR_DEF, __COUNTER__) = \
|
||||
::dgl::runtime::Registry::Register(OpName)
|
||||
|
||||
/**
|
||||
* @brief Macro to register extension type.
|
||||
* This must be registered in a cc file
|
||||
* after the trait extension_class_info is defined.
|
||||
*/
|
||||
#define DGL_REGISTER_EXT_TYPE(T) \
|
||||
DGL_STR_CONCAT(DGL_TYPE_REG_VAR_DEF, __COUNTER__) = \
|
||||
::dgl::runtime::ExtTypeVTable::Register_<T>()
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
#endif // DGL_RUNTIME_REGISTRY_H_
|
||||
@@ -0,0 +1,52 @@
|
||||
/**
|
||||
* Copyright (c) 2017 by Contributors
|
||||
* @file dgl/runtime/serializer.h
|
||||
* @brief Serializer extension to support DGL data types
|
||||
* Include this file to enable serialization of DGLDataType, DGLContext
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_SERIALIZER_H_
|
||||
#define DGL_RUNTIME_SERIALIZER_H_
|
||||
|
||||
#include <dmlc/io.h>
|
||||
#include <dmlc/serializer.h>
|
||||
|
||||
#include "c_runtime_api.h"
|
||||
#include "smart_ptr_serializer.h"
|
||||
|
||||
namespace dmlc {
|
||||
namespace serializer {
|
||||
|
||||
template <>
|
||||
struct Handler<DGLDataType> {
|
||||
inline static void Write(Stream *strm, const DGLDataType &dtype) {
|
||||
Handler<uint8_t>::Write(strm, dtype.code);
|
||||
Handler<uint8_t>::Write(strm, dtype.bits);
|
||||
Handler<uint16_t>::Write(strm, dtype.lanes);
|
||||
}
|
||||
inline static bool Read(Stream *strm, DGLDataType *dtype) {
|
||||
if (!Handler<uint8_t>::Read(strm, &(dtype->code))) return false;
|
||||
if (!Handler<uint8_t>::Read(strm, &(dtype->bits))) return false;
|
||||
if (!Handler<uint16_t>::Read(strm, &(dtype->lanes))) return false;
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Handler<DGLContext> {
|
||||
inline static void Write(Stream *strm, const DGLContext &ctx) {
|
||||
int32_t device_type = static_cast<int32_t>(ctx.device_type);
|
||||
Handler<int32_t>::Write(strm, device_type);
|
||||
Handler<int32_t>::Write(strm, ctx.device_id);
|
||||
}
|
||||
inline static bool Read(Stream *strm, DGLContext *ctx) {
|
||||
int32_t device_type = 0;
|
||||
if (!Handler<int32_t>::Read(strm, &(device_type))) return false;
|
||||
ctx->device_type = static_cast<DGLDeviceType>(device_type);
|
||||
if (!Handler<int32_t>::Read(strm, &(ctx->device_id))) return false;
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace serializer
|
||||
} // namespace dmlc
|
||||
#endif // DGL_RUNTIME_SERIALIZER_H_
|
||||
@@ -0,0 +1,91 @@
|
||||
/**
|
||||
* Copyright (c) 2017 by Contributors
|
||||
* @file dgl/runtime/ndarray.h
|
||||
* @brief shared memory management.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_SHARED_MEM_H_
|
||||
#define DGL_RUNTIME_SHARED_MEM_H_
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <windows.h>
|
||||
#endif // _WIN32
|
||||
#include <string>
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
/**
|
||||
* @brief This class owns shared memory.
|
||||
*
|
||||
* When the object is gone, the shared memory will also be destroyed.
|
||||
* When the shared memory is destroyed, the file corresponding to
|
||||
* the shared memory is removed.
|
||||
*/
|
||||
class SharedMemory {
|
||||
/**
|
||||
* @brief whether the shared memory is owned by the object.
|
||||
*
|
||||
* If shared memory is created in the object, it'll be owned by the object
|
||||
* and will be responsible for deleting it when the object is destroyed.
|
||||
*/
|
||||
bool own_;
|
||||
|
||||
/* @brief the file descripter of the shared memory. */
|
||||
#ifndef _WIN32
|
||||
int fd_;
|
||||
#else // !_WIN32
|
||||
HANDLE handle_;
|
||||
#endif // _WIN32
|
||||
/* @brief the address of the shared memory. */
|
||||
void *ptr_;
|
||||
/* @brief the size of the shared memory. */
|
||||
size_t size_;
|
||||
|
||||
/**
|
||||
* @brief the name of the object.
|
||||
*
|
||||
* In Unix, shared memory is identified by a file. Thus, `name` is actually
|
||||
* the file name that identifies the shared memory.
|
||||
*/
|
||||
std::string name;
|
||||
|
||||
public:
|
||||
/* @brief Get the filename of shared memory file
|
||||
*/
|
||||
std::string GetName() const { return name; }
|
||||
|
||||
/**
|
||||
* @brief constructor of the shared memory.
|
||||
* @param name The file corresponding to the shared memory.
|
||||
*/
|
||||
explicit SharedMemory(const std::string &name);
|
||||
/**
|
||||
* @brief destructor of the shared memory.
|
||||
* It deallocates the shared memory and removes the corresponding file.
|
||||
*/
|
||||
~SharedMemory();
|
||||
/**
|
||||
* @brief create shared memory.
|
||||
* It creates the file and shared memory.
|
||||
* @param sz the size of the shared memory.
|
||||
* @return the address of the shared memory
|
||||
*/
|
||||
void *CreateNew(size_t sz);
|
||||
/**
|
||||
* @brief allocate shared memory that has been created.
|
||||
* @param sz the size of the shared memory.
|
||||
* @return the address of the shared memory
|
||||
*/
|
||||
void *Open(size_t sz);
|
||||
|
||||
/**
|
||||
* @brief check if the shared memory exist.
|
||||
* @param name the name of the shared memory.
|
||||
* @return a boolean value to indicate if the shared memory exists.
|
||||
*/
|
||||
static bool Exist(const std::string &name);
|
||||
};
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
#endif // DGL_RUNTIME_SHARED_MEM_H_
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* Copyright (c) 2017 by Contributors
|
||||
* @file dgl/runtime/serializer.h
|
||||
* @brief Serializer extension to support DGL data types
|
||||
* Include this file to enable serialization of DGLDataType, DGLContext
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_SMART_PTR_SERIALIZER_H_
|
||||
#define DGL_RUNTIME_SMART_PTR_SERIALIZER_H_
|
||||
|
||||
#include <dgl/graph_serializer.h>
|
||||
#include <dmlc/io.h>
|
||||
#include <dmlc/serializer.h>
|
||||
|
||||
#include <memory>
|
||||
|
||||
namespace dmlc {
|
||||
namespace serializer {
|
||||
|
||||
//! \cond Doxygen_Suppress
|
||||
template <typename T>
|
||||
struct Handler<std::shared_ptr<T>> {
|
||||
inline static void Write(Stream *strm, const std::shared_ptr<T> &data) {
|
||||
Handler<T>::Write(strm, *data.get());
|
||||
}
|
||||
inline static bool Read(Stream *strm, std::shared_ptr<T> *data) {
|
||||
// When read, the default initialization behavior of shared_ptr is
|
||||
// shared_ptr<T>(), which is holding a nullptr. Here we need to manually
|
||||
// reset to a real object for further loading
|
||||
if (!(*data)) {
|
||||
data->reset(dgl::Serializer::new_object<T>());
|
||||
}
|
||||
return Handler<T>::Read(strm, data->get());
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct Handler<std::unique_ptr<T>> {
|
||||
inline static void Write(Stream *strm, const std::unique_ptr<T> &data) {
|
||||
Handler<T>::Write(strm, *data.get());
|
||||
}
|
||||
inline static bool Read(Stream *strm, std::unique_ptr<T> *data) {
|
||||
// When read, the default initialization behavior of unique_ptr is
|
||||
// unique_ptr<T>(), which is holding a nullptr. Here we need to manually
|
||||
// reset to a real object for further loading
|
||||
if (!(*data)) {
|
||||
data->reset(dgl::Serializer::new_object<T>());
|
||||
}
|
||||
return Handler<T>::Read(strm, data->get());
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace serializer
|
||||
} // namespace dmlc
|
||||
#endif // DGL_RUNTIME_SMART_PTR_SERIALIZER_H_
|
||||
@@ -0,0 +1,281 @@
|
||||
/**
|
||||
* 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_
|
||||
@@ -0,0 +1,86 @@
|
||||
/**
|
||||
* Copyright (c) 2018 by Contributors
|
||||
* @file dgl/runtime/threading_backend.h
|
||||
* @brief Utilities for manipulating thread pool threads.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_THREADING_BACKEND_H_
|
||||
#define DGL_RUNTIME_THREADING_BACKEND_H_
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
namespace threading {
|
||||
|
||||
/**
|
||||
* @brief A platform-agnostic abstraction for managing a collection of
|
||||
* thread pool threads.
|
||||
*/
|
||||
class ThreadGroup {
|
||||
public:
|
||||
class Impl;
|
||||
|
||||
/**
|
||||
* @brief Creates a collection of threads which run a provided function.
|
||||
*
|
||||
* @param num_workers The total number of worker threads in this group.
|
||||
Includes main thread if `exclude_worker0 = true`
|
||||
* @param worker_callback A callback which is run in its own thread.
|
||||
Receives the worker_id as an argument.
|
||||
* @param exclude_worker0 Whether to use the main thread as a worker.
|
||||
* If `true`, worker0 will not be launched in a new thread and
|
||||
* `worker_callback` will only be called for values >= 1. This
|
||||
* allows use of the main thread as a worker.
|
||||
*/
|
||||
ThreadGroup(
|
||||
int num_workers, std::function<void(int)> worker_callback,
|
||||
bool exclude_worker0 = false);
|
||||
~ThreadGroup();
|
||||
|
||||
/**
|
||||
* @brief Blocks until all non-main threads in the pool finish.
|
||||
*/
|
||||
void Join();
|
||||
|
||||
enum AffinityMode : int {
|
||||
kBig = 1,
|
||||
kLittle = -1,
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief configure the CPU id affinity
|
||||
*
|
||||
* @param mode The preferred CPU type (1 = big, -1 = little).
|
||||
* @param nthreads The number of threads to use (0 = use all).
|
||||
* @param exclude_worker0 Whether to use the main thread as a worker.
|
||||
* If `true`, worker0 will not be launched in a new thread and
|
||||
* `worker_callback` will only be called for values >= 1. This
|
||||
* allows use of the main thread as a worker.
|
||||
*
|
||||
* @return The number of workers to use.
|
||||
*/
|
||||
int Configure(AffinityMode mode, int nthreads, bool exclude_worker0);
|
||||
|
||||
private:
|
||||
Impl* impl_;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief Platform-agnostic no-op.
|
||||
*/
|
||||
// This used to be Yield(), renaming to YieldThread() because windows.h defined
|
||||
// it as a macro in later SDKs.
|
||||
void YieldThread();
|
||||
|
||||
/**
|
||||
* @return the maximum number of effective workers for this system.
|
||||
*/
|
||||
int MaxConcurrency();
|
||||
|
||||
} // namespace threading
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
|
||||
#endif // DGL_RUNTIME_THREADING_BACKEND_H_
|
||||
@@ -0,0 +1,53 @@
|
||||
/**
|
||||
* Copyright (c) 2017 by Contributors
|
||||
* @file dgl/runtime/util.h
|
||||
* @brief Useful runtime util.
|
||||
*/
|
||||
#ifndef DGL_RUNTIME_UTIL_H_
|
||||
#define DGL_RUNTIME_UTIL_H_
|
||||
|
||||
#include "c_runtime_api.h"
|
||||
|
||||
namespace dgl {
|
||||
namespace runtime {
|
||||
|
||||
/**
|
||||
* @brief Check whether type matches the given spec.
|
||||
* @param t The type
|
||||
* @param code The type code.
|
||||
* @param bits The number of bits to be matched.
|
||||
* @param lanes The number of lanes sin the type.
|
||||
*/
|
||||
inline bool TypeMatch(DGLDataType t, int code, int bits, int lanes = 1) {
|
||||
return t.code == code && t.bits == bits && t.lanes == lanes;
|
||||
}
|
||||
} // namespace runtime
|
||||
} // namespace dgl
|
||||
// Forward declare the intrinsic id we need
|
||||
// in structure fetch to enable stackvm in runtime
|
||||
namespace dgl {
|
||||
namespace ir {
|
||||
namespace intrinsic {
|
||||
/** @brief The kind of structure field info used in intrinsic */
|
||||
enum DGLStructFieldKind : int {
|
||||
// array head address
|
||||
kArrAddr,
|
||||
kArrData,
|
||||
kArrShape,
|
||||
kArrStrides,
|
||||
kArrNDim,
|
||||
kArrTypeCode,
|
||||
kArrTypeBits,
|
||||
kArrTypeLanes,
|
||||
kArrByteOffset,
|
||||
kArrDeviceId,
|
||||
kArrDeviceType,
|
||||
kArrKindBound_,
|
||||
// DGLValue field
|
||||
kDGLValueContent,
|
||||
kDGLValueKindBound_
|
||||
};
|
||||
} // namespace intrinsic
|
||||
} // namespace ir
|
||||
} // namespace dgl
|
||||
#endif // DGL_RUNTIME_UTIL_H_
|
||||
Reference in New Issue
Block a user