02-AllReduce 算法实现

#include <mpi.h>
#include <cuda_runtime.h>
#include <vector>

// Ring AllReduce 实现
template<typename T>
class RingAllReduce {
public:
    RingAllReduce(int rank, int worldSize, size_t count)
        : rank_(rank), worldSize_(worldSize), count_(count) {
        // 计算每个块的大小
        chunkSize_ = (count + worldSize - 1) / worldSize;

        // 分配发送/接收缓冲区
        cudaMalloc(&sendBuff_, chunkSize_ * sizeof(T));
        cudaMalloc(&recvBuff_, chunkSize_ * sizeof(T));

        // 计算前驱和后继
        prevRank_ = (rank - 1 + worldSize) % worldSize;
        nextRank_ = (rank + 1) % worldSize;
    }

    ~RingAllReduce() {
        cudaFree(sendBuff_);
        cudaFree(recvBuff_);
    }

    void execute(T* data, cudaStream_t stream) {
        // Phase 1: Reduce-Scatter
        reduceScatter(data, stream);

        // Phase 2: AllGather
        allGather(data, stream);
    }

private:
    void reduceScatter(T* data, cudaStream_t stream) {
        for (int step = 0; step < worldSize_ - 1; step++) {
            // 计算发送和接收的块索引
            int sendChunk = (rank_ - step + worldSize_) % worldSize_;
            int recvChunk = (rank_ - step - 1 + worldSize_) % worldSize_;

            size_t sendOffset = sendChunk * chunkSize_;
            size_t recvOffset = recvChunk * chunkSize_;
            size_t thisChunkSize = getChunkSize(sendChunk);

            // 复制数据到发送缓冲区
            cudaMemcpyAsync(sendBuff_, data + sendOffset,
                           thisChunkSize * sizeof(T),
                           cudaMemcpyDeviceToDevice, stream);

            // 异步发送
            MPI_Request sendReq, recvReq;
            MPI_Isend(sendBuff_, thisChunkSize * sizeof(T), MPI_BYTE,
                     nextRank_, 0, MPI_COMM_WORLD, &sendReq);

            // 异步接收
            MPI_Irecv(recvBuff_, getChunkSize(recvChunk) * sizeof(T), MPI_BYTE,
                     prevRank_, 0, MPI_COMM_WORLD, &recvReq);

            // 等待接收完成
            MPI_Wait(&recvReq, MPI_STATUS_IGNORE);

            // 归约：data[recvChunk] += recvBuff
            reduceKernel<<<gridSize_, blockSize_, 0, stream>>>(
                data + recvOffset, recvBuff_, getChunkSize(recvChunk));

            // 等待发送完成
            MPI_Wait(&sendReq, MPI_STATUS_IGNORE);

            cudaStreamSynchronize(stream);
        }
    }

    void allGather(T* data, cudaStream_t stream) {
        for (int step = 0; step < worldSize_ - 1; step++) {
            // 计算发送和接收的块索引
            int sendChunk = (rank_ - step + 1 + worldSize_) % worldSize_;
            int recvChunk = (rank_ - step + worldSize_) % worldSize_;

            size_t sendOffset = sendChunk * chunkSize_;
            size_t recvOffset = recvChunk * chunkSize_;

            // 复制数据到发送缓冲区
            cudaMemcpyAsync(sendBuff_, data + sendOffset,
                           getChunkSize(sendChunk) * sizeof(T),
                           cudaMemcpyDeviceToDevice, stream);

            MPI_Request sendReq, recvReq;
            MPI_Isend(sendBuff_, getChunkSize(sendChunk) * sizeof(T), MPI_BYTE,
                     nextRank_, 1, MPI_COMM_WORLD, &sendReq);

            MPI_Irecv(recvBuff_, getChunkSize(recvChunk) * sizeof(T), MPI_BYTE,
                     prevRank_, 1, MPI_COMM_WORLD, &recvReq);

            MPI_Wait(&recvReq, MPI_STATUS_IGNORE);

            // 直接复制（不需要归约）
            cudaMemcpyAsync(data + recvOffset, recvBuff_,
                           getChunkSize(recvChunk) * sizeof(T),
                           cudaMemcpyDeviceToDevice, stream);

            MPI_Wait(&sendReq, MPI_STATUS_IGNORE);
            cudaStreamSynchronize(stream);
        }
    }

