3. 性能分析 Profiling#

Overview

Profiling Firedrake programs with PETSc -log_view, and generating flame graphs from the log data.

Firedrake 程序的时间往往不是花在自己写的那几行代码里, 而是花在网格生成、 代码生成与编译、矩阵组装和线性方程组求解这些底层步骤上. Firedrake 建立在 PETSc 之上, 而 PETSc 自带一套性能日志: 它会记录每个内部步骤的调用次数、耗时和浮点运算数, 我们也可以用 PETSc.Log.Event 把自己关心的代码块标注成事件, 一并统计. 本章介绍如何打开这套日志、如何阅读它的输出, 并给出一个可以直接运行的算例.

参考资料:

  1. Optimising Firedrake performance: Firedrake 官方的性能优化指南, 讲火焰图的生成与解读、自定义事件的写法以及几个常见 的性能问题, 是本章的主要依据.

  2. PETSc: Profiling: PETSc 手册的性能分析 一章. -log_view 的各种输出格式、输出表格每一列的含义都在这里.

  3. PetscInitialize: PETSc 初始化阶段识别的命令行选项一览, 各个 -log_* 选项都列在其中.

  4. PetscLogView: 把日志写到指定 viewer 的接口. 本章在 Notebook 中打印完整表格用的 PETSc.Log.view 就是它的 Python 版本.

3.1. log_view#

-log_view 是 PETSc 的命令行选项, 加在脚本后面即可, 不需要改动代码:

python3 your_script.py -log_view

程序结束时 PETSc 会把整个运行过程的性能摘要打印到标准输出. 冒号后面还可以指定输出 文件和格式, 常用的几种是:

选项

作用

-log_view

把性能摘要打印到标准输出, 最常用

-log_view :profile.txt

同上, 但写入文件 profile.txt

-log_view :profile.txt:ascii_flamegraph

输出火焰图格式, 见下面的”火焰图”一节

-log_view :profile.xml:ascii_xml

输出为按调用关系嵌套的 XML

-log_view_memory

在事件表中额外给出内存分配的列

此外还有 -log_all (每个进程各写一个 Log.rank 文件)、-log_trace (实时打印每一步 在做什么) 等选项, 完整列表见 PETSc 手册的 Profiling 一章.

3.1.1. 读懂 -log_view 的输出#

-log_view 的输出分三段.

第一段是全局摘要, 给出整个运行过程的墙钟时间 Time (sec)、浮点运算总数 Flops、 MPI 消息条数和全局归约次数等. 每行有 MaxMax/MinAvgTotal 四列, 分别是 各进程上的最大值、最大值与最小值之比、平均值和总和. 其中 Max/Min 接近 1 说明各 进程负载均衡, 这个比值越大负载越不均衡; 串行运行时非零量的比值恒为 1 (某一项本身 为 0 时 PETSc 直接打成 0.0), 没有参考价值.

第二段 Summary of Stages 按阶段 (stage) 汇总. 不特意划分阶段时, 所有内容都记在 Main Stage 里.

第三段是事件表, 每个事件占一行, 既包括 KSPSolveMatMult 这样的 PETSc 内部事件, 也包括自己用 PETSc.Log.Event 标注的事件. 主要几列的含义是:

  • Count: 该事件被调用的次数, 和 Time/Flop 一样是最大值加最大最小比值;

  • Time (sec): 各进程耗时的最大值, 后面跟着最大值与最小值之比;

  • Flop: 浮点运算数, 同样是最大值加比值;

  • MessAvgLenReduct: MPI 消息条数、平均消息长度 (字节) 和全局归约次数;

  • %T %F %M %L %R: 该事件占总时间、总浮点运算数、总消息条数、总消息长度、总归约 次数的百分比. --- Global --- 是相对整个程序, --- Stage ---- 是相对所在阶段;

  • Mflop/s: 各进程浮点运算数之和除以最大耗时.

