PETSc中MPI并行矩阵仅按行/列拆分的实现问题

在使用PETSc进行矩阵乘法时，默认的mpiaij矩阵格式会同时对行和列进行分块，导致增加进程数后局部矩阵的行列维度都按比例缩小。要实现仅按行或仅按列拆分，需修改矩阵初始化和元素赋值的逻辑。

原代码 #

static char help[] = "Matrix multiplication program.\n"; 
#include <petsc.h> 
#define ROOT 0
int main(int argc,char *argv[]) {
    Mat C,CCt;
    PetscMPIInt numNode, myNode;
    PetscInt i, j, iStart, iEnd, jStart, jEnd,  m = 8, n = 8;
    PetscScalar v;
    PetscRandom rnd;

    PetscCall(PetscInitialize(&argc,&argv,NULL,help));
    PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&numNode));
    PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&myNode));
    
    //random gen
    PetscCall(PetscRandomCreate(PETSC_COMM_WORLD,&rnd));
    PetscCall(PetscRandomSetSeed(rnd,12));

    //ininitalize Matrix
    PetscCall(MatCreate(PETSC_COMM_WORLD, &C));
    PetscCall(MatSetSizes(C, PETSC_DECIDE, PETSC_DECIDE, m, n));
    PetscCall(MatSetFromOptions(C));
    PetscCall(MatSetUp(C));
    
    PetscCall(MatGetOwnershipRange(C, &iStart, &iEnd)); //return the range of row within local
    PetscCall(MatGetOwnershipRangeColumn(C,&jStart, &jEnd));
    PetscCall(PetscPrintf(PETSC_COMM_SELF,"Hello from %d local j range is %d to %d\n",myNode,jStart,jEnd));
    PetscCall(PetscPrintf(PETSC_COMM_SELF,"Hello from %d local i range is %d to %d\n",myNode,iStart,iEnd));
    
    //set values to matrix
    for (i = iStart; i < iEnd; i++) { //loop through all rows
        for (j = jStart; j < jEnd; j++){
            //PetscCall(PetscRandomGetValue(rnd,&v));
            v = 1;
            PetscCall(MatSetValues(C, 1, &i, 1, &j, &v, INSERT_VALUES));
        } 
    }
    
    PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
    
    //matrix mult
    PetscCall(MatMatTransposeMult(C,C,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&CCt));

    //view matrix
    PetscCall(MatView(C,PETSC_VIEWER_STDOUT_WORLD));
    PetscCall(MatView(CCt,PETSC_VIEWER_STDOUT_WORLD));
    
    // Clean up
    PetscCall(MatDestroy(&C));
    PetscCall(MatDestroy(&CCt));
    PetscCall(PetscFinalize());
    return 0; 
}

问题现象 #

• 单进程运行mpiexec -n 1 ./program时，矩阵为8x8全1矩阵，乘法结果为8x8全8矩阵，符合预期。
• 4进程运行mpiexec -n 4 ./program时，每个进程仅持有2行2列的局部矩阵，乘法结果为8x8矩阵，每个元素仅为2，不符合预期。

解决方案 #

1. 仅按行拆分（每个进程持有完整列，仅行分片） #

核心是指定局部列数等于全局列数，让每个进程拥有矩阵的全部列，仅行按进程数拆分。

修改点：

• 矩阵初始化时，将局部列数设为全局列数n，行数用PETSC_DECIDE自动分配：

  PetscCall(MatSetSizes(C, PETSC_DECIDE, n, m, n));
  

• 元素赋值时，遍历所有列（而非局部列范围jStart到jEnd），确保每个进程的局部行包含所有列的元素：

  for (i = iStart; i < iEnd; i++) {
      for (j = 0; j < n; j++){
          v = 1;
          PetscCall(MatSetValues(C, 1, &i, 1, &j, &v, INSERT_VALUES));
      } 
  }
  

修改后完整代码：

