近期的工作需要熟悉scalapack的数据格式，做个简单记录。

采用了直接通过apt安装的方式：

1
sudo apt-get install libscalapack-openmpi-dev

简单编写的测试程序的cmake，注意到直接链接`$SCALAPACK_LIBRARIES会失败。这个包的.cmake没配置这个变量，也没有头文件和相应的头文件路径变量，链接的时候直接链接scalapack-openmpi就可以了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
cmake_minimum_required(VERSION 3.12)
project(ScalapackTest LANGUAGES C)

find_package(MPI REQUIRED)

find_package(ScaLAPACK REQUIRED)  
if(ScaLAPACK_FOUND)
    message(STATUS "ScaLAPACK found at: ${ScaLAPACK_DIR}")  # 配置文件路径
    message(STATUS "ScaLAPACK libraries: ${SCALAPACK_LIBRARIES}")
    message(STATUS "ScaLAPACK include dirs: ${ScaLAPACK_INCLUDE_DIRS}")
endif()

# 设置C++标准
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

add_executable(test_scalapack src/test_scalapack.c)

target_link_libraries(test_scalapack PRIVATE
    scalapack-openmpi
    MPI::MPI_C
)

target_include_directories(test_scalapack PRIVATE ${MPI_CXX_INCLUDE_DIRS})

简单的测试程序主要是为了理解scalapack的数据规格，调用简单选择了矩阵乘，其接口如下：

1
2
3
4
5
6
extern void pdgemm_(
    char *transa, char *transb, int *m, int *n, int *k,
    double *alpha, double *a, int *ia, int *ja, int *desca,
    double *b, int *ib, int *jb, int *descb,
    double *beta, double *c, int *ic, int *jc, int *descc
);

由于scalapack底层是fortran实现，所以首先要注意数据的起始索引是1。scalapack中，是否转置等是以char的方式传入，用N表示转置，T表示转置，C表示共轭转置。ia，ja等变量表示的是子矩阵在全局矩阵中的起始行列索引，索引从1开始。desc是数组，指向包含矩阵全局分布信息的描述符。

进程网格与矩阵描述符

scalapack通过内部的BLACS来管理进程网格和通信。通常情况下进入程序需要创建一个BLACS进程通信上下文，使用以下接口来创建：

1
extern void Cblacs_get(int *context, int *scope, int *ictxt);

第一个参数表示初始上下文标识符，-1表示创建新的上下文。第二个参数表示关联通信域，0表示默认的MPI_COMM_WORLD，最后一个参数用于返回新创建的上下文句柄。

创建上下文后，用以下接口创建进程网格，第二个参数用于控制进程分布的顺序，可以取Row和Col。

1
extern void Cblacs_gridinit(int *ictxt, char *order, int nprow, int npcol);

Cblacs_gridinfo接口用来获取当前进程在进程网格中的信息：

1
2
3
4
5
6
7
extern void Cblacs_gridinfo(
    int ictxt,        // 输入：BLACS上下文句柄
    int *nprow,       // 输出：进程网格行数
    int *npcol,       // 输出：进程网格列数
    int *myrow,       // 输出：当前进程行坐标
    int *mycol        // 输出：当前进程列坐标
);

上述三个接口可以创建和获取进程网格的信息，接下来就要设置矩阵的划分信息了，首先获取本地维度：

1
2
3
4
5
6
7
extern int numroc_(
    int *n,           // 全局维度（行或列）
    int *nb,          // 块大小（行或列）
    int *iproc,       // 当前进程在网格中的坐标（行或列）
    int *isrcproc,    // 起始进程坐标（通常为0）
    int *nprocs       // 进程网格中该方向的总进程数
);

有了本地维度之后，需要初始化矩阵描述符：

1
2
3
4
5
6
7
8
9
10
11
12
extern void descinit_(
    int *desc,        // 输出：9元素描述符数组
    int *m,           // 全局矩阵行数
    int *n,           // 全局矩阵列数
    int *mb,          // 行块大小
    int *nb,          // 列块大小
    int *irsrc,       // 起始进程的行坐标（通常为0）
    int *icsrc,       // 起始进程的列坐标（通常为0）
    int *ictxt,       // BLACS进程网格上下文
    int *lld,         // 本地矩阵的行维度（由numroc_计算）
    int *info         // 输出：错误状态码（0表示成功）
);

之后就可以用矩阵描述符传入计算接口来完成分布式计算了，scalapack的设计中，矩阵描述符和网格等信息和矩阵本身的存储都是分离的，所以要独立分配local_row*local_col的空间作为本地的矩阵数据存储空间，传入计算函数，需要注意大小和创建网格以及创建描述符时一致，并注意资源释放，要分别释放矩阵数据和进程网格，进程网格通过下面的接口释放：

1
Cblacs_gridexit(int *ictxt);

矩阵乘接口调用

以下是利用上面所述接口创建矩阵并调用scalapack矩阵乘的程序：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>

// ScaLAPACK函数声明（Fortran接口）
extern void pdgemm_(char *transa, char *transb, int *m, int *n, int *k,
                    double *alpha, double *a, int *ia, int *ja, int *desca,
                    double *b, int *ib, int *jb, int *descb,
                    double *beta, double *c, int *ic, int *jc, int *descc);

extern void descinit_(int *desc, int *m, int *n, int *mb, int *nb,
                     int *irsrc, int *icsrc, int *ictxt, int *lld, int *info);

extern int numroc_(int *n, int *nb, int *iproc, int *isrcproc, int *nprocs);

int main(int argc, char **argv) {
    MPI_Init(&argc, &argv);
    
    // 基本参数设置
    int ZERO = 0, ONE = 1;
    int N = 1024;          // 全局矩阵维度
    int BLOCK_SIZE = 64;   // 块划分大小
    int NPROW = 2, NPCOL = 2; // 2x2进程网格

    // MPI进程信息
    int myrank, nprocs;
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

    // BLACS进程网格初始化
    int ictxt, myrow, mycol;
    Cblacs_get(-1, 0, &ictxt);
    Cblacs_gridinit(&ictxt, "Row", NPROW, NPCOL);
    Cblacs_gridinfo(ictxt, &NPROW, &NPCOL, &myrow, &mycol);

    // 计算本地矩阵维度
    int local_rows = numroc_(&N, &BLOCK_SIZE, &myrow, &ZERO, &NPROW);
    int local_cols = numroc_(&N, &BLOCK_SIZE, &mycol, &ZERO, &NPCOL);
    local_rows = (local_rows > 0) ? local_rows : 1;
    local_cols = (local_cols > 0) ? local_cols : 1;
    printf("local_rows=%d, local_cols=%d\n", local_rows, local_cols);

    // 矩阵描述符初始化[6,7]
    int descA[9], descB[9], descC[9];
    int info;
    descinit_(descA, &N, &N, &BLOCK_SIZE, &BLOCK_SIZE, &ZERO, &ZERO, 
             &ictxt, &local_rows, &info);
    descinit_(descB, &N, &N, &BLOCK_SIZE, &BLOCK_SIZE, &ZERO, &ZERO,
             &ictxt, &local_rows, &info);
    descinit_(descC, &N, &N, &BLOCK_SIZE, &BLOCK_SIZE, &ZERO, &ZERO,
             &ictxt, &local_rows, &info);

    // 分配本地矩阵内存（列优先存储）
    double *A = (double*)malloc(local_rows * local_cols * sizeof(double));
    double *B = (double*)malloc(local_rows * local_cols * sizeof(double));
    double *C = (double*)calloc(local_rows * local_cols, sizeof(double));

    // 初始化本地矩阵
    for(int i=0; i<local_rows*local_cols; i++) {
        A[i] = (double)rand() / RAND_MAX;
        B[i] = (double)rand() / RAND_MAX;
    }

    // 执行矩阵乘法 C = A*B 
    char trans = 'N';  // 不转置
    double alpha = 1.0, beta = 0.0;
    pdgemm_(&trans, &trans, &N, &N, &N, &alpha,
           A, &ONE, &ONE, descA,
           B, &ONE, &ONE, descB,
           &beta, C, &ONE, &ONE, descC);

    // 资源释放
    free(A); free(B); free(C);
    Cblacs_gridexit(ictxt);
    MPI_Finalize();
    return 0;
}

2D循环块存储

scalapack存储矩阵的方式是2D块循环，前面的测试代码只是简单了解下scalapack怎么使用，并没有体现数据的存储，因为数据都是直接随机生成的。接下来用一个测试程序看看2D块循环。

首先可以通过文档中的例子理解2D块循环的存储方式，下图是一个5x5矩阵划分为2x2矩阵块并存放到2x2进程网格上的示意。每个进程持有的行/列起始索引的差值都是BLOCKSIZE*nprocs，所有的自己负责的块统一按照列有限的顺序存储，这里需要注意的是本地存储不是按照块存储的，而是按照列来存储的，例如图中的划分后的矩阵，0号进程所持有的数据是顺序的a11 a21 a51 a12 a22 a52 a15 a25 a55，而不是按照划分的2x2块存储。

image-20250420155004342

scalapack提供了一个矩阵重分发函数，我们使用这个函数来测试2D块循环：

1
2
3
4
5
6
7
8
extern void pdgemr2d_(
    int *m, int *n,              // 全局矩阵维度
    double *a, int *ia, int *ja, // 源矩阵描述（连续存储）
    int *desca,                  // 源矩阵描述符（描述连续布局）
    double *b, int *ib, int *jb, // 目标矩阵描述（块循环存储）
    int *descb,                  // 目标矩阵描述符（需预先初始化块循环参数）
    int *ictxt                   // BLACS上下文
);

这个函数会把原矩阵描述符描述的矩阵分发到新的描述符中的各个进程。所以可以先在单进程上准备一个完整的矩阵，然后进行分发，检查各个进程上获取的2D块循环的数据是否和上面的描述是一致的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>

extern void Cblacs_get(int, int, int*);
extern void Cblacs_gridinit(int*, const char*, int, int);
extern void Cblacs_gridinfo(int, int*, int*, int*, int*);
extern void Cblacs_gridexit(int);
extern void Cblacs_exit(int);
extern int numroc_(int*, int*, int*, int*, int*);
extern void descinit_(int*, int*, int*, int*, int*, int*, int*, int*, int*, int*);
extern void pdgemr2d_(int*, int*, double*, int*, int*, int*, double*, int*, int*, int*, int*);



int main(int argc, char **argv) {
    MPI_Init(&argc, &argv);
    
    // ============== 基础参数设置 ==============
    int N = 8;          // 全局矩阵维度
    int BLOCK_SIZE = 2;   // 块循环分块大小
    int NPROW = 2, NPCOL = 2; // 2x2进程网格
    int ZERO = 0, ONE = 1;

    // ============== MPI进程信息 ==============
    int myrank, nprocs;
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

    // ============== BLACS进程网格初始化 ==============
    int ictxt, myrow, mycol;
    Cblacs_get(-1, 0, &ictxt);
    Cblacs_gridinit(&ictxt, "Row", NPROW, NPCOL);
    Cblacs_gridinfo(ictxt, &NPROW, &NPCOL, &myrow, &mycol);

    // ============== 本地矩阵维度计算 ==============
    int local_rows = numroc_(&N, &BLOCK_SIZE, &myrow, &ZERO, &NPROW);
    int local_cols = numroc_(&N, &BLOCK_SIZE, &mycol, &ZERO, &NPCOL);
    local_rows = (local_rows > 0) ? local_rows : 1;
    local_cols = (local_cols > 0) ? local_cols : 1;

    // ============== 矩阵描述符初始化 ==============
    int descA[9], descB[9], info;
    int lda = N; 
    descinit_(descA, &N, &N, &N, &N, &ZERO, &ZERO, &ictxt, &lda, &info);
    descinit_(descB, &N, &N, &BLOCK_SIZE, &BLOCK_SIZE, 
             &ZERO, &ZERO, &ictxt, &local_rows, &info);

    // ============== 数据分配与初始化 ==============
    double *global_A = NULL, *distributed_B = NULL;
    if (myrank == 0) {
        global_A = (double*)malloc(N * N * sizeof(double));
        for (int i = 0; i < N*N; i++) {
            global_A[i] = i; 
        }
        printf("-----------------------------------------------------------\n");
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < N; j++) {
                printf("%lf ", global_A[j * N + i]);  // 按列优先打印
            }
            printf("\n");
        }
        printf("-----------------------------------------------------------\n");
    }
    distributed_B = (double*)calloc(local_rows * local_cols, sizeof(double));

    // ============== 调用pdgemr2d分发数据 ==============
    int ia = 1, ja = 1, ib = 1, jb = 1; // 基于1的索引
    pdgemr2d_(&N, &N, global_A, &ia, &ja, descA, 
             distributed_B, &ib, &jb, descB, &ictxt);

    // ============== 基础验证 ==============
    if (myrow == 0 && mycol == 0) {
        for(int i=0;i<local_cols*local_rows;i++) {
            printf("%lf\n", distributed_B[i]);
        }
    }

    // ============== 资源清理 ==============
    if (myrank == 0) free(global_A);
    free(distributed_B);
    Cblacs_gridexit(ictxt);
    MPI_Finalize();
    return 0;
}

输出如下，红框标出了进程0持有的块，下面按照顺序输出了进程0的矩阵存储，显然是先将所有数据拼到一块，然后再按列存储，而不是按块存储。

image-20250420160323232

这种2D块存储的方式是为了将块划分的足够小，这样对于对角矩阵等局部稠密的矩阵能够尽可能每个进程的负载均衡。要将其转为一般的稀疏矩阵格式也可以通过上述的分发接口实现，首先固定块的大小让每个进程只持有一个块的数据，然后再重分发，在进程内部调整为稀疏格式。但是对于不对齐或者无法均分的情况，还要额外做个填充。