    size_t getChunkSize(int chunk) {
        if (chunk == worldSize_ - 1) {
            return count_ - chunk * chunkSize_;
        }
        return chunkSize_;
    }

private:
    int rank_, worldSize_;
    int prevRank_, nextRank_;
    size_t count_, chunkSize_;
    T *sendBuff_, *recvBuff_;
    int gridSize_ = 256, blockSize_ = 256;
};

// 归约 Kernel
template<typename T>
__global__ void reduceKernel(T* dst, const T* src, size_t count) {
    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < count) {
        dst[idx] += src[idx];
    }
}

// Tree AllReduce 实现
template<typename T>
class TreeAllReduce {
public:
    TreeAllReduce(int rank, int worldSize, size_t count)
        : rank_(rank), worldSize_(worldSize), count_(count) {
        buildTree();
        cudaMalloc(&buffer_, count * sizeof(T));
    }

    void execute(T* data, cudaStream_t stream) {
        // Phase 1: Reduce (up)
        reduce(data, stream);

        // Phase 2: Broadcast (down)
        broadcast(data, stream);
    }

private:
    void buildTree() {
        // 构建二叉树
        // 对于 rank r:
        // - 父节点: (r - 1) / 2 (if r > 0)
        // - 左子节点: 2r + 1
        // - 右子节点: 2r + 2

        if (rank_ == 0) {
            parent_ = -1;
        } else {
            parent_ = (rank_ - 1) / 2;
        }

        int leftChild = 2 * rank_ + 1;
        int rightChild = 2 * rank_ + 2;

        if (leftChild < worldSize_) {
            children_.push_back(leftChild);
        }
        if (rightChild < worldSize_) {
            children_.push_back(rightChild);
        }

        // 计算树的深度
        depth_ = 0;
        int nodes = 1;
        while (nodes < worldSize_) {
            depth_++;
            nodes = nodes * 2 + 1;
        }
    }

