From 68017fa03ea2cafaf19a8169d684ca2fe2156c61 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 14 Jan 2025 15:07:19 +0000 Subject: [PATCH 01/18] DO NOT MERGE --- .github/workflows/build.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0eb616c24d..8861f83fe4 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -84,6 +84,7 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ + --package-branch fiat pbrubeck/finat-boundary-quadrature \ || (cat firedrake-install.log && /bin/false) - name: Install test dependencies run: | From f850f7a9dcfc69ca90c6cc73d03d282c9a340ad1 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Fri, 17 Jan 2025 13:44:18 +0000 Subject: [PATCH 02/18] Improve kernel caching Each parloop invocation was doing more than necessary leading to quite poor performance. --- firedrake/scripts/firedrake_clean.py | 7 +- pyop2/caching.py | 3 +- pyop2/global_kernel.py | 111 ++++++++++++++++----------- 3 files changed, 71 insertions(+), 50 deletions(-) diff --git a/firedrake/scripts/firedrake_clean.py b/firedrake/scripts/firedrake_clean.py index f411d498c3..678a3c35fa 100755 --- a/firedrake/scripts/firedrake_clean.py +++ b/firedrake/scripts/firedrake_clean.py @@ -4,10 +4,7 @@ from firedrake.configuration import setup_cache_dirs from pyop2.compilation import clear_compiler_disk_cache as pyop2_clear_cache from firedrake.tsfc_interface import clear_cache as tsfc_clear_cache -try: - import platformdirs as appdirs -except ImportError: - import appdirs +import platformdirs def main(): @@ -20,7 +17,7 @@ def main(): print(f"Removing cached PyOP2 code from {os.environ.get('PYOP2_CACHE_DIR', '???')}") pyop2_clear_cache() - pytools_cache = appdirs.user_cache_dir("pytools", "pytools") + pytools_cache = platformdirs.user_cache_dir("pytools", "pytools") print(f"Removing cached pytools files from {pytools_cache}") if os.path.exists(pytools_cache): shutil.rmtree(pytools_cache, ignore_errors=True) diff --git a/pyop2/caching.py b/pyop2/caching.py index 2948ddede7..7a827d90b0 100644 --- a/pyop2/caching.py +++ b/pyop2/caching.py @@ -547,7 +547,8 @@ def wrapper(*args, **kwargs): value = local_cache.get(key, CACHE_MISS) if value is CACHE_MISS: - value = func(*args, **kwargs) + with PETSc.Log.Event("pyop2: handle cache miss"): + value = func(*args, **kwargs) return local_cache.setdefault(key, value) return wrapper diff --git a/pyop2/global_kernel.py b/pyop2/global_kernel.py index 7edfed0771..433a2992f4 100644 --- a/pyop2/global_kernel.py +++ b/pyop2/global_kernel.py @@ -8,12 +8,15 @@ import loopy as lp import numpy as np import pytools +from loopy.codegen.result import process_preambles from petsc4py import PETSc from pyop2 import mpi +from pyop2.caching import parallel_cache, serial_cache from pyop2.compilation import add_profiling_events, load from pyop2.configuration import configuration from pyop2.datatypes import IntType, as_ctypes +from pyop2.codegen.rep2loopy import generate from pyop2.types import IterationRegion, Constant, READ from pyop2.utils import cached_property, get_petsc_dir @@ -326,8 +329,7 @@ def __call__(self, comm, *args): :arg comm: Communicator the execution is collective over. :*args: Arguments to pass to the compiled kernel. """ - # It is unnecessary to cache this call as it is cached in pyop2/compilation.py - func = self.compile(comm) + func = _compile_global_kernel(self, comm) func(*args) @property @@ -364,48 +366,7 @@ def builder(self): @cached_property def code_to_compile(self): """Return the C/C++ source code as a string.""" - from pyop2.codegen.rep2loopy import generate - - with PETSc.Log.Event("GlobalKernel: generate loopy"): - wrapper = generate(self.builder) - - with PETSc.Log.Event("GlobalKernel: generate device code"): - code = lp.generate_code_v2(wrapper) - - if self.local_kernel.cpp: - from loopy.codegen.result import process_preambles - preamble = "".join(process_preambles(getattr(code, "device_preambles", []))) - device_code = "\n\n".join(str(dp.ast) for dp in code.device_programs) - return preamble + "\nextern \"C\" {\n" + device_code + "\n}\n" - return code.device_code() - - @PETSc.Log.EventDecorator() - @mpi.collective - def compile(self, comm): - """Compile the kernel. - - :arg comm: The communicator the compilation is collective over. - :returns: A ctypes function pointer for the compiled function. - """ - extension = "cpp" if self.local_kernel.cpp else "c" - cppargs = ( - tuple("-I%s/include" % d for d in get_petsc_dir()) - + tuple("-I%s" % d for d in self.local_kernel.include_dirs) - + ("-I%s" % os.path.abspath(os.path.dirname(__file__)),) - ) - ldargs = ( - tuple("-L%s/lib" % d for d in get_petsc_dir()) - + tuple("-Wl,-rpath,%s/lib" % d for d in get_petsc_dir()) - + ("-lpetsc", "-lm") - + tuple(self.local_kernel.ldargs) - ) - - dll = load(self.code_to_compile, extension, cppargs=cppargs, ldargs=ldargs, comm=comm) - add_profiling_events(dll, self.local_kernel.events) - fn = getattr(dll, self.name) - fn.argtypes = self.argtypes - fn.restype = ctypes.c_int - return fn + return _generate_code_from_global_kernel(self) @cached_property def argtypes(self): @@ -427,3 +388,65 @@ def num_flops(self, iterset): elif region not in {IterationRegion.TOP, IterationRegion.BOTTOM}: size = layers - 1 return size * self.local_kernel.num_flops + + @cached_property + def _cppargs(self): + cppargs = [f"-I{d}/include" for d in get_petsc_dir()] + cppargs.extend(f"-I{d}" for d in self.local_kernel.include_dirs) + cppargs.append(f"-I{os.path.abspath(os.path.dirname(__file__))}") + return tuple(cppargs) + + @cached_property + def _ldargs(self): + ldargs = [f"-L{d}/lib" for d in get_petsc_dir()] + ldargs.extend(f"-Wl,-rpath,{d}/lib" for d in get_petsc_dir()) + ldargs.extend(["-lpetsc", "-lm"]) + ldargs.extend(self.local_kernel.ldargs) + return tuple(ldargs) + + +@serial_cache(hashkey=lambda knl: knl.cache_key) +def _generate_code_from_global_kernel(kernel): + with PETSc.Log.Event("GlobalKernel: generate loopy"): + wrapper = generate(kernel.builder) + + with PETSc.Log.Event("GlobalKernel: generate device code"): + code = lp.generate_code_v2(wrapper) + + if kernel.local_kernel.cpp: + preamble = "".join(process_preambles(getattr(code, "device_preambles", []))) + device_code = "\n\n".join(str(dp.ast) for dp in code.device_programs) + return preamble + "\nextern \"C\" {\n" + device_code + "\n}\n" + + return code.device_code() + + +@parallel_cache(hashkey=lambda knl, _: knl.cache_key) +@mpi.collective +def _compile_global_kernel(kernel, comm): + """Compile the kernel. + + Parameters + ---------- + kernel : + The global kernel to generate code for. + comm : + The communicator the compilation is collective over. + + Returns + ------- + A ctypes function pointer for the compiled function. + + """ + dll = load( + kernel.code_to_compile, + "cpp" if kernel.local_kernel.cpp else "c", + cppargs=kernel._cppargs, + ldargs=kernel._ldargs, + comm=comm, + ) + add_profiling_events(dll, kernel.local_kernel.events) + fn = getattr(dll, kernel.name) + fn.argtypes = kernel.argtypes + fn.restype = ctypes.c_int + return fn From 7c153ad687cdc5af00685f00b0c08c4e8d643013 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Fri, 17 Jan 2025 13:47:06 +0000 Subject: [PATCH 03/18] fixup --- firedrake/preconditioners/patch.py | 5 +++-- pyop2/global_kernel.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/firedrake/preconditioners/patch.py b/firedrake/preconditioners/patch.py index 5e9d0d4fa0..d79a03b8a1 100644 --- a/firedrake/preconditioners/patch.py +++ b/firedrake/preconditioners/patch.py @@ -27,6 +27,7 @@ from pyop2.codegen.builder import Pack, MatPack, DatPack from pyop2.codegen.representation import Comparison, Literal from pyop2.codegen.rep2loopy import register_petsc_function +from pyop2.global_kernel import compile_global_kernel __all__ = ("PatchPC", "PlaneSmoother", "PatchSNES") @@ -222,7 +223,7 @@ def matrix_funptr(form, state): 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)) + kernels.append(CompiledKernel(compile_global_kernel(mod, iterset.comm), kinfo)) return cell_kernels, int_facet_kernels @@ -316,7 +317,7 @@ def residual_funptr(form, state): 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)) + kernels.append(CompiledKernel(compile_global_kernel(mod, iterset.comm), kinfo)) return cell_kernels, int_facet_kernels diff --git a/pyop2/global_kernel.py b/pyop2/global_kernel.py index 433a2992f4..aac7d30c3d 100644 --- a/pyop2/global_kernel.py +++ b/pyop2/global_kernel.py @@ -329,7 +329,7 @@ def __call__(self, comm, *args): :arg comm: Communicator the execution is collective over. :*args: Arguments to pass to the compiled kernel. """ - func = _compile_global_kernel(self, comm) + func = compile_global_kernel(self, comm) func(*args) @property @@ -423,7 +423,7 @@ def _generate_code_from_global_kernel(kernel): @parallel_cache(hashkey=lambda knl, _: knl.cache_key) @mpi.collective -def _compile_global_kernel(kernel, comm): +def compile_global_kernel(kernel, comm): """Compile the kernel. Parameters From 45e660d92cab17662644018047427427a2a43eac Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Fri, 17 Jan 2025 15:02:23 +0000 Subject: [PATCH 04/18] test fixups --- tests/pyop2/test_caching.py | 39 ++++++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/tests/pyop2/test_caching.py b/tests/pyop2/test_caching.py index cfd9e6ce7f..5579377ca8 100644 --- a/tests/pyop2/test_caching.py +++ b/tests/pyop2/test_caching.py @@ -309,6 +309,13 @@ def cache(self): int_comm.Set_attr(comm_cache_keyval, _cache_collection) return _cache_collection[default_cache_name] + def code_cache_len_equals(self, expected): + # We need to do this check because different things also get + # put into self.cache + return sum( + 1 for key in self.cache if key[1] == "compile_global_kernel" + ) == expected + @pytest.fixture def a(cls, diterset): return op2.Dat(diterset, list(range(nelems)), numpy.uint32, "a") @@ -328,14 +335,14 @@ def test_same_args(self, iterset, iter2ind1, x, a): a(op2.WRITE), x(op2.READ, iter2ind1)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) op2.par_loop(op2.Kernel(kernel_cpy, "cpy"), iterset, a(op2.WRITE), x(op2.READ, iter2ind1)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) def test_diff_kernel(self, iterset, iter2ind1, x, a): self.cache.clear() @@ -348,7 +355,7 @@ def test_diff_kernel(self, iterset, iter2ind1, x, a): a(op2.WRITE), x(op2.READ, iter2ind1)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) kernel_cpy = "static void cpy(unsigned int* DST, unsigned int* SRC) { *DST = *SRC; }" @@ -357,7 +364,7 @@ def test_diff_kernel(self, iterset, iter2ind1, x, a): a(op2.WRITE), x(op2.READ, iter2ind1)) - assert len(self.cache) == 2 + assert self.code_cache_len_equals(2) def test_invert_arg_similar_shape(self, iterset, iter2ind1, x, y): self.cache.clear() @@ -377,14 +384,14 @@ def test_invert_arg_similar_shape(self, iterset, iter2ind1, x, y): x(op2.RW, iter2ind1), y(op2.RW, iter2ind1)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) op2.par_loop(op2.Kernel(kernel_swap, "swap"), iterset, y(op2.RW, iter2ind1), x(op2.RW, iter2ind1)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) def test_dloop_ignore_scalar(self, iterset, a, b): self.cache.clear() @@ -404,14 +411,14 @@ def test_dloop_ignore_scalar(self, iterset, a, b): a(op2.RW), b(op2.RW)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) op2.par_loop(op2.Kernel(kernel_swap, "swap"), iterset, b(op2.RW), a(op2.RW)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) def test_vector_map(self, iterset, x2, iter2ind2): self.cache.clear() @@ -431,13 +438,13 @@ def test_vector_map(self, iterset, x2, iter2ind2): iterset, x2(op2.RW, iter2ind2)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) op2.par_loop(op2.Kernel(kernel_swap, "swap"), iterset, x2(op2.RW, iter2ind2)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) def test_same_iteration_space_works(self, iterset, x2, iter2ind2): self.cache.clear() @@ -447,12 +454,12 @@ def test_same_iteration_space_works(self, iterset, x2, iter2ind2): op2.par_loop(k, iterset, x2(op2.INC, iter2ind2)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) op2.par_loop(k, iterset, x2(op2.INC, iter2ind2)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) def test_change_dat_dtype_matters(self, iterset, diterset): d = op2.Dat(diterset, list(range(nelems)), numpy.uint32) @@ -463,12 +470,12 @@ def test_change_dat_dtype_matters(self, iterset, diterset): op2.par_loop(k, iterset, d(op2.WRITE)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) d = op2.Dat(diterset, list(range(nelems)), numpy.int32) op2.par_loop(k, iterset, d(op2.WRITE)) - assert len(self.cache) == 2 + assert self.code_cache_len_equals(2) def test_change_global_dtype_matters(self, iterset, diterset): g = op2.Global(1, 0, dtype=numpy.uint32, comm=COMM_WORLD) @@ -479,12 +486,12 @@ def test_change_global_dtype_matters(self, iterset, diterset): op2.par_loop(k, iterset, g(op2.INC)) - assert len(self.cache) == 1 + assert self.code_cache_len_equals(1) g = op2.Global(1, 0, dtype=numpy.float64, comm=COMM_WORLD) op2.par_loop(k, iterset, g(op2.INC)) - assert len(self.cache) == 2 + assert self.code_cache_len_equals(2) class TestSparsityCache: From a778d6800a7cd9d83358675e945b53febf8f9bc6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 22 Jan 2025 09:58:59 +0000 Subject: [PATCH 05/18] DO NOT MERGE --- .github/workflows/build.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0eb616c24d..7f108a77ec 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -84,6 +84,8 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ + --package-branch fiat pbrubeck/simplify-indexed \ + --package-branch ufl pbrubeck/remove-component-tensors \ || (cat firedrake-install.log && /bin/false) - name: Install test dependencies run: | From 1749644549c302f14deb7fb48a91a6b25009a049 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 26 Jan 2025 14:52:51 +0000 Subject: [PATCH 06/18] DO NOT MERGE --- .github/workflows/build.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 0f17161362..77739f82f9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -87,6 +87,8 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ + --package-branch fiat pbrubeck/simplify-indexed \ + --package-branch ufl pbrubeck/remove-component-tensors \ || (cat firedrake-install.log && /bin/false) - name: Install test dependencies From 729339dfedde3a8c6a0b32880562583b3f4a04e7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 26 Jan 2025 15:02:52 +0000 Subject: [PATCH 07/18] Lazy basis transformation --- tsfc/fem.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tsfc/fem.py b/tsfc/fem.py index 10c10dc97d..e20244a0a1 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -641,15 +641,13 @@ def fiat_to_ufl(fiat_dict, order): @translate.register(Argument) def translate_argument(terminal, mt, ctx): - argument_multiindex = ctx.argument_multiindices[terminal.number()] - sigma = tuple(gem.Index(extent=d) for d in mt.expr.ufl_shape) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) def callback(entity_id): finat_dict = ctx.basis_evaluation(element, mt, entity_id) # Filter out irrelevant derivatives - filtered_dict = {alpha: table - for alpha, table in finat_dict.items() + filtered_dict = {alpha: finat_dict[alpha] + for alpha in finat_dict if sum(alpha) == mt.local_derivatives} # Change from FIAT to UFL arrangement @@ -658,13 +656,16 @@ def callback(entity_id): # A numerical hack that FFC used to apply on FIAT tables still # lives on after ditching FFC and switching to FInAT. return ffc_rounding(square, ctx.epsilon) + table = ctx.entity_selector(callback, mt.restriction) if ctx.use_canonical_quadrature_point_ordering: quad_multiindex = ctx.quadrature_rule.point_set.indices quad_multiindex_permuted = _make_quad_multiindex_permuted(mt, ctx) mapper = gem.node.MemoizerArg(gem.optimise.filtered_replace_indices) table = mapper(table, tuple(zip(quad_multiindex, quad_multiindex_permuted))) - return gem.ComponentTensor(gem.Indexed(table, argument_multiindex + sigma), sigma) + + argument_multiindex = ctx.argument_multiindices[terminal.number()] + return gem.partial_indexed(table, argument_multiindex) @translate.register(TSFCConstantMixin) From 8a05b3ad39c0c30540a7af807beb69b2abbec0c3 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 14 Mar 2025 17:07:46 +0000 Subject: [PATCH 08/18] Extract quad degree from Quadrature elements --- .../test_projection_symmetric_tensor.py | 2 +- tsfc/kernel_interface/common.py | 21 ++++++++++++------- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/tests/firedrake/regression/test_projection_symmetric_tensor.py b/tests/firedrake/regression/test_projection_symmetric_tensor.py index bc0e551b73..ebec1c9657 100644 --- a/tests/firedrake/regression/test_projection_symmetric_tensor.py +++ b/tests/firedrake/regression/test_projection_symmetric_tensor.py @@ -32,7 +32,7 @@ def test_projection_symmetric_tensor(mesh, degree, family, tdim): sp = {"mat_type": "matfree", "ksp_type": "preonly", "pc_type": "jacobi"} - fcp = {"quadrature_degree": Nq} + fcp = {} else: Q = TensorFunctionSpace(mesh, family, degree=degree, shape=shape, symmetry=None) Qs = TensorFunctionSpace(mesh, family, degree=degree, shape=shape, symmetry=True) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 3f7daa5383..a34c5025e3 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -295,21 +295,26 @@ def create_context(self): def set_quad_rule(params, cell, integral_type, functions): - # Check if the integral has a quad degree attached, otherwise use - # the estimated polynomial degree attached by compute_form_data + # Check if the integral has a quad degree or quad element attached, + # otherwise use the estimated polynomial degree attached by compute_form_data + quad_rule = params.get("quadrature_rule", "default") try: quadrature_degree = params["quadrature_degree"] except KeyError: - quadrature_degree = params["estimated_polynomial_degree"] - function_degrees = [f.ufl_function_space().ufl_element().degree() - for f in functions] - if all((asarray(quadrature_degree) > 10 * asarray(degree)).all() - for degree in function_degrees): + elements = [f.ufl_function_space().ufl_element() for f in functions] + try: + quadrature_degree, = set(e.degree() for e in elements + if e.family() + in ["Quadrature", "Boundary Quadrature"]) + quad_rule, = set(e.quadrature_scheme() or "default" for e in elements) + except ValueError: + quadrature_degree = params["estimated_polynomial_degree"] + if all((asarray(quadrature_degree) > 10 * asarray(e.degree())).all() + for e in elements): logger.warning("Estimated quadrature degree %s more " "than tenfold greater than any " "argument/coefficient degree (max %s)", quadrature_degree, max_degree(function_degrees)) - quad_rule = params.get("quadrature_rule", "default") if isinstance(quad_rule, str): scheme = quad_rule fiat_cell = as_fiat_cell(cell) From d8238dc4f116b8b9249c44f8fb55a1d40e00a05a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 17 Mar 2025 12:10:06 +0000 Subject: [PATCH 09/18] Projection: support for Boundary Quadrature element --- firedrake/output/vtk_output.py | 2 +- firedrake/preconditioners/fdm.py | 9 ++---- firedrake/projection.py | 23 +++++++++----- tests/firedrake/regression/test_projection.py | 31 +++++++++++++++++++ tests/firedrake/regression/test_quadrature.py | 24 ++++++++++++-- tsfc/driver.py | 4 +-- tsfc/kernel_interface/common.py | 2 +- 7 files changed, 74 insertions(+), 21 deletions(-) diff --git a/firedrake/output/vtk_output.py b/firedrake/output/vtk_output.py index 23072c5019..7f71a16cbe 100644 --- a/firedrake/output/vtk_output.py +++ b/firedrake/output/vtk_output.py @@ -81,7 +81,7 @@ def is_dg(V): :arg V: A FunctionSpace. """ - return V.finat_element.entity_dofs() == V.finat_element.entity_closure_dofs() + return V.finat_element.is_dg() def is_linear(V): diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index d0039a3139..1b6aa9c9aa 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -720,10 +720,7 @@ def _element_kernels(self): @cached_property def insert_mode(self): - is_dg = {} - for Vsub in self.V: - element = Vsub.finat_element - is_dg[Vsub] = element.entity_dofs() == element.entity_closure_dofs() + is_dg = {Vsub: Vsub.finat_element.is_dg() for Vsub in self.V} insert_mode = {} for Vrow, Vcol in product(self.V, self.V): @@ -1935,9 +1932,7 @@ def assemble_reference_tensor(self, V): degree = max(e.degree() for e in line_elements) eta = float(self.appctx.get("eta", degree*(degree+1))) - element = V.finat_element - is_dg = element.entity_dofs() == element.entity_closure_dofs() - + is_dg = V.finat_element.is_dg() Afdm = [] # sparse interval mass and stiffness matrices for each direction Dfdm = [] # tabulation of normal derivatives at the boundary for each direction bdof = [] # indices of point evaluation dofs for each direction diff --git a/firedrake/projection.py b/firedrake/projection.py index 4da5a0bb1b..fdc9809084 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -12,7 +12,7 @@ from firedrake import functionspaceimpl from firedrake import function from firedrake.adjoint_utils import annotate_project -from finat import HDivTrace +from finat import HDivTrace, QuadratureElement __all__ = ['project', 'Projector'] @@ -162,8 +162,7 @@ def __init__( self.bcs = bcs self.constant_jacobian = constant_jacobian try: - element = self.target.function_space().finat_element - is_dg = element.entity_dofs() == element.entity_closure_dofs() + is_dg = self.target.function_space().finat_element.is_dg() is_variable_layers = self.target.function_space().mesh().variable_layers except AttributeError: # Mixed space @@ -172,13 +171,22 @@ def __init__( self.use_slate_for_inverse = (use_slate_for_inverse and is_dg and not is_variable_layers and (not complex_mode or SLATE_SUPPORTS_COMPLEX)) + F = self.target.function_space() + needs_trace = False + if type(F.ufl_element()) is not finat.ufl.MixedElement: + if isinstance(F.finat_element, HDivTrace): + needs_trace = True + elif isinstance(F.finat_element, QuadratureElement): + tdim = F.mesh().topological_dimension() + needs_trace = F.finat_element._rule.point_set.dimension == tdim-1 + self.needs_trace = needs_trace + @cached_property def A(self): + F = self.target.function_space() u = firedrake.TrialFunction(self.target.function_space()) v = firedrake.TestFunction(self.target.function_space()) - F = self.target.function_space() - mixed = isinstance(F.ufl_element(), finat.ufl.MixedElement) - if not mixed and isinstance(F.finat_element, HDivTrace): + if self.needs_trace: if F.extruded: a = ( firedrake.inner(u, v)*firedrake.ds_t @@ -244,8 +252,7 @@ class BasicProjector(ProjectorBase): def rhs_form(self): v = firedrake.TestFunction(self.target.function_space()) F = self.target.function_space() - mixed = isinstance(F.ufl_element(), finat.ufl.MixedElement) - if not mixed and isinstance(F.finat_element, HDivTrace): + if self.needs_trace: # Project onto a trace space by supplying the respective form on the facets. # The measures on the facets differ between extruded and non-extruded mesh. # FIXME The restrictions of cg onto the facets is also a trace space, diff --git a/tests/firedrake/regression/test_projection.py b/tests/firedrake/regression/test_projection.py index 176b56c93f..20168d0fcb 100644 --- a/tests/firedrake/regression/test_projection.py +++ b/tests/firedrake/regression/test_projection.py @@ -391,3 +391,34 @@ def run_extr_trace_projection(x, degree=1, family='DGT'): + area * (w - ref) * (w - ref) * ds_b + area * (w('+') - ref('+')) * (w('+') - ref('+')) * dS_h + area * (w('+') - ref('+')) * (w('+') - ref('+')) * dS_v)) + + +def test_project_scalar_boundary_quadrature(): + nx = 2 + msh = UnitSquareMesh(2 ** nx, 2 ** nx) + x = SpatialCoordinate(msh) + f = x[0]*(2-x[0])*x[1]*(2-x[1]) + + T = FunctionSpace(msh, "BQ", 8) + w = Function(T) + w.project(f, solver_parameters={'mat_type': 'matfree', 'ksp_type': 'preonly', 'pc_type': 'jacobi'}) + + err = w - f + err2 = inner(err, err) + error = sqrt(assemble(err2 * ds + err2('+') * dS)) + assert error < 1E-8 + + +def test_project_vector_boundary_quadrature(): + nx = 2 + msh = UnitCubeMesh(2 ** nx, 2 ** nx, 2 ** nx) + n = FacetNormal(msh) + + T = VectorFunctionSpace(msh, "BQ", 0) + w = Function(T) + w.project(n, solver_parameters={'mat_type': 'matfree', 'ksp_type': 'preonly', 'pc_type': 'jacobi'}) + + err = w - n + err2 = inner(err, err) + error = sqrt(assemble(err2 * ds)) + assert error < 1E-8 diff --git a/tests/firedrake/regression/test_quadrature.py b/tests/firedrake/regression/test_quadrature.py index 6ec5f8b855..9c738b888b 100644 --- a/tests/firedrake/regression/test_quadrature.py +++ b/tests/firedrake/regression/test_quadrature.py @@ -1,8 +1,13 @@ from firedrake import * +import pytest -def test_hand_specified_quadrature(): - mesh = UnitSquareMesh(5, 5) +@pytest.fixture +def mesh(request): + return UnitSquareMesh(5, 5) + + +def test_hand_specified_quadrature(mesh): V = FunctionSpace(mesh, 'CG', 2) v = TestFunction(V) @@ -12,3 +17,18 @@ def test_hand_specified_quadrature(): a_q2 = assemble(a, form_compiler_parameters={'quadrature_degree': 2}) assert not np.allclose(a_q0.dat.data, a_q2.dat.data) + + +@pytest.mark.parametrize("diagonal", [False, True]) +@pytest.mark.parametrize("mat_type", ["matfree", "aij"]) +@pytest.mark.parametrize("family", ["Quadrature", "Boundary Quadrature"]) +def test_quadrature_element(mesh, family, mat_type, diagonal): + V = FunctionSpace(mesh, family, 2) + v = TestFunction(V) + u = TrialFunction(V) + if family == "Boundary Quadrature": + a = inner(u, v) * ds + inner(u('+'), v('+')) * dS + else: + a = inner(u, v) * dx + + assemble(a, mat_type=mat_type, diagonal=diagonal) diff --git a/tsfc/driver.py b/tsfc/driver.py index 38a780a0ad..9e518292bc 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -95,10 +95,10 @@ def compile_integral(integral_data, form_data, prefix, parameters, interface, *, interface = firedrake_interface_loopy.KernelBuilder scalar_type = parameters["scalar_type"] integral_type = integral_data.integral_type - if integral_type.startswith("interior_facet") and diagonal: - raise NotImplementedError("Sorry, we can't assemble the diagonal of a form for interior facet integrals") mesh = integral_data.domain arguments = form_data.preprocessed_form.arguments() + if integral_type.startswith("interior_facet") and diagonal and any(a.function_space().finat_element.is_dg() for a in arguments): + raise NotImplementedError("Sorry, we can't assemble the diagonal of a form for interior facet integrals") kernel_name = f"{prefix}_{integral_type}_integral" # Dict mapping domains to index in original_form.ufl_domains() domain_numbering = form_data.original_form.domain_numbering() diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index a34c5025e3..90d4d8d73a 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -314,7 +314,7 @@ def set_quad_rule(params, cell, integral_type, functions): logger.warning("Estimated quadrature degree %s more " "than tenfold greater than any " "argument/coefficient degree (max %s)", - quadrature_degree, max_degree(function_degrees)) + quadrature_degree, max_degree([e.degree() for e in elements])) if isinstance(quad_rule, str): scheme = quad_rule fiat_cell = as_fiat_cell(cell) From 8047496e23c17fe68bd09dd4e5472de23ba7c9f8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 17 Mar 2025 14:28:50 +0000 Subject: [PATCH 10/18] Projector: refactoring --- firedrake/projection.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/firedrake/projection.py b/firedrake/projection.py index fdc9809084..b19ec87102 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -161,31 +161,30 @@ def __init__( self.form_compiler_parameters = form_compiler_parameters self.bcs = bcs self.constant_jacobian = constant_jacobian - try: - is_dg = self.target.function_space().finat_element.is_dg() - is_variable_layers = self.target.function_space().mesh().variable_layers - except AttributeError: - # Mixed space - is_dg = False - is_variable_layers = True - self.use_slate_for_inverse = (use_slate_for_inverse and is_dg and not is_variable_layers - and (not complex_mode or SLATE_SUPPORTS_COMPLEX)) - F = self.target.function_space() needs_trace = False - if type(F.ufl_element()) is not finat.ufl.MixedElement: + F = self.target.function_space() + if type(F.ufl_element()) is finat.ufl.MixedElement: + use_slate_for_inverse = False + else: + if use_slate_for_inverse: + use_slate_for_inverse = (F.finat_element.is_dg() + and not F.mesh().variable_layers + and (not complex_mode or SLATE_SUPPORTS_COMPLEX)) if isinstance(F.finat_element, HDivTrace): needs_trace = True elif isinstance(F.finat_element, QuadratureElement): tdim = F.mesh().topological_dimension() needs_trace = F.finat_element._rule.point_set.dimension == tdim-1 + + self.use_slate_for_inverse = use_slate_for_inverse self.needs_trace = needs_trace @cached_property def A(self): F = self.target.function_space() - u = firedrake.TrialFunction(self.target.function_space()) - v = firedrake.TestFunction(self.target.function_space()) + u = firedrake.TrialFunction(F) + v = firedrake.TestFunction(F) if self.needs_trace: if F.extruded: a = ( @@ -256,7 +255,7 @@ def rhs_form(self): # Project onto a trace space by supplying the respective form on the facets. # The measures on the facets differ between extruded and non-extruded mesh. # FIXME The restrictions of cg onto the facets is also a trace space, - # but we only cover DGT. + # but we only cover DGT and BQ. if F.extruded: form = ( firedrake.inner(self.source, v)*firedrake.ds_t From 09590ed66813f9a218d5435e8edbbba80e2bce7a Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 17 Mar 2025 15:32:34 +0000 Subject: [PATCH 11/18] clean up --- firedrake/projection.py | 20 +++++--------- tests/firedrake/regression/test_projection.py | 6 ++--- tsfc/kernel_interface/common.py | 26 ++++++++++--------- 3 files changed, 23 insertions(+), 29 deletions(-) diff --git a/firedrake/projection.py b/firedrake/projection.py index b19ec87102..88dd720db9 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -162,22 +162,16 @@ def __init__( self.bcs = bcs self.constant_jacobian = constant_jacobian - needs_trace = False F = self.target.function_space() if type(F.ufl_element()) is finat.ufl.MixedElement: - use_slate_for_inverse = False + slate_supported = False + needs_trace = False else: - if use_slate_for_inverse: - use_slate_for_inverse = (F.finat_element.is_dg() - and not F.mesh().variable_layers - and (not complex_mode or SLATE_SUPPORTS_COMPLEX)) - if isinstance(F.finat_element, HDivTrace): - needs_trace = True - elif isinstance(F.finat_element, QuadratureElement): - tdim = F.mesh().topological_dimension() - needs_trace = F.finat_element._rule.point_set.dimension == tdim-1 - - self.use_slate_for_inverse = use_slate_for_inverse + slate_supported = (F.finat_element.is_dg() and not F.mesh().variable_layers + and (not complex_mode or SLATE_SUPPORTS_COMPLEX)) + needs_trace = F.ufl_element().family() in {"HDiv Trace", "Boundary Quadrature"} + + self.use_slate_for_inverse = use_slate_for_inverse and slate_supported self.needs_trace = needs_trace @cached_property diff --git a/tests/firedrake/regression/test_projection.py b/tests/firedrake/regression/test_projection.py index 20168d0fcb..1424f5ffde 100644 --- a/tests/firedrake/regression/test_projection.py +++ b/tests/firedrake/regression/test_projection.py @@ -394,8 +394,7 @@ def run_extr_trace_projection(x, degree=1, family='DGT'): def test_project_scalar_boundary_quadrature(): - nx = 2 - msh = UnitSquareMesh(2 ** nx, 2 ** nx) + msh = UnitSquareMesh(4, 4) x = SpatialCoordinate(msh) f = x[0]*(2-x[0])*x[1]*(2-x[1]) @@ -410,8 +409,7 @@ def test_project_scalar_boundary_quadrature(): def test_project_vector_boundary_quadrature(): - nx = 2 - msh = UnitCubeMesh(2 ** nx, 2 ** nx, 2 ** nx) + msh = UnitCubeMesh(4, 4, 4) n = FacetNormal(msh) T = VectorFunctionSpace(msh, "BQ", 0) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 90d4d8d73a..1131c1f49e 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -302,19 +302,21 @@ def set_quad_rule(params, cell, integral_type, functions): quadrature_degree = params["quadrature_degree"] except KeyError: elements = [f.ufl_function_space().ufl_element() for f in functions] - try: - quadrature_degree, = set(e.degree() for e in elements - if e.family() - in ["Quadrature", "Boundary Quadrature"]) - quad_rule, = set(e.quadrature_scheme() or "default" for e in elements) - except ValueError: + quad_data = set((e.degree(), e.quadrature_scheme() or "default") for e in elements + if e.family() in {"Quadrature", "Boundary Quadrature"}) + if len(quad_data) == 0: quadrature_degree = params["estimated_polynomial_degree"] - if all((asarray(quadrature_degree) > 10 * asarray(e.degree())).all() - for e in elements): - logger.warning("Estimated quadrature degree %s more " - "than tenfold greater than any " - "argument/coefficient degree (max %s)", - quadrature_degree, max_degree([e.degree() for e in elements])) + if all((asarray(quadrature_degree) > 10 * asarray(e.degree())).all() for e in elements): + logger.warning("Estimated quadrature degree %s more " + "than tenfold greater than any " + "argument/coefficient degree (max %s)", + quadrature_degree, max_degree([e.degree() for e in elements])) + else: + try: + (quadrature_degree, quad_rule), = quad_data + except ValueError: + raise ValueError("The quadrature rule cannot be infered from multiple Quadrature elements") + if isinstance(quad_rule, str): scheme = quad_rule fiat_cell = as_fiat_cell(cell) From 36561765d5a71ca734e9730d7970ae6d380119e6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 17 Mar 2025 21:55:22 +0000 Subject: [PATCH 12/18] Revert pyproject.toml --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8b6c0561e1..dfbc156e8c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dependencies = [ "scipy", "sympy", "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git", - "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/finat-boundary-quadrature", + "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git", "pyadjoint-ad @ git+https://github.com/dolfin-adjoint/pyadjoint.git", "loopy @ git+https://github.com/firedrakeproject/loopy.git@main", ] From c28e86bde811b6e4b7a2f0e945ad3abb7fc7605d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 17 Mar 2025 21:56:16 +0000 Subject: [PATCH 13/18] checkout branches --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 92e385f27b..8c2c7d5eb9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -86,7 +86,7 @@ jobs: --extra-index-url https://download.pytorch.org/whl/cpu \ './firedrake-repo[ci]' pip install "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git@pbrubeck/remove-component-tensors" - pip install "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplify-indexed" + pip install "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/finat-boundary-quadrature" firedrake-clean pip list From 572a99b668c3412552003df3d79e325d9337c459 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 17 Mar 2025 21:57:49 +0000 Subject: [PATCH 14/18] lint --- firedrake/projection.py | 1 - 1 file changed, 1 deletion(-) diff --git a/firedrake/projection.py b/firedrake/projection.py index 88dd720db9..6c5b0c1908 100644 --- a/firedrake/projection.py +++ b/firedrake/projection.py @@ -12,7 +12,6 @@ from firedrake import functionspaceimpl from firedrake import function from firedrake.adjoint_utils import annotate_project -from finat import HDivTrace, QuadratureElement __all__ = ['project', 'Projector'] From 16eec724eda480a16cfc83ad7bc486f7949f10a9 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 18 Mar 2025 11:34:05 +0000 Subject: [PATCH 15/18] force reinstall --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 8c2c7d5eb9..777e4a7438 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -85,8 +85,8 @@ jobs: --no-binary h5py \ --extra-index-url https://download.pytorch.org/whl/cpu \ './firedrake-repo[ci]' - pip install "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git@pbrubeck/remove-component-tensors" - pip install "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/finat-boundary-quadrature" + pip install -I "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git@pbrubeck/remove-component-tensors" + pip install -I "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/finat-boundary-quadrature" firedrake-clean pip list From 2cc0110f8375d0453d8d8a8ce344696fdda482fe Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 18 Mar 2025 15:11:05 +0000 Subject: [PATCH 16/18] Checkout branch --- .github/workflows/build.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 777e4a7438..1ef9d4a3f0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -85,7 +85,6 @@ jobs: --no-binary h5py \ --extra-index-url https://download.pytorch.org/whl/cpu \ './firedrake-repo[ci]' - pip install -I "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git@pbrubeck/remove-component-tensors" pip install -I "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/finat-boundary-quadrature" firedrake-clean pip list From 0d3776f9d9c2a34c1c55b2190f3cd314e3fea1b5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 27 Mar 2025 10:05:48 +0000 Subject: [PATCH 17/18] Apply suggestions from code review --- tsfc/fem.py | 2 -- tsfc/kernel_interface/common.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tsfc/fem.py b/tsfc/fem.py index e20244a0a1..d86be550db 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -656,14 +656,12 @@ def callback(entity_id): # A numerical hack that FFC used to apply on FIAT tables still # lives on after ditching FFC and switching to FInAT. return ffc_rounding(square, ctx.epsilon) - table = ctx.entity_selector(callback, mt.restriction) if ctx.use_canonical_quadrature_point_ordering: quad_multiindex = ctx.quadrature_rule.point_set.indices quad_multiindex_permuted = _make_quad_multiindex_permuted(mt, ctx) mapper = gem.node.MemoizerArg(gem.optimise.filtered_replace_indices) table = mapper(table, tuple(zip(quad_multiindex, quad_multiindex_permuted))) - argument_multiindex = ctx.argument_multiindices[terminal.number()] return gem.partial_indexed(table, argument_multiindex) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 1131c1f49e..667dab75a8 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -315,7 +315,7 @@ def set_quad_rule(params, cell, integral_type, functions): try: (quadrature_degree, quad_rule), = quad_data except ValueError: - raise ValueError("The quadrature rule cannot be infered from multiple Quadrature elements") + raise ValueError("The quadrature rule cannot be inferred from multiple Quadrature elements") if isinstance(quad_rule, str): scheme = quad_rule From 1c9c122b08518a2089f68a7e7698cdda0fd5b520 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 28 Mar 2025 12:06:50 +0000 Subject: [PATCH 18/18] Update .github/workflows/build.yml --- .github/workflows/build.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 1ef9d4a3f0..b8a5582c5b 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -85,7 +85,6 @@ jobs: --no-binary h5py \ --extra-index-url https://download.pytorch.org/whl/cpu \ './firedrake-repo[ci]' - pip install -I "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/finat-boundary-quadrature" firedrake-clean pip list