7. 矩阵和向量的组装 Assembly of Matrices and Vectors#
Firedrake 通过 TSFC [HMLH18] 生成局部组装内核 (如单元刚度矩阵的组装), 然后使用 PyOP2 构建全局内核. TSFC 调用 FInAT [KM19] 生成基函数的求值公式, 并使用 loopy 生成代码.
Overview
How Firedrake assembles matrices and vectors: inspecting the generated C code, and generating local and global assembly kernels with TSFC and PyOP2 (internals).
在读下面的代码之前, 先说清楚本章要展示的是什么.
用户写下的是 UFL 形式 (如 u*v*dx), 但真正组装矩阵和向量时, 用 Python 逐个单元做积分太慢. 因此 Firedrake 会先把形式编译成 C 代码, 再编译成机器码执行. 生成的代码分两层:
局部内核 (local kernel) 由 TSFC 生成, 只处理一个单元 (面积分则是一个面或一对相邻单元): 输入是该单元的坐标自由度 (直边网格上就是顶点坐标) 和各系数在该单元上的自由度, 输出是单元刚度矩阵或单元右端向量.
全局内核 (global kernel) 由 PyOP2 生成, 是套在局部内核外面的一层循环: 遍历网格单元, 按单元的节点映射 (
cell_node_map) 把全局数组里的数据取到单元上, 调用局部内核, 再把结果累加回全局矩阵或向量.
下一节把这两层代码直接打印出来, 分矩阵和右端项两半, 两半的收尾方式不一样:
矩阵那一半: 局部内核是
form00_cell_integral, 外面的全局内核是wrap_form00_cell_integral, 它末尾用 PETSc 的MatSetValuesBlockedLocal把单元矩阵加进全局矩阵;右端项那一半: 局部内核是
form0_cell_integral(只有一个0), 全局内核是wrap_form0_cell_integral, 它末尾没有 PETSc 调用, 而是按dat0[map0[...]] = dat0[map0[...]] + ...直接累加回全局向量.
读本章需要熟悉前面章节里 FunctionSpace、TrialFunction/TestFunction、assemble 的基本用法; 能大致看懂 C 代码会更轻松, 但不是必需的.
本章有两处贴的是 Firedrake 与 TSFC 的内部实现: “生成向量和矩阵组装代码的代码”一节, 以及”生成局部内核”及其之后的全部小节. 这两处的目的是让读者知道每一步归哪一层管, 不要求逐行读懂, 了解即可; 建议先读每节开头的说明, 代码当对照看. 另外要注意, 这些内部接口不属于公开 API, 名字和签名会随 Firedrake 版本变动.
夹在中间的”双线性形式”一节则不能跳过: 它是普通的用户级代码, 之后每个代码单元用的都是那里定义的 form.
7.1. 查看生成的 C 代码#
from firedrake import *
from textwrap import indent
mesh = RectangleMesh(10, 10, 1, 1)
V = FunctionSpace(mesh, 'CG', 1)
u, v = TrialFunction(V), TestFunction(V)
x, y = SpatialCoordinate(mesh)
f = Function(V, name='f').interpolate(sin(x))
f_tilde = conditional(f > 0.8, f, 0)
a = u*v*dx - f_tilde*v*dx
uh = Function(V, name='u_h')
prob = LinearVariationalProblem(lhs(a), rhs(a), uh)
solver = LinearVariationalSolver(prob)
solver.solve()
7.1.1. 组装矩阵的 C 代码#
from pyop2.global_kernel import _generate_code_from_global_kernel
print("Code for mass matrix:\n")
assembler = solver._ctx._assemble_jac.__self__
for parloop in assembler.parloops(solver._ctx._jac):
kernel = parloop.global_kernel
code = _generate_code_from_global_kernel(kernel, parloop.comm)
print(indent(code, " "))
Code for mass matrix:
#include <complex.h>
#include <math.h>
#include <petsc.h>
#include <petsc.h>
#include <stdint.h>
#include <stdbool.h>
#include <math.h>
static void form00_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0);
static void form00_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0)
{
double t0[3 * 3] = { 0.6666666666666666, 0.16666666666666669, 0.16666666666666674, 0.16666666666666663, 0.16666666666666663, 0.6666666666666665, 0.16666666666666663, 0.6666666666666665, 0.16666666666666663 };
double t1;
double t2;
double t3;
double t4[3] = { 0.16666666666666666, 0.16666666666666666, 0.16666666666666666 };
double t5;
double t6;
t1 = -1.0 * coords_0[0];
t2 = -1.0 * coords_0[1];
t3 = fabs((t1 + coords_0[2]) * (t2 + coords_0[5]) + -1.0 * (t1 + coords_0[4]) * (t2 + coords_0[3]));
for (int32_t ip = 0; ip <= 2; ++ip)
{
t5 = t4[ip] * t3;
for (int32_t j = 0; j <= 2; ++j)
{
t6 = t0[ip + 3 * j] * t5;
for (int32_t k = 0; k <= 2; ++k)
A[3 * j + k] = A[3 * j + k] + t0[ip + 3 * k] * t6;
}
}
}
void wrap_form00_cell_integral(int32_t const start, int32_t const end, Mat const mat0, double const *__restrict__ dat0, int32_t const *__restrict__ map0)
{
double t0[3 * 3];
double t1[3 * 2];
for (int32_t n = start; n <= -1 + end; ++n)
{
{
int32_t const i15 = 0;
for (int32_t i16 = 0; i16 <= 2; ++i16)
{
int32_t const i17 = 0;
{
int32_t const i18 = 0;
for (int32_t i19 = 0; i19 <= 2; ++i19)
{
int32_t const i20 = 0;
t0[3 * i16 + i19] = (double) (0.0);
}
}
}
}
{
int32_t const i21 = 0;
for (int32_t i22 = 0; i22 <= 2; ++i22)
for (int32_t i23 = 0; i23 <= 1; ++i23)
t1[2 * i22 + i23] = dat0[2 * map0[3 * n + i22] + i23];
}
form00_cell_integral(&(t0[0]), &(t1[0]));
MatSetValuesBlockedLocal(mat0, 3, &(map0[3 * n]), 3, &(map0[3 * n]), &(t0[0]), ADD_VALUES);
}
}
7.1.2. 组装右端项的 C 代码#
print("\n\nCode for right hand:\n")
residual_assembler = solver._ctx._assemble_residual.__self__
for parloop in residual_assembler.parloops(solver._ctx._F):
kernel = parloop.global_kernel
code = _generate_code_from_global_kernel(kernel, parloop.comm)
print(indent(code, " "))
Code for right hand:
#include <complex.h>
#include <math.h>
#include <petsc.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdbool.h>
#include <math.h>
static void form0_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0, double const *__restrict__ w_0, double const *__restrict__ w_1);
static void form0_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0, double const *__restrict__ w_0, double const *__restrict__ w_1)
{
double t0[3 * 3] = { 0.6666666666666666, 0.16666666666666669, 0.16666666666666674, 0.16666666666666663, 0.16666666666666663, 0.6666666666666665, 0.16666666666666663, 0.6666666666666665, 0.16666666666666663 };
double t1[3] = { 0.16666666666666663, 0.6666666666666665, 0.16666666666666663 };
double t2[3] = { 0.16666666666666663, 0.16666666666666663, 0.6666666666666665 };
double t3[3] = { 0.6666666666666666, 0.16666666666666669, 0.16666666666666674 };
double t4;
double t5;
double t6;
double t7[3] = { 0.16666666666666666, 0.16666666666666666, 0.16666666666666666 };
double t8;
double t9;
t4 = -1.0 * coords_0[0];
t5 = -1.0 * coords_0[1];
t6 = fabs((t4 + coords_0[2]) * (t5 + coords_0[5]) + -1.0 * (t4 + coords_0[4]) * (t5 + coords_0[3]));
for (int32_t ip = 0; ip <= 2; ++ip)
{
t8 = t3[ip] * w_0[0] + t2[ip] * w_0[1] + t1[ip] * w_0[2];
t9 = t7[ip] * t6 * (-1.0 * ((t8 > 0.8) ? t8 : (double) (0.0)) + t3[ip] * w_1[0] + t2[ip] * w_1[1] + t1[ip] * w_1[2]);
for (int32_t j = 0; j <= 2; ++j)
A[j] = A[j] + t0[ip + 3 * j] * t9;
}
}
void wrap_form0_cell_integral(int32_t const start, int32_t const end, double *__restrict__ dat0, double const *__restrict__ dat1, double const *__restrict__ dat2, double const *__restrict__ dat3, int32_t const *__restrict__ map0)
{
double t0[3];
double t1[3 * 2];
double t2[3];
double t3[3];
for (int32_t n = start; n <= -1 + end; ++n)
{
{
int32_t const i15 = 0;
for (int32_t i16 = 0; i16 <= 2; ++i16)
{
int32_t const i17 = 0;
t0[i16] = (double) (0.0);
}
}
for (int32_t i19 = 0; i19 <= 2; ++i19)
{
{
int32_t const i18 = 0;
for (int32_t i20 = 0; i20 <= 1; ++i20)
t1[2 * i19 + i20] = dat1[2 * map0[3 * n + i19] + i20];
}
{
int32_t const i21 = 0;
{
int32_t const i22 = 0;
t2[i19] = dat2[map0[3 * n + i19]];
}
}
{
int32_t const i23 = 0;
{
int32_t const i24 = 0;
t3[i19] = dat3[map0[3 * n + i19]];
}
}
}
form0_cell_integral(&(t0[0]), &(t1[0]), &(t2[0]), &(t3[0]));
for (int32_t i12 = 0; i12 <= 2; ++i12)
{
int32_t const i13 = 0;
{
int32_t const i14 = 0;
dat0[map0[3 * n + i12]] = dat0[map0[3 * n + i12]] + t0[i12];
}
}
}
}
7.2. 生成向量和矩阵组装代码的代码#
上一节看到的 C 代码, 是组装器在背后自动生成的 (那里取的是求解器内部的组装器, 直接调 assemble 走的是同一套机制). 相关代码在 firedrake/assemble.py, 按调用顺序是:
ParloopFormAssembler.local_kernels: 调用tsfc_interface.compile_form, 把形式编译成局部内核;ParloopFormAssembler.parloops: 为每个局部内核建一个ParloopBuilder, 由它的build方法产出可以执行的op2.Parloop;_GlobalKernelBuilder.build: 经_make_global_kernel被ParloopBuilder.build调用, 产出op2.GlobalKernel, 也就是上一节里那层wrap_循环.
firedrake/preconditioners/patch.py 的 matrix_funptr 把”从形式到可调用内核”的整条链路写在了一个函数里, 比 assemble.py 里分散的实现更易读, 因此这里拿它当例子 (残差的版本是同一文件里的 residual_funptr). 它先用 compile_form 编译出局部内核, 再包一层 op2.GlobalKernel 并编译成 C 函数指针, 供 PCPatch 在每个小块上直接调用. 本章最后一节”生成全局内核”就是照着它手工走一遍.
下面是删去细节后的骨架:
def get_map(V, base_mesh, base_integral_type):
return V.topological.entity_node_map(base_mesh.topology, base_integral_type, None, None)
def matrix_funptr(form, state):
from firedrake.tsfc_interface import compile_form
test, trial = map(operator.methodcaller("function_space"), form.arguments())
if test != trial:
raise NotImplementedError("Only for matching test and trial spaces")
if state is not None:
dont_split = (state, )
else:
dont_split = ()
kernels = compile_form(form, "subspace_form", split=False, dont_split=dont_split)
all_meshes = extract_domains(form)
cell_kernels = []
int_facet_kernels = []
ext_facet_kernels = []
for kernel in kernels:
kinfo = kernel.kinfo
mesh = all_meshes[kinfo.domain_number] # integration domain
integral_type = kinfo.integral_type
# 这里省略对 subdomain_id 和 integral_type 的检查
# OK, now we've validated the kernel, let's build the callback
args = []
if integral_type == "cell":
kernels = cell_kernels
elif integral_type == "interior_facet":
kernels = int_facet_kernels
elif integral_type == "exterior_facet":
kernels = ext_facet_kernels
# 这里省略 args 的填充: 局部矩阵、坐标、单元定向、各系数与常数等
wrapper_knl_args = tuple(a.global_kernel_arg for a in args)
mod = op2.GlobalKernel(kinfo.kernel, wrapper_knl_args, subset=True)
kernels.append(CompiledKernel(compile_global_kernel(mod, iterset.comm), kinfo))
return cell_kernels, int_facet_kernels, ext_facet_kernels
7.3. 双线性形式#
从这一节起换一个更复杂的例子: 定义在混合空间 V1*V2 上的双线性形式, 系数里既有混合空间上的函数也有普通函数, 导入方式也从 from firedrake import * 换成了 import firedrake as fd. 它与第一节那个标量质量矩阵算例互不相干: 第一节的内核叫 form00_cell_integral, 而这一节之后的内核按调用时给的前缀命名, 会看到 subspace_form_cell_integral 和 test_form_cell_integral 两种. 不过后面每一个小节用的都是这里定义的 form, 因此这一节不能跳过.
7.3.1. 创建双线性形式#
import firedrake as fd
mesh = fd.RectangleMesh(10, 10, 1, 1)
V1 = fd.FunctionSpace(mesh, 'CG', 1)
V2 = fd.FunctionSpace(mesh, 'CG', 2)
U = V1*V2
u, v = fd.TrialFunction(U), fd.TestFunction(U)
g = fd.Function(U)
h = fd.Function(V1)
r = fd.Function(V2)
form = fd.inner(fd.grad(u[0]*h*g[0]), fd.grad(v[0]))*fd.dx + fd.inner(u[1]*g[1]*r, v[1])*fd.dx
7.3.2. 查看 Form 表达式#
from ufl.formatting import ufl2unicode
ustr = ufl2unicode.ufl2unicode(form)
print(ustr)
∬[rest of domain] v⃗[1] w₂₆ u⃗[1] w⃗₂₂[1] + ∑[i₀]((𝐠𝐫𝐚𝐝 v⃗)[0,i₀] [w⃗₂₂[0] [u⃗[0] (𝐠𝐫𝐚𝐝 w₂₄)[i₂] + (𝐠𝐫𝐚𝐝 u⃗)[0,i₂] w₂₄ ∀ i₂][i₁] + (𝐠𝐫𝐚𝐝 w⃗₂₂)[0,i₁] u⃗[0] w₂₄ ∀ i₁][i₀]) 𝐝𝐱
ufl2unicode 给出的是紧凑的数学写法, 适合快速确认形式有没有写错. 下面的 tree_format 则打印完整的表达式树, 可以看到 UFL 内部真正的节点 (Argument、Coefficient、Grad、Indexed 等), 以及这个形式被拆成了几个 Integral. TSFC 接下来处理的正是这棵树.
from ufl.utils.formatting import tree_format
print(tree_format(form))
Form:
Integral:
integral type: cell
subdomain id: everywhere
integrand:
Conj
Inner
(
Grad
Indexed
(
Argument(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), MeshSequence((Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)))), 0, None)
MultiIndex((FixedIndex(0),))
)
Grad
Product
(
Indexed
(
Coefficient(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), MeshSequence((Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)))), 22)
MultiIndex((FixedIndex(0),))
)
Product
(
Indexed
(
Argument(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), MeshSequence((Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)))), 1, None)
MultiIndex((FixedIndex(0),))
)
Coefficient(WithGeometry(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 1), name=None), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)), 24)
)
)
)
Integral:
integral type: cell
subdomain id: everywhere
integrand:
Product
(
Product
(
Coefficient(WithGeometry(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 2), name=None), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)), 26)
Product
(
Indexed
(
Argument(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), MeshSequence((Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)))), 1, None)
MultiIndex((FixedIndex(1),))
)
Indexed
(
Coefficient(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), MeshSequence((Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)))), 22)
MultiIndex((FixedIndex(1),))
)
)
)
Conj
Indexed
(
Argument(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f3f1f16ecc0>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), MeshSequence((Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 8)))), 0, None)
MultiIndex((FixedIndex(1),))
)
)
7.4. 生成局部内核#
相关代码: firedrake/tsfc_interface.py 中的 compile_form.
下面几个小节把”局部内核是如何造出来的”一层层拆开, 每一节再往里展开一层:
“具体构造过程”: Firedrake 的
compile_form循环体, 算出系数与常数的编号、拼出内核名前缀, 最后交给TSFCKernel;“
TSFCKernel.__init__”: 它调用 TSFC 的compile_form, 再把结果包装成KernelInfo;“
tsfc.driver.compile_form”: 做完预处理后, 对每个IntegralData调用compile_integral;“
tsfc.driver.compile_integral”: 最里面一层, 由KernelBuilder完成从 UFL 到 GEM 再到 loopy 的转换. 这一节只有文字说明, 没有代码单元.
这几个小节的代码单元把外层的循环、缓存和部分校验略去了, 直接对”双线性形式”一节里那个 form 跑一遍.
from firedrake.tsfc_interface import compile_form
kernels = compile_form(form, "subspace_form", split=False)
loopy 负责把内核降到 C 代码. 下面这个单元只做转换、不打印结果 (输出有一百多行), 因此跑完看不到任何输出是正常的; 若想看, 把最后一行的注释去掉即可.
import loopy
idx, kinfo = kernels[0]
code = loopy.generate_code_v2(kinfo[0].code)
# print the code
# print(code.device_code())
7.4.1. 具体构造过程#
这一段是 firedrake/tsfc_interface.py 里 compile_form 循环体的展开: 取出各系数与常数在整个形式中的编号, 拼出内核名前缀, 再交给 TSFCKernel.
from firedrake.parameters import parameters as default_parameters
from firedrake.tsfc_interface import TSFCKernel
from tsfc.ufl_utils import extract_firedrake_constants
parameters = default_parameters["form_compiler"].copy()
nargs = len(form.arguments())
iterable = ([(None, )*nargs, form], )
idx, f = iterable[0]
name = "test_form"
numbering = form.terminal_numbering()
coefficient_numbers = tuple(
numbering[c] for c in form.coefficients()
)
constant_numbers = tuple(
numbering[c] for c in extract_firedrake_constants(f)
)
prefix = name + "".join(map(str, (i for i in idx if i is not None)))
kinfos = TSFCKernel(f, prefix, parameters,
domain_number_map=(0,), # single-mesh form
coefficient_numbers=coefficient_numbers,
constant_numbers=constant_numbers,
dont_split_numbers=(), diagonal=False).kernels
7.4.2. TSFCKernel.__init__#
相关代码: firedrake/tsfc_interface.py 中的 TSFCKernel.__init__.
from tsfc import compile_form as tsfc_compile_form
from firedrake.tsfc_interface import as_pyop2_local_kernel, KernelInfo
tree = tsfc_compile_form(form=f, prefix=name, parameters=parameters,
diagonal=False)
_kernels = []
for kernel in tree:
coefficient_numbers_per_kernel = tuple(
(coefficient_numbers[index], subindices)
for index, subindices in kernel.coefficient_numbers
)
constant_numbers_per_kernel = constant_numbers
events = (kernel.event,)
pyop2_kernel = as_pyop2_local_kernel(kernel.ast, kernel.name,
len(kernel.arguments),
flop_count=kernel.flop_count,
events=events)
_kernels.append(KernelInfo(kernel=pyop2_kernel,
integral_type=kernel.integral_type,
subdomain_id=kernel.subdomain_id,
domain_number=kernel.domain_number,
active_domain_numbers=kernel.active_domain_numbers,
coefficient_numbers=coefficient_numbers_per_kernel,
constant_numbers=constant_numbers_per_kernel,
needs_cell_facets=False,
pass_layer_arg=False,
arguments=kernel.arguments,
events=events))
7.4.3. tsfc.driver.compile_form#
相关代码: tsfc/driver.py 中的 compile_form.
from tsfc.driver import compile_integral
from tsfc.parameters import default_parameters as tsfc_default_parameters, is_complex
from tsfc import fem, ufl_utils
import finat.ufl
complex_mode = parameters and is_complex(parameters.get("scalar_type"))
# Preprocess UFL form in a format suitable for TSFC.
# Mixed coefficients must be split into components, otherwise the generated
# kernel takes fewer arguments than the real one.
form_data = ufl_utils.compute_form_data(
form,
coefficients_to_split=tuple(
c for c in form.coefficients()
if type(c.ufl_element()) == finat.ufl.MixedElement
),
complex_mode=complex_mode,
)
kernels = []
for integral_data in form_data.integral_data:
kernel = compile_integral(integral_data, form_data, prefix=name, parameters=parameters, diagonal=False)
if kernel is not None:
kernels.append(kernel)
compute_form_data 是 UFL 的预处理入口: 展开求导, 把混合空间的系数按分量拆开 (coefficients_to_split), 再按 (网格, 积分类型, 子区域) 给积分分组, 同一组合并成一个 IntegralData. 上面 form 里的两个 dx 项就被并成了一个.
下面打印出第一个 IntegralData 的类型, 以及它里面那个积分的 metadata(). 其中的 estimated_polynomial_degree 是 UFL 在 compute_form_data 里估出的被积函数多项式次数 (由 attach_estimated_degrees 挂上去), 之后 TSFC 的 set_quad_rule 会用它来挑数值积分公式.
# form_data.preprocessed_form.arguments()
print(type(form_data.integral_data[0]))
print(form_data.integral_data[0].integrals[0].metadata())
<class 'ufl.algorithms.domain_analysis.IntegralData'>
{'estimated_polynomial_degree': 8}
7.4.4. tsfc.driver.compile_integral#
相关代码: tsfc/driver.py 中的 compile_integral.
这一层把一个 IntegralData (上一步分好组的一组同类积分) 编译成一个组装内核, 具体工作交给 KernelBuilder:
from tsfc.kernel_interface.firedrake_loopy import KernelBuilder
compile_integral 先建好 KernelBuilder, 把坐标、单元定向、单元尺寸、系数和常数等信息交给它, 然后对这组里的每个积分依次做三件事:
builder.compile_integrand调用tsfc/kernel_interface/common.py的set_quad_rule定下数值积分公式, 再用tsfc.fem.compile_ufl把ufl表达式转成GEM[HMLH18]. 这一步会调用FInAT生成基函数在积分点处的求值公式;builder.construct_integrals应用数值积分公式, 生成积分对应的GEM表达式;builder.stash_integrals把结果暂存进上下文.
这组积分都处理完后, builder.construct_kernel 把 GEM 降到 loopy, 生成的就是本章开头看到的那种 C 代码.
7.5. 生成全局内核#
前面得到的局部内核只算一个单元, 要真正跑起来还得套上 PyOP2 的外层循环. 这一节照着 matrix_funptr 的做法手工走一遍: 先准备参数, 再构建 op2.GlobalKernel. 用到的 kinfo 来自”生成局部内核”一节开头那次调用 (本章出现了两个同名的 compile_form, 这里指 firedrake.tsfc_interface 的那个).
from pyop2 import op2
from firedrake.utils import IntType
from firedrake.preconditioners.patch import LocalDat, LocalMat
import operator
import numpy
test, trial = map(operator.methodcaller("function_space"), form.arguments())
准备全局内核的参数列表: 局部矩阵 (LocalMat)、单元坐标, 以及形式里用到的各个系数. 本章的 form 只有单元积分, 所以取的是 cell_node_map; 注意现在的 patch.py 里用的是更通用的 entity_node_map, 这里用 cell_node_map 是为了写得直白些.
args = []
if kinfo.integral_type == "cell":
get_map = operator.methodcaller("cell_node_map")
elif kinfo.integral_type == "interior_facet":
get_map = operator.methodcaller("interior_facet_node_map")
else:
get_map = None
toset = op2.Set(1, comm=test.comm)
dofset = op2.DataSet(toset, 1)
arity = sum(m.arity*s.cdim
for m, s in zip(get_map(test),
test.dof_dset))
iterset = get_map(test).iterset
entity_node_map = op2.Map(iterset,
toset, arity,
values=numpy.zeros(iterset.total_size*arity, dtype=IntType))
mat = LocalMat(dofset)
arg = mat(op2.INC, (entity_node_map, entity_node_map))
args.append(arg)
mesh = form.ufl_domains()[kinfo.domain_number]
arg = mesh.coordinates.dat(op2.READ, get_map(mesh.coordinates))
args.append(arg)
for n, indices in kinfo.coefficient_numbers:
c = form.coefficients()[n]
for ind in indices:
c_ = c.subfunctions[ind]
map_ = get_map(c_)
arg = c_.dat(op2.READ, map_)
args.append(arg)
if kinfo.integral_type == "interior_facet":
arg = test.ufl_domain().interior_facets.local_facet_dat(op2.READ)
args.append(arg)
wrapper_knl_args = tuple(a.global_kernel_arg for a in args)
构建全局组装内核
mod = op2.GlobalKernel(kinfo.kernel, wrapper_knl_args, subset=True)
把生成的代码打印出来. 这段没有作为代码单元执行, 因为输出有两百多行; 若需要, 可以复制到自己的环境里跑.
from pyop2.global_kernel import _generate_code_from_global_kernel
print(_generate_code_from_global_kernel(mod, iterset.comm))
References
Miklós Homolya, Lawrence Mitchell, Fabio Luporini, and David A. Ham. TSFC: a structure-preserving form compiler. SIAM Journal on Scientific Computing, 40(3):C401–C428, jan 2018. doi:10.1137/17m1130642.
Robert C. Kirby and Lawrence Mitchell. Code generation for generally mapped finite elements. ACM Transactions on Mathematical Software, 45(4):1–23, dec 2019. doi:10.1145/3361745.