cuda编程笔记（32）--OpenMPI的使用

下载

如果你有权限，可以用sudo apt等方法安装，下面我们介绍通过源码生成（能用cuda）的方式

① 下载源码

wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.5.tar.gz
tar -xzf openmpi-4.1.5.tar.gz -C ~/software/openmpi

② 配置编译选项

重点是 --prefix 和 --with-cuda：

./configure --prefix=$HOME/software/openmpi \
            --with-cuda=/usr/local/cuda \
            --enable-mpi-cxx

• --prefix：告诉编译系统安装到 ~/software/openmpi

• --with-cuda：启用 CUDA 支持（前提：系统已有 /usr/local/cuda）

• --enable-mpi-cxx：允许 C++ MPI 程序（mpic++）

③ 编译并安装

make -j$(nproc)
make install

安装完成后，你的可执行文件、头文件、库会分别位于：

~/software/openmpi/bin/
~/software/openmpi/include/
~/software/openmpi/lib/

④ 设置环境变量

在 ~/.bashrc 最后加上以下几行（或使用 vim ~/.bashrc）：

export PATH=$HOME/software/openmpi/bin:$PATH
export LD_LIBRARY_PATH=$HOME/software/openmpi/lib:$LD_LIBRARY_PATH

source ~/.bashrc

验证是否生效：

mpirun --version

mpirun (Open MPI) 4.1.5

⑤ 检查是否支持 CUDA

ompi_info | grep -i cuda

  Configure command line: '--prefix=/home/huangxy/software/openmpi' '--with-cuda=/usr/local/cuda' '--enable-mpi-cxx'
          MPI extensions: affinity, cuda, pcollreq
                 MCA btl: smcuda (MCA v2.1.0, API v3.1.0, Component v4.1.5)
                MCA coll: cuda (MCA v2.1.0, API v2.0.0, Component v4.1.5)

配置环境

用cmake配置环境

# -----------------------------
# OpenMPI 路径
# -----------------------------
set(MPI_HOME "$ENV{HOME}/software/openmpi")
set(MPI_CXX_COMPILER "${MPI_HOME}/bin/mpicxx")
set(MPI_C_COMPILER "${MPI_HOME}/bin/mpicc")
# 用 cmake 自带的 find_package 检测 MPI
find_package(MPI REQUIRED)
if (MPI_FOUND)
    message(STATUS "MPI found: ${MPI_CXX_COMPILER}")
    message(STATUS "MPI include dir: ${MPI_INCLUDE_PATH}")
    message(STATUS "MPI library: ${MPI_LIBRARIES}")
else()
    message(FATAL_ERROR "MPI not found! Please check your MPI installation.")
endif()

# -----------------------------
# 指定头文件路径
# -----------------------------
target_include_directories(my_cuda_program
    PRIVATE
    ${CUDNN_INCLUDE_DIR}
    /usr/local/cuda/include
    ${NCCL_INCLUDE_DIR}
    ${TENSORRT_INCLUDE_DIR}
    ${CNPY_INCLUDE_DIR}
    ${MPI_INCLUDE_PATH}#MPI的头文件路径
)

# -----------------------------
# 链接库
# -----------------------------
target_link_libraries(my_cuda_program PRIVATE
    ${CUDNN_LIBRARY}
    cublasLt
    cublas
    cuda
    cudart
    ${NCCL_LIBRARY}
    nvinfer
    nvinfer_plugin
    nvonnxparser
    ${CNPY_LIBRARY}
    ${MPI_LIBRARIES}#MPI的库路径
)

基础API

MPI_Init

int MPI_Init(int *argc, char ***argv);

功能：
初始化 MPI 环境，所有 MPI 程序必须先调用它（相当于“MPI 世界的 main 函数入口”）。

参数：

• argc, argv：命令行参数的指针（和 main 的参数一样），MPI 可能会解析其中的信息。
  你也可以传 NULL, NULL，但建议原样传递。

返回值：

• MPI_SUCCESS 表示成功

• 其他错误码表示初始化失败（极少见）

说明：

• 所有进程都必须在调用其他 MPI 函数前执行它。

• 内部会建立进程间通信通道。

MPI_Comm_rank

int MPI_Comm_rank(MPI_Comm comm, int *rank);

功能：
获取当前进程在通信器 comm 中的编号（rank）。

参数：

• comm：通信器（一般用 MPI_COMM_WORLD 表示整个进程组）

