样例代码

#include <iostream>
#include <fstream>
#include <vector>
#include <memory>
#include <thread>
#include <chrono>
#include <atomic>
#include <string.h>
#include <pthread.h>
#include "hccl/hccl.h"
#include "hccl/hccl_types.h"

using namespace std;
struct ThreadContext {
    HcclComm comm;
    int32_t device;
};
#define ACLCHECK(ret) do {\
    if(ret != ACL_SUCCESS)\
    {\
        printf("acl interface return err %s:%d, retcode: %d \n", __FILE__, __LINE__, ret);\
        return ret;\
    }\
} while(0)
#define HCCLCHECK(ret) do {\
    if(ret != HCCL_SUCCESS)\
    {\
        printf("hccl interface return err %s:%d, retcode: %d \n", __FILE__, __LINE__, ret);\
        return ret;\
    }\
} while(0)
// 参数平面位置
enum class NICDeployment {
    NIC_DEPLOYMENT_HOST = 0,
    NIC_DEPLOYMENT_DEVICE,
    NIC_DEPLOYMENT_RESERVED
};
 
using HcclRootHandle = struct HcclRootHandleDef {
    char ip[64];
    uint32_t port;
    NICDeployment nicDeploy;
    char identifier[128];
};
 
// rootInfo转HcclRootHandle，并打印
void printHcclRootHandle(const HcclRootInfo& rootInfo) {
    HcclRootHandle rootHandle;
    memcpy(&rootHandle, rootInfo.internal, sizeof(HcclRootHandle));
 
    // 打印HcclRootHandle的内容
    printf("IP: %s\t", rootHandle.ip);
    printf("Port: %u\t", rootHandle.port);
    printf("NIC Deployment: %d\t", static_cast<int>(rootHandle.nicDeploy));
    printf("Identifier: %s\n", rootHandle.identifier);
}
void allreduce_operation(void *sendBuf, void *recvBuf, uint64_t count, HcclDataType dataType,
    HcclReduceOp op, HcclComm comm, aclrtStream stream, int devId)
{
    pthread_setname_np(pthread_self(), "allreduce");
    // 指定集合通信操作使用的设备
    aclrtSetDevice(devId);
    // 尝试执行AllReduce操作
    HcclAllReduce(sendBuf, recvBuf, count, dataType, op, comm, stream);
    // 等待stream中集合通信任务执行完成
    aclrtSynchronizeStream(stream);
}
int Sample(void *arg)
{
    ThreadContext* ctx = (ThreadContext *)arg;
    void* host_buf = nullptr;
    void* send_buff = nullptr;
    void* recv_buff = nullptr;
    uint64_t count = 10;
    int malloc_kSize = count * sizeof(float);
    aclrtEvent start_event, end_event;
    aclrtStream stream;
    ACLCHECK(aclrtCreateStream(&stream));
    ACLCHECK(aclrtCreateEvent(&start_event));
    ACLCHECK(aclrtCreateEvent(&end_event));
    // 申请集合通信操作的内存
    ACLCHECK(aclrtMalloc((void**)&send_buff, malloc_kSize, ACL_MEM_MALLOC_HUGE_FIRST));
    ACLCHECK(aclrtMalloc((void**)&recv_buff, malloc_kSize, ACL_MEM_MALLOC_HUGE_FIRST));
    // 初始化输入内存
    ACLCHECK(aclrtMallocHost((void**)&host_buf, malloc_kSize));
    float* tmpHostBuff = static_cast<float*>(host_buf);
    for (uint32_t i = 0; i < count; ++i) {
        tmpHostBuff[i] = 20000;
    }
    ACLCHECK(aclrtMemcpy((void*)send_buff, malloc_kSize, (void*)host_buf, malloc_kSize, ACL_MEMCPY_HOST_TO_DEVICE));
    // 启动AllReduce操作的线程
    thread allreduce_thread(allreduce_operation, (void *)send_buff, (void*)recv_buff, count,
        HCCL_DATA_TYPE_FP32, HCCL_REDUCE_SUM, ctx->comm, stream, ctx->device);
    // 等待Allreduce线程结束
    allreduce_thread.join();
    cout << "Rank " << ctx->device << " AllReduce operation completed successfully." << endl;
    if (ctx->device < 8) {
        void* resultBuff;
        ACLCHECK(aclrtMallocHost((void**)&resultBuff, malloc_kSize));
        ACLCHECK(aclrtMemcpy((void*)resultBuff, malloc_kSize, (void*)recv_buff, malloc_kSize, ACL_MEMCPY_DEVICE_TO_HOST));
        ACLCHECK(aclrtFreeHost(resultBuff));
    }
    ACLCHECK(aclrtFree(send_buff));
    ACLCHECK(aclrtFree(recv_buff));
    ACLCHECK(aclrtFreeHost(host_buf));
    // 销毁任务流
    ACLCHECK(aclrtDestroyStream(stream));
    ACLCHECK(aclrtDestroyEvent(start_event));
    ACLCHECK(aclrtDestroyEvent(end_event));
    return 0;
}
int multiGetRootInfo(int commNum)
{
    std::ofstream out("rootinfo.bin", std::ios::binary);
    for (int rank = 0; rank < commNum; ++rank) {
        HcclRootInfo rootInfo;
        ACLCHECK(aclrtSetDevice(rank));
        HCCLCHECK(HcclGetRootInfo(&rootInfo));
        out.write(rootInfo.internal, HCCL_ROOT_INFO_BYTES);
        printHcclRootHandle(rootInfo);
        std::cout << "get rank done : " << rank << std::endl;
        
        // 重置设备
        ACLCHECK(aclrtResetDevice(rank));
    }
    out.close();
    return HCCL_SUCCESS;
}
HcclRootInfo loadRootInfo(int commNum, int commId)
{
    std::ifstream in("rootinfo.bin", std::ios::binary);
    for (int i = 0; i < commNum; ++i) {
        HcclRootInfo rootInfo{};
        in.read(rootInfo.internal, HCCL_ROOT_INFO_BYTES);
        if (i == commId) {
            printHcclRootHandle(rootInfo);
            return rootInfo;
        }
    }
    HcclRootInfo rootInfo_{};
    return rootInfo_;
}
int run(int commNum, int commId, int rankNum, int devId, int logicalRankId)
{
    auto rootInfo = loadRootInfo(commNum, commId);
    printHcclRootHandle(rootInfo);
    ACLCHECK(aclrtSetDevice(devId));
    // 创建并初始化通信域配置项
    HcclCommConfig config;
    HcclCommConfigInit(&config);
    // 根据需要修改通信域配置
    config.hcclBufferSize = 50;
    strcpy(config.hcclCommName, "comm_1");
    // 初始化集合通信域
    HcclComm hcclComm;
    HCCLCHECK(HcclCommInitRootInfoConfig(rankNum, &rootInfo, logicalRankId, &config, &hcclComm));
    // 创建任务stream
    struct ThreadContext args;
    args.comm = hcclComm;
    args.device = devId;
    Sample((void *)&args);
    // 销毁集合通信域，只有在未超时操作的情况下销毁
    HCCLCHECK(HcclCommDestroy(hcclComm));
    return 0;
}
int main(int argc, char*argv[])
{
    int devId = std::stoi(argv[1]); // dev id
    int logicalRank = std::stoi(argv[2]); // rank id
    int commId = std::stoi(argv[3]); // commId
    int commNum = std::stoi(argv[4]); // number of comm
    int rankNum = std::stoi(argv[5]); // rank size
    int devCount = rankNum;
    int logicalRankId = logicalRank;
    // 设备资源初始化
    ACLCHECK(aclInit(NULL));
    // 指定集合通信操作使用的设备
    ACLCHECK(aclrtSetDevice(devId));
    // 在rank0循环获取多个rank的rootInfo
    int32_t rootRank = 0;
    std::vector<std::thread> tmp;
    if (devId == rootRank) {
        std::vector<int> rootDevIds;
        std::vector<int> rootLogicalIds;
        for (int i=0; i<commNum; i++) {
            rootDevIds.emplace_back(i); // 每个通信域默认root device = commID
            rootLogicalIds.emplace_back(0);
        }
        HCCLCHECK(multiGetRootInfo(commNum));
        // rank0进程起多个线程，为对应的root调用HcclCommInitRootInfo
        for (int i=0; i<commNum; i++) {
            tmp.emplace_back(std::thread([=]() { run(commNum, i, rankNum, rootDevIds[i], rootLogicalIds[i]); }));
        }
    } else {
        run(commNum, commId, rankNum, devId, logicalRankId);
    }
    for (int i=0; i<tmp.size(); i++) {
        tmp[i].join();
    }
    // 重置设备
    ACLCHECK(aclrtResetDevice(devId));
    // 设备去初始化
    ACLCHECK(aclFinalize());
    return 0;
}