    void reduce(T* data, cudaStream_t stream) {
        // 从叶子节点开始，自底向上
        // 叶子节点没有子节点，直接发送给父节点

        // 接收所有子节点的数据并归约
        for (int child : children_) {
            MPI_Recv(buffer_, count_ * sizeof(T), MPI_BYTE,
                    child, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

            // 归约
            reduceKernel<<<256, 256, 0, stream>>>(data, buffer_, count_);
            cudaStreamSynchronize(stream);
        }

        // 发送给父节点（非根节点）
        if (parent_ >= 0) {
            cudaMemcpy(buffer_, data, count_ * sizeof(T), cudaMemcpyDeviceToHost);
            MPI_Send(buffer_, count_ * sizeof(T), MPI_BYTE,
                    parent_, 0, MPI_COMM_WORLD);
        }
    }

    void broadcast(T* data, cudaStream_t stream) {
        // 从根节点开始，自顶向下

        // 接收父节点的数据（非根节点）
        if (parent_ >= 0) {
            MPI_Recv(buffer_, count_ * sizeof(T), MPI_BYTE,
                    parent_, 1, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            cudaMemcpy(data, buffer_, count_ * sizeof(T), cudaMemcpyHostToDevice);
        }

        // 发送给所有子节点
        cudaMemcpy(buffer_, data, count_ * sizeof(T), cudaMemcpyDeviceToHost);
        for (int child : children_) {
            MPI_Send(buffer_, count_ * sizeof(T), MPI_BYTE,
                    child, 1, MPI_COMM_WORLD);
        }
    }

private:
    int rank_, worldSize_;
    size_t count_;
    int parent_;
    std::vector<int> children_;
    int depth_;
    T* buffer_;
};

template<typename T>
class RecursiveHalvingDoublingAllReduce {
public:
    RecursiveHalvingDoublingAllReduce(int rank, int worldSize, size_t count)
        : rank_(rank), worldSize_(worldSize), count_(count) {
        // 检查是否是 2 的幂
        assert((worldSize & (worldSize - 1)) == 0);

        cudaMalloc(&sendBuff_, count * sizeof(T));
        cudaMalloc(&recvBuff_, count * sizeof(T));
    }

    void execute(T* data, cudaStream_t stream) {
        // Phase 1: Reduce-Scatter with Recursive Halving
        recursiveHalving(data, stream);

        // Phase 2: AllGather with Recursive Doubling
        recursiveDoubling(data, stream);
    }

private:
    void recursiveHalving(T* data, cudaStream_t stream) {
        size_t currentCount = count_;

        // log2(worldSize) 轮
        for (int mask = worldSize_ >> 1; mask > 0; mask >>= 1) {
            int peer = rank_ ^ mask;

            // 确定发送/接收哪半边
            size_t halfCount = currentCount / 2;
            size_t sendOffset, recvOffset;

            if (rank_ < peer) {
                // 发送后半部分，接收前半部分的归约结果
                sendOffset = halfCount;
                recvOffset = 0;
            } else {
                // 发送前半部分，接收后半部分的归约结果
                sendOffset = 0;
                recvOffset = halfCount;
            }

            // 发送
            cudaMemcpyAsync(sendBuff_, data + sendOffset,
                           halfCount * sizeof(T),
                           cudaMemcpyDeviceToDevice, stream);
            cudaStreamSynchronize(stream);

            MPI_Request sendReq, recvReq;
            MPI_Isend(sendBuff_, halfCount * sizeof(T), MPI_BYTE,
                     peer, 0, MPI_COMM_WORLD, &sendReq);
            MPI_Irecv(recvBuff_, halfCount * sizeof(T), MPI_BYTE,
                     peer, 0, MPI_COMM_WORLD, &recvReq);

            MPI_Wait(&recvReq, MPI_STATUS_IGNORE);

            // 归约到接收区域
            reduceKernel<<<256, 256, 0, stream>>>(
                data + recvOffset, recvBuff_, halfCount);
            cudaStreamSynchronize(stream);

            MPI_Wait(&sendReq, MPI_STATUS_IGNORE);

            // 更新有效数据区域
            if (rank_ < peer) {
                // 有效数据在前半部分
                currentCount = halfCount;
            } else {
                // 有效数据在后半部分，移动到开头
                cudaMemcpyAsync(data, data + halfCount,
                               halfCount * sizeof(T),
                               cudaMemcpyDeviceToDevice, stream);
                currentCount = halfCount;
            }
        }
    }

    void recursiveDoubling(T* data, cudaStream_t stream) {
        size_t currentCount = count_ / worldSize_;

        // log2(worldSize) 轮
        for (int mask = 1; mask < worldSize_; mask <<= 1) {
            int peer = rank_ ^ mask;

            // 计算发送/接收位置
            size_t sendOffset = 0;
            size_t recvOffset = currentCount;

            if (rank_ > peer) {
                std::swap(sendOffset, recvOffset);
            }

            // 交换数据
            cudaMemcpyAsync(sendBuff_, data + sendOffset,
                           currentCount * sizeof(T),
                           cudaMemcpyDeviceToDevice, stream);
            cudaStreamSynchronize(stream);

            MPI_Request sendReq, recvReq;
            MPI_Isend(sendBuff_, currentCount * sizeof(T), MPI_BYTE,
                     peer, 1, MPI_COMM_WORLD, &sendReq);
            MPI_Irecv(recvBuff_, currentCount * sizeof(T), MPI_BYTE,
                     peer, 1, MPI_COMM_WORLD, &recvReq);

            MPI_Wait(&recvReq, MPI_STATUS_IGNORE);

            cudaMemcpyAsync(data + recvOffset, recvBuff_,
                           currentCount * sizeof(T),
                           cudaMemcpyDeviceToDevice, stream);
            cudaStreamSynchronize(stream);

            MPI_Wait(&sendReq, MPI_STATUS_IGNORE);

            currentCount *= 2;
        }
    }

private:
    int rank_, worldSize_;
    size_t count_;
    T *sendBuff_, *recvBuff_;
};

template<typename T>
class HierarchicalAllReduce {
public:
    HierarchicalAllReduce(
        int globalRank,
        int worldSize,
        int localRank,
        int localSize,
        size_t count
    ) : globalRank_(globalRank),
        worldSize_(worldSize),
        localRank_(localRank),
        localSize_(localSize),
        count_(count) {

        nodeId_ = globalRank / localSize;
        numNodes_ = worldSize / localSize;

        // 创建节点内通信域
        MPI_Comm_split(MPI_COMM_WORLD, nodeId_, localRank, &localComm_);

        // 创建节点间通信域（仅 local rank 0 参与）
        int interColor = (localRank == 0) ? 0 : MPI_UNDEFINED;
        MPI_Comm_split(MPI_COMM_WORLD, interColor, nodeId_, &interComm_);

        cudaMalloc(&buffer_, count * sizeof(T));
    }

    void execute(T* data, cudaStream_t stream) {
        // Phase 1: 节点内 AllReduce
        intraNodeAllReduce(data, stream);

        // Phase 2: 节点间 AllReduce (仅 local rank 0)
        if (localRank_ == 0) {
            interNodeAllReduce(data, stream);
        }

        // Phase 3: 节点内 Broadcast
        intraNodeBroadcast(data, stream);
    }

private:
    void intraNodeAllReduce(T* data, cudaStream_t stream) {
        // 使用 NCCL 进行节点内 AllReduce
        // 利用 NVLink 高带宽
        ncclAllReduce(data, data, count_, getNcclDataType<T>(),
                     ncclSum, localNcclComm_, stream);
        cudaStreamSynchronize(stream);
    }

    void interNodeAllReduce(T* data, cudaStream_t stream) {
        // 节点间使用 MPI 或网络 NCCL
        cudaMemcpy(buffer_, data, count_ * sizeof(T), cudaMemcpyDeviceToHost);

        // 使用 Ring AllReduce
        T* sendBuff = buffer_;
        T* recvBuff = new T[count_];

        int prevNode = (nodeId_ - 1 + numNodes_) % numNodes_;
        int nextNode = (nodeId_ + 1) % numNodes_;
        size_t chunkSize = count_ / numNodes_;

        // Reduce-Scatter
        for (int step = 0; step < numNodes_ - 1; step++) {
            int sendChunk = (nodeId_ - step + numNodes_) % numNodes_;
            int recvChunk = (nodeId_ - step - 1 + numNodes_) % numNodes_;

            MPI_Request sendReq, recvReq;
            MPI_Isend(sendBuff + sendChunk * chunkSize,
                     chunkSize * sizeof(T), MPI_BYTE,
                     nextNode * localSize_, 0, MPI_COMM_WORLD, &sendReq);
            MPI_Irecv(recvBuff, chunkSize * sizeof(T), MPI_BYTE,
                     prevNode * localSize_, 0, MPI_COMM_WORLD, &recvReq);

            MPI_Wait(&recvReq, MPI_STATUS_IGNORE);

            // 归约
            for (size_t i = 0; i < chunkSize; i++) {
                sendBuff[recvChunk * chunkSize + i] += recvBuff[i];
            }

            MPI_Wait(&sendReq, MPI_STATUS_IGNORE);
        }

        // AllGather (类似)
        // ...

        cudaMemcpy(data, buffer_, count_ * sizeof(T), cudaMemcpyHostToDevice);
        delete[] recvBuff;
    }

    void intraNodeBroadcast(T* data, cudaStream_t stream) {
        // 使用 NCCL Broadcast
        ncclBroadcast(data, data, count_, getNcclDataType<T>(),
                     0, localNcclComm_, stream);
        cudaStreamSynchronize(stream);
    }

private:
    int globalRank_, worldSize_;
    int localRank_, localSize_;
    int nodeId_, numNodes_;
    size_t count_;
    MPI_Comm localComm_, interComm_;
    ncclComm_t localNcclComm_;
    T* buffer_;
};