PETSc4py多进程创建分布式稀疏块矩阵内存分配失败问题排查

我尝试用SciPy组装本地矩阵，再通过PETSc4py的createAIJWithArrays创建分布式稀疏块矩阵，单进程运行正常，但双进程运行时出现内存分配失败的错误。

原始代码 #

import numpy as np
import petsc4py
from petsc4py import PETSc
import scipy.sparse as sp

def main():
    # Initialize parallel PETSc4py
    petsc4py.init()
    
    comm = PETSc.COMM_WORLD
    
    comm_rank = comm.getRank()
    comm_size = comm.getSize()
    
    # Create local matrices
    n = 4
    n_local = n
    n_global = int(comm_size * n)
    
    A_local = sp.random(n_local, n_local, density = 0.5, format = 'csr') \
        + sp.identity(n)
    A_local = (A_local + A_local.transpose()) / 2.

    # Get CSRSpec for createAIJWithArrays
    [II, JJ, VV] = sp.find(A_local)
    [_, II] = np.unique(II, return_index = True)
    II = np.array(np.append(II, len(II) - 1), dtype = np.int32)
    
    # Create PETSc sparse matrix
    A = PETSc.Mat()
    A.createAIJWithArrays(size = [n_global, n_global], csr = [II, JJ, VV],
                          bsize = [n_local, n_local], comm = comm)
    
    # Communicate off-rank values and setup internal data structures for
    # performing parallel operations
    A.assemblyBegin()
    A.assemblyEnd()
    
    PETSc.Sys.Print(A)

    
if __name__ == '__main__':
    main()

错误信息 #

Traceback (most recent call last):
  File "/mnt/c/Users/jltor/Education and Research/misc/PETSc4py/matmpibaij/main.py", line 43, in <module>
    main()
  File "/mnt/c/Users/jltor/Education and Research/misc/PETSc4py/matmpibaij/main.py", line 31, in main
    A.createAIJWithArrays(size = [n_global, n_global], csr = [II, JJ, VV],
  File "petsc4py/PETSc/Mat.pyx", line 449, in petsc4py.PETSc.Mat.createAIJWithArrays
petsc4py.PETSc.Error: error code 55
[0] MatCreateMPIAIJWithArrays() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:4216
[0] MatMPIAIJSetPreallocationCSR() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:4013
[0] MatMPIAIJSetPreallocationCSR_MPIAIJ() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:3936
[0] MatMPIAIJSetPreallocation() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:4148
[0] MatMPIAIJSetPreallocation_MPIAIJ() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:2946
[0] MatSeqAIJSetPreallocation() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/seq/aij.c:3949
[0] MatSeqAIJSetPreallocation_SeqAIJ() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/seq/aij.c:4015
[0] PetscMallocA() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/sys/memory/mal.c:411
[0] PetscMallocAlign() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/sys/memory/mal.c:53
[0] Out of memory. Allocated: 0, Used by process: 80764928
[0] Memory requested 18446744073709551616
[1]PETSC ERROR: ------------------------------------------------------------------------
[1]PETSC ERROR: Caught signal number 11 SEGV: Segmentation Violation, probably memory access out of range
[1]PETSC ERROR: Try option -start_in_debugger or -on_error_attach_debugger
[1]PETSC ERROR: or see https://petsc.org/release/faq/#valgrind and https://petsc.org/release/faq/
[1]PETSC ERROR: configure using --with-debugging=yes, recompile, link, and run
[1]PETSC ERROR: to get more information on the crash.
[1]PETSC ERROR: Run with -malloc_debug to check if memory corruption is causing the crash.

问题原因分析 #

1. CSR行指针数组构造错误：原始代码手动构造的II数组不符合PETSc要求的CSR行指针格式。CSR行指针数组长度应为n_local + 1，其中每个元素表示对应行第一个非零元素在值数组中的偏移量，最后一个元素是总非零元素数。而原始代码通过np.unique和np.append构造的数组完全错误，导致PETSc计算内存需求时出现整数溢出，请求了超大内存（18446744073709551616是无符号64位整数最大值）。
2. 全局列索引未偏移：多进程下，每个进程的本地列索引需要转换为全局索引，否则不同进程的列索引会重叠，导致矩阵结构混乱。

修复后的代码 #

import numpy as np
import petsc4py
from petsc4py import PETSc
import scipy.sparse as sp

def main():
    # Initialize parallel PETSc4py
    petsc4py.init()
    
    comm = PETSc.COMM_WORLD
    
    comm_rank = comm.getRank()
    comm_size = comm.getSize()
    
    # Create local matrices
    n = 4
    n_local = n
    n_global = int(comm_size * n)
    
    # 构造本地对称稀疏矩阵
    A_local = sp.random(n_local, n_local, density=0.5, format='csr') + sp.identity(n)
    A_local = (A_local + A_local.transpose()) / 2.

    # 直接从SciPy CSR矩阵提取正确的行指针、列索引和值
    row_ptr = A_local.indptr.astype(np.int32)
    col_idx = A_local.indices.astype(np.int32)
    values = A_local.data.astype(np.float64)
    
    # 将本地列索引转换为全局列索引
    col_idx += comm_rank * n_local
    
    # Create PETSc sparse block matrix
    A = PETSc.Mat()
    A.createAIJWithArrays(
        size=[n_global, n_global],
        csr=[row_ptr, col_idx, values],
        bsize=[n_local, n_local],
        comm=comm
    )
    
    # 完成矩阵组装
    A.assemblyBegin()
    A.assemblyEnd()
    
    PETSc.Sys.Print(A)

if __name__ == '__main__':
    main()

修复要点说明 #

• 直接使用SciPy CSR矩阵的indptr、indices、data属性，避免手动构造错误的CSR格式数组。
• 对列索引加上comm_rank * n_local偏移量，确保全局索引的唯一性。
• 保持数组类型与PETSc要求一致（行指针和列索引用int32，值用float64）。

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