diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index ddd9d583b1..2f86083a46 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -167,6 +167,14 @@ runs: firedrake-clean pip list + - name: DROP BEFORE MERGE - install the fiat#268 stack + shell: bash + run: | + . venv/bin/activate + pip install --ignore-installed --no-deps --no-build-isolation \ + git+https://github.com/firedrakeproject/fiat.git@pbrubeck/fix/dual-enriched + firedrake-clean + - name: Run firedrake-check shell: bash run: | diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 16bed26f60..586d1707fe 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -76,11 +76,26 @@ def __iter__(self): yield self yield from itertools.chain(*self.bcs) - def function_space(self): + def function_space(self, parent=False): '''The :class:`.FunctionSpace` on which this boundary condition should - be applied.''' - - return self._function_space + be applied. + + Parameters + ---------- + parent : bool + If ``True``, walk up through any indexed or component subspaces + and return the top-level function space instead. + + Returns + ------- + firedrake.functionspaceimpl.WithGeometry + The function space. + ''' + V = self._function_space + if parent: + while V.parent is not None: + V = V.parent + return V def function_space_index(self): fs = self._function_space diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 58823f5451..b9b1d9fcf1 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -351,9 +351,16 @@ def assemble( self._check_mat_type(mat_type) if mat_type == "matfree" and self.rank == 2: + Vrow, Vcol = (arg.function_space() for arg in self.interpolate_args) + Vcol = Vcol.dual() + if bcs is None: + bcs = () + row_bcs = [bc for bc in bcs if bc.function_space(parent=True).topological == Vrow.topological] + col_bcs = [bc for bc in bcs if bc.function_space(parent=True).topological == Vcol.topological] ctx = ImplicitMatrixContext( - self.ufl_interpolate, row_bcs=bcs, col_bcs=bcs, + self.ufl_interpolate, row_bcs=row_bcs, col_bcs=col_bcs, ) + ctx.on_diag = Vrow == Vcol return ImplicitMatrix(self.ufl_interpolate, ctx, bcs=bcs) result = self._get_callable(tensor=tensor, bcs=bcs, mat_type=mat_type, sub_mat_type=sub_mat_type)() diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 89040eda7a..dc13582bc7 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -4,17 +4,14 @@ from firedrake.petsc import PETSc from firedrake.preconditioners.base import PCBase from firedrake.preconditioners.patch import bcdofs -from firedrake.preconditioners.pmg import (prolongation_matrix_matfree, - evaluate_dual, - get_permutation_to_nodal_elements, - cache_generate_code) from firedrake.preconditioners.facet_split import restricted_dofs, split_dofs from firedrake.formmanipulation import ExtractSubBlock from firedrake.functionspace import FunctionSpace, MixedFunctionSpace from firedrake.function import Function from firedrake.cofunction import Cofunction from firedrake.parloops import par_loop -from firedrake.ufl_expr import TestFunction, TestFunctions, TrialFunctions +from firedrake.ufl_expr import TestFunction, TrialFunction, TestFunctions, TrialFunctions +from firedrake.interpolation import interpolate from ufl.algorithms.ad import expand_derivatives from ufl.algorithms.expand_indices import expand_indices from finat.element_factory import create_element @@ -22,6 +19,7 @@ from pyop2.mpi import COMM_SELF from pyop2.sparsity import get_preallocation from pyop2.utils import as_tuple +from pyop2.caching import serial_cache from pyop2 import op2 from tsfc.ufl_utils import extract_firedrake_constants from firedrake.tsfc_interface import compile_form @@ -32,8 +30,11 @@ import finat.ufl import FIAT import finat +import loopy import numpy import ctypes +import os +import tempfile __all__ = ("FDMPC", "PoissonFDMPC") @@ -66,6 +67,7 @@ class FDMPC(PCBase): @PETSc.Log.EventDecorator("FDMInit") def initialize(self, pc): + from firedrake.assemble import assemble petsctools.cite(self._citation) self.comm = pc.comm Amat, Pmat = pc.getOperators() @@ -115,19 +117,15 @@ def initialize(self, pc): else: # Reconstruct Jacobian and bcs with variant element J_fdm = J(*(t.reconstruct(function_space=V_fdm) for t in J.arguments())) - bcs_fdm = [] - for bc in bcs: - W = V_fdm - for index in bc._indices: - W = W.sub(index) - bcs_fdm.append(bc.reconstruct(V=W, g=0)) + bcs_fdm = [bc.reconstruct(V=V_fdm, g=0, indices=bc._indices) for bc in bcs] # Create a new _SNESContext in the variant space self._ctx_ref = self.new_snes_ctx(pc, J_fdm, bcs_fdm, mat_type, fcp=fcp, options_prefix=options_prefix) # Construct interpolation from variant to original spaces - self.fdm_interp = prolongation_matrix_matfree(V_fdm, V, bcs_fdm, []) + interp = interpolate(TrialFunction(V_fdm), V) + self.fdm_interp = assemble(interp, bcs=bcs_fdm, mat_type="matfree").petscmat self.work_vec_x = Amat.createVecLeft() self.work_vec_y = Amat.createVecRight() if use_amat: @@ -2536,3 +2534,213 @@ def vector_map(bsize, ibase, e, result=None): ibase = numpy.arange(bsize, dtype=node_map.values.dtype) return partial(vector_map, bsize, ibase), nel + + +def cache_generate_code(kernel, comm): + _cachedir = os.environ.get('PYOP2_CACHE_DIR', + os.path.join(tempfile.gettempdir(), + 'pyop2-cache-uid%d' % os.getuid())) + + key = kernel.cache_key[0] + shard, disk_key = key[:2], key[2:] + filepath = os.path.join(_cachedir, shard, disk_key) + if os.path.exists(filepath): + with open(filepath, 'r') as f: + code = f.read() + else: + code = loopy.generate_code_v2(kernel.code).device_code() + if comm.rank == 0: + os.makedirs(os.path.join(_cachedir, shard), exist_ok=True) + with open(filepath, 'w') as f: + f.write(code) + comm.barrier() + return code + + +def hash_fiat_element(element): + """FIAT elements are not hashable, + this is not the best way to create a hash""" + restriction = None + e = element + if isinstance(e, FIAT.DiscontinuousElement): + # this hash does not care about inter-element continuity + e = e._element + if isinstance(e, FIAT.RestrictedElement): + restriction = tuple(e._indices) + e = e._element + if len(restriction) == e.space_dimension(): + restriction = None + family = e.__class__.__name__ + degree = e.order + return (family, element.ref_el, degree, restriction) + + +def generate_key_evaluate_dual(source, target, derivative=None): + return hash_fiat_element(source) + hash_fiat_element(target) + (derivative,) + + +@serial_cache(hashkey=generate_key_evaluate_dual) +def compare_element(e1, e2): + """Numerically compare two :class:`FIAT.elements`. + Equality is satisfied if e2.dual_basis(e1.primal_basis) == identity.""" + if e1 is e2: + return True + if e1.space_dimension() != e2.space_dimension(): + return False + B = evaluate_dual(e1, e2) + return numpy.allclose(B, numpy.eye(B.shape[0]), rtol=1E-14, atol=1E-14) + + +def expand_element(ele): + """Expand a FiniteElement as an EnrichedElement of TensorProductElements, + discarding modifiers.""" + if isinstance(ele, finat.FlattenedDimensions): + return expand_element(ele.product) + elif isinstance(ele, (finat.HDivElement, finat.HCurlElement)): + return expand_element(ele.wrappee) + elif isinstance(ele, finat.DiscontinuousElement): + return expand_element(ele.element) + elif isinstance(ele, finat.EnrichedElement): + terms = list(map(expand_element, ele.elements)) + return finat.EnrichedElement(terms) + elif isinstance(ele, finat.TensorProductElement): + factors = list(map(expand_element, ele.factors)) + terms = [tuple()] + for e in factors: + new_terms = [] + for f in e.elements if isinstance(e, finat.EnrichedElement) else [e]: + f_factors = tuple(f.factors) if isinstance(f, finat.TensorProductElement) else (f,) + new_terms.extend(t_factors + f_factors for t_factors in terms) + terms = new_terms + terms = list(map(finat.TensorProductElement, terms)) + return finat.EnrichedElement(terms) + else: + return ele + + +@serial_cache(hashkey=lambda V: V.ufl_element()) +@PETSc.Log.EventDecorator("GetLineElements") +def get_permutation_to_nodal_elements(V): + """Find DOF permutation to factor out the EnrichedElement expansion + into common TensorProductElements. + + This routine exposes structure to e.g vectorize + prolongation of NCE or NCF accross vector components, by permuting all + components into a common TensorProductElement. + + This is temporary while we wait for dual evaluation of :class:`finat.EnrichedElement`. + + Parameters + ---------- + V : + A :class:`.FunctionSpace`. + + Returns + ------- + A 3-tuple of the DOF permutation, the unique terms in expansion + as a list of tuples of :class:`FIAT.FiniteElements`, and the cyclic + permutations of the axes to form the element given by their shifts + in list of `int` tuples + """ + finat_element = V.finat_element + expansion = expand_element(finat_element) + if expansion.space_dimension() != finat_element.space_dimension(): + raise ValueError("Failed to decompose %s into tensor products" % V.ufl_element()) + + nodal_elements = [] + terms = expansion.elements if hasattr(expansion, "elements") else [expansion] + for term in terms: + factors = term.factors if hasattr(term, "factors") else (term,) + fiat_factors = tuple(e.fiat_equivalent for e in reversed(factors)) + if not all(e.is_nodal() for e in fiat_factors): + raise ValueError("Failed to decompose %s into nodal elements" % V.ufl_element()) + nodal_elements.append(fiat_factors) + + shapes = [tuple(e.space_dimension() for e in factors) for factors in nodal_elements] + sizes = list(map(numpy.prod, shapes)) + dof_ranges = numpy.cumsum([0] + sizes) + + dof_perm = [] + unique_nodal_elements = [] + shifts = [] + + visit = [False for e in nodal_elements] + while False in visit: + base = nodal_elements[visit.index(False)] + tdim = len(base) + pshape = tuple(e.space_dimension() for e in base) + unique_nodal_elements.append(base) + + axes_shifts = tuple() + for shift in range(tdim): + if finat_element.formdegree != 2: + shift = (tdim - shift) % tdim + + perm = base[shift:] + base[:shift] + for i, term in enumerate(nodal_elements): + if not visit[i]: + is_perm = all(e1.space_dimension() == e2.space_dimension() + for e1, e2 in zip(perm, term)) + if is_perm: + is_perm = all(compare_element(e1, e2) for e1, e2 in zip(perm, term)) + + if is_perm: + axes_shifts += ((tdim - shift) % tdim, ) + dofs = numpy.arange(*dof_ranges[i:i+2], dtype=PETSc.IntType).reshape(pshape) + dofs = numpy.transpose(dofs, axes=numpy.roll(numpy.arange(tdim), -shift)) + assert dofs.shape == shapes[i] + dof_perm.append(dofs.flat) + visit[i] = True + + shifts.append(axes_shifts) + + dof_perm = get_readonly_view(numpy.concatenate(dof_perm)) + return dof_perm, unique_nodal_elements, shifts + + +def get_readonly_view(arr): + result = arr.view() + result.flags.writeable = False + return result + + +@serial_cache(hashkey=generate_key_evaluate_dual) +def evaluate_dual(source, target, derivative=None): + """Evaluate the action of a set of dual functionals of the target element + on the (derivative of the) basis functions of the source element. + + Parameters + ---------- + source : + A :class:`FIAT.CiarletElement` to interpolate. + target : + A :class:`FIAT.CiarletElement` defining the interpolation space. + derivative : ``str`` or ``None`` + An optional differential operator to apply on the source expression, + either "grad", "curl", or "div". + + Returns + ------- + A read-only :class:`numpy.ndarray` with the evaluation of the target + dual basis on the (derivative of the) source primal basis. + """ + primal = source.get_nodal_basis() + dual = target.get_dual_set() + A = dual.to_riesz(primal) + B = primal.get_coeffs() + if derivative in ("grad", "curl", "div"): + dmats = primal.get_dmats() + B = numpy.tensordot(B, dmats, axes=(-1, -1)) + if derivative == "curl": + d = B.shape[1] + idx = ((i, j) for i in reversed(range(d)) for j in reversed(range(i+1, d))) + B = numpy.stack([((-1)**k) * (B[:, i, j, :] - B[:, j, i, :]) + for k, (i, j) in enumerate(idx)], axis=1) + elif derivative == "div": + B = numpy.trace(B, axis1=1, axis2=2) + elif derivative is not None: + raise ValueError(f"Invalid derivative type {derivative}.") + + B = B.reshape(-1, *A.shape[1:]) + V = numpy.tensordot(A, B, axes=(range(1, A.ndim), range(1, B.ndim))) + return get_readonly_view(V) diff --git a/firedrake/preconditioners/pmg.py b/firedrake/preconditioners/pmg.py index 5fee3d9429..bd2e646a53 100644 --- a/firedrake/preconditioners/pmg.py +++ b/firedrake/preconditioners/pmg.py @@ -1,28 +1,16 @@ -from functools import cached_property, partial -from itertools import chain +from functools import partial from firedrake.dmhooks import (attach_hooks, get_appctx, push_appctx, pop_appctx, add_hook, get_parent, push_parent, pop_parent, get_function_space, set_function_space) from firedrake.petsc import PETSc from firedrake.preconditioners.base import PCBase, SNESBase, PCSNESBase from firedrake.solving_utils import _SNESContext -from firedrake.tsfc_interface import extract_numbered_coefficients -from firedrake.utils import IntType_c -from tsfc import compile_expression_dual_evaluation -from pyop2 import op2 -from pyop2.caching import serial_cache from pyop2.utils import as_tuple import firedrake import finat -import FIAT import ufl import finat.ufl -import loopy -import numpy -import os -import tempfile -import weakref __all__ = ("PMGPC", "PMGSNES") @@ -57,8 +45,6 @@ class PMGBase(PCSNESBase): """ _prefix = "pmg_" - # This is parallel-safe because the keys are ids of a collective objects - _transfer_cache = weakref.WeakKeyDictionary() def coarsen_element(self, ele: finat.ufl.FiniteElementBase) -> finat.ufl.FiniteElementBase: """Coarsen a given element to form the next problem down in the p-hierarchy. @@ -337,22 +323,13 @@ def coarsen_quadrature(metadata: dict | None, fdeg: int, cdeg: int) -> dict | No def create_transfer(self, mat_type, cctx, fctx, cbcs, fbcs): """Create a transfer operator""" - cache = self._transfer_cache.setdefault(fctx, {}) - key = (mat_type, cctx, cbcs, fbcs) - try: - return cache[key] - except KeyError: - if mat_type == "matfree": - construct_mat = prolongation_matrix_matfree - elif mat_type == "aij": - construct_mat = prolongation_matrix_aij - else: - raise ValueError("Unknown matrix type") - cV = cctx._problem.u_restrict.function_space() - fV = fctx._problem.u_restrict.function_space() - cbcs = tuple(cctx._problem.bcs) if cbcs else tuple() - fbcs = tuple(fctx._problem.bcs) if fbcs else tuple() - return cache.setdefault(key, construct_mat(cV, fV, cbcs, fbcs)) + cV = cctx._problem.u_restrict.function_space() + fV = fctx._problem.u_restrict.function_space() + cbcs = tuple(cctx._problem.bcs) if cbcs else tuple() + fbcs = tuple(fctx._problem.bcs) if fbcs else tuple() + bcs = cbcs + fbcs + interp = firedrake.interpolate(firedrake.TrialFunction(cV), fV) + return firedrake.assemble(interp, bcs=bcs, mat_type=mat_type).petscmat def create_interpolation(self, dmc, dmf): prefix = dmc.getOptionsPrefix() @@ -450,1035 +427,3 @@ def step(self, snes, x, f, y): y.aypx(-1, x) snes.setConvergedReason(self.ppc.getConvergedReason()) pop_appctx(self.ppc.dm) - - -def prolongation_transfer_kernel_action(Vf, expr): - kernel = compile_expression_dual_evaluation(expr, Vf.ufl_element()) - coefficients = extract_numbered_coefficients(expr, kernel.coefficient_numbers) - if kernel.needs_external_coords: - coefficients = [Vf.mesh().coordinates] + coefficients - - return op2.Kernel(kernel.ast, kernel.name, - requires_zeroed_output_arguments=True, - flop_count=kernel.flop_count, - events=(kernel.event,)), coefficients - - -def expand_element(ele): - """Expand a FiniteElement as an EnrichedElement of TensorProductElements, - discarding modifiers.""" - if isinstance(ele, finat.FlattenedDimensions): - return expand_element(ele.product) - elif isinstance(ele, (finat.HDivElement, finat.HCurlElement)): - return expand_element(ele.wrappee) - elif isinstance(ele, finat.DiscontinuousElement): - return expand_element(ele.element) - elif isinstance(ele, finat.EnrichedElement): - terms = list(map(expand_element, ele.elements)) - return finat.EnrichedElement(terms) - elif isinstance(ele, finat.TensorProductElement): - factors = list(map(expand_element, ele.factors)) - terms = [tuple()] - for e in factors: - new_terms = [] - for f in e.elements if isinstance(e, finat.EnrichedElement) else [e]: - f_factors = tuple(f.factors) if isinstance(f, finat.TensorProductElement) else (f,) - new_terms.extend(t_factors + f_factors for t_factors in terms) - terms = new_terms - terms = list(map(finat.TensorProductElement, terms)) - return finat.EnrichedElement(terms) - else: - return ele - - -def hash_fiat_element(element): - """FIAT elements are not hashable, - this is not the best way to create a hash""" - restriction = None - e = element - if isinstance(e, FIAT.DiscontinuousElement): - # this hash does not care about inter-element continuity - e = e._element - if isinstance(e, FIAT.RestrictedElement): - restriction = tuple(e._indices) - e = e._element - if len(restriction) == e.space_dimension(): - restriction = None - family = e.__class__.__name__ - degree = e.order - return (family, element.ref_el, degree, restriction) - - -def generate_key_evaluate_dual(source, target, derivative=None): - return hash_fiat_element(source) + hash_fiat_element(target) + (derivative,) - - -def get_readonly_view(arr): - result = arr.view() - result.flags.writeable = False - return result - - -@serial_cache(hashkey=generate_key_evaluate_dual) -def evaluate_dual(source, target, derivative=None): - """Evaluate the action of a set of dual functionals of the target element - on the (derivative of the) basis functions of the source element. - - Parameters - ---------- - source : - A :class:`FIAT.CiarletElement` to interpolate. - target : - A :class:`FIAT.CiarletElement` defining the interpolation space. - derivative : ``str`` or ``None`` - An optional differential operator to apply on the source expression, - either "grad", "curl", or "div". - - Returns - ------- - A read-only :class:`numpy.ndarray` with the evaluation of the target - dual basis on the (derivative of the) source primal basis. - """ - primal = source.get_nodal_basis() - dual = target.get_dual_set() - A = dual.to_riesz(primal) - B = primal.get_coeffs() - if derivative in ("grad", "curl", "div"): - dmats = primal.get_dmats() - B = numpy.tensordot(B, dmats, axes=(-1, -1)) - if derivative == "curl": - d = B.shape[1] - idx = ((i, j) for i in reversed(range(d)) for j in reversed(range(i+1, d))) - B = numpy.stack([((-1)**k) * (B[:, i, j, :] - B[:, j, i, :]) - for k, (i, j) in enumerate(idx)], axis=1) - elif derivative == "div": - B = numpy.trace(B, axis1=1, axis2=2) - elif derivative is not None: - raise ValueError(f"Invalid derivative type {derivative}.") - - B = B.reshape(-1, *A.shape[1:]) - V = numpy.tensordot(A, B, axes=(range(1, A.ndim), range(1, B.ndim))) - return get_readonly_view(V) - - -@serial_cache(hashkey=generate_key_evaluate_dual) -def compare_element(e1, e2): - """Numerically compare two :class:`FIAT.elements`. - Equality is satisfied if e2.dual_basis(e1.primal_basis) == identity.""" - if e1 is e2: - return True - if e1.space_dimension() != e2.space_dimension(): - return False - B = evaluate_dual(e1, e2) - return numpy.allclose(B, numpy.eye(B.shape[0]), rtol=1E-14, atol=1E-14) - - -@serial_cache(hashkey=lambda V: V.ufl_element()) -@PETSc.Log.EventDecorator("GetLineElements") -def get_permutation_to_nodal_elements(V): - """Find DOF permutation to factor out the EnrichedElement expansion - into common TensorProductElements. - - This routine exposes structure to e.g vectorize - prolongation of NCE or NCF accross vector components, by permuting all - components into a common TensorProductElement. - - This is temporary while we wait for dual evaluation of :class:`finat.EnrichedElement`. - - Parameters - ---------- - V : - A :class:`.FunctionSpace`. - - Returns - ------- - A 3-tuple of the DOF permutation, the unique terms in expansion - as a list of tuples of :class:`FIAT.FiniteElements`, and the cyclic - permutations of the axes to form the element given by their shifts - in list of `int` tuples - """ - finat_element = V.finat_element - expansion = expand_element(finat_element) - if expansion.space_dimension() != finat_element.space_dimension(): - raise ValueError("Failed to decompose %s into tensor products" % V.ufl_element()) - - nodal_elements = [] - terms = expansion.elements if hasattr(expansion, "elements") else [expansion] - for term in terms: - factors = term.factors if hasattr(term, "factors") else (term,) - fiat_factors = tuple(e.fiat_equivalent for e in reversed(factors)) - if not all(e.is_nodal() for e in fiat_factors): - raise ValueError("Failed to decompose %s into nodal elements" % V.ufl_element()) - nodal_elements.append(fiat_factors) - - shapes = [tuple(e.space_dimension() for e in factors) for factors in nodal_elements] - sizes = list(map(numpy.prod, shapes)) - dof_ranges = numpy.cumsum([0] + sizes) - - dof_perm = [] - unique_nodal_elements = [] - shifts = [] - - visit = [False for e in nodal_elements] - while False in visit: - base = nodal_elements[visit.index(False)] - tdim = len(base) - pshape = tuple(e.space_dimension() for e in base) - unique_nodal_elements.append(base) - - axes_shifts = tuple() - for shift in range(tdim): - if finat_element.formdegree != 2: - shift = (tdim - shift) % tdim - - perm = base[shift:] + base[:shift] - for i, term in enumerate(nodal_elements): - if not visit[i]: - is_perm = all(e1.space_dimension() == e2.space_dimension() - for e1, e2 in zip(perm, term)) - if is_perm: - is_perm = all(compare_element(e1, e2) for e1, e2 in zip(perm, term)) - - if is_perm: - axes_shifts += ((tdim - shift) % tdim, ) - dofs = numpy.arange(*dof_ranges[i:i+2], dtype=PETSc.IntType).reshape(pshape) - dofs = numpy.transpose(dofs, axes=numpy.roll(numpy.arange(tdim), -shift)) - assert dofs.shape == shapes[i] - dof_perm.append(dofs.flat) - visit[i] = True - - shifts.append(axes_shifts) - - dof_perm = get_readonly_view(numpy.concatenate(dof_perm)) - return dof_perm, unique_nodal_elements, shifts - - -def get_permuted_map(V): - """ - Return a PermutedMap with the same tensor product shape for - every component of H(div) or H(curl) tensor product elements - """ - indices, _, _ = get_permutation_to_nodal_elements(V) - if numpy.all(indices[:-1] < indices[1:]): - return V.cell_node_map() - return op2.PermutedMap(V.cell_node_map(), indices) - - -# Common kernel to compute y = kron(A3, kron(A2, A1)) * x -# Vector and tensor field generalization from Deville, Fischer, and Mund section 8.3.1. -kronmxv_code = """ -#include -#include - -static inline void kronmxv_inplace(PetscBLASInt tflag, - PetscBLASInt mx, PetscBLASInt my, PetscBLASInt mz, - PetscBLASInt nx, PetscBLASInt ny, PetscBLASInt nz, PetscBLASInt nel, - PetscScalar *A1, PetscScalar *A2, PetscScalar *A3, - PetscScalar **x, PetscScalar **y){ - -/* -Kronecker matrix-vector product - -y = op(A) * x, A = kron(A3, kron(A2, A1)) - -where: -op(A) = transpose(A) if tflag>0 else A -op(A1) is mx-by-nx, -op(A2) is my-by-ny, -op(A3) is mz-by-nz, -x is (nx*ny*nz)-by-nel, -y is (mx*my*mz)-by-nel. - -Important notes: -This routine is in-place: the input data in x and y are destroyed in the process. -Need to allocate nel*max(mx, nx)*max(my, ny)*max(mz, nz) memory for both x and y. -*/ - -PetscScalar *ptr[2] = {*x, *y}; -PetscScalar zero = 0.0E0, one = 1.0E0; -PetscBLASInt m, n, k, s, p, lda; -PetscBLASInt ires = 0; - -char tran = 'T', notr = 'N'; -char TA1 = tflag ? tran : notr; -char TA2 = tflag ? notr : tran; - -if(A1){ - m = mx; k = nx; n = ny*nz*nel; - lda = tflag ? nx : mx; - BLASgemm_(&TA1, ¬r, &m, &n, &k, &one, A1, &lda, ptr[ires], &k, &zero, ptr[!ires], &m); - ires = !ires; -} -if(A2){ - p = 0; s = 0; - m = mx; k = ny; n = my; - lda = tflag ? ny : my; - for(PetscBLASInt i=0; i 3: - raise ValueError("More than three factors are not supported") - - # Declare array shapes to be used as literals inside the kernels - nscal = psize*len(shift) - fshape = [e.space_dimension() for e in felem] - cshape = [e.space_dimension() for e in celem] - fshapes.append((nscal,) + tuple(fshape)) - cshapes.append((nscal,) + tuple(cshape)) - - J = [identity_filter(evaluate_dual(ce, fe)).T for ce, fe in zip(celem, felem)] - if any(Jk.size and numpy.isclose(Jk, 0.0E0).all() for Jk in J): - prolong_code.append(f""" - for({IntType_c} i=0; i<{nscal*numpy.prod(fshape)}; i++) {t_out}[i+{fskip}] = 0.0E0; - """) - restrict_code.append(f""" - for({IntType_c} i=0; i<{nscal*numpy.prod(cshape)}; i++) {t_in}[i+{cskip}] = 0.0E0; - """) - else: - Jsize = numpy.cumsum([Jlen] + [Jk.size for Jk in J]) - Jptrs = ["%s+%d" % (mat_name, Jsize[k]) if J[k].size else "NULL" for k in range(len(J))] - Jmats.extend(J) - Jlen = Jsize[-1] - - # The Kronecker product routines assume 3D shapes, so in 1D and 2D we pass NULL instead of J - Jargs = ", ".join(Jptrs+["NULL"]*(3-len(Jptrs))) - fargs = ", ".join(map(str, fshape+[1]*(3-len(fshape)))) - cargs = ", ".join(map(str, cshape+[1]*(3-len(cshape)))) - if in_place: - prolong_code.append(f""" - kronmxv_inplace(0, {fargs}, {cargs}, {nscal}, {Jargs}, &{t_in}, &{t_out}); - """) - restrict_code.append(f""" - kronmxv_inplace(1, {cargs}, {fargs}, {nscal}, {Jargs}, &{t_out}, &{t_in}); - """) - elif shifts == fshifts: - if has_code and psize > 1: - raise ValueError("Single tensor product to many tensor products not implemented for vectors") - # Single tensor product to many - prolong_code.append(f""" - kronmxv(0, {fargs}, {cargs}, {nscal}, {Jargs}, {t_in}+{cskip}, {t_out}+{fskip}, {scratch}, {t_out}+{fskip}); - """) - restrict_code.append(f""" - kronmxv(1, {cargs}, {fargs}, {nscal}, {Jargs}, {t_out}+{fskip}, {t_in}+{cskip}, {t_out}+{fskip}, {scratch}); - """) - else: - # Many tensor products to single tensor product - if has_code: - raise ValueError("Many tensor products to single tensor product not implemented") - fskip = 0 - prolong_code.append(f""" - kronmxv(0, {fargs}, {cargs}, {nscal}, {Jargs}, {t_in}+{cskip}, {t_out}+{fskip}, {t_in}+{cskip}, {t_out}+{fskip}); - """) - restrict_code.append(f""" - kronmxv(1, {cargs}, {fargs}, {nscal}, {Jargs}, {t_out}+{fskip}, {t_in}+{cskip}, {t_out}+{fskip}, {t_in}+{cskip}); - """) - has_code = True - fskip += nscal*numpy.prod(fshape) - cskip += nscal*numpy.prod(cshape) - - # Pass the 1D interpolators as a hexadecimal string - Jdata = ", ".join(map(float.hex, chain.from_iterable(Jk.flat for Jk in Jmats))) - operator_decl.append(f""" - PetscScalar {mat_name}[{Jlen}] = {{ {Jdata} }}; - """) - - operator_decl = "".join(operator_decl) - prolong_code = "".join(prolong_code) - restrict_code = "".join(reversed(restrict_code)) - shapes = [tuple(map(max, zip(*fshapes))), tuple(map(max, zip(*cshapes)))] - - if fskip > numpy.prod(shapes[0]): - shapes[0] = (fskip, 1, 1, 1) - if cskip > numpy.prod(shapes[1]): - shapes[1] = (cskip, 1, 1, 1) - return operator_decl, prolong_code, restrict_code, shapes - - -def get_piola_tensor(mapping, domain, inverse=False): - mapping = mapping.lower() - if mapping == "identity": - return None - elif mapping == "contravariant piola": - if inverse: - return ufl.JacobianInverse(domain)*ufl.JacobianDeterminant(domain) - else: - return ufl.Jacobian(domain)/ufl.JacobianDeterminant(domain) - elif mapping == "covariant piola": - if inverse: - return ufl.Jacobian(domain).T - else: - return ufl.JacobianInverse(domain).T - else: - raise ValueError("Mapping %s is not supported" % mapping) - - -def cache_generate_code(kernel, comm): - _cachedir = os.environ.get('PYOP2_CACHE_DIR', - os.path.join(tempfile.gettempdir(), - 'pyop2-cache-uid%d' % os.getuid())) - - key = kernel.cache_key[0] - shard, disk_key = key[:2], key[2:] - filepath = os.path.join(_cachedir, shard, disk_key) - if os.path.exists(filepath): - with open(filepath, 'r') as f: - code = f.read() - else: - code = loopy.generate_code_v2(kernel.code).device_code() - if comm.rank == 0: - os.makedirs(os.path.join(_cachedir, shard), exist_ok=True) - with open(filepath, 'w') as f: - f.write(code) - comm.barrier() - return code - - -def make_mapping_code(Q, cmapping, fmapping, t_in, t_out): - if fmapping == cmapping: - return None - A = get_piola_tensor(cmapping, Q.mesh(), inverse=False) - B = get_piola_tensor(fmapping, Q.mesh(), inverse=True) - tensor = A - if B: - tensor = ufl.dot(B, tensor) if tensor else B - if tensor is None: - tensor = ufl.Identity(Q.value_shape[0]) - - u = ufl.Coefficient(Q) - expr = ufl.dot(tensor, u) - prolong_map_kernel, coefficients = prolongation_transfer_kernel_action(Q, expr) - prolong_map_code = cache_generate_code(prolong_map_kernel, Q.comm) - prolong_map_code = prolong_map_code.replace("void expression_kernel", "static void prolongation_mapping") - coefficients.remove(u) - - expr = ufl.dot(u, tensor) - restrict_map_kernel, coefficients = prolongation_transfer_kernel_action(Q, expr) - restrict_map_code = cache_generate_code(restrict_map_kernel, Q.comm) - restrict_map_code = restrict_map_code.replace("void expression_kernel", "static void restriction_mapping") - restrict_map_code = restrict_map_code.replace("#include ", "") - restrict_map_code = restrict_map_code.replace("#include ", "") - coefficients.remove(u) - - coef_args = "".join([", c%d" % i for i in range(len(coefficients))]) - coef_decl = "".join([", PetscScalar const *restrict c%d" % i for i in range(len(coefficients))]) - qlen = Q.block_size * Q.finat_element.space_dimension() - prolong_code = f""" - for({IntType_c} i=0; i<{qlen}; i++) {t_out}[i] = 0.0E0; - - prolongation_mapping({t_out}{coef_args}, {t_in}); - """ - restrict_code = f""" - for({IntType_c} i=0; i<{qlen}; i++) {t_in}[i] = 0.0E0; - - restriction_mapping({t_in}{coef_args}, {t_out}); - """ - mapping_code = prolong_map_code + restrict_map_code - return coef_decl, prolong_code, restrict_code, mapping_code, coefficients - - -def make_permutation_code(V, vshape, pshape, t_in, t_out, array_name): - _, _, shifts = get_permutation_to_nodal_elements(V) - shift = shifts[0] - if shift != (0,): - ndof = numpy.prod(vshape) - permutation = numpy.reshape(numpy.arange(ndof), pshape) - axes = numpy.arange(len(shift)) - for k in range(permutation.shape[0]): - permutation[k] = numpy.reshape(numpy.transpose(permutation[k], axes=numpy.roll(axes, -shift[k])), pshape[1:]) - nflip = 0 - mapping = V.ufl_element().mapping().lower() - if mapping == "contravariant piola": - # flip the sign of the first component - nflip = ndof//len(shift) - elif mapping == "covariant piola": - # flip the order of reference components - permutation = numpy.flip(permutation, axis=0) - - permutation = numpy.transpose(numpy.reshape(permutation, vshape)) - pdata = ", ".join(map(str, permutation.flat)) - - decl = f""" - PetscInt {array_name}[{ndof}] = {{ {pdata} }}; - """ - prolong = f""" - for({IntType_c} i=0; i<{ndof}; i++) {t_out}[{array_name}[i]] = {t_in}[i]; - for({IntType_c} i=0; i<{nflip}; i++) {t_out}[i] = -{t_out}[i]; - """ - restrict = f""" - for({IntType_c} i=0; i<{nflip}; i++) {t_out}[i] = -{t_out}[i]; - for({IntType_c} i=0; i<{ndof}; i++) {t_in}[i] = {t_out}[{array_name}[i]]; - """ - else: - decl = "" - prolong = f""" - for({IntType_c} j=0; j<{vshape[1]}; j++) - for({IntType_c} i=0; i<{vshape[0]}; i++) - {t_out}[j + {vshape[1]}*i] = {t_in}[i + {vshape[0]}*j]; - """ - restrict = f""" - for({IntType_c} j=0; j<{vshape[1]}; j++) - for({IntType_c} i=0; i<{vshape[0]}; i++) - {t_in}[i + {vshape[0]}*j] = {t_out}[j + {vshape[1]}*i]; - """ - return decl, prolong, restrict - - -def reference_value_space(V): - element = finat.ufl.WithMapping(V.ufl_element(), mapping="identity") - return V.collapse().reconstruct(element=element) - - -class StandaloneInterpolationMatrix(object): - """ - Interpolation matrix for a single standalone space. - """ - - _cache_kernels = {} - _cache_work = {} - - def __init__(self, Vc, Vf, Vc_bcs, Vf_bcs): - self.uc = self.work_function(Vc) - self.uf = self.work_function(Vf) - self.Vc = self.uc.function_space() - self.Vf = self.uf.function_space() - self.Vc_bcs = Vc_bcs - self.Vf_bcs = Vf_bcs - - fmapping = self.Vf.ufl_element().mapping() - cmapping = self.Vc.ufl_element().mapping() - if type(self.Vf.ufl_element()) is not finat.ufl.MixedElement and fmapping != "identity" and fmapping == cmapping: - # Ignore Piola mapping if it is the same for both source and target, and simply transfer reference values. - self.Vc = reference_value_space(self.Vc) - self.Vf = reference_value_space(self.Vf) - self.uc = firedrake.Function(self.Vc, val=self.uc.dat) - self.uf = firedrake.Function(self.Vf, val=self.uf.dat) - self.Vc_bcs = [bc.reconstruct(V=self.Vc, g=0) for bc in self.Vc_bcs] - self.Vf_bcs = [bc.reconstruct(V=self.Vf, g=0) for bc in self.Vf_bcs] - - def work_function(self, V): - if isinstance(V, firedrake.Function): - return V - key = (V.ufl_element(), V.mesh(), V.boundary_set) - try: - return self._cache_work[key] - except KeyError: - return self._cache_work.setdefault(key, firedrake.Function(V)) - - @cached_property - def _weight(self): - cell_set = self.Vf.mesh().topology.unique().cell_set - weight = firedrake.Function(self.Vf) - wsize = self.Vf.finat_element.space_dimension() * self.Vf.block_size - kernel_code = f""" - void multiplicity(PetscScalar *restrict w) {{ - for (PetscInt i=0; i<{wsize}; i++) w[i] += 1; - }}""" - kernel = op2.Kernel(kernel_code, "multiplicity") - op2.par_loop(kernel, cell_set, weight.dat(op2.INC, weight.cell_node_map())) - with weight.dat.vec as w: - w.reciprocal() - return weight - - @cached_property - def _kernels(self): - try: - self.Vf.finat_element.dual_basis - self.Vc.finat_element.dual_basis - native_interpolation_supported = True - except NotImplementedError: - native_interpolation_supported = False - - if native_interpolation_supported: - return self._build_native_interpolators() - else: - return self._build_custom_interpolators() - - def _build_native_interpolators(self): - from firedrake.interpolation import interpolate, get_interpolator - P = get_interpolator(interpolate(self.uc, self.Vf)) - prolong = partial(P.assemble, tensor=self.uf) - - rf = firedrake.Function(self.Vf.dual(), val=self.uf.dat) - rc = firedrake.Function(self.Vc.dual(), val=self.uc.dat) - vc = firedrake.TestFunction(self.Vc) - R = get_interpolator(interpolate(vc, rf)) - restrict = partial(R.assemble, tensor=rc) - return prolong, restrict - - def _build_custom_interpolators(self): - # We generate custom prolongation and restriction kernels because - # dual evaluation of EnrichedElement is not yet implemented in FInAT - uf_map = get_permuted_map(self.Vf) - uc_map = get_permuted_map(self.Vc) - prolong_kernel, restrict_kernel, coefficients = self.make_blas_kernels(self.Vf, self.Vc) - cell_set = self.Vf.mesh().topology.unique().cell_set - prolong_args = [prolong_kernel, cell_set, - self.uf.dat(op2.INC, uf_map), - self.uc.dat(op2.READ, uc_map), - self._weight.dat(op2.READ, uf_map)] - restrict_args = [restrict_kernel, cell_set, - self.uc.dat(op2.INC, uc_map), - self.uf.dat(op2.READ, uf_map), - self._weight.dat(op2.READ, uf_map)] - coefficient_args = [c.dat(op2.READ, c.cell_node_map()) for c in coefficients] - prolong = op2.ParLoop(*prolong_args, *coefficient_args) - restrict = op2.ParLoop(*restrict_args, *coefficient_args) - return prolong, restrict - - def _prolong(self): - with self.uf.dat.vec_wo as uf: - uf.set(0.0E0) - self._kernels[0]() - - def _restrict(self): - with self.uc.dat.vec_wo as uc: - uc.set(0.0E0) - self._kernels[1]() - - def view(self, mat, viewer=None): - if viewer is None: - return - typ = viewer.getType() - if typ != PETSc.Viewer.Type.ASCII: - return - viewer.printfASCII("Firedrake matrix-free prolongator %s\n" % - type(self).__name__) - - def getInfo(self, mat, info=None): - memory = self.uf.dat.nbytes + self.uc.dat.nbytes - if self._weight is not None: - memory += self._weight.dat.nbytes - if info is None: - info = PETSc.Mat.InfoType.GLOBAL_SUM - if info == PETSc.Mat.InfoType.LOCAL: - return {"memory": memory} - elif info == PETSc.Mat.InfoType.GLOBAL_SUM: - gmem = mat.comm.tompi4py().allreduce(memory, op=op2.MPI.SUM) - return {"memory": gmem} - elif info == PETSc.Mat.InfoType.GLOBAL_MAX: - gmem = mat.comm.tompi4py().allreduce(memory, op=op2.MPI.MAX) - return {"memory": gmem} - else: - raise ValueError("Unknown info type %s" % info) - - def make_blas_kernels(self, Vf, Vc): - """ - Interpolation and restriction kernels between CG / DG - tensor product spaces on quads and hexes. - - Works by tabulating the coarse 1D basis functions - as the (fdegree+1)-by-(cdegree+1) matrix Jhat, - and using the fact that the 2D / 3D tabulation is the - tensor product J = kron(Jhat, kron(Jhat, Jhat)) - """ - cache = self._cache_kernels - key = (Vf.ufl_element(), Vc.ufl_element()) - try: - return cache[key] - except KeyError: - pass - felem = Vf.ufl_element() - celem = Vc.ufl_element() - fmapping = felem.mapping().lower() - cmapping = celem.mapping().lower() - - in_place_mapping = False - coefficients = [] - mapping_code = "" - coef_decl = "" - - if fmapping == cmapping: - # interpolate on each direction via Kroncker product - operator_decl, prolong_code, restrict_code, shapes = make_kron_code(Vc, Vf, "t0", "t1", "J0", "t2") - else: - decl = [""]*4 - prolong = [""]*5 - restrict = [""]*5 - # get embedding element for Vf with identity mapping and collocated vector component DOFs - try: - Qf = Vf if felem.mapping() == "identity" else Vf.reconstruct(mapping="identity") - mapping_output = make_mapping_code(Qf, cmapping, fmapping, "t0", "t1") - in_place_mapping = True - except Exception: - qelem = finat.ufl.FiniteElement("DQ", cell=felem.cell, degree=PMGBase.max_degree(felem)) - if Vf.value_shape: - qelem = finat.ufl.TensorElement(qelem, shape=Vf.value_shape, symmetry=felem.symmetry()) - Qf = Vf.reconstruct(element=qelem) - mapping_output = make_mapping_code(Qf, cmapping, fmapping, "t0", "t1") - - qshape = (Qf.block_size, Qf.finat_element.space_dimension()) - # interpolate to embedding fine space - decl[0], prolong[0], restrict[0], shapes = make_kron_code(Vc, Qf, "t0", "t1", "J0", "t2") - - if mapping_output is not None: - # permute to FInAT ordering, and apply the mapping - decl[1], restrict[1], prolong[1] = make_permutation_code(Vc, qshape, shapes[0], "t0", "t1", "perm0") - coef_decl, prolong[2], restrict[2], mapping_code, coefficients = mapping_output - if not in_place_mapping: - # permute to Kronecker-friendly ordering and interpolate to fine space - decl[2], prolong[3], restrict[3] = make_permutation_code(Vf, qshape, shapes[0], "t1", "t0", "perm1") - decl[3], prolong[4], restrict[4], _shapes = make_kron_code(Qf, Vf, "t0", "t1", "J1", "t2") - shapes.extend(_shapes) - - operator_decl = "".join(decl) - prolong_code = "".join(prolong) - restrict_code = "".join(reversed(restrict)) - - # FInAT elements order the component DOFs related to the same node contiguously. - # We transpose before and after the multiplication times J to have each component - # stored contiguously as a scalar field, thus reducing the number of dgemm calls. - - # We could benefit from loop tiling for the transpose, but that makes the code - # more complicated. - - fshape = (Vf.block_size, Vf.finat_element.space_dimension()) - cshape = (Vc.block_size, Vc.finat_element.space_dimension()) - - lwork = numpy.prod([max(*dims) for dims in zip(*shapes)]) - lwork = max(lwork, max(numpy.prod(fshape), numpy.prod(cshape))) - - if cshape[0] == 1: - coarse_read = f"""for({IntType_c} i=0; i<{numpy.prod(cshape)}; i++) t0[i] = x[i];""" - coarse_write = f"""for({IntType_c} i=0; i<{numpy.prod(cshape)}; i++) x[i] += t0[i];""" - else: - coarse_read = f""" - for({IntType_c} j=0; j<{cshape[1]}; j++) - for({IntType_c} i=0; i<{cshape[0]}; i++) - t0[j + {cshape[1]}*i] = x[i + {cshape[0]}*j]; - """ - coarse_write = f""" - for({IntType_c} j=0; j<{cshape[1]}; j++) - for({IntType_c} i=0; i<{cshape[0]}; i++) - x[i + {cshape[0]}*j] += t0[j + {cshape[1]}*i]; - """ - if (fshape[0] == 1) or in_place_mapping: - fine_read = f"""for({IntType_c} i=0; i<{numpy.prod(fshape)}; i++) t1[i] = y[i] * w[i];""" - fine_write = f"""for({IntType_c} i=0; i<{numpy.prod(fshape)}; i++) y[i] += t1[i] * w[i];""" - else: - fine_read = f""" - for({IntType_c} j=0; j<{fshape[1]}; j++) - for({IntType_c} i=0; i<{fshape[0]}; i++) - t1[j + {fshape[1]}*i] = y[i + {fshape[0]}*j] * w[i + {fshape[0]}*j]; - """ - fine_write = f""" - for({IntType_c} j=0; j<{fshape[1]}; j++) - for({IntType_c} i=0; i<{fshape[0]}; i++) - y[i + {fshape[0]}*j] += t1[j + {fshape[1]}*i] * w[i + {fshape[0]}*j]; - """ - kernel_code = f""" - {mapping_code} - - {kronmxv_code} - - void prolongation(PetscScalar *restrict y, const PetscScalar *restrict x, - const PetscScalar *restrict w{coef_decl}){{ - PetscScalar work[3][{lwork}] = {{0.0E0}}; - PetscScalar *t0 = work[0]; - PetscScalar *t1 = work[1]; - PetscScalar *t2 = work[2]; - {operator_decl} - {coarse_read} - {prolong_code} - {fine_write} - return; - }} - - void restriction(PetscScalar *restrict x, const PetscScalar *restrict y, - const PetscScalar *restrict w{coef_decl}){{ - PetscScalar work[3][{lwork}] = {{0.0E0}}; - PetscScalar *t0 = work[0]; - PetscScalar *t1 = work[1]; - PetscScalar *t2 = work[2]; - {operator_decl} - {fine_read} - {restrict_code} - {coarse_write} - return; - }} - """ - from firedrake.slate.slac.compiler import BLASLAPACK_LIB, BLASLAPACK_INCLUDE - prolong_kernel = op2.Kernel(kernel_code, "prolongation", include_dirs=BLASLAPACK_INCLUDE.split(), - ldargs=BLASLAPACK_LIB.split(), requires_zeroed_output_arguments=True) - restrict_kernel = op2.Kernel(kernel_code, "restriction", include_dirs=BLASLAPACK_INCLUDE.split(), - ldargs=BLASLAPACK_LIB.split(), requires_zeroed_output_arguments=True) - return cache.setdefault(key, (prolong_kernel, restrict_kernel, coefficients)) - - def multTranspose(self, mat, rf, rc): - """ - Implement restriction: restrict residual on fine grid rf to coarse grid rc. - """ - with self.uf.dat.vec_wo as uf: - rf.copy(uf) - for bc in self.Vf_bcs: - bc.zero(self.uf) - - self._restrict() - - for bc in self.Vc_bcs: - bc.zero(self.uc) - with self.uc.dat.vec_ro as uc: - uc.copy(rc) - - def mult(self, mat, xc, xf, inc=False): - """ - Implement prolongation: prolong correction on coarse grid xc to fine grid xf. - """ - with self.uc.dat.vec_wo as uc: - xc.copy(uc) - for bc in self.Vc_bcs: - bc.zero(self.uc) - - self._prolong() - - for bc in self.Vf_bcs: - bc.zero(self.uf) - if inc: - with self.uf.dat.vec_ro as uf: - xf.axpy(1.0, uf) - else: - with self.uf.dat.vec_ro as uf: - uf.copy(xf) - - def multAdd(self, mat, x, y, w): - if y.handle == w.handle: - self.mult(mat, x, w, inc=True) - else: - self.mult(mat, x, w) - w.axpy(1.0, y) - - -class MixedInterpolationMatrix(StandaloneInterpolationMatrix): - """ - Interpolation matrix for a mixed finite element space. - """ - @cached_property - def _weight(self): - return None - - @cached_property - def _standalones(self): - standalones = [] - for i, (uc_sub, uf_sub) in enumerate(zip(self.uc.subfunctions, self.uf.subfunctions)): - Vc_sub_bcs = tuple(bc for bc in self.Vc_bcs if bc.function_space().index == i) - Vf_sub_bcs = tuple(bc for bc in self.Vf_bcs if bc.function_space().index == i) - standalone = StandaloneInterpolationMatrix(uc_sub, uf_sub, Vc_sub_bcs, Vf_sub_bcs) - standalones.append(standalone) - return standalones - - @cached_property - def _kernels(self): - prolong = lambda: [s._prolong() for s in self._standalones] - restrict = lambda: [s._restrict() for s in self._standalones] - return prolong, restrict - - def getNestSubMatrix(self, i, j): - if i == j: - s = self._standalones[i] - sizes = (s.uf.dof_dset.layout_vec.getSizes(), s.uc.dof_dset.layout_vec.getSizes()) - M_shll = PETSc.Mat().createPython(sizes, s, comm=s.uf.comm) - M_shll.setUp() - return M_shll - else: - return None - - -def prolongation_matrix_aij(Vc, Vf, Vc_bcs=(), Vf_bcs=()): - if isinstance(Vf, firedrake.Function): - Vf = Vf.function_space() - if isinstance(Vc, firedrake.Function): - Vc = Vc.function_space() - bcs = Vc_bcs + Vf_bcs - interp = firedrake.interpolate(firedrake.TrialFunction(Vc), Vf) - mat_type = "nest" if len(Vc) > 1 or len(Vf) > 1 else None - mat = firedrake.assemble(interp, bcs=bcs, mat_type=mat_type) - return mat.petscmat - - -def prolongation_matrix_matfree(Vc, Vf, Vc_bcs=[], Vf_bcs=[]): - fele = Vf.ufl_element() - if type(fele) is finat.ufl.MixedElement: - ctx = MixedInterpolationMatrix(Vc, Vf, Vc_bcs, Vf_bcs) - else: - ctx = StandaloneInterpolationMatrix(Vc, Vf, Vc_bcs, Vf_bcs) - - sizes = (Vf.dof_dset.layout_vec.getSizes(), Vc.dof_dset.layout_vec.getSizes()) - M_shll = PETSc.Mat().createPython(sizes, ctx, comm=Vf.comm) - M_shll.setUp() - return M_shll diff --git a/tests/firedrake/multigrid/test_p_multigrid.py b/tests/firedrake/multigrid/test_p_multigrid.py index 547c6af6ef..76febe6e19 100644 --- a/tests/firedrake/multigrid/test_p_multigrid.py +++ b/tests/firedrake/multigrid/test_p_multigrid.py @@ -1,5 +1,4 @@ import pytest -import numpy as np from firedrake import * @@ -74,80 +73,6 @@ def test_reconstruct_degree(tp_mesh, mixed_family): assert e == elist[0].reconstruct(degree=degree) -@pytest.mark.parametrize("family", ["Q", "NCE", "NCF", "DQ"]) -def test_prolong_basic(tp_mesh, family): - """ Interpolate a constant function between low-order and high-order spaces - """ - from firedrake.preconditioners.pmg import prolongation_matrix_matfree - if tp_mesh.topological_dimension == 2: - family = family.replace("N", "RT") - - fs = [FunctionSpace(tp_mesh, family, degree) for degree in (1, 2)] - u, v = [Function(V) for V in fs] - - u.assign(1) - P = prolongation_matrix_matfree(u, v).getPythonContext() - P._prolong() - assert np.allclose(v.dat.data, 1) - - -def test_prolong_de_rham(tp_mesh): - """ Interpolate a linear vector function between [H1]^d, HCurl and HDiv spaces - where it can be exactly represented - """ - from firedrake.preconditioners.pmg import prolongation_matrix_matfree - - tdim = tp_mesh.topological_dimension - b = Constant(list(range(tdim))) - if tp_mesh.extruded_periodic: - expr = b - else: - mat = diag(Constant([tdim+1]*tdim)) + Constant([[-1]*tdim]*tdim) - expr = b + dot(mat, SpatialCoordinate(tp_mesh)) - - cell = tp_mesh.ufl_cell() - elems = [VectorElement(FiniteElement("Q", cell=cell, degree=2)), - FiniteElement("NCE" if tdim == 3 else "RTCE", cell=cell, degree=2), - FiniteElement("NCF" if tdim == 3 else "RTCF", cell=cell, degree=2)] - - fs = [FunctionSpace(tp_mesh, e) for e in elems] - us = [Function(V) for V in fs] - us[0].interpolate(expr) - for u in us: - for v in us: - if u != v: - P = prolongation_matrix_matfree(u, v).getPythonContext() - P._prolong() - assert errornorm(expr, v) < 1E-14 - - -def test_prolong_low_order_to_restricted(tp_mesh, tp_family, variant): - """ Interpolate a low-order function to interior and facet high-order spaces - and ensure that the sum of the two high-order functions is equal to the - low-order function - """ - from firedrake.preconditioners.pmg import prolongation_matrix_matfree - - degree = 5 - cell = tp_mesh.ufl_cell() - element = FiniteElement(tp_family, cell=cell, degree=degree, variant=variant) - Vi = FunctionSpace(tp_mesh, RestrictedElement(element, restriction_domain="interior")) - Vf = FunctionSpace(tp_mesh, RestrictedElement(element, restriction_domain="facet")) - Vc = FunctionSpace(tp_mesh, tp_family, degree=1) - - ui = Function(Vi) - uf = Function(Vf) - uc = Function(Vc) - uc.dat.data[0::2] = 0.0 - uc.dat.data[1::2] = 1.0 - - for v in [ui, uf]: - P = prolongation_matrix_matfree(uc, v).getPythonContext() - P._prolong() - - assert norm(ui + uf - uc, "L2") < 1E-13 - - @pytest.fixture(params=["triangles", "quadrilaterals"], scope="module") def mesh(request): if request.param == "triangles": @@ -358,13 +283,6 @@ def test_p_multigrid_mixed(mat_type): ctx_levels += 1 assert ctx_levels == 3 - # test that the cache is parallel-safe - dummy_eq = type(object).__eq__ - cache = PMGPC._transfer_cache - assert len(cache) > 0 - for k in cache: - assert type(k).__eq__ is dummy_eq - def test_p_fas_scalar(): mat_type = "matfree" @@ -559,52 +477,3 @@ def check_coarsen_quadrature(solver): check_coarsen_quadrature(solver_npmg) iter_npmg = solver_npmg.snes.getLinearSolveIterations() assert 2*iter_pfas <= iter_npmg - - -@pytest.fixture -def piola_mesh(): - return UnitDiskMesh(3) - - -@pytest.mark.parametrize("mat_type", ("matfree", "aij")) -@pytest.mark.parametrize("mixed", (False, True), ids=("standalone", "mixed")) -@pytest.mark.parametrize("family, degree", (("CG", 4), ("N2curl", 2), ("N1div", 3))) -def test_pmg_transfer_piola(piola_mesh, family, degree, mixed, mat_type): - """Test prolongation and restriction kernels for piola-mapped elements. - """ - from firedrake.preconditioners.pmg import prolongation_matrix_matfree, prolongation_matrix_aij - Vf = FunctionSpace(piola_mesh, family, degree) - if mixed: - DG = FunctionSpace(Vf.mesh(), "DG", 2) - Vf = Vf * Vf * DG - Vc = Vf.reconstruct(degree=[1, 1, 1]) - else: - Vc = Vf.reconstruct(degree=1) - - Vf_bcs = [DirichletBC(Vf.sub(0), 0, "on_boundary")] - Vc_bcs = [DirichletBC(Vc.sub(0), 0, "on_boundary")] - if mat_type == "matfree": - P = prolongation_matrix_matfree(Vc, Vf, Vc_bcs, Vf_bcs) - else: - P = prolongation_matrix_aij(Vc, Vf, Vc_bcs, Vf_bcs) - - uc = Function(Vc) - uf = Function(Vf) - with uc.dat.vec_wo as xc: - xc.setRandom() - for bc in Vc_bcs: - bc.zero(uc) - with uc.dat.vec_ro as xc, uf.dat.vec as xf: - P.mult(xc, xf) - assert norm(uf - uc) < 1E-12 - - rc = Cofunction(Vc.dual()) - rf = Cofunction(Vf.dual()) - with rf.dat.vec_wo as xf: - xf.setRandom() - for bc in Vf_bcs: - bc.zero(rf) - with rf.dat.vec_ro as xf, rc.dat.vec as xc: - P.multTranspose(xf, xc) - - assert abs(assemble(action(rf, uf)) - assemble(action(rc, uc))) < 1E-11 diff --git a/tests/tsfc/test_dual_evaluation.py b/tests/tsfc/test_dual_evaluation.py index 16d07f172b..5fb65d1ef0 100644 --- a/tests/tsfc/test_dual_evaluation.py +++ b/tests/tsfc/test_dual_evaluation.py @@ -56,3 +56,18 @@ def test_ufl_only_shape_mismatch(): assert to_element.value_shape == (2,) with pytest.raises(ValueError): compile_expression_dual_evaluation(expr, W.ufl_element()) + + +def test_dual_argument_on_concatenated_dual_basis(): + """The dual basis of a facet-restricted element is a Concatenate. + + A Cofunction dual argument sums over the concatenated index, so the + contraction has to be split along the Concatenate before it is formed. + """ + mesh = ufl.Mesh(finat.ufl.VectorElement("Q", ufl.quadrilateral, 1)) + element = finat.ufl.FiniteElement("Q", ufl.quadrilateral, 3) + V = ufl.FunctionSpace(mesh, element["facet"]) + W = ufl.FunctionSpace(mesh, element.reconstruct(degree=2)["facet"]) + expr = ufl.Interpolate(ufl.Argument(V, 0), ufl.Cofunction(W.dual())) + kernel = compile_expression_dual_evaluation(expr, W.ufl_element()) + assert kernel.flop_count > 0 diff --git a/tsfc/driver.py b/tsfc/driver.py index 2c480f9c55..357cece6a4 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -14,6 +14,7 @@ import gem import gem.impero_utils as impero_utils +from gem.unconcatenate import unconcatenate import finat from finat.element_factory import as_fiat_cell @@ -353,10 +354,15 @@ def compile_expression_dual_evaluation(expression, ufl_element, *, gem_dual = builder.coefficient_map[dual_arg] if complex_mode: evaluation = gem.MathFunction('conj', evaluation) - # The dual argument contracts over the nodes, so the basis indices are - # reduction indices like the points, not indices of the return value. - evaluation = evaluation * gem_dual[basis_indices] - quadrature_multiindex += tuple(basis_indices) + # The dual argument contracts over the nodes. Split the dual basis + # along its Concatenate nodes first, as assembly does for coefficient + # evaluation. Each block then sums over its own basis indices, rather + # than over the concatenated index, which nothing can be split along. + var, = gem.optimise.remove_componenttensors([gem_dual[basis_indices]]) + summands = [gem.IndexSum(gem.Product(expr, v), v.index_ordering()) + for v, expr in unconcatenate([(var, evaluation)], + kernel_cfg["index_cache"])] + evaluation = gem.optimise.make_sum(summands) basis_indices = () else: argument_multiindices[dual_arg.number()] = basis_indices