盒子
盒子
文章目录
  1. GPU
    1. DLPack
    2. Matrix coordinate (COO) assembly

PETSc[08] petsc4py中的GPU

2026.01.01 Zhuoyu Chen petsc  热度 ℃

GPU

PETSc 的 GPU 支持，需要注意：

  • Vec, Mat对象的 类型（VecType、MatType）决定了算子实现是否走 CUDA/HIP/Kokkos 等device后端，可以通过 -vec_type {cuda,hip,kokkos} 与 -mat_type {aijcusparse,aijhipsparse,aijkokkos} ，来让PETSc在GPU上创建计算对象
  • 数据在 host 与 device 之间需要通过通信来进行传递，而PETSc 需要维护两侧一致性，任何一次访问都可能触发隐式的数据迁移
  • 使用GPU计算，核心问题在于，如何让数据长期停留在 device，并减少 host/device 之间的数据迁移，以及减少host上的计算

一个常用的命令行形式如下：

1
2
3
4
mpirun -n 4 python test.py \
  -vec_type cuda \
  -mat_type aijcusparse \
  -log_view_gpu_time -log_view

在MINI TUTORIAL: PETSC ON THE GPU, JUNCHAO ZHANG中，有更详细的PETSc+CUDA的介绍。

在GPU上，直接–log_view会获得NaN，需要改用–log_view_gpu_time –log_view

DLPack

petsc4py 为 Vec 提供 toDLPack([mode])，用于导出一个 DLPack PyCapsule 来包装向量数据，或者直接使用createWithDLPack方法。矩阵侧也提供 Mat.toDLPack([mode])，用于导出矩阵数据。在DLPack文档中，有更详细的相关信息。

下面代码给出在GPU上创建数据方法：

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
import torch, torch.utils.dlpack
import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc

PETSC_OFFLOAD = {
    1: "PETSC_OFFLOAD_CPU",
    2: "PETSC_OFFLOAD_GPU",
    3: "PETSC_OFFLOAD_BOTH",
}

x = PETSc.Vec().create()
x.setSizes(1000); x.setType('cuda')
x.set(1.0)

A = PETSc.Mat().create()
A.setSizes(1000); A.setType('aijcusparse')
A.setDiagonal(x)

PETSc.Sys.Print(f"{x.getType()=}")
PETSc.Sys.Print(f"{A.getType()=}")
PETSc.Sys.Print(f"{PETSC_OFFLOAD[x.getOffloadMask()]=}")

y = torch.ones(1000, dtype=torch.float64, device="cuda")
y = torch.utils.dlpack.to_dlpack(y)
y = PETSc.Vec().createWithDLPack(y)
PETSc.Sys.Print(f"{y.getType()=}")
PETSc.Sys.Print(f"{PETSC_OFFLOAD[y.getOffloadMask()]=}")

输出为：

1
2
3
4
5
x.getType()='seqcuda'
A.getType()='seqaijcusparse'
PETSC_OFFLOAD[x.getOffloadMask()]='PETSC_OFFLOAD_BOTH'
y.getType()='seqcuda'
PETSC_OFFLOAD[y.getOffloadMask()]='PETSC_OFFLOAD_GPU'

如果要在GPU上进行PETSc的Vec进行零拷贝的访问，可以考虑下面的包装器：

1
2
3
4
5
6
7
8
9
10
import contextlib, torch
@contextlib.contextmanager
def gpu_access(vec, mode="rw"):
  # High-level wrapper for zero-copy GPU access via DLPack.
  dlpack = vec.toDLPack(mode=mode)
  tensor = torch.utils.dlpack.from_dlpack(dlpack)
  yield tensor
  if mode != "r": torch.cuda.synchronize()
with gpu_access(U, mode="r") as u:
  ...

Matrix coordinate (COO) assembly

在MINI TUTORIAL: PETSC ON THE GPU, JUNCHAO ZHANG中，有更详细的介绍。

在CPU上通过MatSetValues进行矩阵装配方法，并不是在GPU上使用Mat的最佳实践途径，而PETSc加入了矩阵坐标系统(matric coordinate, COO)，来提高装配效率，以实现下面三个步骤：

  • 用户给出非零元(nonzeros)的坐标(coordinate)
  • PETSc使用该数据，创建稀疏矩阵模式(sparsity pattern)，并指定插入/修改计划，以及GPU上的通信计划
  • 用户给出非零元(nonzeros)的值，并且需要保证顺序和之前给出的坐标顺序相同

在PETSc中，这个可以简化成两个语句：

  • MatSetPreallocationCOO(Mat A, PetscCount n, PetscInt i[], PetscInt j[])
  • MatSetValuesCOO(Mat A, const PetscScalar v[], InsertMode imode)

并且可以通过MatDuplicate(A,op,&B)进行矩阵结构的复制，以及MatSetValuesCOO(B, ...)给矩阵分配不同值。

下面是一个示例代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import sys, petsc4py
petsc4py.init(sys.argv)
from petsc4py import PETSc
import numpy as np

n = 4
A = PETSc.Mat().create()
A.setSizes(n)
A.setType('aijcusparse')
A.setUp()

rows = np.arange(n, dtype=np.int32)
cols = np.arange(n, dtype=np.int32)

I = np.repeat(rows, n).astype(np.int32)   # [0,0,0,0, 1,1,1,1, ...]
J = np.tile(cols, n).astype(np.int32)     # [0,1,2,3, 0,1,2,3, ...]
V = (100.0 * I + J).astype(np.float64)    # A[i,j] = 100*i + j

A.setPreallocationCOO(I, J)
A.setValuesCOO(V)

A.assemble()
A.view()

其输出为：

1
2
3
4
5
6
Mat Object: 1 MPI process
  type: seqaijcusparse
row 0: (0, 0.)  (1, 1.)  (2, 2.)  (3, 3.)
row 1: (0, 100.)  (1, 101.)  (2, 102.)  (3, 103.)
row 2: (0, 200.)  (1, 201.)  (2, 202.)  (3, 203.)
row 3: (0, 300.)  (1, 301.)  (2, 302.)  (3, 303.)