文件
rocm-systems/projects/rccl/src/include/core.h
T
Wenkai Du 1e182391ad Add bfloat16 support in RCCL
Preprocessor symbol RCCL_BFLOAT16 is used as feature indicator


[ROCm/rccl commit: 5e109ed400]
2019-11-18 13:45:53 -08:00

67 行
1.6 KiB
C

/*************************************************************************
* Copyright (c) 2015-2019, NVIDIA CORPORATION. All rights reserved.
*
* See LICENSE.txt for license information
************************************************************************/
#ifndef NCCL_CORE_H_
#define NCCL_CORE_H_
#include <pthread.h>
#include <algorithm>
#include "nccl.h"
#include "debug.h"
#include "checks.h"
#include "alloc.h"
#include "transport.h"
#include "devcomm.h"
#include "comm.h"
#include "info.h"
#include "argcheck.h"
#include <cstdio>
#include <unistd.h>
#include <stdlib.h>
#ifdef PROFAPI
#define NCCL_API(ret, func, args...) \
__attribute__ ((visibility("default"))) \
__attribute__ ((alias(#func))) \
ret p##func (args); \
extern "C" \
__attribute__ ((visibility("default"))) \
__attribute__ ((weak)) \
ret func(args)
#else
#define NCCL_API(ret, func, args...) \
extern "C" \
__attribute__ ((visibility("default"))) \
ret func(args)
#endif // end PROFAPI
int ncclCudaCompCap();
ncclResult_t ncclNvlinkGpu(int* nvlink);
int64_t ncclTreeThreshold();
static __inline__ int ncclTypeSize(ncclDataType_t type) {
switch (type) {
case ncclInt8:
case ncclUint8:
return 1;
case ncclFloat16:
case ncclBfloat16:
return 2;
case ncclInt32:
case ncclUint32:
case ncclFloat32:
return 4;
case ncclInt64:
case ncclUint64:
case ncclFloat64:
return 8;
default:
return -1;
}
}
#endif // end include guard