排查性能问题时, 通常先按 %T 找出最费时间的几项, 再看它们的 Mess/Reduct 和 最大最小比值, 判断瓶颈是计算本身还是通信与负载不均衡. 并行情形下各列的详细解释见 Interpreting -log_view Output: Parallel Performance.

3.1.2. 火焰图#

运行脚本时把选项换成 -log_view :profile.txt:ascii_flamegraph, PETSc 会把日志写成 火焰图格式 (命令里的 py/profile_poisson.py 是本章最后一节”写成脚本运行”给出的 脚本, 这里先看命令的形式):

python3 py/profile_poisson.py -log_view :profile.txt:ascii_flamegraph

生成的 profile.txt 每行是一条调用栈和它的耗时 (单位是微秒), 形如

Main Stage;firedrake;MyAssemble;firedrake.assemble.assemble;ParLoopExecute 288

把这个文件上传到在线工具 speedscope 就能看到火焰图; 也可以用 FlameGraph 提供的 flamegraph.pl 生成 SVG, 注意加上 --countname us, 否则鼠标悬停时显示的单位会写成 samples 而不是微秒. 火焰图比表格多出来的信息是事件之间的嵌套关系: Firedrake 自身的 不少函数也带着事件装饰器, 因此能一眼看出时间是花在组装、代码编译还是求解上.

3.1.3. 自定义事件#

-log_view 默认只统计 PETSc 与 Firedrake 内部的事件. 用 PETSc.Log.EventPETSc.Log.EventDecorator 可以把自己关心的代码块也标注成事件, 它的名字和运行时间 会一并出现在上面的表格和火焰图中.

  1. PETSc.Log.Event 用作上下文管理器, 适合标注一段代码:

    from firedrake.petsc import PETSc
    
    with PETSc.Log.Event("foo"):
        do_something_expensive()
    
  2. PETSc.Log.EventDecorator 用作装饰器, 适合标注整个函数:

    from firedrake.petsc import PETSc
    
    @PETSc.Log.EventDecorator("foo")
    def do_something_expensive():
        ...
    

    不传名字时, 事件名取”模块名.函数名”的形式. 事件表里的 firedrake.assemble.assemble 就是这么来的.

事件名会原样出现在事件表的第一列, 而这一列的宽度是 16 个字符, 名字一旦超过 16 个 字符, 这一行后面的所有列都会整体右移, 和表头对不齐, 因此事件名要尽量取短.

3.2. 一个可运行的算例#

下面把上面的内容串起来: 对一个小规模的 Poisson 问题, 给网格生成、矩阵组装和线性 求解分别加上自定义事件, 再把各事件的耗时读出来.

Notebook 里没有命令行, 加不了 -log_view, 所以要用 PETSc.Log.begin() 手动打开 日志, 然后用事件对象的 getPerfInfo()PETSc.Log.view 取结果. 需要火焰图或者 其他命令行选项时, 仍然要把代码写成脚本运行, 见本节末尾.

import petsctools
from firedrake.petsc import PETSc

# 命令行传了 -log_view 时 PETSc 已经自动开启日志, 此时不要再调用 Log.begin():
# 那会额外启动一个扁平日志处理器, 用 ascii_xml / ascii_flamegraph 这类嵌套格式时,
# 程序退出会报 "Trying to end paused event, not allowed".
if 'log_view' not in petsctools.get_commandline_options():
    PETSc.Log.begin()

t_start = PETSc.Log.getTime()   # 记下起点, 后面用来算耗时占比

print('logging active:', PETSc.Log.isActive())
logging active: True

petsctools.get_commandline_options() 返回 PETSc 从命令行认出的选项名. Notebook 的 启动参数里没有 log_view, 于是 PETSc.Log.begin() 会被调用; 把同一段代码放进脚本并带上 -log_view 运行时, 这次调用会被跳过. 这个写法取自 Firedrake 文档的 Optimising 一节.

之所以要加这个判断: 传了 -log_view 时 PETSc 已经自己把日志开好了 (普通格式开的是 扁平日志, :...:ascii_xml:...:ascii_flamegraph 开的是嵌套日志), 这时再调用 一次 PETSc.Log.begin() 会额外启动一个扁平日志处理器; 用后两种嵌套格式时, 程序退出 时会报 Trying to end paused event, not allowed.

