PETSc[03] petsc4py中的Mat

Mat

一个 Mat 的使用，通常要经历以下阶段：

create
setSizes
setType / setFromOptions
preallocate
setValues
assemble
compute / reuse / view / destroy

矩阵Mat的创建是类似于Vec的：

A = PETSc.Mat().create()
A.setSizes(n)                    # create a matrix of size n*n
# A.setSizes((n, n))             # create a matrix of size n*n
# A.setSizes((n, m))             # create a matrix of size n*m
# A.setSizes(((nl, n), (ml, m))) # create a matrix of size n*m with local sizes nl*ml
A.setFromOptions()
A.setUp()

也可以在创建矩阵的时候，指定矩阵的存储格式，具体见Mat.pyx文档，主要的函数有：

createAIJ
createBAIJ
createSBAIJ
createAIJCRL
createAIJWithArrays
createDense
createDenseCUDA
createDiagonal
...

在MatType文档中，给出了详细的-mat_type类型。一般来说，使用-mat_type aij可以设置为稀疏矩阵(CSR)，用-mat_type dense可以设置为稠密矩阵。在CUDA上，则分别使用 -mat_type aijcusparse 和-mat_type densecuda进行设置。

需要注意的是，Mat的内存分配是按照行进行分配的，每个 rank 只拥有一段连续的全局行索引区间，可以通过：

rstart, rend = A.getOwnershipRange()

来查询本 rank 负责的行 [rstart, rend)。

对于稀疏矩阵来说，最好是通过预分配空间，来减少内存分配的时间，比如：

A = PETSc.Mat().createAIJ(
    size=((n_local, n_global), (m_local, m_global)),
    nnz=5,
    comm=comm)

表示每一行最多有 5 个非零元素(nnz)，当然这个指的是d_nz=5, o_nz=5。

或者采用nnz=(d_nz, o_nz)来将对角块(Diagonal Block)、非对角块(Off-diagonal Block)处的非零元素数量进行分配，类似于：
$$
\begin{array}{c c}
\left[
\begin{array}{c|c c|c c}
& \text{col 0} & \text{col 1} & \text{col 2} & \text{col 3} \\
\hline
\text{proc 0, row 0} & \boxed{A_{\mathrm{diag}}} & \boxed{A_{\mathrm{diag}}} & A_{\mathrm{off-diag}} & A_{\mathrm{off-diag}} \\
\text{proc 0, row 1} & \boxed{A_{\mathrm{diag}}} & \boxed{A_{\mathrm{diag}}} & A_{\mathrm{off-diag}} & A_{\mathrm{off-diag}} \\
\hline
\text{proc 1, row 2} & A_{\mathrm{off-diag}} & A_{\mathrm{off-diag}} & \boxed{A_{\mathrm{diag}}} & \boxed{A_{\mathrm{diag}}} \\
\text{proc 1, row 3} & A_{\mathrm{off-diag}} & A_{\mathrm{off-diag}} & \boxed{A_{\mathrm{diag}}} & \boxed{A_{\mathrm{diag}}} \\
\end{array}
\right]
&
\begin{array}{l}
\leftarrow \text{proc 0}\\[2.2em]
\leftarrow \text{proc 1}
\end{array}
\end{array}
$$
考虑一个$10 \times 10$的Mat对象，被分配到了四个处理器上，并且设置了nnz=5，那么平均分配的结果是：

• proc 0占有3行，所以diag大小为$3 \times 3$，off-diag大小为$3 \times 7$

  • 总nnz数量为$3 \times (\min(3,5)+\min(7,5))=24$

• proc 1占有3行，所以diag大小为$3 \times 3$，off-diag大小为$3 \times 7$

  • 总nnz数量为$3 \times (\min(3,5)+\min(7,5))=24$

• proc 2占有2行，所以diag大小为$2 \times 2$，off-diag大小为$2 \times 8$

  • 总nnz数量为$2 \times (\min(2,5)+\min(8,5))=14$

