MPI_Probe

MPI_Probe 用于：等待一条匹配的消息到达，获取消息的信息，但不真正接收消息。

它通常与 MPI_Get_count、MPI_Recv 配合，用来处理：

• 不知道消息长度
• 不知道消息来自哪个进程
• 不知道消息的 tag
• 根据不同 tag 处理不同类型的消息

函数原型

int MPI_Probe(
    int source,
    int tag,
    MPI_Comm comm,
    MPI_Status *status
);

参数含义：

参数	含义
source	指定发送方 rank，也可以使用 MPI_ANY_SOURCE
tag	指定消息标签，也可以使用 MPI_ANY_TAG
comm	通信器，例如 MPI_COMM_WORLD
status	保存探测到的消息信息

基本用法

MPI_Status status;

MPI_Probe(
    MPI_ANY_SOURCE,
    MPI_ANY_TAG,
    MPI_COMM_WORLD,
    &status
);

执行成功后，可以查看消息来源和标签：

int source = status.MPI_SOURCE;
int tag    = status.MPI_TAG;

printf("消息来自 rank %d，tag = %d\n", source, tag);

此时消息仍然在 MPI 消息队列中，必须继续调用 MPI_Recv：

int value;

MPI_Recv(
    &value,
    1,
    MPI_INT,
    status.MPI_SOURCE,
    status.MPI_TAG,
    MPI_COMM_WORLD,
    MPI_STATUS_IGNORE
);

典型流程

MPI_Status status;

// 1. 等待并查看消息
MPI_Probe(
    MPI_ANY_SOURCE,
    MPI_ANY_TAG,
    MPI_COMM_WORLD,
    &status
);

// 2. 获取消息中的元素数量
int count;
MPI_Get_count(&status, MPI_INT, &count);

// 3. 分配接收空间
std::vector<int> data(count);

// 4. 真正接收消息
MPI_Recv(
    data.data(),
    count,
    MPI_INT,
    status.MPI_SOURCE,
    status.MPI_TAG,
    MPI_COMM_WORLD,
    MPI_STATUS_IGNORE
);

关键关系：

MPI_Probe：查看消息，但不取走
MPI_Get_count：获取消息包含多少个元素
MPI_Recv：真正取走消息

完整示例：接收未知长度数组

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

int main(int argc, char **argv)
{
    MPI_Init(&argc, &argv);

    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    if (rank == 0) {
        int data[] = {10, 20, 30, 40, 50};

        MPI_Send(
            data,
            5,
            MPI_INT,
            1,
            100,
            MPI_COMM_WORLD
        );
    }
    else if (rank == 1) {
        MPI_Status status;

        // 只探测，不接收
        MPI_Probe(
            0,
            100,
            MPI_COMM_WORLD,
            &status
        );

        int count = 0;
        MPI_Get_count(&status, MPI_INT, &count);

        std::vector<int> data(count);

        // 真正接收刚才探测到的消息
        MPI_Recv(
            data.data(),
            count,
            MPI_INT,
            status.MPI_SOURCE,
            status.MPI_TAG,
            MPI_COMM_WORLD,
            MPI_STATUS_IGNORE
        );

        printf(
            "收到来自 rank %d 的 %d 个整数：",
            status.MPI_SOURCE,
            count
        );

        for (int value : data) {
            printf("%d ", value);
        }

        printf("\n");
    }

    MPI_Finalize();
    return 0;
}

根据tag决定消息类型

例如定义不同消息：

enum MessageTag {
    TAG_TASK_REQUEST = 1,
    TAG_TASK_DATA    = 2,
    TAG_TASK_DONE    = 3,
    TAG_STOP         = 4
};

接收方先查看下一条消息：

MPI_Status status;

MPI_Probe(
    MPI_ANY_SOURCE,
    MPI_ANY_TAG,
    MPI_COMM_WORLD,
    &status
);

然后根据 tag 分别处理：

switch (status.MPI_TAG) {
case TAG_TASK_REQUEST: {
    int request;

    MPI_Recv(
        &request,
        1,
        MPI_INT,
        status.MPI_SOURCE,
        TAG_TASK_REQUEST,
        MPI_COMM_WORLD,
        MPI_STATUS_IGNORE
    );

    printf("rank %d 请求任务\n", status.MPI_SOURCE);
    break;
}

case TAG_TASK_DONE: {
    int taskId;

    MPI_Recv(
        &taskId,
        1,
        MPI_INT,
        status.MPI_SOURCE,
        TAG_TASK_DONE,
        MPI_COMM_WORLD,
        MPI_STATUS_IGNORE
    );

    printf(
        "rank %d 完成任务 %d\n",
        status.MPI_SOURCE,
        taskId
    );
    break;
}
}

这在"主进程调度多个 worker"的程序中很常见。

MPI_Probe会阻塞

如果没有匹配的消息：

MPI_Probe(...);

当前进程会一直等待。

如果不想阻塞，使用 MPI_Iprobe：

int flag = 0;
MPI_Status status;

MPI_Iprobe(
    MPI_ANY_SOURCE,
    MPI_ANY_TAG,
    MPI_COMM_WORLD,
    &flag,
    &status
);

if (flag) {
    // 有消息
    int value;

    MPI_Recv(
        &value,
        1,
        MPI_INT,
        status.MPI_SOURCE,
        status.MPI_TAG,
        MPI_COMM_WORLD,
        MPI_STATUS_IGNORE
    );
}
else {
    // 当前没有消息，可以继续做其他工作
}

函数	没有匹配消息时	是否取走消息
MPI_Probe	阻塞等待	否
MPI_Iprobe	立即返回，flag=0	否
MPI_Recv	阻塞等待	是

使用时需要注意

1. Probe以后一定要Recv

错误示例：

while (true) {
    MPI_Probe(
        MPI_ANY_SOURCE,
        MPI_ANY_TAG,
        MPI_COMM_WORLD,
        &status
    );

    // 没有MPI_Recv
}

因为消息没有被取走，下一次 MPI_Probe 可能仍然看到同一条消息。

正确写法：

MPI_Probe(..., &status);
MPI_Recv(..., status.MPI_SOURCE, status.MPI_TAG, ...);

2. Recv最好使用Probe得到的来源和tag

推荐：

MPI_Probe(
    MPI_ANY_SOURCE,
    MPI_ANY_TAG,
    comm,
    &status
);

MPI_Recv(
    buffer,
    count,
    datatype,
    status.MPI_SOURCE,
    status.MPI_TAG,
    comm,
    MPI_STATUS_IGNORE
);

这样可以确保接收的是刚刚探测到的那条消息。

3. MPI_Get_count的类型要匹配

MPI_Get_count(&status, MPI_INT, &count);

这里返回的是消息包含多少个 MPI_INT，不是字节数。

如果发送的是：

MPI_Send(data, 5, MPI_INT, ...);

那么：

MPI_Get_count(&status, MPI_INT, &count);

通常得到：

count = 5

一句话总结：

MPI_Probe 就像先查看快递单，知道是谁发的、是什么类型、有多大；MPI_Recv 才是真正把快递取走。