MeshHierarchy: assemble prolongation into MatAIJ

firedrake/mg/embedded.py

@@ -5,6 +5,7 @@
 from enum import IntEnum
 from firedrake.petsc import PETSc
 from firedrake.embedding import get_embedding_dg_element
+from .interface import assemble_prolongation_aij
 from finat.element_factory import create_element
 
 __all__ = ("TransferManager", )
@@ -32,7 +33,7 @@ def __init__(self, ufl_element, value_shape):
             self._work_vec = {}
             self._V_dof_weights = {}
 
-    def __init__(self, *, native_transfers=None, use_averaging=True):
+    def __init__(self, *, native_transfers=None, use_averaging=True, mat_type="matfree"):
         """
         An object for managing transfers between levels in a multigrid
         hierarchy (possibly via embedding in DG spaces).
@@ -43,10 +44,13 @@ def __init__(self, *, native_transfers=None, use_averaging=True):
         :arg use_averaging: Use averaging to approximate the
            projection out of the embedded DG space? If False, a global
            L2 projection will be performed.
+        :arg mat_type: The matrix assembly type for prolongation/restriction.
         """
         self.native_transfers = native_transfers or {}
         self.use_averaging = use_averaging
         self.caches = {}
+        self.mat_type = mat_type
+        self._mat_cache = {}
 
     def is_native(self, element, gdim, op):
         if element in self.native_transfers:
@@ -78,7 +82,12 @@ def _native_transfer(self, element, gdim, op):
             return self.native_transfers[element][op]
         except KeyError:
             if self.is_native(element, gdim, op):
-                ops = firedrake.prolong, firedrake.restrict, firedrake.inject
+                if self.mat_type == "aij":
+                    ops = self.prolong_aij, self.restrict_aij, firedrake.inject
+                elif self.mat_type == "matfree":
+                    ops = firedrake.prolong, firedrake.restrict, firedrake.inject
+                else:
+                    raise ValueError(f"Unsupported mat_type {self.mat_type}")
                 return self.native_transfers.setdefault(element, ops)[op]
         return None
 
@@ -362,3 +371,25 @@ def restrict(self, source, target):
                 self.DG_inv_mass(VDGt).mult(dgv, dgwork)
                 self.V_DG_mass(Vt, VDGt).multTranspose(dgwork, t)
         self.cache_dat_versions(Vs_star, Op.RESTRICT, source, target)
+
+    def prolongation_matrix(self, Vc, Vf):
+        key = (Vc, Vf)
+        try:
+            return self._mat_cache[key]
+        except KeyError:
+            P = assemble_prolongation_aij(Vc, Vf)
+            return self._mat_cache.setdefault(key, P)
+
+    def prolong_aij(self, uc, uf):
+        Vc = uc.function_space()
+        Vf = uf.function_space()
+        P = self.prolongation_matrix(Vc, Vf)
+        with uc.dat.vec_ro as x, uf.dat.vec_wo as y:
+            P.petscmat.mult(x, y)
+
+    def restrict_aij(self, rf, rc):
+        Vc = rc.function_space().dual()
+        Vf = rf.function_space().dual()
+        P = self.prolongation_matrix(Vc, Vf)
+        with rf.dat.vec_ro as x, rc.dat.vec_wo as y:
+            P.petscmat.multTranspose(x, y)

firedrake/mg/interface.py

@@ -1,15 +1,19 @@
 from pyop2 import op2
+from fractions import Fraction
 
 from firedrake import ufl_expr, dmhooks
+from firedrake.assemble import assemble
+from firedrake.interpolation import interpolate
 from firedrake.function import Function
 from firedrake.cofunction import Cofunction
+from firedrake.matrix import AssembledMatrix
 from firedrake.petsc import PETSc
 from ufl.duals import is_dual
 from . import utils
 from . import kernels
 
 
-__all__ = ["prolong", "restrict", "inject"]
+__all__ = ["prolong", "restrict", "inject", "assemble_prolongation_aij"]
 
 
 def check_arguments(coarse, fine, needs_dual=False):
