8. 组装内核的构建

Firedrake 通过 TSFC [HMLH18] 生成局部组装内核 (如单元刚度矩阵的组装), 然后使用 PyOP2 构建全局内核. TSFC 调用 FInAT [KM19] 生成基函数的求值公式, 并使用 loopy 生成代码.

以下代码来自 firedrake/patch.py, 它的作用为生成局部内核, 即单元上矩阵或残差的组装代码.

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:
        interface = make_builder(dont_split=(state, ))
    else:
        interface = None

    kernels = compile_form(form, "subspace_form", split=False, interface=interface)

    cell_kernels = []
    int_facet_kernels = []
    for kernel in kernels:
        kinfo = kernel.kinfo

        # OK, now we've validated the kernel, let's build the callback
        args = []

        if kinfo.integral_type == "cell":
            get_map = operator.methodcaller("cell_node_map")
            kernels = cell_kernels
        elif kinfo.integral_type == "interior_facet":
            get_map = operator.methodcaller("interior_facet_node_map")
            kernels = int_facet_kernels
        else:
            get_map = None

        # build 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(mod.compile(iterset.comm), kinfo))

    return cell_kernels, int_facet_kernels

8.1. 双线性形式

8.1.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

8.1.2. 查看 Form 表达式

1. 使用 可视化包 graphviz 和 ufl 提供的 ufl2dot

   import graphviz
   from ufl.formatting import ufl2dot
   
   source = ufl2dot.ufl2dot(form, labeling="compact")[0]
   dot = graphviz.Source(source)
   dot
   

2. 使用 ufl 中的 ufl2unicode 模块

   from ufl.formatting import ufl2unicode
   
   ustr = ufl2unicode.ufl2unicode(form)
   
   # The unicode string can not show correctly without space!
   for i in ustr:
       print(i, end=' ')
   

3. 使用 ufl 中的 printing.tree_format

   from ufl.formatting import printing
   
   printing.tree_format(form)
   

import graphviz
from ufl.formatting import ufl2dot

source = ufl2dot.ufl2dot(form, labeling="compact")[0]
dot = graphviz.Source(source)
filename = dot.render('figures/form', format='png')

Fig. 8.1 DAG of form

8.2. 生成局部内核

from firedrake.tsfc_interface import compile_form

kernels = compile_form(form, "subspace_form", split=False, interface=None)

通过 loopy 生成代码

import loopy

idx, kinfo = kernels[0]
code = loopy.generate_code_v2(kinfo[0].code)

输出代码

print(code.device_code())

8.2.1. 具体构造过程

from firedrake.parameters import parameters as default_parameters
from firedrake.tsfc_interface import TSFCKernel

nargs = len(form.arguments())

iterable = ([(None, )*nargs, form], )

idx, f = iterable[0]
name = "test_form"

coefficient_numbers = form.coefficient_numbering()
number_map = tuple(coefficient_numbers[c] for c in f.coefficients())
parameters = default_parameters["form_compiler"].copy()

kinfos = TSFCKernel(form=f, name=name, parameters=parameters, number_map=number_map, interface=None, diagonal=False).kernels

8.2.2. TSFCKernel 的初始化

from tsfc import compile_form as tsfc_compile_form
from firedrake.tsfc_interface import as_pyop2_local_kernel

import collections

KernelInfo = collections.namedtuple("KernelInfo",
                                    ["kernel",
                                     "integral_type",
                                     "oriented",
                                     "subdomain_id",
                                     "domain_number",
                                     "coefficient_map",
                                     "needs_cell_facets",
                                     "pass_layer_arg",
                                     "needs_cell_sizes",
                                     "arguments",
                                     "events"])

tree = tsfc_compile_form(form=f, prefix=name, parameters=parameters,
                                 interface=None,
                                 diagonal=False)

_kernels = []
for kernel in tree:
    # Unwind coefficient numbering
    numbers = tuple((number_map[number], indices) for number, indices in kernel.coefficient_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,
                              oriented=kernel.oriented,
                              subdomain_id=kernel.subdomain_id,
                              domain_number=kernel.domain_number,
                              coefficient_map=numbers,
                              needs_cell_facets=False,
                              pass_layer_arg=False,
                              needs_cell_sizes=kernel.needs_cell_sizes,
                              arguments=kernel.arguments,
                              events=events))

8.2.3. tsfc.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

complex_mode = parameters and is_complex(parameters.get("scalar_type"))

# Preprocess UFL form in a format suitable for TSFC
fd = ufl_utils.compute_form_data(form, complex_mode=complex_mode)

kernels = []
for integral_data in fd.integral_data:
    kernel = compile_integral(integral_data, fd, prefix=name, parameters=parameters, interface=None, diagonal=False)
    if kernel is not None:
        kernels.append(kernel)

fd.preprocessed_form.arguments()

(Argument(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f17b438c370>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f17b438c370>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 1)), 0, None),
 Argument(WithGeometry(MixedFunctionSpace(IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f17b438c370>, FiniteElement('Lagrange', triangle, 1), name=None, index=0, component=None), IndexedProxyFunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f17b438c370>, FiniteElement('Lagrange', triangle, 2), name=None, index=1, component=None), name='None_None'), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 1)), 1, None))

# print(fd.integral_data[0].integrals[0])
# print(tree_format(intd.integrals[0]))

8.2.4. tsfc 中的 compile_integral

编译 ufl 积分表达式为组装核

interface 为 None 时, 默认使用 KernelBuilder 进行构建.

from tsfc.kernel_interface.firedrake_loopy import KernelBuilder

builder.compile_integrand 调用 tsfc.fem.compile_ufl 转化 ufl 表达式为 GEM [HMLH18]. 这里会调用 FInAT 生成基函数在积分点处的求值公式.

builder.construct_integrals 应用数值积分公式生成相应表达式.

8.3. 生成全局内核

from pyop2 import op2
from firedrake.utils import cached_property, complex_mode, IntType
from firedrake.preconditioners.patch import LocalDat, LocalMat

import operator
import numpy

test, trial = map(operator.methodcaller("function_space"), form.arguments())

准备全局参数

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_map:
    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)

输出代码

print(mod.code_to_compile)

References

[HMLH18] (1,2)

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.

[KM19]

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.