static char help[] = "Matrix multiplication program.\n"; 
#include <petsc.h> 
#define ROOT 0
int main(int argc,char *argv[]) {
    Mat C,CCt;
    PetscMPIInt numNode, myNode;
    PetscInt i, j, iStart, iEnd, m = 8, n = 8;
    PetscScalar v;
    PetscRandom rnd;

    PetscCall(PetscInitialize(&argc,&argv,NULL,help));
    PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&numNode));
    PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&myNode));
    
    PetscCall(PetscRandomCreate(PETSC_COMM_WORLD,&rnd));
    PetscCall(PetscRandomSetSeed(rnd,12));

    // 仅按行拆分：局部列数=全局列数n
    PetscCall(MatCreate(PETSC_COMM_WORLD, &C));
    PetscCall(MatSetSizes(C, PETSC_DECIDE, n, m, n));
    PetscCall(MatSetFromOptions(C));
    PetscCall(MatSetUp(C));
    
    PetscCall(MatGetOwnershipRange(C, &iStart, &iEnd));
    PetscCall(PetscPrintf(PETSC_COMM_SELF,"进程%d的本地行范围：%d 到 %d\n",myNode,iStart,iEnd));
    
    // 遍历所有列赋值
    for (i = iStart; i < iEnd; i++) {
        for (j = 0; j < n; j++){
            v = 1;
            PetscCall(MatSetValues(C, 1, &i, 1, &j, &v, INSERT_VALUES));
        } 
    }
    
    PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
    
    PetscCall(MatMatTransposeMult(C,C,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&CCt));

    PetscCall(MatView(C,PETSC_VIEWER_STDOUT_WORLD));
    PetscCall(MatView(CCt,PETSC_VIEWER_STDOUT_WORLD));
    
    // 清理资源（补充销毁随机数生成器）
    PetscCall(MatDestroy(&C));
    PetscCall(MatDestroy(&CCt));
    PetscCall(PetscRandomDestroy(&rnd));
    PetscCall(PetscFinalize());
    return 0; 
}

运行效果：

4进程运行时，每个进程持有2行完整的8列，矩阵乘法结果为8x8全8矩阵，与单进程结果一致。

2. 仅按列拆分（每个进程持有完整行，仅列分片） #

核心是指定局部行数等于全局行数，让每个进程拥有矩阵的全部行，仅列按进程数拆分。

修改点：

• 矩阵初始化时，将局部行数设为全局行数m，列数用PETSC_DECIDE自动分配：

  PetscCall(MatSetSizes(C, m, PETSC_DECIDE, m, n));
  

• 元素赋值时，遍历所有行（而非局部行范围iStart到iEnd），确保每个进程的局部列包含所有行的元素：

  for (i = 0; i < m; i++) {
      for (j = jStart; j < jEnd; j++){
          v = 1;
          PetscCall(MatSetValues(C, 1, &i, 1, &j, &v, INSERT_VALUES));
      } 
  }
  

修改后完整代码：

static char help[] = "Matrix multiplication program.\n"; 
#include <petsc.h> 
#define ROOT 0
int main(int argc,char *argv[]) {
    Mat C,CCt;
    PetscMPIInt numNode, myNode;
    PetscInt i, j, jStart, jEnd, m = 8, n = 8;
    PetscScalar v;
    PetscRandom rnd;

    PetscCall(PetscInitialize(&argc,&argv,NULL,help));
    PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD,&numNode));
    PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD,&myNode));
    
    PetscCall(PetscRandomCreate(PETSC_COMM_WORLD,&rnd));
    PetscCall(PetscRandomSetSeed(rnd,12));

    // 仅按列拆分：局部行数=全局行数m
    PetscCall(MatCreate(PETSC_COMM_WORLD, &C));
    PetscCall(MatSetSizes(C, m, PETSC_DECIDE, m, n));
    PetscCall(MatSetFromOptions(C));
    PetscCall(MatSetUp(C));
    
    PetscCall(MatGetOwnershipRangeColumn(C,&jStart, &jEnd));
    PetscCall(PetscPrintf(PETSC_COMM_SELF,"进程%d的本地列范围：%d 到 %d\n",myNode,jStart,jEnd));
    
    // 遍历所有行赋值
    for (i = 0; i < m; i++) {
        for (j = jStart; j < jEnd; j++){
            v = 1;
            PetscCall(MatSetValues(C, 1, &i, 1, &j, &v, INSERT_VALUES));
        } 
    }
    
    PetscCall(MatAssemblyBegin(C, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(C, MAT_FINAL_ASSEMBLY));
    
    PetscCall(MatMatTransposeMult(C,C,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&CCt));

    PetscCall(MatView(C,PETSC_VIEWER_STDOUT_WORLD));
    PetscCall(MatView(CCt,PETSC_VIEWER_STDOUT_WORLD));
    
    // 清理资源
    PetscCall(MatDestroy(&C));
    PetscCall(MatDestroy(&CCt));
    PetscCall(PetscRandomDestroy(&rnd));
    PetscCall(PetscFinalize());
    return 0; 
}

运行效果：

4进程运行时，每个进程持有8行完整的2列，矩阵乘法结果同样为8x8全8矩阵，与单进程结果一致。

内容的提问来源于stack exchange，提问作者HLTran