@@ -290,3 +294,74 @@ def inject(fine, coarse):
             coarse = new_coarse.interpolate(coarse)
         fine = coarse
     return coarse
+
+
+def _bc_matches_space(bc, V):
+    fs = bc.function_space()
+    while fs.component is not None and fs.parent is not None:
+        fs = fs.parent
+    return fs == V
+
+
+@PETSc.Log.EventDecorator()
+def assemble_prolongation_aij(Vc, Vf, bcs=None):
+    if len(Vc) > 1 or len(Vc) > 1:
+        raise NotImplementedError("Mixed spaces are handled through TransferManager")
+    arguments = (ufl_expr.TestFunction(Vf.dual()), ufl_expr.TrialFunction(Vc))
+    Vtarget = Vf
+    if needs_quadrature := not Vf.finat_element.has_pointwise_dual_basis:
+        # Introduce an intermediate quadrature target space
+        Vf = Vf.quadrature_space()
+    Vrow, Vcol = Vf, Vc
+
+    hierarchy, levelc = utils.get_level(Vcol.mesh())
+    hierarchyf, levelf = utils.get_level(Vrow.mesh())
+    if hierarchy is None or hierarchy is not hierarchyf:
+        raise ValueError("Mismatching hierarchies")
+    if levelc + Fraction(1, hierarchy.refinements_per_level) != levelf:
+        raise ValueError("Only implemented on consecutive levels")
+
+    row_map = utils.identity_node_map(Vrow)
+    col_map = utils.fine_node_to_coarse_node_map(Vrow, Vcol)
+    sparsity = op2.Sparsity((Vrow.dof_dset, Vcol.dof_dset),
+                            [(row_map, col_map, None)],
+                            name=f"{Vrow.name}_{Vcol.name}_hierarchy_interpolation_sparsity",
+                            nest=False,
+                            block_sparse=False)
+    mat = op2.Mat(sparsity)
+
+    lgmaps = None
+    if bcs:
+        row_bcs = [bc for bc in bcs if _bc_matches_space(bc, Vrow)]
+        col_bcs = [bc for bc in bcs if _bc_matches_space(bc, Vcol)]
+        if row_bcs or col_bcs:
+            lgmaps = [(Vrow.local_to_global_map(row_bcs), Vcol.local_to_global_map(col_bcs))]
+
+    kernel = kernels.prolong_matrix_kernel(Vcol, Vrow)
+    node_locations = utils.physical_node_locations(Vrow)
+    source_mesh = Vcol.mesh()
+    source_coords = source_mesh.coordinates
+    compose_map = lambda u: utils.fine_node_to_coarse_node_map(Vrow, u.function_space())
+    kernel_args = [
+        mat(op2.INC, (row_map, col_map), lgmaps=lgmaps),
+        node_locations.dat(op2.READ),
+        source_coords.dat(op2.READ, compose_map(source_coords)),
+    ]
+    if kernel.oriented:
+        co = source_mesh.cell_orientations()
+        kernel_args.append(co.dat(op2.READ, compose_map(co)))
+    if kernel.needs_cell_sizes:
+        cs = source_mesh.cell_sizes
+        kernel_args.append(cs.dat(op2.READ, compose_map(cs)))
+    source_coords.dat.global_to_local_begin(op2.READ)
+    source_coords.dat.global_to_local_end(op2.READ)
+    op2.par_loop(kernel, Vrow.node_set, *kernel_args)
+    mat.assemble()
+    result = mat.handle
+
+    if needs_quadrature:
+        interp = interpolate(ufl_expr.TrialFunction(Vf), Vtarget)
+        Q = assemble(interp, bcs=bcs, mat_type="aij").petscmat
+        result = Q.matMult(result)
+
+    return AssembledMatrix(arguments, result, bcs=bcs)

firedrake/mg/kernels.py

@@ -253,6 +253,119 @@ def prolong_kernel(expression, Vf):
         return cache.setdefault(key, transfer_kernel)
 
 