• proc 3占有2行，所以diag大小为$2 \times 2$，off-diag大小为$2 \times 8$

  • 总nnz数量为$2 \times (\min(2,5)+\min(8,5))=14$

因此整体的nnz数量为$24+24+14+14=76$，可以用下面代码进行验证：

A = PETSc.Mat().createAIJ(size=10, nnz=5, comm=comm)
A.assemble()
PETSc.Sys.Print(A.getInfo()["nz_allocated"])
# Run:    mpirun -np 4 python test.py
# Output: 76.0

区分 diag / off-diag 是很有必要的，在SpMV中，列对应的向量元素是不是local(diag)的，决定了要不要通信。

比如说，对于计算下面的这个矩阵乘法：
$$
\begin{pmatrix}
a_{11} & a_{12} \quad (\text{proc 0}) \\
a_{21} & a_{22} \quad (\text{proc 1})
\end{pmatrix}
\begin{pmatrix}
x_1 \quad(\text{proc 0}) \\
x_2 \quad(\text{proc 1})
\end{pmatrix}=
\begin{pmatrix}
\boxed{a_{11} x_1}\quad (\text{same proc})+\boxed{a_{12} x_2}\quad (\text{diff proc}) \\
\boxed{a_{21} x_1}\quad (\text{diff proc})+\boxed{a_{22} x_2}\quad (\text{same proc})
\end{pmatrix}
$$

可以看到，在同一个处理器上的部分，可以快速计算，也就是$a_{11}x_{1}, a_{22}x_{2}$都是在处理器内部直接计算的，不需要与其他处理器进行通信。然而，在proc 0上计算$a_{12}\cdot x_{2}$，或者在 proc 1上计算$a_{21}\cdot x_{1}$，都需要从另一个处理器上获取对应的Vec的值，因此PETSc可以根据区分diag/off-diag，来预分配需要通信获取的值：

1. 先根据 off-diag 的列索引，一次性通信所需的 $x_j$
2. 把这些远程 $x_j$ 放到一个连续的本地缓冲区 x_off
3. 计算 y_diag += A_diag * x_diag
4. 计算 y_diag += A_off * x_off

从而，可以获得 diag上的Vec的值。

在Mat中进行赋值操作，是类似于Vec的，注意要使用assemble()来保证矩阵成功组装：

A.zeroEntries() # clear any preallocated values
A.setValue(i, j, value)
A.setValues(rows, cols, values)
A.assemble()    # equal to assemblyBegin+assemblyEnd

或者通过getOwnershipRange()获取所在行，进行赋值：

rstart, rend = A.getOwnershipRange()
rows = np.arange(rstart, rend, dtype=np.int32)
cols = np.array([1, 8], dtype=np.int32)
dats = np.ones((rend-rstart, cols.shape[0]), dtype=float)
A.setValues(rows, cols, dats)
A.assemble()

注意，更新完之后，有A[rows[i], cols[j]] = dats[i, j]。如果需要逻辑列的范围，可以用A.getOwnershipRangeColumn()获取，此时$[rstart, rend) \times [cstart, cend)$构成了diag部分的元素，其余的$[rstart, rend)$行构成了off-diag部分的元素。

当然，也可以直接通过A[i, j]来进行矩阵的坐标访问，这个通过mat_getitem和mat_setitem实现。更一般的，如果有切片$r,c$，同时有一个数据集$B$，其大小为$n_r \times n_c$，则可以通过A[r, c] = B来对相应的元素进行赋值，比如：

A[1:3, 1:3] = np.ones((2, 2))

这里我们再考虑size=((n_local, n_global), (m_local, m_global))这个语句的真实含义，其实是：

• 矩阵的总大小为 n_global*m_global
• 矩阵分配到若干个处理器上，每个处理器管理连续的n_local行
• 每个处理器的diag部分大小为m_local

实际上，可以通过下面的函数，创建与矩阵大小匹配的Vec对象，从而进行合理分配处理器的矩阵向量乘：

createVecs
createVecRight
createVecLeft