7. 矩阵和向量的组装 Assembly of Matrices and Vectors#
Firedrake 通过 TSFC [HMLH18] 生成局部组装内核 (如单元刚度矩阵的组装), 然后使用 PyOP2 构建全局内核. TSFC 调用 FInAT [KM19] 生成基函数的求值公式, 并使用 loopy 生成代码.
Firedrake uses TSFC [HMLH18] to generate the local assembly kernels (the assembly of the element stiffness matrix, for instance), and then builds the global kernel with PyOP2. TSFC calls FInAT [KM19] to generate the evaluation formulas of the basis functions, and uses loopy to generate the code.
概要
Firedrake 如何组装矩阵和向量: 查看生成的 C 代码, 以及用 TSFC 和 PyOP2 生成局部内核与全局内核 (内部实现).
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[...]] + ...直接累加回全局向量.
本章有两处贴的是 Firedrake 与 TSFC 的内部实现: “生成向量和矩阵组装代码的代码”一节, 以及”生成局部内核”及其之后的全部小节, 不要求逐行读懂. 这些内部接口不属于公开 API, 名字和签名会随 Firedrake 版本变动.
夹在中间的”双线性形式”一节则不能跳过: 它是普通的用户级代码, 之后每个代码单元用的都是那里定义的 form.
What the user writes down is a UFL form (such as u*v*dx), but when matrices and vectors are actually assembled, integrating cell by cell in Python is far too slow. Firedrake therefore first compiles the form into C code, which is then compiled into machine code and executed. The generated code has two layers:
The local kernel is generated by TSFC and handles a single cell only (a single facet, or a pair of neighboring cells, for facet integrals): its inputs are the coordinate degrees of freedom of that cell (simply the vertex coordinates on a straight-edged mesh) and the degrees of freedom of each coefficient on that cell, and its output is the element stiffness matrix or the element right-hand side vector.
The global kernel is generated by PyOP2 and is a loop wrapped around the local kernel: it runs over the mesh cells, gathers the data of the global arrays onto each cell through the cell node map (
cell_node_map), calls the local kernel, and then accumulates the result back into the global matrix or vector.
The next section prints these two layers of code directly, in two halves, one for the matrix and one for the right-hand side:
The matrix half: the local kernel is
form00_cell_integraland the global kernel around it iswrap_form00_cell_integral, which at its end uses PETSc’sMatSetValuesBlockedLocalto add the element matrix into the global matrix;The right-hand side half: the local kernel is
form0_cell_integral(with a single0) and the global kernel iswrap_form0_cell_integral, which makes no PETSc call at its end but accumulates back into the global vector directly asdat0[map0[...]] = dat0[map0[...]] + ....
Two places in this chapter quote the internal implementation of Firedrake and TSFC: the section “The code that generates the vector and matrix assembly code”, and the section “Generating the local kernel” together with all the subsections after it; they need not be understood line by line. These internal interfaces are not public API: their names and signatures change with the Firedrake version.
The section “Bilinear form” sandwiched between them, on the other hand, must not be skipped: it is ordinary user-level code, and every code cell after it uses the form defined there.
7.1. 查看生成的 C 代码 Inspecting the generated C code#
例子取 P1 空间上的 a = u*v*dx - f_tilde*v*dx: 左端是质量矩阵, 右端带一个系数 f_tilde (用 conditional 从 f 截断得到). 求解一次之后, 生成的 C 代码就留在求解器内部, 下面两小节分别把矩阵和右端项那份打印出来.
The example is a = u*v*dx - f_tilde*v*dx on a P1 space: the left-hand side is the mass matrix, and the right-hand side carries a coefficient f_tilde (obtained from f by a conditional cutoff). After one solve the generated C code is held inside the solver; the two subsections below print the matrix half and the right-hand-side half.
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 代码 C code for assembling the matrix#
从求解器内部的组装器取出全局内核, 转成 C 源码打印. 输出里 form00_cell_integral 是局部内核, 只以单元坐标为输入; wrap_form00_cell_integral 是外层循环.
The global kernel is taken from the assembler inside the solver and turned back into C source. In the output form00_cell_integral is the local kernel, whose only input is the cell coordinates; wrap_form00_cell_integral is the outer loop.
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 代码 C code for assembling the right-hand side#
做法相同, 换成残差的组装器. 与矩阵那份对照, 局部内核 form0_cell_integral 的参数里多了系数在单元上的自由度.
The same is done with the residual assembler. Compared with the matrix half, the parameter list of the local kernel form0_cell_integral additionally carries the degrees of freedom of the coefficient on the cell.
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. 生成向量和矩阵组装代码的代码 The code that generates the vector and matrix assembly code#
上一节看到的 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_循环.
The C code seen in the previous section is generated automatically by the assembler behind the scenes (the assembler taken there is the one inside the solver; calling assemble directly goes through the same mechanism). The relevant code is in firedrake/assemble.py; in call order:
ParloopFormAssembler.local_kernels: callstsfc_interface.compile_formto compile the form into local kernels;ParloopFormAssembler.parloops: builds oneParloopBuilderper local kernel, whosebuildmethod produces an executableop2.Parloop;_GlobalKernelBuilder.build: called byParloopBuilder.buildthrough_make_global_kernel, produces theop2.GlobalKernel, thewrap_loop of the previous section.
firedrake/preconditioners/patch.py 的 matrix_funptr 把”从形式到可调用内核”的整条链路写在了一个函数里, 比 assemble.py 里分散的实现更易读, 因此这里拿它当例子 (残差的版本是同一文件里的 residual_funptr). 它先用 compile_form 编译出局部内核, 再包一层 op2.GlobalKernel 并编译成 C 函数指针, 供 PCPatch 在每个小块上直接调用. 本章最后一节”生成全局内核”就是照着它手工走一遍.
下面是删去细节后的骨架:
matrix_funptr in firedrake/preconditioners/patch.py writes the whole chain “from the form to a callable kernel” inside a single function, which is easier to read than the implementation scattered over assemble.py, so it is taken as the example here (the residual version is residual_funptr in the same file). It first compiles the local kernel with compile_form, then wraps it in an op2.GlobalKernel and compiles that into a C function pointer, which PCPatch calls directly on each patch. The last section of this chapter, “Generating the global kernel”, walks through it by hand.
Below is the skeleton with the details removed; the two Chinese comments in it mark what has been left out, namely the checks on subdomain_id and integral_type, and the filling of args (the local matrix, the coordinates, the cell orientations, the coefficients and the constants, and so on):
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. 双线性形式 Bilinear form#
从这一节起换一个更复杂的例子: 定义在混合空间 V1*V2 上的双线性形式, 系数里既有混合空间上的函数也有普通函数, 导入方式也从 from firedrake import * 换成了 import firedrake as fd. 它与第一节那个标量质量矩阵算例互不相干: 第一节的内核叫 form00_cell_integral, 而这一节之后的内核按调用时给的前缀命名, 会看到 subspace_form_cell_integral 和 test_form_cell_integral 两种. 不过后面每一个小节用的都是这里定义的 form, 因此这一节不能跳过.
From this section on we switch to a more complicated example: a bilinear form on the mixed space V1*V2, whose coefficients include both a function on the mixed space and an ordinary function; the import also changes from from firedrake import * to import firedrake as fd. It has nothing to do with the scalar mass matrix example of the first section: the kernel there is called form00_cell_integral, whereas the kernels from this section on are named after the prefix given at the call site, so both subspace_form_cell_integral and test_form_cell_integral will show up. Every subsection below does use the form defined here, however, so this section must not be skipped.
7.3.1. 创建双线性形式 Creating the bilinear form#
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 表达式 Inspecting the Form expression#
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 接下来处理的正是这棵树.
ufl2unicode gives a compact mathematical rendering, handy for a quick check that the form has been written correctly. The tree_format below prints the full expression tree instead, showing the nodes UFL really uses internally (Argument, Coefficient, Grad, Indexed, and so on) and how many Integrals the form is split into. This tree is precisely what TSFC processes next.
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 0x7fea271f8cb0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7fea271f8cb0>, 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 0x7fea271f8cb0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7fea271f8cb0>, 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 0x7fea271f8cb0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7fea271f8cb0>, 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 0x7fea271f8cb0>, 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 0x7fea271f8cb0>, 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 0x7fea271f8cb0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7fea271f8cb0>, 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 0x7fea271f8cb0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7fea271f8cb0>, 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 0x7fea271f8cb0>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7fea271f8cb0>, 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. 生成局部内核 Generating the local kernel#
相关代码: 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 跑一遍.
Relevant code: compile_form in firedrake/tsfc_interface.py.
The next few subsections take the construction of the local kernel apart layer by layer:
“The construction in detail”: the body of Firedrake’s
compile_formloop, which works out the numbering of the coefficients and constants, assembles the kernel name prefix, and finally hands over toTSFCKernel;“
TSFCKernel.__init__”: it calls TSFC’scompile_formand wraps the result into aKernelInfo;“
tsfc.driver.compile_form”: once the preprocessing is done, it callscompile_integralfor everyIntegralData;“
tsfc.driver.compile_integral”: the innermost layer, whereKernelBuildercarries out the conversion from UFL to GEM and then to loopy. This subsection has text only, no code cell.
The code cells of these subsections leave out the outer loops, the caching and part of the validation, and simply run on the form of the section “Bilinear form”.
from firedrake.tsfc_interface import compile_form
kernels = compile_form(form, "subspace_form", split=False)
loopy 负责把内核降到 C 代码. 下面这个单元只做转换、不打印结果 (输出有一百多行), 因此跑完看不到任何输出是正常的; 若想看, 把最后一行的注释去掉即可.
loopy is responsible for lowering the kernel to C code. The cell below only performs the conversion and does not print the result (the output is more than a hundred lines), so seeing no output after it has run is normal; to see it, remove the comment on the last line.
import loopy
idx, kinfo = kernels[0]
code = loopy.generate_code_v2(kinfo[0].code)
# print the code
# print(code.device_code())
7.4.1. 具体构造过程 The construction in detail#
这一段是 firedrake/tsfc_interface.py 里 compile_form 循环体的展开: 取出各系数与常数在整个形式中的编号, 拼出内核名前缀, 再交给 TSFCKernel.
This is an expansion of the body of the compile_form loop in firedrake/tsfc_interface.py: it picks out the numbering of each coefficient and constant within the whole form, assembles the kernel name prefix, and then hands over to 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__.
Relevant code: TSFCKernel.__init__ in firedrake/tsfc_interface.py.
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.
Relevant code: compile_form in tsfc/driver.py.
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 会用它来挑数值积分公式.
compute_form_data is UFL’s preprocessing entry point: it expands the derivatives, splits the coefficients on mixed spaces into their components (coefficients_to_split), and then groups the integrals by (mesh, integral type, subdomain), merging each group into one IntegralData. The two dx terms of the form above are merged into one.
Below we print the type of the first IntegralData together with the metadata() of the integral it contains. The estimated_polynomial_degree there is the polynomial degree of the integrand estimated by UFL inside compute_form_data (attached by attach_estimated_degrees); TSFC’s set_quad_rule later uses it to pick the quadrature 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:
Relevant code: compile_integral in tsfc/driver.py.
This layer compiles one IntegralData (a group of integrals of the same kind, formed in the previous step) into one assembly kernel, and delegates the actual work to 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 代码.
compile_integral first builds the KernelBuilder and hands it the coordinates, the cell orientations, the cell sizes, the coefficients and the constants, and then does three things for each integral of the group in turn:
builder.compile_integrandcallsset_quad_ruleintsfc/kernel_interface/common.pyto fix the quadrature rule, and then turns theuflexpression intoGEM[HMLH18] withtsfc.fem.compile_ufl. This step callsFInATto generate the evaluation formulas of the basis functions at the quadrature points;builder.construct_integralsapplies the quadrature rule and produces theGEMexpression of the integral;builder.stash_integralsstashes the result into the context.
Once all the integrals of the group have been processed, builder.construct_kernel lowers the GEM to loopy, and what comes out is the kind of C code seen at the beginning of this chapter.
7.5. 生成全局内核 Generating the global kernel#
前面得到的局部内核只算一个单元, 要真正跑起来还得套上 PyOP2 的外层循环. 这一节照着 matrix_funptr 的做法手工走一遍: 先准备参数, 再构建 op2.GlobalKernel. 用到的 kinfo 来自”生成局部内核”一节开头那次调用 (本章出现了两个同名的 compile_form, 这里指 firedrake.tsfc_interface 的那个).
The local kernel obtained so far computes a single cell only; to make it actually run, PyOP2’s outer loop still has to be wrapped around it. This section walks through the approach of matrix_funptr by hand: first prepare the arguments, then build the op2.GlobalKernel. The kinfo used comes from the call at the beginning of the section “Generating the local kernel” (two functions named compile_form appear in this chapter; the one meant here is the one from 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 是为了写得直白些.
Prepare the argument list of the global kernel: the local matrix (LocalMat), the cell coordinates, and each of the coefficients used in the form. The form of this chapter has cell integrals only, so cell_node_map is taken; note that today’s patch.py uses the more general entity_node_map, and cell_node_map is used here to keep things plain.
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)
构建全局组装内核 Building the global assembly kernel
mod = op2.GlobalKernel(kinfo.kernel, wrapper_knl_args, subset=True)
把生成的代码打印出来. 这段没有作为代码单元执行, 因为输出有两百多行; 若需要, 可以复制到自己的环境里跑.
Printing the generated code. This is not run as a code cell because the output is more than two hundred lines; if needed, it can be copied into your own environment and run there.
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.