+def prolong_matrix_kernel(Vc, Vf):
+    hierarchy, levelf = utils.get_level(Vf.mesh())
+    hierarchy, levelc = utils.get_level(Vc.mesh())
+    if Vc.mesh().extruded:
+        assert Vf.mesh().extruded
+        level_ratio = (Vc.mesh().layers - 1) // (Vf.mesh().layers - 1)
+    else:
+        level_ratio = 1
+    if levelf <= levelc:
+        raise ValueError("Can only build hierarchy interpolation matrices from coarse to fine spaces")
+    ncandidate = hierarchy.fine_to_coarse_cells[levelf].shape[1] * level_ratio
+    coordinates = Vc.mesh().coordinates
+    key = (("prolong_matrix", ncandidate)
+           + (Vf.block_size,)
+           + _make_element_key(Vf.finat_element)
+           + _make_element_key(Vc.finat_element)
+           + _make_element_key(coordinates.function_space().finat_element))
+    cache = hierarchy._shared_data_cache["transfer_kernels"]
+    try:
+        return cache[key]
+    except KeyError:
+        kernel = dual_evaluation_kernel(ufl.TrialFunction(Vc), ufl.TestFunction(Vf.dual()))
+        evaluate_code = lp.generate_code_v2(kernel.ast).device_code()
+        to_reference_kernel = to_reference_coordinates(coordinates.ufl_element())
+        coords_element = create_element(coordinates.ufl_element())
+        element = create_element(Vc.ufl_element())
+        num_verts = len(element.cell.get_vertices())
+        row_dim = Vf.block_size
+        source_cell_inc = element.space_dimension()
+        source_stencil_inc = ncandidate * source_cell_inc
+        local_tensor_size = row_dim * source_stencil_inc
+        cell_tensor_size = row_dim * source_cell_inc
+
+        kernel_code = """#include <petsc.h>
+        %(to_reference)s
+        %(evaluate)s
+        __attribute__((noinline)) /* Clang bug */
+        static void pyop2_kernel_prolong_matrix(PetscScalar *A, const PetscScalar *X, const PetscScalar *Xc
+                                                %(cell_orient)s%(cell_sizes)s)
+        {
+            PetscScalar Xref[%(tdim)d];
+            PetscScalar B[%(cell_tensor_size)d];
+            int cell = -1;
+            int bestcell = -1;
+            double bestdist = 1e10;
+            for (int i = 0; i < %(local_tensor_size)d; i++) {
+                A[i] = 0;
+            }
+            for (int i = 0; i < %(cell_tensor_size)d; i++) {
+                B[i] = 0;
+            }
+            for (int i = 0; i < %(ncandidate)d; i++) {
+                const PetscScalar *Xci = Xc + i*%(Xc_cell_inc)d;
+                double celldist = 2*bestdist;
+                to_reference_coords_kernel(Xref, X, Xci);
+                if (%(inside_cell)s) {
+                    cell = i;
+                    break;
+                }
+
+                celldist = %(celldist_l1_c_expr)s;
+                if (celldist < bestdist) {
+                    bestdist = celldist;
+                    bestcell = i;
+                }
+
+            }
+            if (cell == -1) {
+                /* We didn't find a cell that contained this point exactly.
+                   Did we find one that was close enough? */
+                if (bestdist < 10) {
+                    cell = bestcell;
+                } else {
+                    fprintf(stderr, "Could not identify cell in transfer operator. Point: ");
+                    for (int coord = 0; coord < %(tdim)s; coord++) {
+                      fprintf(stderr, "%%.14e ", X[coord]);
+                    }
+                    fprintf(stderr, "\\n");
+                    fprintf(stderr, "Number of candidates: %%d. Best distance located: %%14e", %(ncandidate)d, bestdist);
+                    abort();
+                }
+            }
+            const PetscScalar *Xci = Xc + cell*%(Xc_cell_inc)d;
+            pyop2_kernel_evaluate(%(kernel_args)s);
+            for (int i = 0; i < %(row_dim)d; i++) {
+                for (int j = 0; j < %(source_cell_inc)d; j++) {
+                    A[i*%(source_stencil_inc)d + cell*%(source_cell_inc)d + j] =
+                        B[i*%(source_cell_inc)d + j];
+                }
+            }
+        }
+        """ % {"to_reference": str(to_reference_kernel),
+               "evaluate": evaluate_code,
+               "cell_orient": ", const PetscScalar *co" if kernel.oriented else "",
+               "cell_sizes": ", const PetscScalar *cs" if kernel.needs_cell_sizes else "",
+               "kernel_args": _make_kernel_args(kernel, element, "B", "co+cell", f"cs+cell*{num_verts}", "Xci", "Xref"),
+               "ncandidate": ncandidate,
+               "row_dim": row_dim,
+               "source_cell_inc": source_cell_inc,
+               "source_stencil_inc": source_stencil_inc,
+               "cell_tensor_size": cell_tensor_size,
+               "local_tensor_size": local_tensor_size,
+               "inside_cell": inside_check(element.cell, eps=1e-8, X="Xref"),
+               "celldist_l1_c_expr": celldist_l1_c_expr(element.cell, X="Xref"),
+               "Xc_cell_inc": coords_element.space_dimension(),
+               "tdim": element.cell.get_spatial_dimension()}
+
+        transfer_kernel = op2.Kernel(kernel_code, name="pyop2_kernel_prolong_matrix")
+        transfer_kernel.oriented = kernel.oriented
+        transfer_kernel.needs_cell_sizes = kernel.needs_cell_sizes
+        return cache.setdefault(key, transfer_kernel)
+
+
 def restrict_kernel(Vf, Vc):
     hierarchy, levelf = utils.get_level(Vf.mesh())
     if Vf.mesh().extruded:

