mpi4py 进阶之 scalapy

本文从本人简书博客同步过来

在上一篇中我们非常简要地介绍了 ScaLAPACK 软件。虽然 ScaLAPACK 在设计上作了很多工作使其方法接口与 LAPACK 尽量保持一致，但是直接使用 Fortran 或 C 语言按照上一篇中介绍的步骤使用 ScaLAPACK 仍然是一件比较麻烦和容易出错的事情，就好比我们使用 numpy.linalg 或 scipy.linalg （在底层调用 BLAS 和 LAPACK）中的相关函数比直接调用 BLAS 和 LAPACK 中的相关例程要容易和方便的多，我们也希望使用一个 Python 包装之后的 ScaLAPACK，下面我们就将介绍这样一个工具 scalapy。

为什么使用 scalapy

numpy.linalg 或 scipy.linalg 中提供了非常丰富且易用的线性代数计算函数，能够满足我们大多数的线性代数运算需求，为什么还要使用 ScaLAPACK 和 scalapy 呢？对规模不是很大的线性代数问题一般没必要使用 ScaLAPACK 和 scalapy，但是对比较大规模甚至很大规模的线性代数问题，比如超过 5000 × 5000 的矩阵运算，ScaLAPACK 和 scalapy 的优势就能显现了。对如此大规模的线性代数问题，你会发现 numpy.linalg 或 scipy.linalg 中的函数可能会计算的比较慢，因为 numpy.linalg 或 scipy.linalg 都不是分布式内存的计算工具，因此其只能利用单台机器的计算能力。同时，受制于单台机器可用内存的大小，所能解决问题的规模也会受到限制，不然就会超过内存的大小而导致程序崩溃。另外如果 numpy.linalg 或 scipy.linalg 使用的是 32 位的 LAPACK/BlAS 库，即使单台机器的内存充足，所能计算的问题的规模也会受到限制，比如矩阵元素的数目不能超过一个 32 位整型数所能表示的最大整数 2147483647，因此可计算的最大方块矩阵大小为 46340 × 46340，超过就会出错。对这些情况，ScaLAPACK 和 scalapy 就能帮上忙了，因为它们是分布式内存的计算工具，因此可以把计算分布到多台机器（如集群或超级计算机上），以充分利用多台机器的强大计算能力和更大的总内存空间来加快问题的求解或者解决某些在单台机器上无法完成的计算任务。

下载与安装 scalapy

软件依赖

• MPI 实现，目前支持 OpenMPI，IntelMPI 和 MPICH；
• ScaLAPACK 库， Netlib 和 Intel MKL 的 ScaLAPACK 都支持；
• Python 2.7， 3.2 或更新版本；
• numpy；
• mpi4py；
• Cython；
• f2py (可单独安装，或使用较新版本的 numpy 中包含的 f2py)。

安装

前往 https://github.com/jrs65/scalapy 下载 scalapy。可以直接下载 .zip 格式的压缩包，并解压，或者使用 git 将该软件 clone 到本地：

$ git clone https://github.com/jrs65/scalapy.git

进入软件顶层目录，使用以下命令进行安装：

$ python setup.py install [--user]

主要方法接口

scalapy 的核心成分 —— core 模块

core 模块提供了 1 个全局函数 initmpi 和 4 个类 ProcessContext，DistributedMatrix，ScalapyException 和 ScalapackException，下面分别进行介绍。

函数 initmpi

全局的 scalapy 初始化函数，一般会首先调用。

initmpi(gridshape, block_shape=[32, 32])

使用 scalapy 之前的初始化工作，会初始化一个全局的进程网格上下文 _context （为我们下面将要介绍的 ProcessContext 对象），和一个全局的块状循环分布的块大小 _block_shape。gridshape 是一个二元数组指定初始化的进程网格的大小为 gridshape[0] × gridshape[1]， block_shape 也是一个二元数组指定全局块状循环分布的块大小，一般设置成 [16, 16]，[32, 32] 或 [64, 64] 能获得比较高的计算性能，块的行和列也可以设置成不同大小。

注意：调用 initmpi 初始化的全局 ProcessContext 对象的通信子为 mpi4py.MPI.COMM_WORLD，如果需要使用其它的通信子，则必须手动使用（下面将会介绍的） ProcessContext 的构造函数来初始化一个 ProcessContext 对象。

类 ProcessContext

ProcessContext 对象存储有关 MPI/BLACS 进程的相关信息，这些信息会被（后面将会介绍的） DistributedMatrix 所使用。

方法

__init__(self, grid_shape, comm=None)

构造函数，以一个指定的进程网格大小 grid_shape 和一个 MPI 通信子 comm 创建一个 ProcessContext 对象。grid_shape 是一个二元数组指定初始化的进程网格的大小为 grid_shape[0] × grid_shape[1]， comm 如果为 None，则会使用 mpi4py.MPI.COMM_WORLD。

属性

grid_shape

所创建的进程网格的 shape，为一个二元 tuple。

grid_position

本进程在进程网格中的位置，为一个二元 tuple。

mpi_comm

此 ProcessContext 对象所使用的 MPI 通信子。

blacs_context

BLACS context 句柄。

all_grid_positions

一个 shape 为 (mpi_comm_size, 2) 的数组，其第 i 行是第 i 个进程在进程网格中的位置。

all_mpi_ranks

一个 shape 为 (grid_shape[0], grid_shape[1]) 的整数数组，all_mpi_ranks[i, j] 给出进程网格位置 (i, j) 上的进程的 rank。

类 DistributedMatrix

一个以块状循环的方式分布在多个 MPI 进程组成的一个进程网格上的分布式矩阵。块状循环分布矩阵是 ScaLAPACK 相关例程内部所使用的数据分布形式。

DistributedMatrix 分布在各个进程上的数据为基本的 numpy （2维）数组，下面是 DistributedMatrix 中使用的 ScaLAPACK 数据类型与 numpy 类型和 mpi4py 类型的对照表：

numpy 类型	MPI 类型	ScaLAPACK 类型	说明
np.float32	MPI.FLOAT	‘S’	单精度浮点数
np.float64	MPI.DOUBLE	‘D’	双精度浮点数