打开日志之后就可以正常写 Firedrake 代码了. 注意 PETSc.Log.begin() 必须放在被计时 的代码之前, 在它之前发生的事情 (例如 import firedrake 本身) 不会被记录.

from firedrake import *

N = 64
params = {'ksp_type': 'cg', 'pc_type': 'jacobi', 'ksp_rtol': 1e-10}

with PETSc.Log.Event('MyMesh'):
    mesh = UnitSquareMesh(N, N)
    V = FunctionSpace(mesh, 'CG', 1)

x, y = SpatialCoordinate(mesh)
u, v = TrialFunction(V), TestFunction(V)
a = inner(grad(u), grad(v))*dx
L = inner(sin(pi*x)*sin(pi*y), v)*dx
bc = DirichletBC(V, 0, 'on_boundary')
u_h = Function(V, name='u_h')

with PETSc.Log.Event('MyAssemble'):
    A = assemble(a, bcs=bc)
    b = assemble(L, bcs=bc)

with PETSc.Log.Event('MySolve'):
    solve(A, u_h, b, solver_parameters=params)

print('number of dofs:', V.dim())
number of dofs: 4225

这里把组装和求解拆成了两步 (先 assemble 得到矩阵 A 和右端项 b, 再解线性方程 组), 这样两个事件的耗时才分得开; 直接写 solve(a == L, u_h, bcs=bc) 时, 组装是在 求解内部完成的, 事件表里就只剩一个数.

PETSc.Log.Event(name) 用同一个名字再取一次, 拿到的还是同一个事件, 因此可以在事后 用 getPerfInfo() 读出它的统计量: 一个字典, 含调用次数 count、累计耗时 time (秒)、浮点运算数 flops, 以及并行时才非零的 numMessagesmessageLengthnumReductions.

def show_events(names):
    """打印各事件的调用次数、累计耗时、耗时占比和浮点运算数."""
    elapsed = PETSc.Log.getTime() - t_start
    print('event'.ljust(14) + 'count'.rjust(6) + 'time (s)'.rjust(10)
          + '%T'.rjust(7) + 'flop'.rjust(12))
    for name in names:
        info = PETSc.Log.Event(name).getPerfInfo()
        count, t, flop = info['count'], info['time'], info['flops']
        print(f'{name:<14}{count:>6}{t:>10.4f}{100*t/elapsed:>6.1f}%{flop:>12.3e}')


show_events(['MyMesh', 'MyAssemble', 'MySolve'])
event          count  time (s)     %T        flop
MyMesh             1    0.2395  77.0%   0.000e+00
MyAssemble         1    0.0352  11.3%   3.635e+06
MySolve            1    0.0233   7.5%   1.125e+07

%T 一列是各事件耗时占 t_start 以来墙钟时间的比例. 它只是对 -log_view 同名列的 粗略模仿: 分母里还包含 Notebook 单元之间的空闲时间, 交互式运行时等得越久这个比例越小, 因此看它的相对大小即可, 不必当作精确值.

具体的秒数取决于机器, 也取决于 Firedrake 的磁盘缓存是否预热, 但可以看出: 三个事件 的耗时都远大于真正的线性代数运算 (后面完整表格里 KSPSolve 只有几毫秒). 原因是 Firedrake 会把变分形式即时编译成 C 代码 (JIT), 第一次组装和第一次求解都包含了代码 生成、编译、稀疏结构构建和求解器 setup 这些一次性开销. 下面把同样的组装和求解再做 一遍就能看出差别.

with PETSc.Log.Event('MyAssemble2'):
    A = assemble(a, bcs=bc)
    b = assemble(L, bcs=bc)

u_h.assign(0)      # 让两次求解从同一个初值出发, 迭代步数才可比
with PETSc.Log.Event('MySolve2'):
    solve(A, u_h, b, solver_parameters=params)