firedrake/mg/utils.py

@@ -8,6 +8,16 @@
 from firedrake.cython import mgimpl as impl
 
 
+def identity_node_map(V):
+    cache = V.mesh()._shared_data_cache["hierarchy_identity_node_map"]
+    key = (V.ufl_element(), V.boundary_set)
+    try:
+        return cache[key]
+    except KeyError:
+        values = numpy.arange(V.node_set.total_size, dtype=IntType).reshape(-1, 1)
+        return cache.setdefault(key, op2.Map(V.node_set, V.node_set, 1, values=values))
+
+
 def fine_node_to_coarse_node_map(Vf, Vc):
     if len(Vf) > 1:
         assert len(Vf) == len(Vc)

firedrake/solving_utils.py

@@ -314,16 +314,28 @@ def transfer_manager(self):
         if self._transfer_manager is None:
             opts = PETSc.Options()
             prefix = self.options_prefix or ""
-            if opts.hasName(prefix + "mg_transfer_manager"):
-                managername = opts[prefix + "mg_transfer_manager"]
-            elif opts.hasName(prefix + "fas_transfer_manager"):
-                managername = opts[prefix + "fas_transfer_manager"]
-            else:
-                managername = None
 
+            def get_transfer_option(mg_name, fas_name, default=None):
+                mg_name = prefix + mg_name
+                fas_name = prefix + fas_name
+                has_mg = opts.hasName(mg_name)
+                has_fas = opts.hasName(fas_name)
+                if has_mg and has_fas:
+                    warning(f"Both '{mg_name}' and '{fas_name}' options were supplied; "
+                            f"ignoring '{fas_name}'.")
+                if has_mg:
+                    return opts[mg_name]
+                elif has_fas:
+                    return opts[fas_name]
+                else:
+                    return default
+
+            managername = get_transfer_option("mg_transfer_manager", "fas_transfer_manager")
             if managername is None:
                 from firedrake import TransferManager
-                transfer = TransferManager(use_averaging=True)
+                mat_type = get_transfer_option("mg_transfer_mat_type", "fas_transfer_mat_type",
+                                               default="matfree")
+                transfer = TransferManager(use_averaging=True, mat_type=mat_type)
             else:
                 (modname, objname) = managername.rsplit('.', 1)
                 mod = __import__(modname)