You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

如何使用MPI_IN_PLACE实现MPI_Allgatherv原地操作以节省内存

用MPI_IN_PLACE实现原地MPI_Allgatherv操作

要实现原地MPI_Allgatherv,核心是利用MPI_IN_PLACE让发送、接收缓冲区复用同一块内存,避免额外的大数组复制。你的代码已经满足原地操作的前提——每个进程的本地数据已经存放在最终数组的对应位置,只需要修改MPI_Allgatherv的调用参数即可。

修改步骤

  • 移除额外接收数组:删掉gathered_array的定义,直接用原array作为接收缓冲区。
  • 指定MPI_IN_PLACE为发送缓冲区:调用MPI_Allgatherv时,第一个参数传MPI_IN_PLACE,此时MPI会忽略发送计数、类型参数(可填0或原数值,不影响结果)。
  • 接收缓冲区复用原array:接收参数直接指向array,保持recvcounts和displs不变,这些参数定义了各进程数据在接收缓冲区中的存放位置。

修改后的完整代码

// mpiexec -np 3 ./a.out
#include <stdlib.h>
#include <stdio.h>
#include <mpi.h>

#define MAX_ELES 100
#define CUR_ELES 6

void print_arr(int arr[], int rank) {
    printf("RANK %d: ", rank);
    for (int i = 0; i < CUR_ELES; i++)
        printf("%d ", arr[i]);
    printf("\n"); 
}

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

    // 初始化数组为0
    int array[MAX_ELES];
    for (int position = 0; position < CUR_ELES; position++)
        array[position] = 0;

    // 计算Allgatherv的位移和数据范围
    int lower_bound, upper_bound;
    int N = CUR_ELES;
    int P = size;

    int *displs = malloc(size*sizeof(int));
    int *recvcounts = malloc(size*sizeof(int));
    int *lower_bounds = malloc(size*sizeof(int));
    int *upper_bounds = malloc(size*sizeof(int));

    for (int r = 0; r < size; r++) { 
        lower_bounds[r] = r*(N/P); 
        upper_bounds[r] = r == P-1 ? N : (r+1)*(N/P) ;
        displs[r] = lower_bounds[r];
        recvcounts[r] = upper_bounds[r] - lower_bounds[r];
    }

    lower_bound = lower_bounds[rank]; // 对应各进程的起始索引:{0, 2, 4}
    upper_bound = upper_bounds[rank]; // 对应各进程的结束索引:{2, 4, 6}

    // 初始化本地负责的数组段
    for (int b = lower_bound; b < upper_bound; b++)
        array[b] = rank+1;

    // 打印初始化后的数组
    print_arr(array, rank);

    // 原地Allgatherv操作:复用array作为接收缓冲区
    MPI_Allgatherv(
        MPI_IN_PLACE, 0, MPI_INT,  // 发送缓冲区设为MPI_IN_PLACE,后续参数被忽略
        array, recvcounts, displs, MPI_INT, comm);

    // 打印聚合结果(每个进程的array都已包含完整数据)
    if (rank == 0) {
        printf("---------\n");
        printf("RANK %d: Gathered\n", rank);
        print_arr(array, rank);
        printf("---------\n");
    }

    // 释放动态分配的内存
    free(displs);
    free(recvcounts);
    free(lower_bounds);
    free(upper_bounds);

    MPI_Finalize();
}

关键说明

  • 原地操作的核心前提是每个进程的本地数据已经位于接收缓冲区的正确位移位置,你的代码中各进程初始化时就把数据写到了array的对应段,正好匹配displs和recvcounts的定义,MPI聚合时不会覆盖未发送的本地数据。
  • 使用MPI_IN_PLACE时,发送端的计数、类型参数会被MPI忽略,填0是为了明确表示这些参数无效。
  • 聚合完成后,每个进程的array都会包含所有进程的完整数据,和原代码中gathered_array的内容完全一致,但无需额外占用内存。

内容的提问来源于stack exchange,提问作者Jared

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.08.04 20:50:29