show_events(['MyAssemble', 'MyAssemble2', 'MySolve', 'MySolve2'])
event          count  time (s)     %T        flop
MyAssemble         1    0.0352  10.7%   3.635e+06
MyAssemble2        1    0.0060   1.8%   3.635e+06
MySolve            1    0.0233   7.1%   1.125e+07
MySolve2           1    0.0079   2.4%   1.125e+07

两次组装的 flop 完全相同, 说明做的是同一件事, 但第二次的耗时低了一个数量级以上; 求解也是同样的情况. 由此可以记住两条经验:

  • 给自己的代码计时之前应当先”热身”, 即先跑一遍把编译和 setup 的开销排除掉, 否则 测到的主要是编译时间;

  • 反过来, 若程序只求解一次就结束, 那么编译开销确实是真实成本, 这时应当关注的是 如何复用缓存, 而不是优化求解器.

Firedrake 会把编译好的代码缓存到磁盘, 因此同一段代码第二次运行时”第一次组装”也会 比冷启动快. 冷热两种情形下上面的对比都成立, 只是差距大小不同.

最后看完整的 -log_view 表格. PETSc.Log.view 接受一个 viewer, 不给参数时写到标准 输出, 但那样会连同主机名、用户名一起打印出来, 因此这里把它写进一个临时文件, 再挑 需要的行打印.

import os
import tempfile

with tempfile.TemporaryDirectory() as tmpdir:
    filename = os.path.join(tmpdir, 'profile.txt')
    viewer = PETSc.Viewer().createASCII(filename, comm=COMM_WORLD)
    PETSc.Log.view(viewer)
    viewer.destroy()            # 先关闭再读, 确保内容已经写入文件
    with open(filename) as f:
        lines = f.read().splitlines()

summary = ('Time (sec):', 'Flops:', 'Event  ', '                   Max Ratio')
events = ('MyMesh', 'MyAssemble', 'MyAssemble2', 'MySolve', 'MySolve2',
          'KSPSolve', 'MatMult', 'PCApply', 'firedrake.assemble.assemble')
for line in lines:
    if line.startswith(summary) or line.split(' ')[0] in events:
        print(line)
Time (sec):           1.431e+00     1.000   1.431e+00
Flops:                2.978e+07     1.000   2.978e+07  2.978e+07
Event                Count      Time (sec)     Flop                              --- Global ---  --- Stage ----  Total
                   Max Ratio  Max     Ratio   Max  Ratio  Mess   AvgLen  Reduct  %T %F %M %L %R  %T %F %M %L %R Mflop/s
MatMult              200 1.0 4.2545e-03 1.0 1.08e+07 1.0 0.0e+00 0.0e+00 0.0e+00  0 36  0  0  0   1 36  0  0  0  2533
PCApply              202 1.0 3.8953e-04 1.0 8.53e+05 1.0 0.0e+00 0.0e+00 0.0e+00  0  3  0  0  0   0  3  0  0  0  2191
KSPSolve               2 1.0 5.9630e-03 1.0 2.18e+07 1.0 0.0e+00 0.0e+00 0.0e+00  0 73  0  0  0   2 73  0  0  0  3657
MyMesh                 1 1.0 2.3945e-01 1.0 0.00e+00 0.0 0.0e+00 0.0e+00 0.0e+00 17  0  0  0  0  71  0  0  0  0     0
MyAssemble             1 1.0 3.5152e-02 1.0 3.63e+06 1.0 0.0e+00 0.0e+00 0.0e+00  2 12  0  0  0  10 12  0  0  0   103
firedrake.assemble.assemble       4 1.0 4.1166e-02 1.0 7.27e+06 1.0 0.0e+00 0.0e+00 0.0e+00  3 24  0  0  0  12 24  0  0  0   177
MySolve                1 1.0 2.3261e-02 1.0 1.13e+07 1.0 0.0e+00 0.0e+00 0.0e+00  2 38  0  0  0   7 38  0  0  0   484
MyAssemble2            1 1.0 6.0468e-03 1.0 3.63e+06 1.0 0.0e+00 0.0e+00 0.0e+00  0 12  0  0  0   2 12  0  0  0   601
MySolve2               1 1.0 7.8624e-03 1.0 1.13e+07 1.0 0.0e+00 0.0e+00 0.0e+00  1 38  0  0  0   2 38  0  0  0  1431

