已合并
[shmem] fix with aclshmem #29045
王超创建于 1月5日
[shmem] fix with aclshmem #29045
已合并
共 7 个文件变更+261-122
| @@ -13,71 +13,18 @@ | |||
| 13 | 13 | ||
| 14 | extern "C" { | 14 | extern "C" { |
| 15 | 15 | ||
| 16 | -/** | ||
| 17 | - * @private | ||
| 18 | -*/ | ||
| 19 | - | ||
| 20 | - | ||
| 21 | -/// \def SHMEM_DEVICE | ||
| 22 | -/// \brief A macro that identifies a function on the device side. | ||
| 23 | - | ||
| 24 | 16 | ||
| 25 | /** | 17 | /** |
| 26 | - * @addtogroup group_enums | 18 | + * @brief Data operation engine type. |
| 27 | - * @{ | ||
| 28 | -*/ | ||
| 29 | - | ||
| 30 | -/** | ||
| 31 | - * @brief Team's index. | ||
| 32 | -*/ | ||
| 33 | -enum shmem_team_index_t{ | ||
| 34 | - SHMEM_TEAM_INVALID = -1, | ||
| 35 | - SHMEM_TEAM_WORLD = 0 | ||
| 36 | -}; | ||
| 37 | - | ||
| 38 | -/** | ||
| 39 | - * @brief Data op engine type. | ||
| 40 | */ | 19 | */ |
| 41 | enum data_op_engine_type_t { | 20 | enum data_op_engine_type_t { |
| 42 | - SHMEM_DATA_OP_MTE = 0x01, | 21 | + ACLSHMEM_DATA_OP_MTE = 0x01, |
| 22 | + ACLSHMEM_DATA_OP_SDMA = 0x02, | ||
| 23 | + ACLSHMEM_DATA_OP_ROCE = 0x04, | ||
| 43 | }; | 24 | }; |
| 44 | 25 | ||
| 45 | -/** | ||
| 46 | - * @brief signal ops, used by signaler in p2p synchronization | ||
| 47 | - */ | ||
| 48 | -enum { | ||
| 49 | - SHMEM_SIGNAL_SET, | ||
| 50 | - SHMEM_SIGNAL_ADD | ||
| 51 | -}; | ||
| 52 | - | ||
| 53 | -/** | ||
| 54 | - * @brief signal compare ops, used by signalee in p2p synchronization | ||
| 55 | - */ | ||
| 56 | -enum { | ||
| 57 | - SHMEM_CMP_EQ, | ||
| 58 | - SHMEM_CMP_NE, | ||
| 59 | - SHMEM_CMP_GT, | ||
| 60 | - SHMEM_CMP_GE, | ||
| 61 | - SHMEM_CMP_LT, | ||
| 62 | - SHMEM_CMP_LE | ||
| 63 | -}; | ||
| 64 | - | ||
| 65 | -/**@} */ // end of group_enums | ||
| 66 | - | ||
| 67 | -/** | ||
| 68 | - * @defgroup group_typedef Typedef | ||
| 69 | - * @{ | ||
| 70 | - | ||
| 71 | -*/ | ||
| 72 | -/** | ||
| 73 | - * @brief A typedef of int | ||
| 74 | -*/ | ||
| 75 | -typedef int shmem_team_t; | ||
| 76 | - | ||
| 77 | -/**@} */ // end of group_typedef | ||
| 78 | - | ||
| 79 | 26 | ||
| 80 | } | 27 | } |
| 81 | 28 | ||
| 82 | 29 | ||
| 83 | -#endif /*SHMEM_TYPES_H*/ | 30 | +#endif // SHMEM_TYPES_H |
| @@ -283,8 +283,8 @@ NpuSysCtrl::SysStatus NpuSysCtrl::Finalize() | |||
| 283 | NPU_CHECK_WARN(c10_npu::DestroyUsedStreams()); | 283 | NPU_CHECK_WARN(c10_npu::DestroyUsedStreams()); |
| 284 | NPU_CHECK_WARN(c10_npu::ResetUsedDevices()); | 284 | NPU_CHECK_WARN(c10_npu::ResetUsedDevices()); |
| 285 | 285 | ||
| 286 | - if (c10d::symmetric_memory::Shmem_finalize_exist()) { | 286 | + if (c10d::symmetric_memory::Aclshmem_finalize_exist()) { |
| 287 | - auto ret = c10d::symmetric_memory::Shmem_finalize(); | 287 | + auto ret = c10d::symmetric_memory::Aclshmem_finalize(); |
| 288 | ASCEND_LOGI("shmem_finalize emd, ret is %d", ret); | 288 | ASCEND_LOGI("shmem_finalize emd, ret is %d", ret); |
| 289 | } | 289 | } |
| 290 | 290 | ||
| @@ -25,32 +25,57 @@ void initialize_npushmem_with_store( | |||
| 25 | 25 | ||
| 26 | logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, rank is %d, world_size is %d.", rank, world_size); | 26 | logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, rank is %d, world_size is %d.", rank, world_size); |
| 27 | 27 | ||
| 28 | - uint32_t status = c10d::symmetric_memory::Shmem_set_conf_store_tls(false, nullptr, 0); | 28 | + uint32_t status = c10d::symmetric_memory::Aclshmemx_set_conf_store_tls(false, nullptr, 0); |
| 29 | TORCH_CHECK(status == 0, "shmem_set_conf_store_tls failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | 29 | TORCH_CHECK(status == 0, "shmem_set_conf_store_tls failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); |
| 30 | 30 | ||
| 31 | - shmem_uniqueid_t unique_id; | ||
| 32 | - if (rank == 0) { | ||
| 33 | - status = c10d::symmetric_memory::Shmem_get_uniqueid(&unique_id); | ||
| 34 | - TORCH_CHECK(status == 0, "shmem_get_uniqueid failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | ||
| 35 | - logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, Shmem_get_uniqueid rank is %d, version %d, internal is %s.", | ||
| 36 | - rank, unique_id.version, unique_id.internal); | ||
| 37 | - } | ||
| 38 | - auto unique_ids = storeExchange.all_gather(store, rank, world_size, unique_id); | ||
| 39 | - logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, unique_id rank is %d, version %d, internal is %s.", | ||
| 40 | - rank, unique_ids[0].version, unique_ids[0].internal); | ||
| 41 | - | ||
| 42 | int64_t init_size = c10_npu::option::OptionsManager::GetShmemSymmetricSize(); | 31 | int64_t init_size = c10_npu::option::OptionsManager::GetShmemSymmetricSize(); |
| 43 | - shmem_init_attr_t* attributes; | 32 | + if (c10d::symmetric_memory::Aclshmemx_get_uniqueid_exist()) { |
| 44 | - logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, start shmem_set_attr rank is %d, world_size is %d, size is %llu.", | 33 | + // gitcode version |
| 45 | - rank, world_size, init_size); | 34 | + aclshmemx_uniqueid_t unique_id; |
| 46 | - status = c10d::symmetric_memory::Shmem_set_attr(rank, world_size, init_size, nullptr, &attributes); | 35 | + if (rank == 0) { |
| 47 | - TORCH_CHECK(status == 0, "shmem_set_attr failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | 36 | + status = c10d::symmetric_memory::Aclshmemx_get_uniqueid(&unique_id); |
| 48 | - logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, end shmem_set_attr rank is %d, world_size is %d, size is %llu.", | 37 | + TORCH_CHECK(status == 0, "aclshmemx_get_uniqueid failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); |
| 49 | - rank, world_size, init_size); | 38 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, aclshmemx_get_uniqueid rank is %d, version %d, internal is %s.", |
| 39 | + rank, unique_id.version, unique_id.internal); | ||
| 40 | + } | ||
| 41 | + auto unique_ids = storeExchange.all_gather(store, rank, world_size, unique_id); | ||
| 42 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, unique_id rank is %d, version %d, internal is %s.", | ||
| 43 | + rank, unique_ids[0].version, unique_ids[0].internal); | ||
| 50 | 44 | ||
| 51 | - status = c10d::symmetric_memory::Shmem_set_attr_uniqueid_args(rank, world_size, &unique_ids[0], attributes); | 45 | + aclshmemx_init_attr_t attr; |
| 52 | - TORCH_CHECK(status == 0, "Shmem_set_attr_uniqueid_args failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | 46 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, start aclshmemx_set_attr_uniqueid_args rank is %d, world_size is %d, size is %llu.", |
| 53 | - logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store success, rank is %d, world_size is %d.", rank, world_size); | 47 | + rank, world_size, init_size); |
| 48 | + | ||
| 49 | + status = c10d::symmetric_memory::Aclshmemx_set_attr_uniqueid_args(rank, world_size, init_size, &unique_ids[0], &attr); | ||
| 50 | + TORCH_CHECK(status == 0, "aclshmemx_set_attr_uniqueid_args failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | ||
| 51 | + | ||
| 52 | + status = c10d::symmetric_memory::Aclshmemx_init_attr(ACLSHMEMX_INIT_WITH_DEFAULT, &attr); | ||
| 53 | + TORCH_CHECK(status == 0, "aclshmemx_init_attr failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | ||
| 54 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store success, rank is %d, world_size is %d.", rank, world_size); | ||
| 55 | + } else { | ||
| 56 | + shmem_uniqueid_t unique_id; | ||
| 57 | + if (rank == 0) { | ||
| 58 | + status = c10d::symmetric_memory::Shmemx_get_uniqueid(&unique_id); | ||
| 59 | + TORCH_CHECK(status == 0, "shmem_get_uniqueid failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | ||
| 60 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, Shmem_get_uniqueid rank is %d, version %d, internal is %s.", | ||
| 61 | + rank, unique_id.version, unique_id.internal); | ||
| 62 | + } | ||
| 63 | + auto unique_ids = storeExchange.all_gather(store, rank, world_size, unique_id); | ||
| 64 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, unique_id rank is %d, version %d, internal is %s.", | ||
| 65 | + rank, unique_ids[0].version, unique_ids[0].internal); | ||
| 66 | + | ||
| 67 | + shmem_init_attr_t* attributes; | ||
| 68 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, start shmem_set_attr rank is %d, world_size is %d, size is %llu.", | ||
| 69 | + rank, world_size, init_size); | ||
| 70 | + status = c10d::symmetric_memory::Shmem_set_attr(rank, world_size, init_size, nullptr, &attributes); | ||
| 71 | + TORCH_CHECK(status == 0, "shmem_set_attr failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | ||
| 72 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store, end shmem_set_attr rank is %d, world_size is %d, size is %llu.", | ||
| 73 | + rank, world_size, init_size); | ||
| 74 | + | ||
| 75 | + status = c10d::symmetric_memory::Shmem_set_attr_uniqueid_args(rank, world_size, &unique_ids[0], attributes); | ||
| 76 | + TORCH_CHECK(status == 0, "shmem_set_attr_uniqueid_args failed, status is ", status, DIST_ERROR(ErrCode::INTERNAL)); | ||
| 77 | + logger->debug("NPUSHMEMSymmetricMemoryAllocator initialize_npushmem_with_store success, rank is %d, world_size is %d.", rank, world_size); | ||
| 78 | + } | ||
| 54 | is_initialized = true; | 79 | is_initialized = true; |
| 55 | } | 80 | } |
| 56 | 81 | ||
| @@ -13,9 +13,16 @@ namespace symmetric_memory { | |||
| 13 | GET_FUNCTION(libshmem, funcName) | 13 | GET_FUNCTION(libshmem, funcName) |
| 14 | 14 | ||
| 15 | REGISTER_LIBRARY(libshmem) | 15 | REGISTER_LIBRARY(libshmem) |
| 16 | +LOAD_FUNCTION(aclshmemx_set_conf_store_tls) | ||
| 17 | +LOAD_FUNCTION(aclshmemx_get_uniqueid) | ||
| 18 | +LOAD_FUNCTION(aclshmemx_set_attr_uniqueid_args) | ||
| 19 | +LOAD_FUNCTION(aclshmemx_init_attr) | ||
| 20 | +LOAD_FUNCTION(aclshmem_malloc) | ||
| 21 | +LOAD_FUNCTION(aclshmem_free) | ||
| 22 | +LOAD_FUNCTION(aclshmem_ptr) | ||
| 23 | +LOAD_FUNCTION(aclshmem_finalize) | ||
| 16 | LOAD_FUNCTION(shmem_set_conf_store_tls) | 24 | LOAD_FUNCTION(shmem_set_conf_store_tls) |
| 17 | LOAD_FUNCTION(shmem_set_attr) | 25 | LOAD_FUNCTION(shmem_set_attr) |
| 18 | -LOAD_FUNCTION(shmem_init_attr) | ||
| 19 | LOAD_FUNCTION(shmem_get_uniqueid) | 26 | LOAD_FUNCTION(shmem_get_uniqueid) |
| 20 | LOAD_FUNCTION(shmem_set_attr_uniqueid_args) | 27 | LOAD_FUNCTION(shmem_set_attr_uniqueid_args) |
| 21 | LOAD_FUNCTION(shmem_malloc) | 28 | LOAD_FUNCTION(shmem_malloc) |
| @@ -23,14 +30,18 @@ LOAD_FUNCTION(shmem_free) | |||
| 23 | LOAD_FUNCTION(shmem_ptr) | 30 | LOAD_FUNCTION(shmem_ptr) |
| 24 | LOAD_FUNCTION(shmem_finalize) | 31 | LOAD_FUNCTION(shmem_finalize) |
| 25 | 32 | ||
| 26 | -int32_t Shmem_set_conf_store_tls(bool enable, const char *tls_info, const uint32_t tls_info_len) | 33 | +int32_t Aclshmemx_set_conf_store_tls(bool enable, const char *tls_info, const uint32_t tls_info_len) |
| 27 | { | 34 | { |
| 28 | typedef int32_t (*ShmemApiFunc)(bool, const char *, const uint32_t); | 35 | typedef int32_t (*ShmemApiFunc)(bool, const char *, const uint32_t); |
| 29 | static ShmemApiFunc shmem_set_conf_store_tls_func = nullptr; | 36 | static ShmemApiFunc shmem_set_conf_store_tls_func = nullptr; |
| 37 | + if (shmem_set_conf_store_tls_func == nullptr) { | ||
| 38 | + shmem_set_conf_store_tls_func = (ShmemApiFunc)GET_FUNC(aclshmemx_set_conf_store_tls); | ||
| 39 | + } | ||
| 30 | if (shmem_set_conf_store_tls_func == nullptr) { | 40 | if (shmem_set_conf_store_tls_func == nullptr) { |
| 31 | shmem_set_conf_store_tls_func = (ShmemApiFunc)GET_FUNC(shmem_set_conf_store_tls); | 41 | shmem_set_conf_store_tls_func = (ShmemApiFunc)GET_FUNC(shmem_set_conf_store_tls); |
| 32 | } | 42 | } |
| 33 | - TORCH_CHECK(shmem_set_conf_store_tls_func, "Failed to find function ", "shmem_set_conf_store_tls", PTA_ERROR(ErrCode::NOT_FOUND)); | 43 | + TORCH_CHECK(shmem_set_conf_store_tls_func, "Failed to find function ", |
| 44 | + "aclshmemx_set_conf_store_tls or shmem_set_conf_store_tls", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 34 | return shmem_set_conf_store_tls_func(enable, tls_info, tls_info_len); | 45 | return shmem_set_conf_store_tls_func(enable, tls_info, tls_info_len); |
| 35 | } | 46 | } |
| 36 | 47 | ||
| @@ -46,25 +57,41 @@ int32_t Shmem_set_attr(int32_t my_rank, int32_t n_ranks, uint64_t local_mem_size | |||
| 46 | return shmem_set_attr_func(my_rank, n_ranks, local_mem_size, ip_port, attributes); | 57 | return shmem_set_attr_func(my_rank, n_ranks, local_mem_size, ip_port, attributes); |
| 47 | } | 58 | } |
| 48 | 59 | ||
| 49 | -int32_t Shmem_init_attr(shmem_init_attr_t *attributes) | 60 | +int Shmemx_get_uniqueid(shmem_uniqueid_t *uid) |
| 50 | -{ | ||
| 51 | - typedef int32_t (*ShmemApiFunc)(shmem_init_attr_t *); | ||
| 52 | - static ShmemApiFunc shmem_init_attr_func = nullptr; | ||
| 53 | - if (shmem_init_attr_func == nullptr) { | ||
| 54 | - shmem_init_attr_func = (ShmemApiFunc)GET_FUNC(shmem_init_attr); | ||
| 55 | - } | ||
| 56 | - TORCH_CHECK(shmem_init_attr_func, "Failed to find function ", "shmem_init_attr", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 57 | - return shmem_init_attr_func(attributes); | ||
| 58 | -} | ||
| 59 | - | ||
| 60 | -int32_t Shmem_get_uniqueid(shmem_uniqueid_t *uid) | ||
| 61 | { | 61 | { |
| 62 | typedef int32_t (*ShmemApiFunc)(shmem_uniqueid_t *); | 62 | typedef int32_t (*ShmemApiFunc)(shmem_uniqueid_t *); |
| 63 | static ShmemApiFunc shmem_get_uniqueid_func = nullptr; | 63 | static ShmemApiFunc shmem_get_uniqueid_func = nullptr; |
| 64 | if (shmem_get_uniqueid_func == nullptr) { | 64 | if (shmem_get_uniqueid_func == nullptr) { |
| 65 | shmem_get_uniqueid_func = (ShmemApiFunc)GET_FUNC(shmem_get_uniqueid); | 65 | shmem_get_uniqueid_func = (ShmemApiFunc)GET_FUNC(shmem_get_uniqueid); |
| 66 | } | 66 | } |
| 67 | - TORCH_CHECK(shmem_get_uniqueid_func, "Failed to find function ", "shmem_get_uniqueid", PTA_ERROR(ErrCode::NOT_FOUND)); | 67 | + TORCH_CHECK(shmem_get_uniqueid_func, "Failed to find function ", |
| 68 | + "shmem_get_uniqueid", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 69 | + return shmem_get_uniqueid_func(uid); | ||
| 70 | +} | ||
| 71 | + | ||
| 72 | +bool Aclshmemx_get_uniqueid_exist() | ||
| 73 | +{ | ||
| 74 | + const static bool shmemApiFuncExist = []() -> bool { | ||
| 75 | + try { | ||
| 76 | + auto func = GET_FUNC(aclshmemx_get_uniqueid); | ||
| 77 | + return func != nullptr; | ||
| 78 | + } catch (...) { | ||
| 79 | + // libshmem.so not exist | ||
| 80 | + return false; | ||
| 81 | + } | ||
| 82 | + }(); | ||
| 83 | + return shmemApiFuncExist; | ||
| 84 | +} | ||
| 85 | + | ||
| 86 | +int32_t Aclshmemx_get_uniqueid(aclshmemx_uniqueid_t *uid) | ||
| 87 | +{ | ||
| 88 | + typedef int32_t (*ShmemApiFunc)(aclshmemx_uniqueid_t *); | ||
| 89 | + static ShmemApiFunc shmem_get_uniqueid_func = nullptr; | ||
| 90 | + if (shmem_get_uniqueid_func == nullptr) { | ||
| 91 | + shmem_get_uniqueid_func = (ShmemApiFunc)GET_FUNC(aclshmemx_get_uniqueid); | ||
| 92 | + } | ||
| 93 | + TORCH_CHECK(shmem_get_uniqueid_func, "Failed to find function ", | ||
| 94 | + "aclshmemx_get_uniqueid or shmem_get_uniqueid", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 68 | return shmem_get_uniqueid_func(uid); | 95 | return shmem_get_uniqueid_func(uid); |
| 69 | } | 96 | } |
| 70 | 97 | ||
| @@ -75,49 +102,88 @@ int Shmem_set_attr_uniqueid_args(int rank_id, int nranks, const shmem_uniqueid_t | |||
| 75 | if (shmem_set_attr_uniqueid_args_func == nullptr) { | 102 | if (shmem_set_attr_uniqueid_args_func == nullptr) { |
| 76 | shmem_set_attr_uniqueid_args_func = (ShmemApiFunc)GET_FUNC(shmem_set_attr_uniqueid_args); | 103 | shmem_set_attr_uniqueid_args_func = (ShmemApiFunc)GET_FUNC(shmem_set_attr_uniqueid_args); |
| 77 | } | 104 | } |
| 78 | - TORCH_CHECK(shmem_set_attr_uniqueid_args_func, "Failed to find function ", "shmem_set_attr_uniqueid_args", PTA_ERROR(ErrCode::NOT_FOUND)); | 105 | + TORCH_CHECK(shmem_set_attr_uniqueid_args_func, "Failed to find function ", |
| 106 | + "shmem_set_attr_uniqueid_args", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 79 | return shmem_set_attr_uniqueid_args_func(rank_id, nranks, uid, attr); | 107 | return shmem_set_attr_uniqueid_args_func(rank_id, nranks, uid, attr); |
| 80 | } | 108 | } |
| 81 | 109 | ||
| 82 | -void *Shmem_malloc(size_t size) | 110 | +int Aclshmemx_set_attr_uniqueid_args(int rank_id, int nranks, int64_t local_mem_size, |
| 111 | + aclshmemx_uniqueid_t *uid, aclshmemx_init_attr_t *aclshmem_attr) | ||
| 112 | +{ | ||
| 113 | + typedef int32_t (*ShmemApiFunc)(int, int, int64_t, aclshmemx_uniqueid_t *, aclshmemx_init_attr_t *); | ||
| 114 | + static ShmemApiFunc aclshmemx_set_attr_uniqueid_args_func = nullptr; | ||
| 115 | + if (aclshmemx_set_attr_uniqueid_args_func == nullptr) { | ||
| 116 | + aclshmemx_set_attr_uniqueid_args_func = (ShmemApiFunc)GET_FUNC(aclshmemx_set_attr_uniqueid_args); | ||
| 117 | + } | ||
| 118 | + TORCH_CHECK(aclshmemx_set_attr_uniqueid_args_func, "Failed to find function ", | ||
| 119 | + "aclshmemx_set_attr_uniqueid_args", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 120 | + return aclshmemx_set_attr_uniqueid_args_func(rank_id, nranks, local_mem_size, uid, aclshmem_attr); | ||
| 121 | +} | ||
| 122 | + | ||
| 123 | +int Aclshmemx_init_attr(aclshmemx_bootstrap_t bootstrap_flags, aclshmemx_init_attr_t *attributes) | ||
| 124 | +{ | ||
| 125 | + typedef int32_t (*ShmemApiFunc)(aclshmemx_bootstrap_t, aclshmemx_init_attr_t *); | ||
| 126 | + static ShmemApiFunc aclshmemx_init_attr_func = nullptr; | ||
| 127 | + if (aclshmemx_init_attr_func == nullptr) { | ||
| 128 | + aclshmemx_init_attr_func = (ShmemApiFunc)GET_FUNC(aclshmemx_init_attr); | ||
| 129 | + } | ||
| 130 | + TORCH_CHECK(aclshmemx_init_attr_func, "Failed to find function ", | ||
| 131 | + "aclshmemx_init_attr", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 132 | + return aclshmemx_init_attr_func(bootstrap_flags, attributes); | ||
| 133 | +} | ||
| 134 | + | ||
| 135 | +void *Aclshmem_malloc(size_t size) | ||
| 83 | { | 136 | { |
| 84 | typedef void* (*ShmemApiFunc)(size_t); | 137 | typedef void* (*ShmemApiFunc)(size_t); |
| 85 | static ShmemApiFunc shmem_malloc_func = nullptr; | 138 | static ShmemApiFunc shmem_malloc_func = nullptr; |
| 139 | + if (shmem_malloc_func == nullptr) { | ||
| 140 | + shmem_malloc_func = (ShmemApiFunc)GET_FUNC(aclshmem_malloc); | ||
| 141 | + } | ||
| 86 | if (shmem_malloc_func == nullptr) { | 142 | if (shmem_malloc_func == nullptr) { |
| 87 | shmem_malloc_func = (ShmemApiFunc)GET_FUNC(shmem_malloc); | 143 | shmem_malloc_func = (ShmemApiFunc)GET_FUNC(shmem_malloc); |
| 88 | } | 144 | } |
| 89 | - TORCH_CHECK(shmem_malloc_func, "Failed to find function ", "shmem_malloc", PTA_ERROR(ErrCode::NOT_FOUND)); | 145 | + TORCH_CHECK(shmem_malloc_func, "Failed to find function ", |
| 146 | + "aclshmem_malloc or shmem_malloc", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 90 | return shmem_malloc_func(size); | 147 | return shmem_malloc_func(size); |
| 91 | } | 148 | } |
| 92 | 149 | ||
| 93 | -void Shmem_free(void *ptr) | 150 | +void Aclshmem_free(void *ptr) |
| 94 | { | 151 | { |
| 95 | typedef void (*ShmemApiFunc)(void *); | 152 | typedef void (*ShmemApiFunc)(void *); |
| 96 | static ShmemApiFunc shmem_free_func = nullptr; | 153 | static ShmemApiFunc shmem_free_func = nullptr; |
| 154 | + if (shmem_free_func == nullptr) { | ||
| 155 | + shmem_free_func = (ShmemApiFunc)GET_FUNC(aclshmem_free); | ||
| 156 | + } | ||
| 97 | if (shmem_free_func == nullptr) { | 157 | if (shmem_free_func == nullptr) { |
| 98 | shmem_free_func = (ShmemApiFunc)GET_FUNC(shmem_free); | 158 | shmem_free_func = (ShmemApiFunc)GET_FUNC(shmem_free); |
| 99 | } | 159 | } |
| 100 | - TORCH_CHECK(shmem_free_func, "Failed to find function ", "shmem_free", PTA_ERROR(ErrCode::NOT_FOUND)); | 160 | + TORCH_CHECK(shmem_free_func, "Failed to find function ", |
| 161 | + "aclshmem_free or shmem_free", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 101 | return shmem_free_func(ptr); | 162 | return shmem_free_func(ptr); |
| 102 | } | 163 | } |
| 103 | 164 | ||
| 104 | -void *Shmem_ptr(void *ptr, int pe) | 165 | +void *Aclshmem_ptr(void *ptr, int pe) |
| 105 | { | 166 | { |
| 106 | typedef void* (*ShmemApiFunc)(void *, int); | 167 | typedef void* (*ShmemApiFunc)(void *, int); |
| 107 | static ShmemApiFunc shmem_ptr_func = nullptr; | 168 | static ShmemApiFunc shmem_ptr_func = nullptr; |
| 169 | + if (shmem_ptr_func == nullptr) { | ||
| 170 | + shmem_ptr_func = (ShmemApiFunc)GET_FUNC(aclshmem_ptr); | ||
| 171 | + } | ||
| 108 | if (shmem_ptr_func == nullptr) { | 172 | if (shmem_ptr_func == nullptr) { |
| 109 | shmem_ptr_func = (ShmemApiFunc)GET_FUNC(shmem_ptr); | 173 | shmem_ptr_func = (ShmemApiFunc)GET_FUNC(shmem_ptr); |
| 110 | } | 174 | } |
| 111 | - TORCH_CHECK(shmem_ptr_func, "Failed to find function ", "shmem_ptr", PTA_ERROR(ErrCode::NOT_FOUND)); | 175 | + TORCH_CHECK(shmem_ptr_func, "Failed to find function ", |
| 176 | + "aclshmem_ptr or shmem_ptr", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 112 | return shmem_ptr_func(ptr, pe); | 177 | return shmem_ptr_func(ptr, pe); |
| 113 | } | 178 | } |
| 114 | 179 | ||
| 115 | -bool Shmem_finalize_exist() | 180 | +bool Aclshmem_finalize_exist() |
| 116 | { | 181 | { |
| 117 | const static bool shmemApiFuncExist = []() -> bool { | 182 | const static bool shmemApiFuncExist = []() -> bool { |
| 118 | try { | 183 | try { |
| 119 | - auto func = GET_FUNC(shmem_finalize) | 184 | + auto func1 = GET_FUNC(aclshmem_finalize); |
| 120 | - return func != nullptr; | 185 | + auto func2 = GET_FUNC(shmem_finalize); |
| 186 | + return func1 != nullptr || func2 != nullptr; | ||
| 121 | } catch (...) { | 187 | } catch (...) { |
| 122 | // libshmem.so not exist | 188 | // libshmem.so not exist |
| 123 | return false; | 189 | return false; |
| @@ -126,14 +192,18 @@ bool Shmem_finalize_exist() | |||
| 126 | return shmemApiFuncExist; | 192 | return shmemApiFuncExist; |
| 127 | } | 193 | } |
| 128 | 194 | ||
| 129 | -int Shmem_finalize(void) | 195 | +int Aclshmem_finalize(void) |
| 130 | { | 196 | { |
| 131 | typedef int (*ShmemApiFunc)(void); | 197 | typedef int (*ShmemApiFunc)(void); |
| 132 | static ShmemApiFunc shmem_finalize_func = nullptr; | 198 | static ShmemApiFunc shmem_finalize_func = nullptr; |
| 199 | + if (shmem_finalize_func == nullptr) { | ||
| 200 | + shmem_finalize_func = (ShmemApiFunc)GET_FUNC(aclshmem_finalize); | ||
| 201 | + } | ||
| 133 | if (shmem_finalize_func == nullptr) { | 202 | if (shmem_finalize_func == nullptr) { |
| 134 | shmem_finalize_func = (ShmemApiFunc)GET_FUNC(shmem_finalize); | 203 | shmem_finalize_func = (ShmemApiFunc)GET_FUNC(shmem_finalize); |
| 135 | } | 204 | } |
| 136 | - TORCH_CHECK(shmem_finalize_func, "Failed to find function ", "shmem_finalize", PTA_ERROR(ErrCode::NOT_FOUND)); | 205 | + TORCH_CHECK(shmem_finalize_func, "Failed to find function ", |
| 206 | + "aclshmem_finalize or shmem_finalize", PTA_ERROR(ErrCode::NOT_FOUND)); | ||
| 137 | return shmem_finalize_func(); | 207 | return shmem_finalize_func(); |
| 138 | } | 208 | } |
| 139 | 209 | ||
| @@ -7,26 +7,32 @@ | |||
| 7 | namespace c10d { | 7 | namespace c10d { |
| 8 | namespace symmetric_memory { | 8 | namespace symmetric_memory { |
| 9 | 9 | ||
| 10 | -int32_t Shmem_set_conf_store_tls(bool enable, const char *tls_info, const uint32_t tls_info_len); | 10 | +int32_t Aclshmemx_set_conf_store_tls(bool enable, const char *tls_info, const uint32_t tls_info_len); |
| 11 | 11 | ||
| 12 | int32_t Shmem_set_attr(int32_t my_rank, int32_t n_ranks, uint64_t local_mem_size, const char *ip_port, | 12 | int32_t Shmem_set_attr(int32_t my_rank, int32_t n_ranks, uint64_t local_mem_size, const char *ip_port, |
| 13 | shmem_init_attr_t **attributes); | 13 | shmem_init_attr_t **attributes); |
| 14 | 14 | ||
| 15 | -int32_t Shmem_init_attr(shmem_init_attr_t *attributes); | 15 | +int Shmemx_get_uniqueid(shmem_uniqueid_t *uid); |
| 16 | 16 | ||
| 17 | -int Shmem_get_uniqueid(shmem_uniqueid_t *uid); | 17 | +bool Aclshmemx_get_uniqueid_exist(); |
| 18 | + | ||
| 19 | +int Aclshmemx_get_uniqueid(aclshmemx_uniqueid_t *uid); | ||
| 18 | 20 | ||
| 19 | int Shmem_set_attr_uniqueid_args(int rank_id, int nranks, const shmem_uniqueid_t *uid, shmem_init_attr_t *attr); | 21 | int Shmem_set_attr_uniqueid_args(int rank_id, int nranks, const shmem_uniqueid_t *uid, shmem_init_attr_t *attr); |
| 20 | 22 | ||
| 21 | -void *Shmem_malloc(size_t size); | 23 | +int Aclshmemx_set_attr_uniqueid_args(int rank_id, int nranks, int64_t local_mem_size, aclshmemx_uniqueid_t *uid, aclshmemx_init_attr_t *aclshmem_attr); |
| 22 | 24 | ||
| 23 | -void Shmem_free(void *ptr); | 25 | +int Aclshmemx_init_attr(aclshmemx_bootstrap_t bootstrap_flags, aclshmemx_init_attr_t *attributes); |
| 24 | 26 | ||
| 25 | -void *Shmem_ptr(void *ptr, int pe); | 27 | +void *Aclshmem_malloc(size_t size); |
| 26 | 28 | ||
| 27 | -bool Shmem_finalize_exist(); | 29 | +void Aclshmem_free(void *ptr); |
| 28 | 30 | ||
| 29 | -int Shmem_finalize(void); | 31 | +void *Aclshmem_ptr(void *ptr, int pe); |
| 32 | + | ||
| 33 | +bool Aclshmem_finalize_exist(); | ||
| 34 | + | ||
| 35 | +int Aclshmem_finalize(void); | ||
| 30 | 36 | ||
| 31 | } // namespace symmetric_memory | 37 | } // namespace symmetric_memory |
| 32 | } // namespace c10d | 38 | } // namespace c10d |
| @@ -23,7 +23,7 @@ NPUSHMEMAllocation::~NPUSHMEMAllocation() | |||
| 23 | auto device = c10::Device(at::DeviceType::PrivateUse1, device_idx); | 23 | auto device = c10::Device(at::DeviceType::PrivateUse1, device_idx); |
| 24 | at::DeviceGuard device_guard(device); | 24 | at::DeviceGuard device_guard(device); |
| 25 | logger->debug("~NPUSHMEMAllocation, start Shmem_free, ptr is %p.", ptr); | 25 | logger->debug("~NPUSHMEMAllocation, start Shmem_free, ptr is %p.", ptr); |
| 26 | - Shmem_free(ptr); // shmem_free has no return value | 26 | + Aclshmem_free(ptr); // shmem_free has no return value |
| 27 | logger->debug("~NPUSHMEMAllocation, end Shmem_free, ptr is %p.", ptr); | 27 | logger->debug("~NPUSHMEMAllocation, end Shmem_free, ptr is %p.", ptr); |
| 28 | } | 28 | } |
| 29 | 29 | ||
| @@ -66,7 +66,7 @@ NPUSHMEMSymmetricMemory::NPUSHMEMSymmetricMemory( | |||
| 66 | TORCH_INTERNAL_ASSERT(!group_info.rank_to_global_rank.empty()); | 66 | TORCH_INTERNAL_ASSERT(!group_info.rank_to_global_rank.empty()); |
| 67 | rank_to_global_rank_ = group_info.rank_to_global_rank; | 67 | rank_to_global_rank_ = group_info.rank_to_global_rank; |
| 68 | for (int r = 0; r < world_size_; ++r) { | 68 | for (int r = 0; r < world_size_; ++r) { |
| 69 | - auto buffer = Shmem_ptr(allocation->ptr, rank_to_global_rank_[r]); | 69 | + auto buffer = Aclshmem_ptr(allocation->ptr, rank_to_global_rank_[r]); |
| 70 | TORCH_CHECK(buffer != nullptr, "shmem_ptr return nullptr with ptr ", allocation->ptr, DIST_ERROR(ErrCode::MEMORY)); | 70 | TORCH_CHECK(buffer != nullptr, "shmem_ptr return nullptr with ptr ", allocation->ptr, DIST_ERROR(ErrCode::MEMORY)); |
| 71 | buffers_.push_back(buffer); | 71 | buffers_.push_back(buffer); |
| 72 | logger->debug("[rank %d] NPUSHMEMSymmetricMemory shmem_ptr, r is %d, rank_to_global_rank is %d, ptr is %p, shmem_ptr is %p.", | 72 | logger->debug("[rank %d] NPUSHMEMSymmetricMemory shmem_ptr, r is %d, rank_to_global_rank is %d, ptr is %p, shmem_ptr is %p.", |
| @@ -205,7 +205,7 @@ void* NPUSHMEMSymmetricMemoryAllocator::alloc( | |||
| 205 | int world_size = group_info.world_size; | 205 | int world_size = group_info.world_size; |
| 206 | npushmem_extension::initialize_npushmem_with_store(store, rank, world_size); | 206 | npushmem_extension::initialize_npushmem_with_store(store, rank, world_size); |
| 207 | 207 | ||
| 208 | - auto ptr = Shmem_malloc(size); | 208 | + auto ptr = Aclshmem_malloc(size); |
| 209 | TORCH_CHECK(ptr != nullptr, "shmem_malloc return nullptr with size ", size, DIST_ERROR(ErrCode::MEMORY)); | 209 | TORCH_CHECK(ptr != nullptr, "shmem_malloc return nullptr with size ", size, DIST_ERROR(ErrCode::MEMORY)); |
| 210 | auto allocation = | 210 | auto allocation = |
| 211 | std::make_shared<NPUSHMEMAllocation>(ptr, size, device_idx); | 211 | std::make_shared<NPUSHMEMAllocation>(ptr, size, device_idx); |


代码一致性: 宏'ACLSHMEM_MAX_IP_PORT_LEN'使用#define定义,而其他类似的长度常量如'SHMEM_UNIQUE_ID_INNER_LEN'和'ACLSHMEM_UNIQUE_ID_INNER_LEN'使用constexpr定义。这种不一致的常量定义方式降低了代码的一致性。
问题类型: 代码一致性 文件路径:
third_party/shmem/include/shmem_host_def.h行号: 74 问题代码:修改建议:
此评论由代码审查工具自动生成