wehub-resource-sync
435 lines
23 KiB
C++
435 lines
23 KiB
C++
/**
|
|
* Copyright (c) 2020 by Contributors
|
|
* @file dgl/aten/macro.h
|
|
* @brief Common macros for aten package.
|
|
*/
|
|
|
|
#ifndef DGL_ATEN_MACRO_H_
|
|
#define DGL_ATEN_MACRO_H_
|
|
|
|
///////////////////////// Dispatchers //////////////////////////
|
|
|
|
/**
|
|
* Dispatch according to device:
|
|
*
|
|
* ATEN_XPU_SWITCH(array->ctx.device_type, XPU, {
|
|
* // Now XPU is a placeholder for array->ctx.device_type
|
|
* DeviceSpecificImplementation<XPU>(...);
|
|
* });
|
|
*/
|
|
#define ATEN_XPU_SWITCH(val, XPU, op, ...) \
|
|
do { \
|
|
if ((val) == kDGLCPU) { \
|
|
constexpr auto XPU = kDGLCPU; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << "Operator " << (op) << " does not support " \
|
|
<< dgl::runtime::DeviceTypeCode2Str(val) << " device."; \
|
|
} \
|
|
} while (0)
|
|
|
|
/**
|
|
* Dispatch according to device:
|
|
*
|
|
* XXX(minjie): temporary macro that allows CUDA operator
|
|
*
|
|
* ATEN_XPU_SWITCH(array->ctx.device_type, XPU, {
|
|
* // Now XPU is a placeholder for array->ctx.device_type
|
|
* DeviceSpecificImplementation<XPU>(...);
|
|
* });
|
|
*
|
|
* We treat pinned memory as normal host memory if we don't want
|
|
* to enable CUDA UVA access for this operator
|
|
*/
|
|
#ifdef DGL_USE_CUDA
|
|
#define ATEN_XPU_SWITCH_CUDA(val, XPU, op, ...) \
|
|
do { \
|
|
if ((val) == kDGLCPU) { \
|
|
constexpr auto XPU = kDGLCPU; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val) == kDGLCUDA) { \
|
|
constexpr auto XPU = kDGLCUDA; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << "Operator " << (op) << " does not support " \
|
|
<< dgl::runtime::DeviceTypeCode2Str(val) << " device."; \
|
|
} \
|
|
} while (0)
|
|
#else // DGL_USE_CUDA
|
|
#define ATEN_XPU_SWITCH_CUDA ATEN_XPU_SWITCH
|
|
#endif // DGL_USE_CUDA
|
|
|
|
/**
|
|
* Dispatch according to integral type (either int32 or int64):
|
|
*
|
|
* ATEN_ID_TYPE_SWITCH(array->dtype, IdType, {
|
|
* // Now IdType is the type corresponding to data type in array.
|
|
* // For instance, one can do this for a CPU array:
|
|
* DType *data = static_cast<DType *>(array->data);
|
|
* });
|
|
*/
|
|
#define ATEN_ID_TYPE_SWITCH(val, IdType, ...) \
|
|
do { \
|
|
CHECK_EQ((val).code, kDGLInt) << "ID must be integer type"; \
|
|
if ((val).bits == 32) { \
|
|
typedef int32_t IdType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 64) { \
|
|
typedef int64_t IdType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << "ID can only be int32 or int64"; \
|
|
} \
|
|
} while (0)
|
|
|
|
/**
|
|
* Dispatch according to bits (either int32 or int64):
|
|
*
|
|
* ATEN_ID_BITS_SWITCH(bits, IdType, {
|
|
* // Now IdType is the type corresponding to data type in array.
|
|
* // For instance, one can do this for a CPU array:
|
|
* DType *data = static_cast<DType *>(array->data);
|
|
* });
|
|
*/
|
|
#define ATEN_ID_BITS_SWITCH(bits, IdType, ...) \
|
|
do { \
|
|
CHECK((bits) == 32 || (bits) == 64) << "bits must be 32 or 64"; \
|
|
if ((bits) == 32) { \
|
|
typedef int32_t IdType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((bits) == 64) { \
|
|
typedef int64_t IdType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << "ID can only be int32 or int64"; \
|
|
} \
|
|
} while (0)
|
|
|
|
/**
|
|
* Dispatch according to float type (either float32 or float64):
|
|
*
|
|
* ATEN_FLOAT_TYPE_SWITCH(array->dtype, FloatType, {
|
|
* // Now FloatType is the type corresponding to data type in array.
|
|
* // For instance, one can do this for a CPU array:
|
|
* FloatType *data = static_cast<FloatType *>(array->data);
|
|
* });
|
|
*/
|
|
#define ATEN_FLOAT_TYPE_SWITCH(val, FloatType, val_name, ...) \
|
|
do { \
|
|
CHECK_EQ((val).code, kDGLFloat) << (val_name) << " must be float type"; \
|
|
if ((val).bits == 32) { \
|
|
typedef float FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 64) { \
|
|
typedef double FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << (val_name) << " can only be float32 or float64"; \
|
|
} \
|
|
} while (0)
|
|
|
|
/**
|
|
* Dispatch according to float type, including 16bits
|
|
* (float16/bfloat16/float32/float64).
|
|
*/
|
|
#ifdef DGL_USE_CUDA
|
|
#if BF16_ENABLED
|
|
#define ATEN_FLOAT_TYPE_SWITCH_16BITS(val, FloatType, XPU, val_name, ...) \
|
|
do { \
|
|
CHECK((val).code == kDGLFloat || (val.code == kDGLBfloat)) \
|
|
<< (val_name) << " must be float type"; \
|
|
if ((val).bits == 32) { \
|
|
typedef float FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 64) { \
|
|
typedef double FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ( \
|
|
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLFloat) { \
|
|
typedef __half FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ( \
|
|
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLBfloat) { \
|
|
typedef __nv_bfloat16 FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ( \
|
|
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLFloat) { \
|
|
LOG(FATAL) << (val_name) << " can't be float16 on CPU"; \
|
|
} else if ( \
|
|
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLBfloat) { \
|
|
typedef BFloat16 FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << (val_name) \
|
|
<< " can only be float16/bfloat16/float32/float64 on GPU"; \
|
|
} \
|
|
} while (0)
|
|
#else // BF16_ENABLED
|
|
#define ATEN_FLOAT_TYPE_SWITCH_16BITS(val, FloatType, XPU, val_name, ...) \
|
|
do { \
|
|
CHECK((val).code == kDGLFloat || (val.code == kDGLBfloat)) \
|
|
<< (val_name) << " must be float type"; \
|
|
if ((val).bits == 32) { \
|
|
typedef float FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 64) { \
|
|
typedef double FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ( \
|
|
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLFloat) { \
|
|
typedef __half FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ( \
|
|
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLBfloat) { \
|
|
LOG(FATAL) << "bfloat16 requires CUDA >= 11.0"; \
|
|
} else if ( \
|
|
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLFloat) { \
|
|
LOG(FATAL) << (val_name) << " can't be float16 on CPU"; \
|
|
} else if ( \
|
|
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLBfloat) { \
|
|
typedef BFloat16 FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << (val_name) \
|
|
<< " can only be float16/float32/float64 on GPU"; \
|
|
} \
|
|
} while (0)
|
|
#endif // BF16_ENABLED
|
|
#else // DGL_USE_CUDA
|
|
#define ATEN_FLOAT_TYPE_SWITCH_16BITS(val, FloatType, XPU, val_name, ...) \
|
|
do { \
|
|
CHECK((val).code == kDGLFloat || (val.code == kDGLBfloat)) \
|
|
<< (val_name) << " must be float type"; \
|
|
if ((val).bits == 32) { \
|
|
typedef float FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 64) { \
|
|
typedef double FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ( \
|
|
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLBfloat) { \
|
|
typedef BFloat16 FloatType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << (val_name) \
|
|
<< " can only be bfloat16/float32/float64 on CPU"; \
|
|
} \
|
|
} while (0)
|
|
#endif // DGL_USE_CUDA
|
|
|
|
/**
|
|
* Dispatch according to data type (int32, int64, float32 or float64):
|
|
*
|
|
* ATEN_DTYPE_SWITCH(array->dtype, DType, {
|
|
* // Now DType is the type corresponding to data type in array.
|
|
* // For instance, one can do this for a CPU array:
|
|
* DType *data = static_cast<DType *>(array->data);
|
|
* });
|
|
*/
|
|
#define ATEN_DTYPE_SWITCH(val, DType, val_name, ...) \
|
|
do { \
|
|
if ((val).code == kDGLInt && (val).bits == 32) { \
|
|
typedef int32_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).code == kDGLInt && (val).bits == 64) { \
|
|
typedef int64_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).code == kDGLFloat && (val).bits == 32) { \
|
|
typedef float DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).code == kDGLFloat && (val).bits == 64) { \
|
|
typedef double DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << (val_name) \
|
|
<< " can only be int32, int64, float32 or float64"; \
|
|
} \
|
|
} while (0)
|
|
|
|
/**
|
|
* Dispatch according to data type (int8, uint8, float32 or float64):
|
|
*
|
|
* ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(array->dtype, DType, {
|
|
* // Now DType is the type corresponding to data type in array.
|
|
* // For instance, one can do this for a CPU array:
|
|
* DType *data = static_cast<DType *>(array->data);
|
|
* });
|
|
*/
|
|
#define ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(val, DType, val_name, ...) \
|
|
do { \
|
|
if ((val).code == kDGLInt && (val).bits == 8) { \
|
|
typedef int8_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).code == kDGLUInt && (val).bits == 8) { \
|
|
typedef uint8_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).code == kDGLFloat && (val).bits == 32) { \
|
|
typedef float DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).code == kDGLFloat && (val).bits == 64) { \
|
|
typedef double DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << (val_name) \
|
|
<< " can only be int8, uint8, float32 or float64"; \
|
|
} \
|
|
} while (0)
|
|
|
|
/**
|
|
* Dispatch data type only based on bit-width (8-bit, 16-bit, 32-bit, 64-bit):
|
|
*
|
|
* ATEN_DTYPE_BITS_ONLY_SWITCH(array->dtype, DType, {
|
|
* // Now DType is the type which has the same bit-width with the
|
|
* // data type in array.
|
|
* // Do not use for computation, but only for read and write.
|
|
* // For instance, one can do this for a CPU array:
|
|
* DType *data = static_cast<DType *>(array->data);
|
|
* });
|
|
*/
|
|
#define ATEN_DTYPE_BITS_ONLY_SWITCH(val, DType, val_name, ...) \
|
|
do { \
|
|
if ((val).bits == 8) { \
|
|
typedef int8_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 16) { \
|
|
typedef int16_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 32) { \
|
|
typedef int32_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).bits == 64) { \
|
|
typedef int64_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << (val_name) \
|
|
<< " can only be 8-bit, 16-bit, 32-bit, or 64-bit"; \
|
|
} \
|
|
} while (0)
|
|
|
|
/**
|
|
* Dispatch according to integral type of CSR graphs.
|
|
* Identical to ATEN_ID_TYPE_SWITCH except for a different error message.
|
|
*/
|
|
#define ATEN_CSR_DTYPE_SWITCH(val, DType, ...) \
|
|
do { \
|
|
if ((val).code == kDGLInt && (val).bits == 32) { \
|
|
typedef int32_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else if ((val).code == kDGLInt && (val).bits == 64) { \
|
|
typedef int64_t DType; \
|
|
{ __VA_ARGS__ } \
|
|
} else { \
|
|
LOG(FATAL) << "CSR matrix data can only be int32 or int64"; \
|
|
} \
|
|
} while (0)
|
|
|
|
// Macro to dispatch according to device context and index type.
|
|
#define ATEN_CSR_SWITCH(csr, XPU, IdType, op, ...) \
|
|
ATEN_XPU_SWITCH((csr).indptr->ctx.device_type, XPU, op, { \
|
|
ATEN_ID_TYPE_SWITCH((csr).indptr->dtype, IdType, {{__VA_ARGS__}}); \
|
|
});
|
|
|
|
// Macro to dispatch according to device context and index type.
|
|
#define ATEN_COO_SWITCH(coo, XPU, IdType, op, ...) \
|
|
ATEN_XPU_SWITCH((coo).row->ctx.device_type, XPU, op, { \
|
|
ATEN_ID_TYPE_SWITCH((coo).row->dtype, IdType, {{__VA_ARGS__}}); \
|
|
});
|
|
|
|
#define CHECK_VALID_CONTEXT(VAR1, VAR2) \
|
|
CHECK( \
|
|
((VAR1)->ctx == (VAR2)->ctx) || (VAR1).IsPinned() || \
|
|
((VAR1).NumElements() == 0)) /* Let empty arrays pass */ \
|
|
<< "Expected " << (#VAR2) << "(" << (VAR2)->ctx << ")" \
|
|
<< " to have the same device " \
|
|
<< "context as " << (#VAR1) << "(" << (VAR1)->ctx << "). " \
|
|
<< "Or " << (#VAR1) << "(" << (VAR1)->ctx << ")" \
|
|
<< " is pinned";
|
|
|
|
/**
|
|
* Macro to dispatch according to the context of array and dtype of csr
|
|
* to enable CUDA UVA ops.
|
|
* Context check is covered here to avoid confusion with CHECK_SAME_CONTEXT.
|
|
* If csr has the same context with array, same behivor as ATEN_CSR_SWITCH_CUDA.
|
|
* If csr is pinned, array's context will conduct the actual operation.
|
|
*/
|
|
#define ATEN_CSR_SWITCH_CUDA_UVA(csr, array, XPU, IdType, op, ...) \
|
|
do { \
|
|
CHECK_VALID_CONTEXT(csr.indices, array); \
|
|
ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, op, { \
|
|
ATEN_ID_TYPE_SWITCH((csr).indptr->dtype, IdType, {{__VA_ARGS__}}); \
|
|
}); \
|
|
} while (0)
|
|
|
|
// Macro to dispatch according to device context (allowing cuda)
|
|
#ifdef DGL_USE_CUDA
|
|
#define ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, op, ...) \
|
|
ATEN_XPU_SWITCH_CUDA((csr).indptr->ctx.device_type, XPU, op, { \
|
|
ATEN_ID_TYPE_SWITCH((csr).indptr->dtype, IdType, {{__VA_ARGS__}}); \
|
|
});
|
|
|
|
// Macro to dispatch according to device context and index type.
|
|
#define ATEN_COO_SWITCH_CUDA(coo, XPU, IdType, op, ...) \
|
|
ATEN_XPU_SWITCH_CUDA((coo).row->ctx.device_type, XPU, op, { \
|
|
ATEN_ID_TYPE_SWITCH((coo).row->dtype, IdType, {{__VA_ARGS__}}); \
|
|
});
|
|
#else // DGL_USE_CUDA
|
|
#define ATEN_CSR_SWITCH_CUDA ATEN_CSR_SWITCH
|
|
#define ATEN_COO_SWITCH_CUDA ATEN_COO_SWITCH
|
|
#endif // DGL_USE_CUDA
|
|
|
|
///////////////////////// Array checks //////////////////////////
|
|
|
|
#define IS_INT32(a) ((a)->dtype.code == kDGLInt && (a)->dtype.bits == 32)
|
|
#define IS_INT64(a) ((a)->dtype.code == kDGLInt && (a)->dtype.bits == 64)
|
|
#define IS_FLOAT32(a) ((a)->dtype.code == kDGLFloat && (a)->dtype.bits == 32)
|
|
#define IS_FLOAT64(a) ((a)->dtype.code == kDGLFloat && (a)->dtype.bits == 64)
|
|
|
|
#define CHECK_IF(cond, prop, value_name, dtype_name) \
|
|
CHECK(cond) << "Expecting " << (prop) << " of " << (value_name) << " to be " \
|
|
<< (dtype_name)
|
|
|
|
#define CHECK_INT32(value, value_name) \
|
|
CHECK_IF(IS_INT32(value), "dtype", value_name, "int32")
|
|
#define CHECK_INT64(value, value_name) \
|
|
CHECK_IF(IS_INT64(value), "dtype", value_name, "int64")
|
|
#define CHECK_INT(value, value_name) \
|
|
CHECK_IF( \
|
|
IS_INT32(value) || IS_INT64(value), "dtype", value_name, \
|
|
"int32 or int64")
|
|
#define CHECK_FLOAT32(value, value_name) \
|
|
CHECK_IF(IS_FLOAT32(value), "dtype", value_name, "float32")
|
|
#define CHECK_FLOAT64(value, value_name) \
|
|
CHECK_IF(IS_FLOAT64(value), "dtype", value_name, "float64")
|
|
#define CHECK_FLOAT(value, value_name) \
|
|
CHECK_IF( \
|
|
IS_FLOAT32(value) || IS_FLOAT64(value), "dtype", value_name, \
|
|
"float32 or float64")
|
|
|
|
#define CHECK_NDIM(value, _ndim, value_name) \
|
|
CHECK_IF((value)->ndim == (_ndim), "ndim", value_name, _ndim)
|
|
|
|
#define CHECK_SAME_DTYPE(VAR1, VAR2) \
|
|
CHECK((VAR1)->dtype == (VAR2)->dtype) \
|
|
<< "Expected " << (#VAR2) << " to be the same type as " << (#VAR1) \
|
|
<< "(" << (VAR1)->dtype << ")" \
|
|
<< ". But got " << (VAR2)->dtype << ".";
|
|
|
|
#define CHECK_SAME_CONTEXT(VAR1, VAR2) \
|
|
CHECK((VAR1)->ctx == (VAR2)->ctx) \
|
|
<< "Expected " << (#VAR2) << " to have the same device context as " \
|
|
<< (#VAR1) << "(" << (VAR1)->ctx << ")" \
|
|
<< ". But got " << (VAR2)->ctx << ".";
|
|
|
|
#define CHECK_NO_OVERFLOW(dtype, val) \
|
|
do { \
|
|
if (sizeof(val) == 8 && (dtype).bits == 32) \
|
|
CHECK_LE((val), 0x7FFFFFFFL) \
|
|
<< "int32 overflow for argument " << (#val) << "."; \
|
|
} while (0);
|
|
|
|
#define CHECK_IS_ID_ARRAY(VAR) \
|
|
CHECK((VAR)->ndim == 1 && (IS_INT32(VAR) || IS_INT64(VAR))) \
|
|
<< "Expected argument " << (#VAR) << " to be an 1D integer array.";
|
|
|
|
#endif // DGL_ATEN_MACRO_H_
|