From a89fad79ed962b4e836db5264a14d2ae887949b6 Mon Sep 17 00:00:00 2001 From: Li Baoming <41820386+baominghelly@users.noreply.github.com> Date: Mon, 7 Sep 2026 10:50:02 +0800 Subject: [PATCH] feat(ascend): add HCCL backend support --- CMakeLists.txt | 84 +++++++++++++++++++- scripts/gen_bridge.py | 3 + src/CMakeLists.txt | 24 ++++++ src/backend.h | 5 ++ src/backend_device_map.h | 4 + src/backends/ccl/hccl/api.h | 90 +++++++++++++++++++++ src/backends/ccl/hccl/ascend/api.h | 15 ++++ src/backends/ccl/hccl/impl/all_reduce.h | 17 ++++ src/backends/ccl/hccl/impl/comm_destroy.h | 17 ++++ src/backends/ccl/hccl/impl/comm_init_all.h | 76 ++++++++++++++++++ src/backends/ccl/hccl/type_map.h | 91 ++++++++++++++++++++++ src/backends/mpi/ompi/impl/comm_init_all.h | 17 +++- src/base/comm_init_all.h | 21 ++--- src/device.h | 5 ++ src/devices/ascend/device_.h | 31 ++++++++ src/devices/ascend/runtime_.h | 85 ++++++++++++++++++++ 16 files changed, 565 insertions(+), 20 deletions(-) create mode 100644 src/backends/ccl/hccl/api.h create mode 100644 src/backends/ccl/hccl/ascend/api.h create mode 100644 src/backends/ccl/hccl/impl/all_reduce.h create mode 100644 src/backends/ccl/hccl/impl/comm_destroy.h create mode 100644 src/backends/ccl/hccl/impl/comm_init_all.h create mode 100644 src/backends/ccl/hccl/type_map.h create mode 100644 src/devices/ascend/device_.h create mode 100644 src/devices/ascend/runtime_.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 4277c1c..609d36a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -12,6 +12,7 @@ option(WITH_ILUVATAR "Enable ILUVATAR GPU support" OFF) option(WITH_METAX "Enable MetaX GPU support" OFF) option(WITH_MOORE "Enable Moore GPU support" OFF) option(WITH_CAMBRICON "Enable Cambricon MLU support" OFF) +option(WITH_ASCEND "Enable Ascend NPU support" OFF) option(WITH_HYGON "Enable Hygon DCU support" OFF) set(WITH_CPU ON CACHE INTERNAL "CPU backend is always enabled") @@ -23,6 +24,7 @@ option(WITH_OMPI "Enable OpenMPI backend" OFF) option(WITH_MPICH "Enable MPICH backend" OFF) option(WITH_NCCL "Enable NCCL backend" OFF) option(WITH_MCCL "Enable MCCL backend" OFF) +option(WITH_HCCL "Enable HCCL backend" OFF) # ========================================================= # --- MISC. BUILD OPTIONS --- @@ -178,6 +180,35 @@ if(AUTO_DETECT_DEVICES) message(STATUS "Cambricon environment not detected.") endif() + # Ascend + set(ASCEND_FOUND FALSE) + + if(DEFINED ENV{ASCEND_HOME_PATH} OR DEFINED ENV{ASCEND_TOOLKIT_HOME} OR DEFINED ENV{ASCEND_HOME}) + set(ASCEND_FOUND TRUE) + elseif(EXISTS "/dev/davinci0" OR EXISTS "/dev/davinci_manager") + set(ASCEND_FOUND TRUE) + else() + find_program(ASCEND_SMI_PATH npu-smi) + if(ASCEND_SMI_PATH) + execute_process( + COMMAND ${ASCEND_SMI_PATH} info + RESULT_VARIABLE ASCEND_SMI_RESULT + OUTPUT_QUIET + ERROR_QUIET + ) + if(ASCEND_SMI_RESULT EQUAL 0) + set(ASCEND_FOUND TRUE) + endif() + endif() + endif() + + if(ASCEND_FOUND) + set(WITH_ASCEND ON) + message(STATUS "Ascend environment detected.") + else() + message(STATUS "Ascend environment not detected.") + endif() + # Hygon DCU if(NOT WITH_HYGON) set(HYGON_FOUND FALSE) @@ -307,10 +338,30 @@ if(AUTO_DETECT_BACKENDS) else() message(STATUS "No suitable device environment, skipping MCCL detection.") endif() + + # Detect HCCL dependencies. + if(WITH_ASCEND) + set(_HCCL_HINTS + "${ASCEND_HOME}" + "$ENV{ASCEND_HOME_PATH}" + "$ENV{ASCEND_TOOLKIT_HOME}" + "$ENV{ASCEND_HOME}" + /usr/local/Ascend/ascend-toolkit/latest + ) + find_path(AUTO_HCCL_INC NAMES hccl/hccl.h HINTS ${_HCCL_HINTS} PATH_SUFFIXES include QUIET) + find_library(AUTO_HCCL_LIB NAMES hccl HINTS ${_HCCL_HINTS} PATH_SUFFIXES lib64 aarch64-linux/lib64 QUIET) + + if(AUTO_HCCL_INC AND AUTO_HCCL_LIB) + set(WITH_HCCL ON) + message(STATUS "Auto-detected HCCL backend.") + else() + message(STATUS "HCCL library/headers not found in Ascend paths.") + endif() + endif() endif() # Fallback: If no backends are enabled or auto-detected, fall back to OpenMPI as the default bootstrap profile. -if(NOT WITH_OMPI AND NOT WITH_MPICH AND NOT WITH_NCCL AND NOT WITH_MCCL) +if(NOT WITH_OMPI AND NOT WITH_MPICH AND NOT WITH_NCCL AND NOT WITH_MCCL AND NOT WITH_HCCL) set(WITH_OMPI ON) message(STATUS "No backend specified or detected. Defaulting to `WITH_OMPI=ON`") endif() @@ -420,6 +471,26 @@ if(WITH_CAMBRICON) find_library(CAMBRICON_PAPI_LIB NAMES cnpapi HINTS "${NEUWARE_HOME}/lib64" REQUIRED) endif() +if(WITH_ASCEND) + set(_ASCEND_HINTS) + if(ASCEND_HOME) + list(APPEND _ASCEND_HINTS "${ASCEND_HOME}") + endif() + foreach(_ascend_env ASCEND_HOME_PATH ASCEND_TOOLKIT_HOME ASCEND_HOME) + if(DEFINED ENV{${_ascend_env}} AND NOT "$ENV{${_ascend_env}}" STREQUAL "") + list(APPEND _ASCEND_HINTS "$ENV{${_ascend_env}}") + endif() + endforeach() + list(APPEND _ASCEND_HINTS + /usr/local/Ascend/ascend-toolkit/latest + ) + + find_path(ASCEND_INC NAMES acl/acl.h HINTS ${_ASCEND_HINTS} PATH_SUFFIXES include REQUIRED) + find_library(ASCENDCL_LIB NAMES ascendcl HINTS ${_ASCEND_HINTS} PATH_SUFFIXES lib64 aarch64-linux/lib64 REQUIRED) + + include_directories(${ASCEND_INC}) +endif() + if(WITH_HYGON) set(HYGON_DTK_ROOT "") foreach(_hygon_env DTKROOT DTK_ROOT ROCM_PATH) @@ -504,6 +575,17 @@ if(WITH_MCCL) include_directories(${MCCL_INC}) endif() +if(WITH_HCCL) + if(NOT WITH_ASCEND) + message(FATAL_ERROR "HCCL backend requires Ascend device support. Please enable `WITH_ASCEND`.") + endif() + + find_library(HCCL_LIB NAMES hccl HINTS ${_ASCEND_HINTS} PATH_SUFFIXES lib64 aarch64-linux/lib64 REQUIRED) + find_path(HCCL_INC NAMES hccl/hccl.h HINTS ${_ASCEND_HINTS} PATH_SUFFIXES include REQUIRED) + + include_directories(${HCCL_INC}) +endif() + # Python is required for code generation. find_package(Python3 REQUIRED) diff --git a/scripts/gen_bridge.py b/scripts/gen_bridge.py index 211080b..2952e58 100644 --- a/scripts/gen_bridge.py +++ b/scripts/gen_bridge.py @@ -32,16 +32,19 @@ "mpich": ["backends/mpi/ompi/impl"], "nccl": ["backends/ccl/nccl/impl"], "mccl": ["backends/ccl/mccl/impl"], + "hccl": ["backends/ccl/hccl/impl"], } BACKEND_COMMON_HEADERS = { "nccl": ["backends/ccl/nccl/type_map.h"], "mccl": ["backends/ccl/mccl/type_map.h"], + "hccl": ["backends/ccl/hccl/type_map.h"], } CCL_PROVIDER_BACKENDS = { "nccl": "backends/ccl/nccl", "mccl": "backends/ccl/mccl", + "hccl": "backends/ccl/hccl", } # ================================================================= diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 42c1171..d8f14f8 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -215,6 +215,20 @@ if(WITH_CAMBRICON) ) endif() +# Ascend +if(WITH_ASCEND) + list(APPEND DEVICE_LIST "ascend") + + file(GLOB_RECURSE ASCEND_SRCS + "devices/ascend/*.cc" + "devices/ascend/*.cpp" + ) + + target_sources(infiniccl PRIVATE ${ASCEND_SRCS}) + target_include_directories(infiniccl PRIVATE ${ASCEND_INC}) + target_link_libraries(infiniccl PRIVATE ${ASCENDCL_LIB}) +endif() + # Hygon if(WITH_HYGON) list(APPEND DEVICE_LIST "hygon") @@ -279,6 +293,16 @@ if(WITH_MCCL) target_link_libraries(infiniccl PRIVATE ${MCCL_LIB}) endif() +# HCCL +if(WITH_HCCL) + list(APPEND BACKEND_LIST "hccl") + file(GLOB_RECURSE HCCL_SRCS "backends/ccl/hccl/*.cc" "backends/ccl/hccl/*.cpp") + + target_sources(infiniccl PRIVATE ${HCCL_SRCS}) + target_include_directories(infiniccl PRIVATE ${HCCL_INC}) + target_link_libraries(infiniccl PRIVATE ${HCCL_LIB}) +endif() + # ========================================================= # --- File Generation --- # ========================================================= diff --git a/src/backend.h b/src/backend.h index 0b02662..84a6725 100644 --- a/src/backend.h +++ b/src/backend.h @@ -63,6 +63,11 @@ struct BackendPriority { static constexpr int value = 10; }; +template <> +struct BackendPriority { + static constexpr int value = 10; +}; + } // namespace infini::ccl #endif // INFINI_CCL_BACKEND_H_ diff --git a/src/backend_device_map.h b/src/backend_device_map.h index 96766c0..02f028c 100644 --- a/src/backend_device_map.h +++ b/src/backend_device_map.h @@ -29,6 +29,10 @@ template <> struct IsSupportedCombination : std::true_type {}; +template <> +struct IsSupportedCombination + : std::true_type {}; + }; // namespace infini::ccl #endif // INFINI_CCL_BACKEND_DEVICE_MAP_H_ diff --git a/src/backends/ccl/hccl/api.h b/src/backends/ccl/hccl/api.h new file mode 100644 index 0000000..5f14957 --- /dev/null +++ b/src/backends/ccl/hccl/api.h @@ -0,0 +1,90 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_HCCL_API_H_ +#define INFINI_CCL_BACKENDS_CCL_HCCL_API_H_ + +#include + +#include +#include + +#include "backends/ccl/common/api.h" +#include "logging.h" +#include "return_status_impl.h" +#include "runtime.h" + +namespace infini::ccl { + +template +struct HcclApi { + static constexpr BackendType kBackendType = BackendType::kHccl; + static constexpr Device::Type kDeviceType = device; + + using Comm = HcclComm; + using Result = HcclResult; + using DataType = HcclDataType; + using RedOp = HcclReduceOp; + using Stream = typename Runtime::Stream; + + struct ThreadLocalStream { + Stream stream{}; + Result status = HCCL_SUCCESS; + + ThreadLocalStream() { + if (Runtime::StreamCreate(&stream) != 0) { + status = HCCL_E_RUNTIME; + } + } + + ~ThreadLocalStream() { + if (stream != nullptr) { + Runtime::StreamDestroy(stream); + } + } + }; + + static ReturnStatus Check(Result result) { + if (result != HCCL_SUCCESS) { + const char* message = HcclGetErrorString(result); + LOG(message ? message : "Unknown HCCL error"); + return ReturnStatus::kSystemError; + } + return ReturnStatus::kSuccess; + } + + static Result CommInitAll(Comm* comms, int n_dev, const int* dev_list, + const int*) { + return HcclCommInitAll( + static_cast(n_dev), + reinterpret_cast(const_cast(dev_list)), comms); + } + + static Result CommDestroy(Comm comm) { return HcclCommDestroy(comm); } + + static Result ResolveStream(Stream requested, Stream* resolved) { + if (requested != nullptr) { + *resolved = requested; + return HCCL_SUCCESS; + } + + thread_local ThreadLocalStream default_stream; + if (default_stream.status != HCCL_SUCCESS) { + return default_stream.status; + } + *resolved = default_stream.stream; + return HCCL_SUCCESS; + } + + static Result AllReduce(const void* send_buff, void* recv_buff, size_t count, + DataType data_type, RedOp op, Comm comm, + Stream stream) { + auto status = ResolveStream(stream, &stream); + if (status != HCCL_SUCCESS) { + return status; + } + return HcclAllReduce(const_cast(send_buff), recv_buff, count, + data_type, op, comm, stream); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_HCCL_API_H_ diff --git a/src/backends/ccl/hccl/ascend/api.h b/src/backends/ccl/hccl/ascend/api.h new file mode 100644 index 0000000..705c4a2 --- /dev/null +++ b/src/backends/ccl/hccl/ascend/api.h @@ -0,0 +1,15 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_HCCL_ASCEND_API_H_ +#define INFINI_CCL_BACKENDS_CCL_HCCL_ASCEND_API_H_ + +#include "backends/ccl/hccl/api.h" +#include "devices/ascend/runtime_.h" + +namespace infini::ccl { + +template <> +struct CclApi + : HcclApi {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_HCCL_ASCEND_API_H_ diff --git a/src/backends/ccl/hccl/impl/all_reduce.h b/src/backends/ccl/hccl/impl/all_reduce.h new file mode 100644 index 0000000..db22008 --- /dev/null +++ b/src/backends/ccl/hccl/impl/all_reduce.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_ALL_REDUCE_H_ +#define INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_ALL_REDUCE_H_ + +#include "backends/ccl/common/impl/all_reduce.h" + +namespace infini::ccl { + +template +class AllReduceImpl + : public CclAllReduceImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_ALL_REDUCE_H_ diff --git a/src/backends/ccl/hccl/impl/comm_destroy.h b/src/backends/ccl/hccl/impl/comm_destroy.h new file mode 100644 index 0000000..7d64786 --- /dev/null +++ b/src/backends/ccl/hccl/impl/comm_destroy.h @@ -0,0 +1,17 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_COMM_DESTROY_H_ +#define INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_COMM_DESTROY_H_ + +#include "backends/ccl/common/impl/comm_destroy.h" + +namespace infini::ccl { + +template +class CommDestroyImpl + : public CclCommDestroyImpl {}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_COMM_DESTROY_H_ diff --git a/src/backends/ccl/hccl/impl/comm_init_all.h b/src/backends/ccl/hccl/impl/comm_init_all.h new file mode 100644 index 0000000..8420dc6 --- /dev/null +++ b/src/backends/ccl/hccl/impl/comm_init_all.h @@ -0,0 +1,76 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_COMM_INIT_ALL_H_ +#define INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_COMM_INIT_ALL_H_ + +#include +#include +#include + +#include "backends/ccl/common/comm_instance.h" +#include "base/comm_init_all.h" +#include "communicator.h" +#include "runtime.h" + +namespace infini::ccl { + +template +class CommInitAllImpl { + public: + static ReturnStatus Apply(void** comm_handles, int n_dev, + const int* dev_list) { + using Api = CclApi; + using CommInstance = CclCommInstance; + using Rt = Runtime; + + if (!comm_handles || !dev_list || n_dev <= 0) { + return ReturnStatus::kInvalidArgument; + } + + auto** comms = reinterpret_cast(comm_handles); + for (int i = 0; i < n_dev; ++i) { + if (comms[i]) { + return ReturnStatus::kInvalidArgument; + } + } + + std::vector> wrappers; + std::vector> instances; + std::vector backend_comms(n_dev); + std::vector rank_list(n_dev); + wrappers.reserve(n_dev); + instances.reserve(n_dev); + std::iota(rank_list.begin(), rank_list.end(), 0); + + // HCCL requires every local device context to be initialized first. + for (int i = n_dev - 1; i >= 0; --i) { + auto status = Rt::Check(Rt::SetDevice(dev_list[i])); + if (status != ReturnStatus::kSuccess) { + return status; + } + } + + auto status = Api::Check(Api::CommInitAll(backend_comms.data(), n_dev, + dev_list, rank_list.data())); + if (status != ReturnStatus::kSuccess) { + return status; + } + + for (int i = 0; i < n_dev; ++i) { + wrappers.emplace_back( + std::make_unique(device, dev_list[i])); + instances.emplace_back(std::make_unique()); + instances.back()->handle = backend_comms[i]; + wrappers.back()->set_world_info(i, n_dev); + wrappers.back()->set_intra_comm(std::move(instances.back())); + comms[i] = wrappers.back().release(); + } + + return ReturnStatus::kSuccess; + } +}; + +template <> +struct BackendEnabled : std::true_type {}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_HCCL_IMPL_COMM_INIT_ALL_H_ diff --git a/src/backends/ccl/hccl/type_map.h b/src/backends/ccl/hccl/type_map.h new file mode 100644 index 0000000..ad15887 --- /dev/null +++ b/src/backends/ccl/hccl/type_map.h @@ -0,0 +1,91 @@ +#ifndef INFINI_CCL_BACKENDS_CCL_HCCL_TYPE_MAP_H_ +#define INFINI_CCL_BACKENDS_CCL_HCCL_TYPE_MAP_H_ + +#include + +#include "backends/ccl/common/api.h" +#include "comm_impl.h" +#include "data_type_impl.h" + +namespace infini::ccl { + +inline bool DataTypeToHcclType(DataType dtype, HcclDataType* hccl_dtype) { + switch (dtype) { + case DataType::kInt8: + *hccl_dtype = HCCL_DATA_TYPE_INT8; + return true; + case DataType::kInt16: + *hccl_dtype = HCCL_DATA_TYPE_INT16; + return true; + case DataType::kInt32: + *hccl_dtype = HCCL_DATA_TYPE_INT32; + return true; + case DataType::kInt64: + *hccl_dtype = HCCL_DATA_TYPE_INT64; + return true; + case DataType::kUInt8: + *hccl_dtype = HCCL_DATA_TYPE_UINT8; + return true; + case DataType::kUInt16: + *hccl_dtype = HCCL_DATA_TYPE_UINT16; + return true; + case DataType::kUInt32: + *hccl_dtype = HCCL_DATA_TYPE_UINT32; + return true; + case DataType::kUInt64: + *hccl_dtype = HCCL_DATA_TYPE_UINT64; + return true; + case DataType::kFloat16: + *hccl_dtype = HCCL_DATA_TYPE_FP16; + return true; + case DataType::kBFloat16: + *hccl_dtype = HCCL_DATA_TYPE_BFP16; + return true; + case DataType::kFloat32: + *hccl_dtype = HCCL_DATA_TYPE_FP32; + return true; + case DataType::kFloat64: + *hccl_dtype = HCCL_DATA_TYPE_FP64; + return true; + default: + return false; + } +} + +inline bool RedOpToHcclOp(ReductionOpType red_op, HcclReduceOp* hccl_op) { + switch (red_op) { + case ReductionOpType::kSum: + *hccl_op = HCCL_REDUCE_SUM; + return true; + case ReductionOpType::kProd: + *hccl_op = HCCL_REDUCE_PROD; + return true; + case ReductionOpType::kMax: + *hccl_op = HCCL_REDUCE_MAX; + return true; + case ReductionOpType::kMin: + *hccl_op = HCCL_REDUCE_MIN; + return true; + default: + return false; + } +} + +template <> +struct CclTypeMap { + using Api = CclApi; + + static bool ToBackendDataType(DataType dtype, + typename Api::DataType* backend_dtype) { + return DataTypeToHcclType(dtype, backend_dtype); + } + + static bool ToBackendRedOp(ReductionOpType red_op, + typename Api::RedOp* backend_op) { + return RedOpToHcclOp(red_op, backend_op); + } +}; + +} // namespace infini::ccl + +#endif // INFINI_CCL_BACKENDS_CCL_HCCL_TYPE_MAP_H_ diff --git a/src/backends/mpi/ompi/impl/comm_init_all.h b/src/backends/mpi/ompi/impl/comm_init_all.h index e834622..7cf320b 100644 --- a/src/backends/mpi/ompi/impl/comm_init_all.h +++ b/src/backends/mpi/ompi/impl/comm_init_all.h @@ -12,19 +12,28 @@ namespace infini::ccl { template class CommInitAllImpl { public: - static ReturnStatus Apply(Communicator *comm, int n_dev, - const int *dev_list) { + static ReturnStatus Apply(void** comm_handles, int n_dev, + const int* dev_list) { constexpr Device::Type kDev = ListGetBest(ActiveDevices{}); using Rt = Runtime; - if (!comm) { + if (!comm_handles) { // TODO(lzm): change to use `glog`. LOG("Failed to initialize OpenMPI communicator: invalid " "communicator pointer."); return ReturnStatus::kInternalError; } + Communicator*& comm = *reinterpret_cast(comm_handles); + if (comm && comm->inter_comm()) { + LOG("Invalid communicator handle for `CommInitAll`."); + return ReturnStatus::kInvalidArgument; + } + if (!comm) { + comm = new Communicator(kDev, 0); + } + int rank, size; auto inst = std::make_unique(); INFINI_CHECK_MPI(MPI_Comm_dup(MPI_COMM_WORLD, &inst->handle)); @@ -35,7 +44,7 @@ class CommInitAllImpl { comm->set_inter_comm(std::move(inst)); int local_rank = 0; - char *local_rank_str = getenv("OMPI_COMM_WORLD_LOCAL_RANK"); + char* local_rank_str = getenv("OMPI_COMM_WORLD_LOCAL_RANK"); if (local_rank_str) { local_rank = atoi(local_rank_str); } diff --git a/src/base/comm_init_all.h b/src/base/comm_init_all.h index e737306..2371e50 100644 --- a/src/base/comm_init_all.h +++ b/src/base/comm_init_all.h @@ -13,25 +13,16 @@ struct CommInitAllImpl; class CommInitAll : public Operation { public: - template - static ReturnStatus Execute(void **comm_handle, Args &&...args) { - Communicator *&comm = *reinterpret_cast(comm_handle); - if (comm && comm->inter_comm()) { - // TODO(lzm): change to use `glog`. + template + static ReturnStatus Execute(void** comm_handles, int n_dev, + const int* dev_list) { + if (!comm_handles || n_dev <= 0) { LOG("Invalid communicator handle for `CommInitAll`."); return ReturnStatus::kInvalidArgument; } - constexpr Device::Type kDev = - ListGetBest(ActiveDevices{}); - - if (!comm) { - comm = new Communicator(kDev, 0); - } - - return CommInitAllImpl::Apply( - comm, std::forward(args)...); + return CommInitAllImpl::Apply(comm_handles, + n_dev, dev_list); } }; diff --git a/src/device.h b/src/device.h index 4809420..1d21db3 100644 --- a/src/device.h +++ b/src/device.h @@ -156,6 +156,11 @@ struct DevicePriority { static constexpr int value = 5; }; +template <> +struct DevicePriority { + static constexpr int value = 5; +}; + template <> struct DevicePriority { static constexpr int value = 5; diff --git a/src/devices/ascend/device_.h b/src/devices/ascend/device_.h new file mode 100644 index 0000000..527605f --- /dev/null +++ b/src/devices/ascend/device_.h @@ -0,0 +1,31 @@ +#ifndef INFINI_CCL_DEVICES_ASCEND_DEVICE_H_ +#define INFINI_CCL_DEVICES_ASCEND_DEVICE_H_ + +#include + +#include "device.h" + +namespace infini::ccl { + +template <> +struct DeviceEnabled : std::true_type {}; + +template <> +MemorySpace GetMemorySpace(const void* ptr) { + if (!ptr) { + return MemorySpace::kHost; + } + + aclrtPtrAttributes attributes{}; + if (aclrtPointerGetAttributes(ptr, &attributes) != ACL_SUCCESS) { + return MemorySpace::kHost; + } + + return attributes.location.type == ACL_MEM_LOCATION_TYPE_DEVICE + ? MemorySpace::kDevice + : MemorySpace::kHost; +} + +} // namespace infini::ccl + +#endif // INFINI_CCL_DEVICES_ASCEND_DEVICE_H_ diff --git a/src/devices/ascend/runtime_.h b/src/devices/ascend/runtime_.h new file mode 100644 index 0000000..6168f87 --- /dev/null +++ b/src/devices/ascend/runtime_.h @@ -0,0 +1,85 @@ +#ifndef INFINI_CCL_DEVICES_ASCEND_RUNTIME_H_ +#define INFINI_CCL_DEVICES_ASCEND_RUNTIME_H_ + +#include + +#include + +#include "devices/ascend/device_.h" +#include "logging.h" +#include "return_status_impl.h" +#include "runtime.h" + +namespace infini::ccl { + +template <> +struct Runtime + : DeviceRuntime> { + using Stream = aclrtStream; + + static constexpr Device::Type kDeviceType = Device::Type::kAscend; + + static constexpr auto Check = + [](auto status, ReturnStatus err_code = ReturnStatus::kSystemError) { + if (status != ACL_SUCCESS) { + const char* message = aclGetRecentErrMsg(); + if (message) { + LOG(message); + } else { + LOG(("Ascend ACL error code: " + + std::to_string(static_cast(status))) + .c_str()); + } + return err_code; + } + return ReturnStatus::kSuccess; + }; + + static constexpr auto Malloc = [](void** ptr, size_t size) { + return aclrtMalloc(ptr, size, ACL_MEM_MALLOC_HUGE_FIRST); + }; + + static constexpr auto Memcpy = [](void* dst, const void* src, size_t count, + aclrtMemcpyKind kind) { + return aclrtMemcpy(dst, count, src, count, kind); + }; + + static constexpr auto Free = aclrtFree; + + static constexpr auto MemcpyHostToDevice = ACL_MEMCPY_HOST_TO_DEVICE; + + static constexpr auto MemcpyDeviceToHost = ACL_MEMCPY_DEVICE_TO_HOST; + + static constexpr auto Memset = [](void* ptr, int value, size_t count) { + return aclrtMemset(ptr, count, value, count); + }; + + static constexpr auto GetDevice = aclrtGetDevice; + + static constexpr auto SetDevice = aclrtSetDevice; + + static aclError EnsureDeviceContext() { + int32_t device_id = 0; + const aclError status = aclrtGetDevice(&device_id); + if (status == ACL_SUCCESS) { + return ACL_SUCCESS; + } + return aclrtSetDevice(0); + } + + static constexpr auto DeviceSynchronize = aclrtSynchronizeDevice; + + static constexpr auto StreamCreate = aclrtCreateStream; + + static constexpr auto StreamDestroy = aclrtDestroyStream; + + static constexpr auto StreamSynchronize = [](aclrtStream stream) { + return stream ? aclrtSynchronizeStream(stream) : aclrtSynchronizeDevice(); + }; +}; + +static_assert(Runtime::Validate()); + +} // namespace infini::ccl + +#endif // INFINI_CCL_DEVICES_ASCEND_RUNTIME_H_