• rank：输出参数，当前进程的 rank（从 0 到 size-1）

返回值：

• MPI_SUCCESS 表示成功。

MPI_Comm_size

int MPI_Comm_size(MPI_Comm comm, int *size);

功能：
获取通信器中进程总数。

参数：

• comm：通信器（通常是 MPI_COMM_WORLD）

• size：输出参数，进程总数

返回值：

• MPI_SUCCESS 表示成功。

MPI_Finalize

int MPI_Finalize();

功能：
结束 MPI 环境，释放所有通信资源。

参数：

• 无。

返回值：

• MPI_SUCCESS 表示成功。

说明：

• 所有进程都必须调用它，且必须在最后。

• 调用后不能再使用任何 MPI 函数。

MPI_Send

int MPI_Send(
    const void *buf, int count, MPI_Datatype datatype,
    int dest, int tag, MPI_Comm comm
);

功能：
向通信器 comm 中编号为 dest 的进程发送消息。

参数：

• buf：要发送的数据首地址(可以发送GPU数据)

• count：发送的数据元素数量

• datatype：元素类型（如 MPI_INT, MPI_FLOAT）

• dest：目标进程的 rank

• tag：消息标识（整型，可区分不同通信）

• comm：通信器（如 MPI_COMM_WORLD）

返回值：

• MPI_SUCCESS 表示成功。

说明：

• 是阻塞发送：直到系统缓冲区可安全释放或对方已接收。

• 可以用非阻塞版本 MPI_Isend。

MPI_Recv

int MPI_Recv(
    void *buf, int count, MPI_Datatype datatype,
    int source, int tag, MPI_Comm comm, MPI_Status *status
);

功能：
从通信器 comm 中的某个进程接收消息。

参数：

• buf：接收数据的缓冲区地址

• count：最多能接收的元素数量

• datatype：元素类型（必须和发送端一致）

• source：发送方的 rank（或 MPI_ANY_SOURCE）

• tag：消息标识（或 MPI_ANY_TAG）

• comm：通信器

• status：输出参数（可传 MPI_STATUS_IGNORE）

返回值：

• MPI_SUCCESS 表示成功。

MPI_Bcast

int MPI_Bcast(void *buffer, int count, MPI_Datatype datatype,
                int root, MPI_Comm comm);

功能：
把一个进程（称为 root）的数据广播给通信组中所有其他进程。所有进程必须同时调用该函数。

返回值：

• MPI_SUCCESS 表示成功。

MPI_Allreduce

int MPI_Allreduce(const void *sendbuf, void *recvbuf, int count,
                  MPI_Datatype datatype, MPI_Op op, MPI_Comm comm);

全规约（Allreduce）操作：
将所有进程的数据按照某种操作（如加和、最大值、最小值）规约后，把结果分发给每个进程。即「所有人贡献 → 所有人获得结果」。

返回值：

• MPI_SUCCESS 表示成功。

代码示例

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

int main(int argc, char** argv) {
    MPI_Init(&argc, &argv);
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    if (size != 2) {
        if(rank==0) std::cout << "This demo requires 2 processes.\n";
        MPI_Finalize();
        return 0;
    }

    // 分配 GPU 内存
    float *d_data;
    cudaSetDevice(rank); // 每个进程用不同 GPU
    cudaMalloc(&d_data, 4 * sizeof(float));

    // 初始化
    float h_data[4] = {rank*1.0f, rank*2.0f, rank*3.0f, rank*4.0f};
    cudaMemcpy(d_data, h_data, 4*sizeof(float), cudaMemcpyHostToDevice);

    // 发送/接收 GPU 数据（CUDA-aware MPI）
    if (rank == 0) {
        MPI_Send(d_data, 4, MPI_FLOAT, 1, 0, MPI_COMM_WORLD);
    } else if (rank == 1) {
        MPI_Recv(d_data, 4, MPI_FLOAT, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);

        // 拷回 CPU 打印
        float h_recv[4];
        cudaMemcpy(h_recv, d_data, 4*sizeof(float), cudaMemcpyDeviceToHost);
        std::cout << "GPU1 received: ";
        for(int i=0;i<4;i++) std::cout << h_recv[i] << " ";
        std::cout << std::endl;
    }

    cudaFree(d_data);
    MPI_Finalize();
    return 0;
}

编译完后，要运行。由于是进程间通信，得启动多个进程

mpirun -np 2 ./mpi_test

GPU1 received: 0 0 0 0