这就是脚本里加 -log_view 时会看到的表格 (完整输出有一百多行, 这里只留了表头和 几个事件). 可以对照前面”读懂 -log_view 的输出”一节逐列看: 自定义事件和 PETSc 内部 事件混在同一张表里, KSPSolveCount 是 2, 正好对应两次求解; 单进程运行时凡是 非零量的比值都是 1.0 (某一项本身为 0 时 PETSc 直接打成 0.0, 例如 MyMesh 那行 Flop 的比值), MessAvgLenReduct 三列全为 0. 表里的 firedrake.assemble.assemble 是 Firedrake 自己标注的事件, 它的名字超过 16 个字符, 因此后面的列整体右移了, 正是上一节说的那种情况.

3.2.1. 写成脚本运行#

命令行选项只能用在脚本上. 脚本 py/profile_poisson.py 是同一 个算例, 但去掉了 PETSc.Log.begin(), 改由 -log_view 自动开启日志, 并且用装饰器 的写法标注组装:

"""用 PETSc 自定义事件标注 Poisson 求解的各个阶段, 配合 -log_view 使用.

在 firedrake 目录下运行:

    python3 py/profile_poisson.py -log_view
    python3 py/profile_poisson.py -log_view :profile.txt:ascii_flamegraph

脚本中没有调用 PETSc.Log.begin(): 命令行传了 -log_view 时 PETSc 已经自己把日志
开好了 (普通格式开的是扁平日志, :...:ascii_xml 和 :...:ascii_flamegraph 开的是
嵌套日志). 这时再调用一次 Log.begin() 会额外启动一个扁平日志处理器; 用后两种
嵌套格式时, 程序退出会报 "Trying to end paused event, not allowed".
"""
from firedrake import *
from firedrake.petsc import PETSc


@PETSc.Log.EventDecorator('MyAssemble')
def build_system(a, L, bc):
    return assemble(a, bcs=bc), assemble(L, bcs=bc)


N = 64
with PETSc.Log.Event('MyMesh'):
    mesh = UnitSquareMesh(N, N)
    V = FunctionSpace(mesh, 'CG', 1)

x, y = SpatialCoordinate(mesh)
u, v = TrialFunction(V), TestFunction(V)
a = inner(grad(u), grad(v))*dx
L = inner(sin(pi*x)*sin(pi*y), v)*dx
bc = DirichletBC(V, 0, 'on_boundary')
u_h = Function(V, name='u_h')

A, b = build_system(a, L, bc)

with PETSc.Log.Event('MySolve'):
    solve(A, u_h, b, solver_parameters={'ksp_type': 'cg',
                                        'pc_type': 'jacobi',
                                        'ksp_rtol': 1e-10})

PETSc.Sys.Print(f'dofs = {V.dim()}, L2 norm of u_h = {norm(u_h):.6e}')

在激活了 Firedrake 环境的终端中, 切换到本笔记本所在的 firedrake 目录后运行

python3 py/profile_poisson.py -log_view

即可看到完整的性能摘要; 换成

python3 py/profile_poisson.py -log_view :profile.txt:ascii_flamegraph

则会在当前目录生成火焰图文件 profile.txt. 用

mpiexec -n 4 python3 py/profile_poisson.py -log_view

并行运行时, 表格中的最大最小比值和 MPI 相关的列才会真正有内容, 前面说的负载均衡 判断也才用得上.

找到瓶颈之后如何改, 可以参考 Firedrake 文档 Optimising 的 Common performance issues 一节, 其中最常见的一条是: 每次调用 solve() 都会新建一个 PETSc 求解器对象, 在循环里反复求解时应当把 NonlinearVariationalProblemNonlinearVariationalSolver 建在循环外面复用.