# SHMEM API 样例 SHMEM包含host和device两类接口。host接口用SHMEM_HOST_API宏标识,device接口用SHMEM_DEVICE宏标识。 此章介绍各类API的常见接口的使用样例。 ## Init API SHMEM的初始化接口样例 ### 初始化状态 ```c++ enum { SHMEM_STATUS_NOT_INITIALIZED = 0, // 未初始化 SHMEM_STATUS_SHM_CREATED, // 完成共享内存堆创建 SHMEM_STATUS_IS_INITIALIZED, // 初始化完成 SHMEM_STATUS_INVALID = INT_MAX, }; ``` ### 初始化所需的attributes ```c++ // 初始化属性 typedef struct { int version; // 版本 int my_rank; // 当前rank int n_ranks; // 总rank数 const char* ip_port; // ip端口 uint64_t local_mem_size; // 本地申请内存大小 shmem_init_optional_attr_t option_attr; // 可选参数 } shmem_init_attr_t; // 可选属性 typedef struct { data_op_engine_type_t data_op_engine_type; // 数据引擎 // timeout uint32_t shm_init_timeout; uint32_t shm_create_timeout; uint32_t control_operation_timeout; } shmem_init_optional_attr_t; ``` ### 初始化样例 ```c++ #include #include #include #include "shmem_api.h" aclInit(nullptr); int status; int device_id = 0; int rank_id = 0; int n_ranks = 8; uint64_t local_mem_size = 1024UL * 1024UL * 1024; const char* test_global_ipport = "tcp://127.0.0.1:8666"; status = aclrtSetDevice(device_id); shmem_init_attr_t* attributes; shmem_set_attr(rank_id, n_ranks, local_mem_size, test_global_ipport, &attributes); // shmem_init_attr_t* attributes = new shmem_init_attr_t{rank_id, n_ranks, test_global_ipport, local_mem_size, {0, SHMEM_DATA_OP_MTE, 120, 120, 120}}; // 自定义attr shmem_init_attr(attributes); // delete attributes; // 销毁自定义attr status = shmem_init_status(); if (status == SHMEM_STATUS_IS_INITIALIZED) { std::cout << "Init success!" << std::endl; } //################你的任务################# //######################################### status = shmem_finalize(); aclrtResetDevice(device_id); aclFinalize(); ``` ### 自定义日志打印 自定义日志打印是可选操作,如果未注册自定义日志打印函数,日志则默认打屏,接口示例如下。 ```c void cpp_logger_example(int level, const char* msg) { // do print here } // set self-defined log int ret = shmem_set_extern_logger(cpp_logger_example); // set log level. 0-debug, 1-info, 2-warn, 3-error ret = shmem_set_log_level(level); ``` ### 注册私钥口令解密函数 shmem的多卡之间业务面TCP通信,通过memfabric_hybrid提供的能力实现,为了保证通信安全,该特性默认开启,使用时传入初始化TLS信息,TLS信息格式参考docs/security.md样例,其中私钥是密文,私钥口令以密文形式存储,需通过注册的解密回调函数进行解密,调用者实现具体的解密逻辑。 ```c int my_key_password_decrypt_handler(const char *cipherText, size_t cipherTextLen, char *plainText, size_t &plainTextLen) { // cipherText: input encrypted key password // plainText: output decrypted key password // plainTextLen: output decrypted key password length // do decrypt here } const char *pk = "xxx"; uint32_t pk_len = strlen(pk); const char *password = "xxxx"; uint32_t pw_len = strlen(password); int ret = shmem_set_config_store_tls_key(pk, pk_len, password, pw_len, my_key_password_decrypt_handler); ``` 如需关闭加密特性,则调用如下接口。关闭后,则无需调用shmem_register_decrypt_handler注册接口。 ```c int ret = shmem_set_conf_store_tls(false, nullptr, 0); ``` ### Team API SHMEM的通信域管理接口样例 ### host侧接口样例 ```c++ // ################### 调用初始化相关接口 ########################### //... // ###################### 子通信域切分 ############################# shmem_team_t team_odd; int start = 1; int stride = 2; int team_size = 4; shmem_team_split_strided(SHMEM_TEAM_WORLD, start, stride, team_size, &team_odd); // ##################### host侧取值 ############################### if (rank_id & 1) { // shmem_team_n_pes(team_odd): Returns the number of PEs in the team. int team_n_pes = shmem_team_n_pes(team_odd); // team_n_pes == team_size // shmem_team_my_pe(team_odd): Returns the number of the calling PE in the specified team. int team_my_pe = shmem_team_my_pe(team_odd); // team_my_pe == rank_id / stride // shmem_n_pes(): Returns the number of PEs running in the program. int n_pes = shmem_n_pes(); // n_pes == n_ranks // shmem_my_pe(): Returns the PE number of the local PE int my_pe = shmem_my_pe(); // my_pe == rank_id } // #################### 相关资源释放 ################################ shmem_team_destroy(team_odd); // ################## 调用去初始化相关接口 ########################### //... ``` ### device侧接口样例 ```c++ class kernel_state_test { public: __aicore__ inline kernel_state_test() {} __aicore__ inline void Init(GM_ADDR gva, shmem_team_t team_id) { gva_gm = (__gm__ int *)gva; team_idx= team_id; rank = smem_shm_get_global_rank(); // 获取当前rank rank_size = smem_shm_get_global_rank_size(); // 获取总rank数 } __aicore__ inline void Process() { AscendC::PipeBarrier(); // ##################### device侧取值 ############################### // shmem_int32_p 是RMA功能提供的接口,此处可简易理解为在device存储第二个入参的函数的结果。 // shmem_n_pes(): Returns the number of PEs running in the program. shmem_int32_p(gva_gm, shmem_n_pes(), rank); // shmem_my_pe(): Returns the PE number of the local PE shmem_int32_p(gva_gm + 1, shmem_my_pe(), rank); // shmem_team_my_pe(team_idx): Returns the number of the calling PE in the specified team. shmem_int32_p(gva_gm + 2, shmem_team_my_pe(team_idx), rank); // shmem_team_n_pes(team_idx): Returns the number of PEs in the team. shmem_int32_p(gva_gm + 3, shmem_team_n_pes(team_idx), rank); // shmem_team_translate_pe(team_idx, 1, SHMEM_TEAM_WORLD): Translate a given PE number in one team into the corresponding PE number in another team. shmem_int32_p(gva_gm + 4, shmem_team_translate_pe(team_idx, 1, SHMEM_TEAM_WORLD), rank); } private: __gm__ int *gva_gm; shmem_team_t team_idx; int64_t rank; int64_t rank_size; }; extern "C" __global__ __aicore__ void device_state_test(GM_ADDR gva, int team_id) { kernel_state_test op; op.Init(gva, (shmem_team_t)team_id); op.Process(); } void get_device_state(uint32_t block_dim, void* stream, uint8_t* gva, shmem_team_t team_id) { device_state_test<<>>(gva, (int)team_id); } ``` ## Mem API SHMEM的内存管理接口样例 ```c++ // ################### 调用初始化相关接口 ########################### //... // ################## 内存管理接口调用 ########################### // 分配1024 bytes,返回所分配内存的指针ptr. void *ptr = shmem_malloc(1024); // 释放ptr对应的被分配的内存空间. shmem_free(ptr); // ################## 调用去初始化相关接口 ########################### //... ``` ## RMA API SHMEM的远端内存访问接口样例 ```c++ class kernel_p { public: __aicore__ inline kernel_p() {} __aicore__ inline void Init(GM_ADDR gva, float val) { gva_gm = (__gm__ float *)gva; value = val; rank = smem_shm_get_global_rank(); // 获取当前rank rank_size = smem_shm_get_global_rank_size(); // 获取总rank数 } __aicore__ inline void Process() { // 把value的值put到共享内存gva_gm在(rank + 1) % rank_size中的对应位置。 shmem_float_p(gva_gm, value, (rank + 1) % rank_size); } private: __gm__ float *gva_gm; float value; int64_t rank; int64_t rank_size; }; extern "C" __global__ __aicore__ void p_num_test(GM_ADDR gva, float val) { kernel_p op; op.Init(gva, val); op.Process(); } void put_one_num_do(uint32_t block_dim, void* stream, uint8_t* gva, float val) { p_num_test<<>>(gva, val); } ``` ## Sync API SHMEM的同步管理接口样例 ```c++ // 任务1 // ... // 阻塞直到所有任务完成。 shmem_barrier_all(); // 任务2 // ... ```