From 5357fa71e584246be2eefd886aa977cf08f0f85b Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Mon, 22 Jun 2026 11:20:20 +0100 Subject: [PATCH 01/24] Fix submesh deadlock Don't branch on local size. --- .github/workflows/core.yml | 6 ++++-- firedrake/cython/dmcommon.pyx | 4 ++-- firedrake/extrusion_utils.py | 3 ++- firedrake/mesh.py | 4 ++-- pyop3/axis_tree/parallel.py | 1 + pyop3/expr/tensor/dat.py | 24 ++++++++++++++++++++---- 6 files changed, 31 insertions(+), 11 deletions(-) diff --git a/.github/workflows/core.yml b/.github/workflows/core.yml index 7a759d5b68..22e257b883 100644 --- a/.github/workflows/core.yml +++ b/.github/workflows/core.yml @@ -210,7 +210,8 @@ jobs: if: success() || steps.install.conclusion == 'success' run: | . venv/bin/activate - git clone --depth 1 https://github.com/thetisproject/thetis.git thetis-repo --branch ${{ inputs.base_ref }} + : # UNDO ME + git clone --depth 1 https://github.com/thetisproject/thetis.git thetis-repo --branch connorjward/pyop3 pip install --verbose ./thetis-repo python -m pytest -n 8 --verbose thetis-repo/test_adjoint/test_swe_adjoint.py timeout-minutes: 10 @@ -228,7 +229,8 @@ jobs: if: success() || steps.install.conclusion == 'success' run: | . venv/bin/activate - git clone --depth 1 https://github.com/g-adopt/g-adopt.git g-adopt-repo --branch ${{ inputs.base_ref }} + : # UNDO ME + git clone --depth 1 https://github.com/connorjward/g-adopt.git g-adopt-repo --branch connorjward/pyop3 pip install --verbose ./g-adopt-repo make -C g-adopt-repo/demos/mantle_convection/base_case check timeout-minutes: 5 diff --git a/firedrake/cython/dmcommon.pyx b/firedrake/cython/dmcommon.pyx index 4c786e7830..706a3cbae2 100644 --- a/firedrake/cython/dmcommon.pyx +++ b/firedrake/cython/dmcommon.pyx @@ -4626,8 +4626,8 @@ def filter_is(is_: PETSc.IS, start: IntType, end: IntType) -> PETSc.IS: cdef: PETSc.IS filtered_is - # empty ISes remain empty - if not is_ or is_.size == 0: + # cast null ISes to empty ones + if not is_: return PETSc.IS().createGeneral(np.empty(0, dtype=IntType)) filtered_is = is_.duplicate() diff --git a/firedrake/extrusion_utils.py b/firedrake/extrusion_utils.py index 0210a5b2b7..2424808adc 100644 --- a/firedrake/extrusion_utils.py +++ b/firedrake/extrusion_utils.py @@ -69,7 +69,8 @@ def make_extruded_coords(extruded_topology, base_coords, ext_coords, layer_height = numpy.cumsum(numpy.concatenate(([0], layer_height))) layer_heights = layer_height.size - layer_height = op3.Dat.from_array(layer_height) + # NOTE: Not sure about this in parallel, needs an SF? + layer_height = op3.Dat.from_array(layer_height, comm=extruded_topology.comm) if kernel is not None: raise NotImplementedError diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 03e1cf0707..d684cdb2db 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -1268,7 +1268,7 @@ def _facet_support_dat(self, facet_type: Literal["exterior"] | Literal["interior [ op3.Subset( facet_support_dat.axes.root.component.label, - op3.Dat.from_array(selected_facets), + op3.Dat.from_array(selected_facets, comm=self.comm), label=facet_axis.component.label, ) ], @@ -1961,7 +1961,7 @@ def interior_facets(self) -> op3.IndexedAxisTree: # Maybe doesn't have to be a method either def _facet_subset(self, plex_indices_is: PETSc.IS, component_renumbering: PETSc.Section, component_label) -> op3.Slice: subset_indices = dmcommon.section_offsets(component_renumbering, plex_indices_is, sort=True) - subset_dat = op3.Dat.from_array(subset_indices.indices) + subset_dat = op3.Dat.from_array(subset_indices.indices, comm=self.comm) return op3.Slice(self.name, [op3.Subset(component_label, subset_dat)]) @cached_property diff --git a/pyop3/axis_tree/parallel.py b/pyop3/axis_tree/parallel.py index ca370f816a..f0131eab68 100644 --- a/pyop3/axis_tree/parallel.py +++ b/pyop3/axis_tree/parallel.py @@ -83,6 +83,7 @@ def _collect_sf_graphs_rec(axis_tree: AbstractNonUnitAxisTree, path: ConcretePat raise NotImplementedError("This will be very inefficient") # FIXME: Only need to call the inner bit once and repeatedly add? + # FIXME: This is not parallel safe for point in range(component.local_size): sfs.extend( _collect_sf_graphs_rec(axis_tree, path_, indices | {axis.label: point}) diff --git a/pyop3/expr/tensor/dat.py b/pyop3/expr/tensor/dat.py index a130eafdc9..98c76865e6 100644 --- a/pyop3/expr/tensor/dat.py +++ b/pyop3/expr/tensor/dat.py @@ -16,6 +16,7 @@ from petsc4py import PETSc import pyop3.arrayref +import pyop3.axis_tree import pyop3.device import pyop3.record from pyop3 import utils @@ -240,9 +241,24 @@ def null(cls, axes, dtype=AbstractBuffer.DEFAULT_DTYPE, *, buffer_kwargs=idict() return cls(axes, buffer=buffer, **kwargs) @classmethod - def from_array(cls, array: np.ndarray, *, buffer_kwargs=None, **kwargs) -> Dat: + def from_array( + cls, + array: np.ndarray, + *, + sf: StarForest | None = None, + comm: MPI.Comm = MPI.COMM_SELF, + buffer_kwargs=None, + **kwargs, + ) -> Dat: from pyop3 import Scalar + if sf and sf.comm != comm: + raise pyop3.exceptions.CommMismatchException + + # If no SF is provided then we assume no overlap + if sf is None: + sf = pyop3.sf.local_sf(array.size, comm=comm) + buffer_kwargs = buffer_kwargs or {} name = utils.maybe_generate_name(kwargs.pop("name", None), kwargs.pop("prefix", None), cls.DEFAULT_PREFIX) @@ -252,9 +268,9 @@ def from_array(cls, array: np.ndarray, *, buffer_kwargs=None, **kwargs) -> Dat: if "name" not in buffer_kwargs: buffer_kwargs["name"] = f"{name}_buffer" - # NOTE: Should this size *always* be a Scalar? - axes = Axis(Scalar(array.size)) - buffer = ArrayBuffer(array, **buffer_kwargs) + # NOTE: Should this size *always* be a Scalar? maybe not if rank_constant is True... + axes = pyop3.axis_tree.Axis(pyop3.axis_tree.AxisComponent(Scalar(array.size), sf=sf)) + buffer = ArrayBuffer(array, sf=sf, **buffer_kwargs) return cls(axes, buffer=buffer, **kwargs) @classmethod From 2dd00a8cdc51056d21d910563ccc082176005247 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Mon, 22 Jun 2026 17:41:17 +0100 Subject: [PATCH 02/24] Parallel fixes for ensemble, avoid weakref and unordered sets AT ALL COSTS --- firedrake/mesh.py | 12 ++--- firedrake/variational_solver.py | 4 +- pyop3/axis_tree/parallel.py | 11 ++--- pyop3/axis_tree/tree.py | 82 ++++++++++++++------------------ pyop3/axis_tree/visitors/size.py | 10 ++-- pyop3/buffer.py | 8 ++-- pyop3/cache.py | 17 +++++-- pyop3/collections.py | 11 +++++ pyop3/expr/base.py | 20 ++++++++ pyop3/expr/buffer.py | 3 +- pyop3/expr/tensor/base.py | 3 +- pyop3/expr/tensor/mat.py | 3 +- pyop3/expr/visitors/__init__.py | 1 + pyop3/index_tree/tree.py | 4 ++ pyop3/insn/base.py | 41 ++++++++-------- pyop3/mpi.py | 59 ++++++++++++++++++++++- pyop3/obj.py | 6 +++ pyop3/sf.py | 29 +---------- pyop3/utils.py | 47 ------------------ pyop3/visitors.py | 64 +++++++++++++++++++++++++ 20 files changed, 262 insertions(+), 173 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index d684cdb2db..32a0d9d16b 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2466,7 +2466,7 @@ def submesh_map_child_parent(self, source_integral_type, source_subset_points, r # or to a mix of facet types). # We don't know a priori what this type is so we instead detect it here. target_exterior_facets = dmcommon.intersect_is( - PETSc.IS().createGeneral(target_subset_points), + PETSc.IS().createGeneral(target_subset_points, comm=MPI.COMM_SELF), target._exterior_facet_plex_indices, ) with temp_internal_comm(self.comm) as icomm: @@ -2475,7 +2475,7 @@ def submesh_map_child_parent(self, source_integral_type, source_subset_points, r ) target_interior_facets = dmcommon.intersect_is( - PETSc.IS().createGeneral(target_subset_points), + PETSc.IS().createGeneral(target_subset_points, comm=MPI.COMM_SELF), target._interior_facet_plex_indices, ) with temp_internal_comm(self.comm) as icomm: @@ -6895,13 +6895,13 @@ def get_mesh_topologies(expr) -> frozenset[AbstractMeshTopology]: # FIXME: This isn't valid for certain inputs (e.g. ZeroBaseForm) but this # is a very heavy-handed way to fix that try: - return frozenset({d.topology for d in extract_domains(expr)}) + return op3.collections.OrderedFrozenSet([d_.topology for d in extract_domains(expr) for d_ in d]) except: - return frozenset() + return () def extract_mesh_topologies(mesh) -> frozenset[MeshTopology]: if isinstance(mesh, MeshSequenceGeometry): - return frozenset({m.topology for m in mesh}) + return op3.collections.OrderedFrozenSet([m.topology for m in mesh]) else: - return frozenset({mesh.topology}) + return op3.collections.OrderedFrozenSet([mesh.topology]) diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py index 101b666f6f..cda0f2f69c 100644 --- a/firedrake/variational_solver.py +++ b/firedrake/variational_solver.py @@ -146,9 +146,9 @@ def _mesh_topologies(self) -> frozenset: # TODO: This breaks for certain inputs (e.g. FormSum) but this # is a very heavy-handed way to fix that try: - return frozenset({d.topology for d in extract_domains(self.F)}) + return op3.collections.OrderedFrozenSet([d.topology for d in extract_domains(self.F)]) except: - return frozenset() + return () @staticmethod def compute_bc_lifting(J: ufl.BaseForm | slate.TensorBase, diff --git a/pyop3/axis_tree/parallel.py b/pyop3/axis_tree/parallel.py index f0131eab68..6e3c0b57dc 100644 --- a/pyop3/axis_tree/parallel.py +++ b/pyop3/axis_tree/parallel.py @@ -9,6 +9,7 @@ from mpi4py import MPI from pyop3.axis_tree.tree import AbstractNonUnitAxisTree +import pyop3.visitors from pyop3 import utils from pyop3.dtypes import IntType, as_numpy_dtype from pyop3.sf import StarForest, create_petsc_section_sf @@ -81,10 +82,8 @@ def _collect_sf_graphs_rec(axis_tree: AbstractNonUnitAxisTree, path: ConcretePat elif subaxis := axis_tree.node_map.get(path_): if isinstance(size := component.size, numbers.Integral) and size > 1: raise NotImplementedError("This will be very inefficient") - - # FIXME: Only need to call the inner bit once and repeatedly add? - # FIXME: This is not parallel safe - for point in range(component.local_size): + # have to use .size here because .local_size means we deadlock + for point in range(component.size): sfs.extend( _collect_sf_graphs_rec(axis_tree, path_, indices | {axis.label: point}) ) @@ -163,7 +162,7 @@ def concatenate_star_forests(star_forests: Sequence[StarForest]) -> StarForest: ilocal = np.concatenate(local_leaf_indicess) iremote = np.concatenate(remote_leaf_indicess) - comm = utils.single_comm(star_forests, "comm") + comm = pyop3.visitors.single_comm(*star_forests) return StarForest.from_graph(total_size, ilocal, iremote, comm) # old code, it fixed some bug or other... @@ -187,7 +186,7 @@ def concatenate_star_forests(star_forests: Sequence[StarForest]) -> StarForest: ilocal = np.concatenate(local_leaf_indicess) iremote = np.concatenate(remote_leaf_indicess) - comm = utils.single_comm(star_forests, "comm") + comm = pyop3.visitors.single_comm(*star_forests) return StarForest.from_graph(size, ilocal, iremote, comm) # because ghost points are already at the back? diff --git a/pyop3/axis_tree/tree.py b/pyop3/axis_tree/tree.py index 7e394be266..fae86031ce 100644 --- a/pyop3/axis_tree/tree.py +++ b/pyop3/axis_tree/tree.py @@ -34,7 +34,7 @@ from pyop3.constants import PYOP3_DECIDE from pyop3.dtypes import IntType from pyop3.exceptions import InvalidIndexTargetException, Pyop3Exception -from pyop3.sf import DistributedObject, AbstractStarForest, NullStarForest, ParallelAwareObject, StarForest, local_sf, single_star_sf +from pyop3.sf import AbstractStarForest, NullStarForest, StarForest, local_sf, single_star_sf from pyop3.mpi import collective, temp_internal_comm from pyop3 import utils from pyop3.labeled_tree import ( @@ -230,6 +230,10 @@ def get_disk_cache_key(self, visitor) -> Hashable: get_instruction_executor_cache_key = get_disk_cache_key + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.get_comm(self.size) + def __init__(self, size, label=None): from pyop3 import as_linear_buffer_expression, Tensor @@ -258,13 +262,6 @@ def __post_init__(self) -> None: # }}} - @property - def comm(self) -> MPI.Comm: - if isinstance(self.size, numbers.Integral): - return MPI.COMM_SELF - else: - return self.size.comm - def __str__(self) -> str: if self.label is None: return str(self.size) @@ -385,6 +382,10 @@ def get_disk_cache_key(self, visitor) -> Hashable: get_instruction_executor_cache_key = get_disk_cache_key + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self.regions, self.size, self.sf) + def __init__( self, regions, @@ -493,10 +494,6 @@ def _all_regions(self) -> tuple[AxisComponentRegion]: def has_non_trivial_regions(self) -> bool: return len(self.regions) > 1 or utils.just_one(self.regions).label is not None - @property - def comm(self) -> MPI.Comm | None: - return self.sf.comm if self.sf else None - @property def region_labels(self) -> tuple[ComponentRegionLabelT]: return tuple(r.label for r in self.regions) @@ -557,7 +554,7 @@ def regionless(self) -> AxisComponent: @pyop3.record.frozenrecord() -class Axis(LoopIterable, MultiComponentLabelledNode, ParallelAwareObject): +class Axis(LoopIterable, MultiComponentLabelledNode): # {{{ instance attrs @@ -572,6 +569,10 @@ def get_disk_cache_key(self, visitor) -> Hashable: get_instruction_executor_cache_key = get_disk_cache_key + @cached_property + def comm(self) -> MPI.Comm | None: + return pyop3.visitors.common_comm(*self.components) + def __init__( self, components, @@ -651,10 +652,6 @@ def component_index(self, component) -> int: def matching_component(self, component_label: ComponentLabelT) -> AxisComponent: return self.components[self.component_index(component_label)] - @property - def comm(self) -> MPI.Comm | None: - return utils.single_comm(self.components, "comm", allow_undefined=True) - @property def size(self): return self._tree.size @@ -894,7 +891,7 @@ def __contains__(self, obj: Any, /) -> bool: node_map = idict({idict(): None}) -class AbstractNonUnitAxisTree(AbstractAxisTree, ContextFreeLoopIterable, LabelledTree, DistributedObject): +class AbstractNonUnitAxisTree(AbstractAxisTree, ContextFreeLoopIterable, LabelledTree): """Base class for non-unit axis trees.""" # {{{ abstract methods @@ -1424,13 +1421,6 @@ def index(self) -> LoopIndex: return LoopIndex(self) - @property - def comm(self): - from pyop3.debug import warn_todo - warn_todo("This comm choice is unsafe") - return MPI.COMM_SELF - - UNIT_AXIS_TREE = _UnitAxisTree() """Placeholder value for an axis tree that is guaranteed to have a single entry. @@ -1463,6 +1453,10 @@ def get_disk_cache_key(self, visitor) -> Hashable: get_instruction_executor_cache_key = get_disk_cache_key + @cached_property + def comm(self): + return pyop3.visitors.common_comm(*self._node_map.values()) + def __init__(self, node_map: Mapping[PathT, Node] | None | None = None) -> None: object.__setattr__(self, "_node_map", as_node_map(node_map)) @@ -1513,10 +1507,6 @@ def blocked(self, block_shape: Sequence[int, ...] | int) -> AxisTree: else: return self[self._block_indices(block_shape)].materialize() - @property - def comm(self): - return utils.single_comm(self.nodes, "comm", allow_undefined=True) or MPI.COMM_SELF - # TODO: rename to local_section def section(self, path: PathT, component: ComponentT, indices=idict()) -> PETSc.Section: # NOTE: This is the same as indexedaxistree but offsets are known to increase linearly @@ -1740,6 +1730,10 @@ def get_instruction_executor_cache_key(self, visitor) -> Hashable: return (type(self), node_map_key, visitor(self._unindexed), targets_key) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self._node_map.values(), self._unindexed) + # TODO: where to put *, and order? def __init__( self, @@ -1944,10 +1938,6 @@ def section(self, path: PathT, component: ComponentT, indices=idict()) -> PETSc. # }}} - @property - def comm(self): - return self.unindexed.comm - @property def layouts(self): return self.unindexed.layouts @@ -2083,6 +2073,10 @@ def get_instruction_executor_cache_key(self, visitor) -> Hashable: return (type(self), visitor(self._unindexed), targets_key) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self._node_maps.values(), self.unindexed) + def __init__( self, unindexed: AxisTree | None, @@ -2123,10 +2117,6 @@ def targets(self) -> tuple[idict[ConcretePathT, tuple[AxisTarget, ...]], ...]: idict(): self._targets[idict()] + self.materialize().targets[idict()] }) - @property - def comm(self) -> MPI.Comm: - return self.unindexed.comm - def materialize(self): return UNIT_AXIS_TREE @@ -2291,7 +2281,11 @@ class AxisForest(AbstractAxisTreeLike): _trees: tuple def get_instruction_executor_cache_key (self, visitor) -> Hashable: - return (type(self), tuple(map(visitor, self.trees))) + return (type(self), tuple(map(visitor, self._trees))) + + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self._trees) def __new__( cls, @@ -2440,10 +2434,6 @@ def getitem(self, indices, *, strict=False): else: return AxisForest(indexed_trees_) - @property - def comm(self) -> MPI.Comm: - return utils.common_comm(self.trees, "comm") - def materialize(self) -> AxisForest: return type(self)((tree.materialize() for tree in self.trees)) @@ -2521,6 +2511,10 @@ def get_instruction_executor_cache_key(self, visitor) -> Hashable: trees_key = idict(trees_key) return (type(self), trees_key) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self.trees.values()) + def __init__(self, trees: Mapping): trees = idict(trees) @@ -2536,10 +2530,6 @@ def __post_init__(self) -> None: def context_map(self): # old alias return self.trees - @property - def comm(self) -> MPI.Comm: - return utils.single_comm(self.context_map.values(), "comm") - def __repr__(self) -> str: return f"{type(self).__name__}({self.context_map!r})" diff --git a/pyop3/axis_tree/visitors/size.py b/pyop3/axis_tree/visitors/size.py index b6444d4a06..98edc8cac4 100644 --- a/pyop3/axis_tree/visitors/size.py +++ b/pyop3/axis_tree/visitors/size.py @@ -42,6 +42,7 @@ def _axis_tree_size_rec(axis_tree: AxisTree, path): # TODO: just be a cached method? Or globally cache? @cached_on(lambda tree, *a, **kw: tree, lambda tree, path, label: (path, label)) def compute_axis_tree_component_size(axis_tree: AbstractNonUnitAxisTree, path: PathT, component_label: ComponentLabelT): + import pyop3 from pyop3 import Scalar from pyop3.expr.visitors import replace_terminals, replace @@ -126,9 +127,10 @@ def compute_axis_tree_component_size(axis_tree: AbstractNonUnitAxisTree, path: P component_size = replace(component_size, axis_to_loop_var_replace_map) subtree_size_expr = replace(subtree_size_tmp, axis_to_loop_var_replace_map) - Loop(i, - component_size.iassign(subtree_size_expr) - )() + pyop3.loop(i, + component_size.iassign(subtree_size_expr), + eager=True, + ) if component_size_axes is UNIT_AXIS_TREE: # ick way to make sure that if we have sizes wrapped up into Scalars that this @@ -147,5 +149,3 @@ def compute_axis_tree_component_size(axis_tree: AbstractNonUnitAxisTree, path: P return replace(XXX, axis_to_loop_var_replace_map) else: return component.size * subtree_size - - diff --git a/pyop3/buffer.py b/pyop3/buffer.py index f1c3387eba..a13a234f37 100644 --- a/pyop3/buffer.py +++ b/pyop3/buffer.py @@ -22,7 +22,7 @@ from pyop3.cache import cached_method from pyop3.collections import OrderedFrozenSet from pyop3.dtypes import IntType, ScalarType, DTypeT -from pyop3.sf import DistributedObject, NullStarForest, StarForest, local_sf +from pyop3.sf import NullStarForest, StarForest, local_sf from pyop3.utils import UniqueNameGenerator, as_tuple, deprecated, maybe_generate_name, readonly from pyop3.device import ( Device, @@ -898,7 +898,7 @@ def _make_petsc_mat( submat = cls._make_petsc_mat(submat_spec, preallocator=preallocator) submats[i, j] = submat - comm = utils.single_comm(submats.flatten(), "comm") + comm = pyop3.visitors.single_comm(*submats.flatten()) return PETSc.Mat().createNest(submats, comm=comm) else: assert isinstance(mat_spec, FullPetscMatBufferSpec) @@ -916,7 +916,7 @@ def _make_non_nested_petsc_mat(cls, mat_spec: FullPetscMatBufferSpec, *, preallo row_axes = row_spec column_axes = column_spec - comm = utils.single_comm([row_axes, column_axes], "comm") + comm = pyop3.visitors.single_comm(row_axes, column_axes) if mat_type == "rvec": mode = "row" @@ -932,7 +932,7 @@ def _make_non_nested_petsc_mat(cls, mat_spec: FullPetscMatBufferSpec, *, preallo if preallocator: mat_type = PETSc.Mat.Type.PREALLOCATOR - comm = utils.single_comm([row_spec.lgmap, column_spec.lgmap], "comm") + comm = pyop3.visitors.single_comm(row_spec.lgmap, column_spec.lgmap) mat = PETSc.Mat().create(comm) mat.setType(mat_type) diff --git a/pyop3/cache.py b/pyop3/cache.py index b91bdc3eb5..6c14c3fa81 100644 --- a/pyop3/cache.py +++ b/pyop3/cache.py @@ -542,11 +542,17 @@ def make_instrumented_cache(): comm_caches = get_comm_caches(comm) if heavy: if cache_id not in comm_caches: - comm_caches[cache_id] = weakref.WeakKeyDictionary() + # comm_caches[cache_id] = weakref.WeakKeyDictionary() + comm_caches[cache_id] = {} caches = [] cache_type = None for lifetime_obj in _heavy_caches: + assert lifetime_obj.comm.size >= comm.size + if pyop3.config.spmd_strict: + # cache access must be collective over lifetime objects + lifetime_obj.comm.barrier() + try: cache = comm_caches[cache_id][lifetime_obj] except KeyError: @@ -657,7 +663,9 @@ def decorator(func): return decorator -_heavy_caches = weakref.WeakSet() +# NOTE: It is safe to reference the object because it is only accessible +# via a context manager that pops the reference after use. +_heavy_caches = pyop3.collections.OrderedSet() class heavy_caches: @@ -675,7 +683,7 @@ def __init__(self, objs: Any) -> None: self._objs = objs # keep track of the objects we inserted ourselves, if they were already # there then we don't want to remove them! - self._added_objs = set() + self._added_objs = pyop3.collections.OrderedSet() def __enter__(self) -> None: for obj in self._objs: @@ -694,11 +702,12 @@ def with_heavy_caches(get_obj: Callable) -> Callable: def decorator(func): def wrapper(*args, **kwargs): obj = get_obj(*args, **kwargs) + assert pyop3.collections.is_ordered_sequence(obj) with heavy_caches(obj): return func(*args, **kwargs) return wrapper return decorator -with_self_heavy_cache = with_heavy_caches(lambda self, *a, **kw: {self}) +with_self_heavy_cache = with_heavy_caches(lambda self, *a, **kw: (self,)) """Method decorator that sets ``self`` as a heavy cache.""" diff --git a/pyop3/collections.py b/pyop3/collections.py index 872d53162d..ff1171908e 100644 --- a/pyop3/collections.py +++ b/pyop3/collections.py @@ -121,6 +121,17 @@ def update(self, /, *others): for item in other: self.add(item) + def remove(self, value): + try: + index = self._values.index(value) + except ValueError: + raise KeyError + else: + self._values.pop(index) + + def clear(self): + self._values.clear() + class OrderedFrozenSet(AbstractOrderedSet): diff --git a/pyop3/expr/base.py b/pyop3/expr/base.py index 752a380ad3..e63993a3ed 100644 --- a/pyop3/expr/base.py +++ b/pyop3/expr/base.py @@ -171,6 +171,14 @@ def operands(self) -> tuple[ExpressionT, ...]: # }}} + # {{{ interface impls + + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self.operands) + + # }}} + @pyop3.record.frozenrecord() class UnaryOperator(Operator, metaclass=abc.ABCMeta): @@ -535,6 +543,10 @@ def get_disk_cache_key(self, visitor) -> Hashable: get_instruction_executor_cache_key = get_disk_cache_key + @property + def comm(self) -> MPI.Comm: + return self.axis.comm + def __init__(self, axis: Axis) -> None: assert len(axis.components) == 1 assert axis.component.sf is None @@ -576,6 +588,10 @@ def disk_cache_key(self, renamer): def instruction_executor_cache_key(self, renamer): return (type(self),) + @property + def comm(self) -> MPI.Comm: + return MPI.COMM_SELF + @property def local_max(self) -> NoReturn: raise TypeError @@ -616,6 +632,10 @@ def get_disk_cache_key(self, visitor) -> Hashable: get_instruction_executor_cache_key = get_disk_cache_key + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(self.loop_index, self.axis) + def __init__(self, loop_index, axis) -> None: from pyop3 import LoopIndex diff --git a/pyop3/expr/buffer.py b/pyop3/expr/buffer.py index 03242e7ebb..66090d97db 100644 --- a/pyop3/expr/buffer.py +++ b/pyop3/expr/buffer.py @@ -13,7 +13,6 @@ from pyop3.labeled_tree import is_subpath from pyop3.axis_tree import UNIT_AXIS_TREE from pyop3.buffer import AbstractBuffer, ArrayBuffer -from pyop3.sf import DistributedObject from pyop3.collections import OrderedFrozenSet from .base import Expression, as_str @@ -21,7 +20,7 @@ # TODO: Should inherit from Terminal (but Terminal has odd attrs) -class BufferExpression(Expression, DistributedObject, metaclass=abc.ABCMeta): +class BufferExpression(Expression, metaclass=abc.ABCMeta): # {{{ abstract methods diff --git a/pyop3/expr/tensor/base.py b/pyop3/expr/tensor/base.py index 0f5610709b..1c61683b94 100644 --- a/pyop3/expr/tensor/base.py +++ b/pyop3/expr/tensor/base.py @@ -17,7 +17,6 @@ from pyop3.expr.base import ExpressionT import pyop3.record from pyop3 import utils -from pyop3.sf import DistributedObject from pyop3.axis_tree import ContextAware from pyop3.axis_tree.tree import AbstractNonUnitAxisTree from pyop3.expr import TerminalExpression @@ -28,7 +27,7 @@ import pyop3.insn.exec -class Tensor(ContextAware, TerminalExpression, DistributedObject, abc.ABC): +class Tensor(ContextAware, TerminalExpression, abc.ABC): DEFAULT_PREFIX: ClassVar[str] = "array" diff --git a/pyop3/expr/tensor/mat.py b/pyop3/expr/tensor/mat.py index 9fa8b6c9a9..fc157ec24b 100644 --- a/pyop3/expr/tensor/mat.py +++ b/pyop3/expr/tensor/mat.py @@ -18,6 +18,7 @@ import pyop3.dtypes import pyop3.index_tree import pyop3.record +import pyop3.visitors from pyop3 import utils from pyop3.cache import cached_method from .base import Tensor, ReshapeTensorTransform, TensorTransform @@ -357,7 +358,7 @@ def nest_labels(self) -> tuple[tuple[int, int], ...]: def make_full_mat_buffer_spec(partial_spec: PetscMatBufferSpec, row_axes: AbstractNonUnitAxisTree, column_axes: AbstractNonUnitAxisTree) -> FullMatBufferSpec: if isinstance(partial_spec, NonNestedPetscMatBufferSpec): - comm = utils.common_comm((row_axes, column_axes), "comm") + comm = pyop3.visitors.common_comm(row_axes, column_axes) if partial_spec.mat_type in {"rvec", "cvec"}: row_spec = row_axes diff --git a/pyop3/expr/visitors/__init__.py b/pyop3/expr/visitors/__init__.py index abf3d11212..bce71ec7e6 100644 --- a/pyop3/expr/visitors/__init__.py +++ b/pyop3/expr/visitors/__init__.py @@ -878,6 +878,7 @@ def _(dat, /) -> OrderedFrozenSet: @memory_cache(heavy=True) +@pyop3.mpi.collective def materialize_composite_dat(composite_dat: pyop3.expr.CompositeDat, comm: MPI.Comm) -> pyop3.expr.LinearDatBufferExpression: axes = composite_dat.axis_tree diff --git a/pyop3/index_tree/tree.py b/pyop3/index_tree/tree.py index b53b2e5e3f..2fce6cc9a1 100644 --- a/pyop3/index_tree/tree.py +++ b/pyop3/index_tree/tree.py @@ -401,6 +401,10 @@ def get_instruction_executor_cache_key(self, visitor): visitor.renamer.add(self.id, "LoopIndex"), ) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.get_comm(self.iterset) + def __init__(self, iterset: AbstractNonUnitAxisTree, *, id=utils.PYOP3_DECIDE): id = id if id is not utils.PYOP3_DECIDE else self.unique_label() diff --git a/pyop3/insn/base.py b/pyop3/insn/base.py index 701558f2b3..2139f4d02c 100644 --- a/pyop3/insn/base.py +++ b/pyop3/insn/base.py @@ -26,6 +26,7 @@ import pyop3.compile import pyop3.expr import pyop3.record +import pyop3.visitors from pyop3 import utils from pyop3.cache import with_heavy_caches, with_self_heavy_cache, memory_cache, cached_method from pyop3.collections import OrderedFrozenSet, OrderedSet, is_ordered_mapping @@ -33,7 +34,6 @@ from pyop3.axis_tree import AxisTree from pyop3.axis_tree.tree import UNIT_AXIS_TREE, AxisForest, ContextFree, ContextSensitive, axis_tree_is_valid_subset, matching_axis_tree from pyop3.expr import BufferExpression, Tensor, Scalar, Dat, Mat -from pyop3.sf import DistributedObject from pyop3.dtypes import dtype_limits from pyop3.exceptions import Pyop3Exception from pyop3.utils import ( @@ -98,7 +98,7 @@ class UnprocessedExpressionException(Pyop3Exception): """Exception raised when pyop3 expected a preprocessed expression.""" -class Instruction(Node, DistributedObject, abc.ABC): +class Instruction(Node, abc.ABC): def __init__(self) -> None: object.__setattr__(self, "_hit_executor_cache", True) @@ -185,6 +185,10 @@ def get_instruction_executor_cache_key(self, visitor) -> Hashable: tuple(map(visitor, self.statements)), ) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(self.index, *self.statements) + def __init__( self, index: LoopIndex, @@ -205,11 +209,6 @@ def __init__( def global_arguments(self) -> OrderedFrozenSet[Tensor]: return OrderedFrozenSet().union(*(stmt.global_arguments for stmt in self.statements)) - @property - def comm(self) -> MPI.Comm: - # TODO: check iterset - return utils.common_comm(self.statements, "comm") - # }}} def __str__(self) -> str: @@ -238,6 +237,10 @@ def get_disk_cache_key(self, visitor): def collect_buffers(self, visitor): return OrderedFrozenSet().union(*(map(visitor, self.instructions))) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self.instructions) + def __init__(self, instructions: Iterable[Instruction]) -> None: instructions = tuple(instructions) object.__setattr__(self, "instructions", instructions) @@ -248,10 +251,6 @@ def __init__(self, instructions: Iterable[Instruction]) -> None: child_attrs = ("instructions",) - @property - def comm(self) -> MPI.Comm: - return utils.common_comm(self.instructions, "comm") - @property def global_arguments(self) -> OrderedFrozenSet[Tensor]: return OrderedFrozenSet().union(*(insn.global_arguments for insn in self.instructions)) @@ -538,10 +537,6 @@ def argument_shapes(self): for arg in self.function.code.default_entrypoint.args ) - @property - def comm(self) -> MPI.Comm: - return utils.common_comm(self.arguments, "comm", allow_undefined=True) or MPI.COMM_SELF - @pyop3.record.frozenrecord() class CalledFunction(AbstractCalledFunction): @@ -558,6 +553,10 @@ def get_instruction_executor_cache_key(self, visitor) -> Hashable: tuple(map(visitor, self._arguments)), ) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self._arguments) + def __init__(self, function: Function, arguments: Iterable): arguments = tuple(arguments) @@ -615,6 +614,10 @@ def get_disk_cache_key(self, visitor): def collect_buffers(self, visitor): return OrderedFrozenSet().union(*(map(visitor, self._arguments))) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(*self._arguments) + def __init__(self, function: Function, arguments: Iterable): arguments = tuple(arguments) @@ -766,6 +769,10 @@ def get_instruction_executor_cache_key(self, visitor) -> Hashable: self._assignment_type, ) + @cached_property + def comm(self) -> MPI.Comm: + return pyop3.visitors.common_comm(self._assignee, self._expression) + def __init__(self, assignee: Any, expression: Any, assignment_type: AssignmentType | str) -> None: assignment_type = AssignmentType(assignment_type) @@ -786,10 +793,6 @@ def __post_init__(self) -> None: expression: ClassVar[property] = pyop3.record.attr("_expression") assignment_type: ClassVar[property] = pyop3.record.attr("_assignment_type") - @property - def comm(self) -> MPI.Comm: - return utils.common_comm([self.assignee, self.expression], "comm", allow_undefined=True) or MPI.COMM_SELF - # NOTE: Wrong type here... @property def shape(self) -> tuple[AxisTree, ...]: diff --git a/pyop3/mpi.py b/pyop3/mpi.py index 277da2b158..9b453bbbd8 100644 --- a/pyop3/mpi.py +++ b/pyop3/mpi.py @@ -36,6 +36,7 @@ """PyOP2 MPI communicator.""" +from collections.abc import Iterable from typing import Any, Callable from petsc4py import PETSc from mpi4py import MPI # noqa @@ -50,6 +51,7 @@ import weakref import pyop3.config +from pyop3 import utils from pyop3.exceptions import CompilationException from pyop3.log import debug, LOGGER, DEBUG @@ -163,7 +165,7 @@ def collective(fn): extra = """\ This function is logically collective over MPI ranks, it is an error to call it on fewer than all the ranks in MPI communicator. -PYOP2_SPMD_STRICT=1 is in your environment and function calls will be +PYOP3_SPMD_STRICT=1 is in your environment and function calls will be guarded by a barrier where possible.""" @wraps(fn) @@ -556,6 +558,61 @@ def compilation_comm(comm, obj): return comp_comm +def comm_is_subset(comm_small: MPI.Comm, comm_large: MPI.Comm) -> bool: + """Return if a communicator is a subset of another. + + This is a local operation. No communication is required. + + """ + # optimisation: if the small comm only has one rank in it then + # it is necessarily a subset of the bigger comm + if comm_small.size == 1: + return True + + group_small = comm_small.Get_group() + group_large = comm_large.Get_group() + group_intersect = MPI.Group.Intersect(group_small, group_large) + # If the intersection group doesn't match the small group then there + # are elements in the small group that aren't in the large group and + # therefore the comms aren't compatible. + return MPI.Group.Compare(group_intersect, group_small) != MPI.UNEQUAL + + +@collective +def common_comm(*comms: Iterable[MPI.Comm]) -> MPI.Comm: + """Return a communicator valid for all objects. + + The valid communicator is defined as the one with the largest size. + + Parameters + ---------- + comms + A collection of communicators. + + Returns + ------- + MPI.Comm + A communicator that the provided objects are safely collective over. + + """ + shared_comm = None + for comm in utils.iterflat(comms): + if shared_comm is None: + shared_comm = comm + continue + + if comm.size > shared_comm.size: + big_comm = comm + small_comm = shared_comm + else: + big_comm = shared_comm + small_comm = comm + assert comm_is_subset(small_comm, big_comm) + shared_comm = big_comm + assert shared_comm is not None + return shared_comm + + def finalize_safe_debug(): ''' Return function for debug output. diff --git a/pyop3/obj.py b/pyop3/obj.py index f7eb38314d..104423c183 100644 --- a/pyop3/obj.py +++ b/pyop3/obj.py @@ -31,3 +31,9 @@ def get_disk_cache_key(self, renamer) -> Hashable: f"'get_disk_cache_key' not implemented for '{type(self).__qualname__}'" ) + @property + def comm(self) -> MPI.Comm: + """The communicator over which this object is collective.""" + raise NotImplementedError( + f"'comm' not implemented for '{type(self).__qualname__}'" + ) diff --git a/pyop3/sf.py b/pyop3/sf.py index adb49f8890..3f3d45339c 100644 --- a/pyop3/sf.py +++ b/pyop3/sf.py @@ -23,38 +23,11 @@ from ._sf_cy import filter_petsc_sf, create_petsc_section_sf, renumber_petsc_sf # noqa: F401 -class ParallelAwareObject(abc.ABC): - """Abstract class for objects that know about communicators. - - Unlike `DistributedObject`s, it is allowed for objects inheriting from - this class to have `None` for communicator values. - - """ - - @property - @abc.abstractmethod - def comm(self) -> MPI.Comm | None: - pass - - -class DistributedObject(ParallelAwareObject, metaclass=abc.ABCMeta): - """Abstract class for objects that have a parallel execution context. - - The expected usage is for classes to implement the attribute `user_comm`. - - """ - - @property - @abc.abstractmethod - def comm(self) -> MPI.Comm: - pass - - class BufferSizeMismatchException(Exception): pass -class AbstractStarForest(DistributedObject, abc.ABC): +class AbstractStarForest(abc.ABC): # {{{ abstract methods diff --git a/pyop3/utils.py b/pyop3/utils.py index 4833af50a7..0a056d78fb 100644 --- a/pyop3/utils.py +++ b/pyop3/utils.py @@ -540,53 +540,6 @@ def _(hashable: Hashable) -> Hashable: return hashable -def single_comm(objects, /, comm_attr: str, *, allow_undefined: bool = False) -> MPI.Comm | None: - assert len(objects) > 0 - - comm = None - for item in iterflat(objects): - item_comm = getattr(item, comm_attr, None) - - if item_comm is None: - if allow_undefined: - continue - else: - raise CommNotFoundException("Object does not have an associated communicator") - - if comm is None: - comm = item_comm - elif item_comm != comm: - raise CommMismatchException("Multiple communicators found") - return comm - - -@collective -def common_comm(objects, /, comm_attr: str, *, allow_undefined: bool = False) -> MPI.Comm | None: - """Return a communicator valid for all objects. - - This is defined as the communicator with the largest size. I *think* that - this is the right way to think about this. - - """ - assert len(objects) > 0 - - selected_comm = None - for item in iterflat(objects): - item_comm = getattr(item, comm_attr, None) - - if item_comm is None: - if allow_undefined: - continue - else: - raise CommNotFoundException("Object does not have an associated communicator") - - if selected_comm is None or item_comm.size > selected_comm.size: - selected_comm = item_comm - if not allow_undefined: - assert selected_comm is not None - return selected_comm - - def iterflat(iterable): if isinstance(iterable, np.ndarray): iterable = iterable.flatten() diff --git a/pyop3/visitors.py b/pyop3/visitors.py index 006b107825..e9afc03e2d 100644 --- a/pyop3/visitors.py +++ b/pyop3/visitors.py @@ -11,9 +11,12 @@ from typing import Any from immutabledict import immutabledict as idict +from mpi4py import MPI +from petsc4py import PETSc import pyop3.node import pyop3.obj +import pyop3.sf from pyop3 import utils from pyop3.collections import OrderedFrozenSet @@ -195,3 +198,64 @@ def get_instruction_executor_cache_key(obj: pyop3.obj.Pyop3Object) -> Hashable: the buffers at the top-level. """ return InstructionExecutorCacheKeyGetter()(obj) + + +@functools.singledispatch +def get_comm(obj: Any, /) -> MPI.Comm: + """Return the communicator associated with an object. + + If no communicator is available (e.g. trying to get the comm of an integer) + then ``COMM_SELF`` is used. + + """ + utils.raise_visitor_type_error(obj) + + +@get_comm.register +def _(comm: MPI.Comm, /) -> MPI.Comm: + return comm + + +@get_comm.register +def _(obj: PETSc.Object, /) -> MPI.Comm: + return obj.comm + + +@get_comm.register +def _(sf: pyop3.sf.AbstractStarForest, /) -> MPI.Comm: + return sf.comm + + +@get_comm.register +def _(obj: pyop3.obj.Pyop3Object, /) -> MPI.Comm: + return obj.comm + + +@get_comm.register +def _(num: numbers.Number | types.NoneType, /) -> MPI.Comm: + return MPI.COMM_SELF + + +def common_comm(*objects: Iterable[Any]) -> MPI.Comm: + """Return a communicator valid for all objects. + + The valid communicator is defined as the one with the largest size. + + Parameters + ---------- + objects + Communicator-carrying objects to inspect. All object must define + a ``comm`` attribute. + + Returns + ------- + MPI.Comm + A communicator that the provided objects are safely collective over. + + """ + return pyop3.mpi.common_comm(*(map(get_comm, objects))) + + +def single_comm(*objects: Iterable[Any]) -> MPI.Comm: + """Return the single comm shared by all objects.""" + return utils.single_valued(map(get_comm, objects)) From 6283656d4fe3aca11e99a9f00bb078f04eaf6019 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Mon, 22 Jun 2026 17:57:44 +0100 Subject: [PATCH 03/24] fixup --- firedrake/extrusion_utils.py | 2 +- pyop3/axis_tree/tree.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/extrusion_utils.py b/firedrake/extrusion_utils.py index 2424808adc..31ec3978da 100644 --- a/firedrake/extrusion_utils.py +++ b/firedrake/extrusion_utils.py @@ -16,7 +16,7 @@ @PETSc.Log.EventDecorator() -@with_heavy_caches(lambda extr_top, *a, **kw: {extr_top}) +@with_heavy_caches(lambda extr_top, *a, **kw: [extr_top]) def make_extruded_coords(extruded_topology, base_coords, ext_coords, layer_height, extrusion_type='uniform', kernel=None): """ diff --git a/pyop3/axis_tree/tree.py b/pyop3/axis_tree/tree.py index fae86031ce..9df344a454 100644 --- a/pyop3/axis_tree/tree.py +++ b/pyop3/axis_tree/tree.py @@ -1147,7 +1147,7 @@ def owned(self): if self.comm.size == 1: return self else: - return self.with_region_label(OWNED_REGION_LABEL) + return self.with_region_label(OWNED_REGION_LABEL, allow_missing=True) @cached_property def unconstrained(self): From 5d222afe2473bc882afb43ad9966ee887c57f99e Mon Sep 17 00:00:00 2001 From: India Marsden Date: Tue, 23 Jun 2026 11:01:07 +0100 Subject: [PATCH 04/24] Remove use of FUSE env var --- .github/actions/install/action.yml | 2 +- firedrake/functionspace.py | 3 +-- firedrake/mesh.py | 27 +++++++++++++++++---------- firedrake/utility_meshes.py | 6 ++++++ 4 files changed, 25 insertions(+), 13 deletions(-) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 560a80cd3c..ca4790e37b 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -156,7 +156,7 @@ runs: with: repository: firedrakeproject/fiat path: fiat-repo - ref: refs/heads/indiamai/integrate_fuse + ref: refs/heads/indiamai/fuse_cell_change - name: Install checked out FIAT shell: bash diff --git a/firedrake/functionspace.py b/firedrake/functionspace.py index 9310989b0d..f3fae74c9a 100644 --- a/firedrake/functionspace.py +++ b/firedrake/functionspace.py @@ -6,7 +6,6 @@ """ import itertools import ufl -from ufl.cell import as_cell import finat.ufl from pyop3.pyop2_utils import flatten @@ -69,7 +68,7 @@ def make_scalar_element(mesh, family, degree, vfamily, vdegree, variant, quad_sc quad_scheme=quad_scheme) # If second element was passed in, use it lb = finat.ufl.FiniteElement(vfamily, - cell=ufl.interval, + cell=finat.ufl.as_cell("interval"), degree=vdegree, variant=variant, quad_scheme=quad_scheme) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 8f4f998f3d..d001f63f08 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -328,7 +328,7 @@ class AbstractMeshTopology(abc.ABC): """A representation of an abstract mesh topology without a concrete PETSc DM implementation""" - def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=None): + def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=None, use_fuse=False): """Initialise a mesh topology. Parameters @@ -357,7 +357,8 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, Communicator. submesh_parent: AbstractMeshTopology Submesh parent. - + use_fuse: bool + Flag to determine if the mesh is based on a UFC cell or a FUSE cell """ if comm.size == 1: # in serial the point sf isn't initialised @@ -481,6 +482,7 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, # To set, do e.g. # target_mesh._parallel_compatible = {weakref.ref(source_mesh)} self._parallel_compatible = None + self._use_fuse = use_fuse layers = None """No layers on unstructured mesh""" @@ -1692,6 +1694,7 @@ def __init__( permutation_name=None, submesh_parent=None, comm=COMM_WORLD, + use_fuse=False, ): """Initialise a mesh topology. @@ -1723,6 +1726,8 @@ def __init__( Submesh parent. comm : mpi4py.MPI.Comm Communicator. + use_fuse: bool + Flag to determine if the underlying cell is a UFC cell or FUSE cell. """ if distribution_parameters is None: @@ -1740,7 +1745,7 @@ def __init__( # Disable auto distribution and reordering before setFromOptions is called. plex.distributeSetDefault(False) plex.reorderSetDefault(PETSc.DMPlex.ReorderDefaultFlag.FALSE) - super().__init__(plex, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=submesh_parent) + super().__init__(plex, name, reorder, sfXB, perm_is, distribution_name, permutation_name, comm, submesh_parent=submesh_parent, use_fuse=use_fuse) @cached_property def _entity_indices(self): @@ -1858,7 +1863,7 @@ def _ufl_cell(self): # represent a mesh topology (as here) have geometric dimension # equal their topological dimension. This is reflected in the # corresponding UFL mesh. - return as_cell(_cells[tdim][nfacets]) + return as_cell(_cells[tdim][nfacets], self._use_fuse) @cached_property def _ufl_mesh(self): @@ -1927,7 +1932,7 @@ def entity_orientations_dat(self): cell_axis = self.cells.root # # so instead we do # cell_axis = op3.Axis([self.points.root.components[0]], self.points.root.label) - if bool(os.environ.get("FIREDRAKE_USE_FUSE", 0)): + if self._use_fuse: return self.entity_orientations_dat_fuse # TODO: This is quite a funky way of getting this. We should be able to get @@ -2652,9 +2657,11 @@ def __init__(self, mesh, layers, periodic=False, name=None): self._shared_data_cache = defaultdict(dict) self._max_work_functions = {} + self._use_fuse = mesh._use_fuse + @cached_property def _ufl_cell(self): - return ufl.TensorProductCell(self._base_mesh.ufl_cell(), as_cell("interval")) + return ufl.TensorProductCell(self._base_mesh.ufl_cell(), as_cell("interval", self._use_fuse)) @cached_property def _ufl_mesh(self): @@ -3485,7 +3492,7 @@ def _mark_entity_classes(self): @cached_property def _ufl_cell(self): - return as_cell(_cells[0][0]) + return as_cell(_cells[0][0], self._use_fuse) @cached_property def _ufl_mesh(self): @@ -4849,7 +4856,7 @@ def Mesh(meshfile, **kwargs): distribution_name=kwargs.get("distribution_name"), permutation_name=kwargs.get("permutation_name"), submesh_parent=submesh_parent.topology if submesh_parent else None, - comm=user_comm) + comm=user_comm, use_fuse=kwargs.get("use_fuse")) mesh = make_mesh_from_mesh_topology(topology, name) if from_netgen: @@ -4990,9 +4997,9 @@ def ExtrudedMesh(mesh, layers, layer_height=None, extrusion_type='uniform', peri if extrusion_type == 'radial_hedgehog': helement = helement.reconstruct(family="DG", variant="equispaced") if periodic: - velement = finat.ufl.FiniteElement("DP", as_cell("interval"), 1, variant="equispaced") + velement = finat.ufl.FiniteElement("DP", as_cell("interval", self._use_fuse), 1, variant="equispaced") else: - velement = finat.ufl.FiniteElement("Lagrange", as_cell("interval"), 1) + velement = finat.ufl.FiniteElement("Lagrange", as_cell("interval", self._use_fuse), 1) element = finat.ufl.TensorProductElement(helement, velement) if gdim is None: diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index 12868ab0e7..3f9a7ba12c 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -670,6 +670,7 @@ def RectangleMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a rectangular mesh @@ -730,6 +731,7 @@ def RectangleMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse, ) @@ -845,6 +847,7 @@ def SquareMesh( name: str = DEFAULT_MESH_NAME, distribution_name: str | None = None, permutation_name: str | None = None, + use_fuse: bool = False, ): """Generate a square mesh. @@ -905,6 +908,7 @@ def SquareMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse, ) @@ -920,6 +924,7 @@ def UnitSquareMesh( name: str = DEFAULT_MESH_NAME, distribution_name: str | None = None, permutation_name: str | None = None, + use_fuse: bool = False, ): """Generate a unit square mesh. @@ -977,6 +982,7 @@ def UnitSquareMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse, ) From 1ca1d1b2d4c1e6a93bd119c4d3e2e7b05fe1e16d Mon Sep 17 00:00:00 2001 From: India Marsden Date: Tue, 23 Jun 2026 11:34:52 +0100 Subject: [PATCH 05/24] update all utility meshes --- firedrake/utility_meshes.py | 130 +++++++++++++++++++++++++++++++++++- 1 file changed, 127 insertions(+), 3 deletions(-) diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index 3f9a7ba12c..be87907bc0 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -75,7 +75,7 @@ reorder_noop = False -def _postprocess_periodic_mesh(coords, comm, distribution_parameters, reorder, name, distribution_name, permutation_name): +def _postprocess_periodic_mesh(coords, comm, distribution_parameters, reorder, name, distribution_name, permutation_name, use_fuse=False): dm = coords.function_space().mesh().topology.topology_dm dm.removeLabel("firedrake_is_ghost") dm.removeLabel("exterior_facets") @@ -95,6 +95,7 @@ def _postprocess_periodic_mesh(coords, comm, distribution_parameters, reorder, n name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) @PETSc.Log.EventDecorator() @@ -107,6 +108,7 @@ def OneTetMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """ Mesh with only two tets in @@ -128,6 +130,8 @@ def OneTetMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The left hand boundary point has boundary marker 1, while the right hand point has marker 2. @@ -149,6 +153,7 @@ def OneTetMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) return m @@ -162,6 +167,7 @@ def TwoTetMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """ Mesh with only two tets in @@ -183,6 +189,8 @@ def TwoTetMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The left hand boundary point has boundary marker 1, while the right hand point has marker 2. @@ -203,6 +211,7 @@ def TwoTetMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) return m @@ -219,6 +228,7 @@ def IntervalMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """ Generate a uniform mesh of an interval. @@ -240,6 +250,8 @@ def IntervalMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The left hand boundary point has boundary marker 1, while the right hand point has marker 2. @@ -274,6 +286,7 @@ def IntervalMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) @@ -286,6 +299,7 @@ def UnitIntervalMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """ Generate a uniform mesh of the interval [0,1]. @@ -302,6 +316,8 @@ def UnitIntervalMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The left hand (:math:`x=0`) boundary point has boundary marker 1, while the right hand (:math:`x=1`) point has marker 2. @@ -328,6 +344,7 @@ def PeriodicIntervalMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a periodic mesh of an interval. @@ -344,6 +361,8 @@ def PeriodicIntervalMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. """ plex = PETSc.DMPlex().createBoxMesh( (ncells,), @@ -353,7 +372,7 @@ def PeriodicIntervalMesh( periodic=True, interpolate=True, sparseLocalize=False, - comm=comm + comm=comm, ) _mark_mesh_boundaries(plex) @@ -365,6 +384,7 @@ def PeriodicIntervalMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse, ) @@ -377,6 +397,7 @@ def PeriodicUnitIntervalMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a periodic mesh of the unit interval. @@ -392,6 +413,8 @@ def PeriodicUnitIntervalMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. """ return PeriodicIntervalMesh( ncells, @@ -402,6 +425,7 @@ def PeriodicUnitIntervalMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=False, ) @@ -415,6 +439,7 @@ def OneElementThickMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """ Generate a rectangular mesh in the domain with corners [0,0] @@ -433,6 +458,8 @@ def OneElementThickMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. """ left = np.arange(ncells, dtype=np.int32) @@ -565,6 +592,7 @@ def OneElementThickMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=False, ) topverts = Vc.cell_node_list[:, 1::2].flatten() mash.coordinates.dat.data_with_halos[topverts, 1] = Ly @@ -599,6 +627,7 @@ def UnitTriangleMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a mesh of the reference triangle @@ -613,6 +642,8 @@ def UnitTriangleMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. """ coords = [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]] cells = [[0, 1, 2]] @@ -651,6 +682,7 @@ def UnitTriangleMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=False, ) @@ -695,6 +727,8 @@ def RectangleMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The boundary edges in this mesh are numbered as follows: @@ -746,6 +780,7 @@ def TensorRectangleMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a rectangular mesh @@ -760,6 +795,8 @@ def TensorRectangleMesh( from bottom left to top right (``"right"``), or top left to bottom right (``"left"``), or put in both diagonals (``"crossed"``). Default is ``"left"``. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The boundary edges in this mesh are numbered as follows: @@ -831,6 +868,7 @@ def TensorRectangleMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=False ) @@ -878,6 +916,9 @@ def SquareMesh( permutation_name The name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + use_fuse + Flag indicting if mesh should be created using a fuse cell + Default: False, mesh is based on the ufc cell Returns ------- @@ -953,6 +994,9 @@ def UnitSquareMesh( permutation_name The name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + use_fuse + Flag indicting if mesh should be created using a fuse cell + Default: False, mesh is based on the ufc cell Returns ------- @@ -1001,6 +1045,7 @@ def PeriodicRectangleMesh( name: str = DEFAULT_MESH_NAME, distribution_name: str | None = None, permutation_name: str | None = None, + use_fuse=False, ) -> MeshGeometry: """Generate a periodic rectangular mesh. @@ -1035,6 +1080,9 @@ def PeriodicRectangleMesh( permutation_name The name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + use_fuse + Flag indicting if mesh should be created using a fuse cell + Default: False, mesh is based on the ufc cell Returns ------- @@ -1094,6 +1142,7 @@ def PeriodicRectangleMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=False ) @@ -1111,6 +1160,7 @@ def PeriodicSquareMesh( name: str = DEFAULT_MESH_NAME, distribution_name: str | None = None, permutation_name: str | None = None, + use_fuse=False, ): """Generate a periodic square mesh. @@ -1143,6 +1193,9 @@ def PeriodicSquareMesh( permutation_name The name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + use_fuse + Flag indicting if mesh should be created using a fuse cell + Default: False, mesh is based on the ufc cell Returns ------- @@ -1176,6 +1229,7 @@ def PeriodicSquareMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse, ) @@ -1192,6 +1246,7 @@ def PeriodicUnitSquareMesh( name: str = DEFAULT_MESH_NAME, distribution_name: str | None = None, permutation_name: str | None = None, + use_fuse: bool = False, ): """Generate a periodic unit square mesh. @@ -1222,6 +1277,9 @@ def PeriodicUnitSquareMesh( permutation_name The name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + use_fuse + Flag indicting if mesh should be created using a fuse cell + Default: False, mesh is based on the ufc cell Returns ------- @@ -1317,6 +1375,7 @@ def CircleManifoldMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse, ) if degree > 1: new_coords = Function(VectorFunctionSpace(m, "CG", degree)) @@ -1331,6 +1390,7 @@ def CircleManifoldMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) m._radius = radius return m @@ -1360,6 +1420,8 @@ def UnitDiskMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ vertices = np.array( [[0, 0], [1, 0], [1, 1], [0, 1], [-1, 1], [-1, 0], [-1, -1], [0, -1], [1, -1]], @@ -1412,6 +1474,7 @@ def UnitDiskMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) return m @@ -1440,6 +1503,8 @@ def UnitBallMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ vertices = np.array( [ @@ -1499,6 +1564,7 @@ def UnitBallMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) return m @@ -1509,6 +1575,7 @@ def UnitTetrahedronMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a mesh of the reference tetrahedron. @@ -1520,6 +1587,8 @@ def UnitTetrahedronMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ coords = [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]] cells = [[0, 1, 2, 3]] @@ -1533,6 +1602,7 @@ def UnitTetrahedronMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) return m @@ -1680,6 +1750,7 @@ def TensorBoxMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) return m @@ -1749,6 +1820,7 @@ def BoxMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) else: xcoords = np.linspace(0, Lx, nx + 1, dtype=np.double) @@ -1781,6 +1853,7 @@ def CubeMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a mesh of a cube @@ -1800,6 +1873,8 @@ def CubeMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell The boundary surfaces are numbered as follows: @@ -1839,6 +1914,7 @@ def UnitCubeMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a mesh of a unit cube @@ -1857,6 +1933,8 @@ def UnitCubeMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell The boundary surfaces are numbered as follows: @@ -1898,6 +1976,7 @@ def PeriodicBoxMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a periodic mesh of a 3D box. @@ -1933,6 +2012,9 @@ def PeriodicBoxMesh( permutation_name : str or None Name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + use_fuse + Flag indicting if mesh should be created using a fuse cell + Default: False, mesh is based on the ufc cell Returns ------- @@ -1975,7 +2057,8 @@ def PeriodicBoxMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, - comm=comm) + comm=comm, + use_fuse=use_fuse) else: # TODO: When hexahedra -> simplex refinement is implemented this can go away. if tuple(directions) != (True, True, True): @@ -2020,6 +2103,7 @@ def PeriodicBoxMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) new_coordinates = Function( @@ -2070,6 +2154,7 @@ def PeriodicUnitCubeMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a periodic mesh of a unit cube @@ -2099,6 +2184,9 @@ def PeriodicUnitCubeMesh( permutation_name : str or None Name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + use_fuse + Flag indicting if mesh should be created using a fuse cell + Default: False, mesh is based on the ufc cell Returns ------- @@ -2149,6 +2237,7 @@ def IcosahedralSphereMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False ): """Generate an icosahedral approximation to the surface of the sphere. @@ -2176,6 +2265,8 @@ def IcosahedralSphereMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ if refinement_level < 0 or refinement_level % 1: raise RuntimeError("Number of refinements must be a non-negative integer") @@ -2248,6 +2339,7 @@ def IcosahedralSphereMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) if degree > 1: new_coords = Function(VectorFunctionSpace(m, "CG", degree)) @@ -2261,6 +2353,7 @@ def IcosahedralSphereMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) m._radius = radius return m @@ -2276,6 +2369,7 @@ def UnitIcosahedralSphereMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate an icosahedral approximation to the unit sphere. @@ -2294,6 +2388,8 @@ def UnitIcosahedralSphereMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ return IcosahedralSphereMesh( 1.0, @@ -2324,6 +2420,7 @@ def OctahedralSphereMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate an octahedral approximation to the surface of the sphere. @@ -2348,6 +2445,8 @@ def OctahedralSphereMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ if refinement_level < 0 or refinement_level % 1: raise ValueError("Number of refinements must be a non-negative integer") @@ -2406,6 +2505,7 @@ def OctahedralSphereMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) if degree > 1: # use it to build a higher-order mesh @@ -2416,6 +2516,7 @@ def OctahedralSphereMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) # remap to a cone @@ -2478,6 +2579,7 @@ def UnitOctahedralSphereMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False ): """Generate an octahedral approximation to the unit sphere. @@ -2500,6 +2602,8 @@ def UnitOctahedralSphereMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ return OctahedralSphereMesh( 1.0, @@ -2660,6 +2764,7 @@ def CubedSphereMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate an cubed approximation to the surface of the sphere. @@ -2679,6 +2784,8 @@ def CubedSphereMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ if refinement_level < 0 or refinement_level % 1: raise RuntimeError("Number of refinements must be a non-negative integer") @@ -2700,6 +2807,7 @@ def CubedSphereMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) if degree > 1: @@ -2713,6 +2821,7 @@ def CubedSphereMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) m._radius = radius return m @@ -2728,6 +2837,7 @@ def UnitCubedSphereMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a cubed approximation to the unit sphere. @@ -2745,6 +2855,8 @@ def UnitCubedSphereMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ return CubedSphereMesh( 1.0, @@ -2790,6 +2902,8 @@ def TorusMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell """ if nR < 3 or nr < 3: @@ -2847,6 +2961,7 @@ def TorusMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) return m @@ -2879,6 +2994,8 @@ def AnnulusMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if ``None``, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. """ if nt < 3: raise ValueError("Must have at least 3 cells in the circumferential direction") @@ -2931,6 +3048,8 @@ def SolidTorusMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if ``None``, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. """ if nR < 3: raise ValueError("Must have at least 3 cells in the major direction") @@ -2995,6 +3114,8 @@ def CylinderMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The boundary edges in this mesh are numbered as follows: @@ -3122,6 +3243,7 @@ def CylinderMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, + use_fuse=use_fuse ) @@ -3162,6 +3284,8 @@ def PartiallyPeriodicRectangleMesh( :kwarg permutation_name: the name of entity permutation (reordering) used when checkpointing; if `None`, the name is automatically generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. The boundary edges in this mesh are numbered as follows: From caf6c5362d5c1e7df8fcf093fe7ab757a8bb4646 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Tue, 23 Jun 2026 14:37:18 +0100 Subject: [PATCH 06/24] Fix for submesh --- firedrake/mesh.py | 170 ++++++++++-------- firedrake/mg/opencascade_mh.py | 4 +- firedrake/parloops.py | 17 +- firedrake/preconditioners/bddc.py | 4 +- pyop3/axis_tree/tree.py | 22 ++- pyop3/buffer.py | 2 +- pyop3/exceptions.py | 7 + pyop3/expr/buffer.py | 20 ++- pyop3/expr/tensor/dat.py | 7 + pyop3/expr/visitors/__init__.py | 4 +- pyop3/index_tree/tree.py | 9 +- pyop3/labeled_tree.py | 16 +- pyop3/record.py | 10 +- pyop3/utils.py | 1 + .../firedrake/regression/test_exodus_mesh.py | 2 +- 15 files changed, 168 insertions(+), 127 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 32a0d9d16b..48bc467091 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -657,12 +657,12 @@ def _local_facet_orientation_dat(self, facet_type: Literal["exterior", "interior if facet_type == "exterior": local_facet_numbers_dat = self._exterior_facet_local_numbers_dat arity = 1 - facet_to_cell_map = self._facet_support_dat("exterior").data_ro + facet_to_cell_map = self._facet_support_dats("exterior")[1].data_ro else: assert facet_type == "interior" local_facet_numbers_dat = self._interior_facet_local_numbers_dat arity = 2 - facet_to_cell_map = self._facet_support_dat("interior").data_ro + facet_to_cell_map = self._facet_support_dats("interior")[1].data_ro facet_to_cell_map = facet_to_cell_map.reshape((-1, arity)) @@ -1159,34 +1159,95 @@ def support(self, index): def _support(self): supports = {} - # 1-tuple here because in theory support(facet) could map to other valid things (like points) - exterior_facets_axis = self.exterior_facets.owned.as_axis() - supports[idict({exterior_facets_axis.label: exterior_facets_axis.component.label})] = ( - ( - op3.TabulatedMapComponent( - self.name, - self.cell_label, - self._facet_support_dat("exterior"), - label=0, - ), - ), - ) + # Facet supports: + for facet_type in ["exterior", "interior"]: + # We have seperate (but identical) maps mapping from either only + # owned points or all points. + for only_owned in [False, True]: + if facet_type == "exterior": + facets_axes = self.exterior_facets + else: + facets_axes = self.interior_facets + if only_owned: + facets_axes = facets_axes.owned + facets_axis = facets_axes.as_axis() + + support_dat = self._facet_support_dat(facet_type, only_owned=only_owned) + + # 1-tuple here because in theory support(facet) could map to other valid things (like points) + supports[idict({facets_axis.label: facets_axis.component.label})] = ( + ( + op3.TabulatedMapComponent( + self.name, + self.cell_label, + support_dat, + label=0, + ), + ), + ) - interior_facets_axis = self.interior_facets.owned.as_axis() - supports[idict({interior_facets_axis.label: interior_facets_axis.component.label})] = ( - ( - op3.TabulatedMapComponent( - self.name, - self.cell_label, - self._facet_support_dat("interior"), - label=0, - ), - ), + return op3.Map(supports, name="support") + + @cached_method() + def _facet_support_dat( + self, + facet_type: Literal["exterior", "interior"], + *, + only_owned: bool, + ) -> op3.Dat: + """Return a dat encoding the support of exterior/interior facets.""" + # Get the support map for *all* facets in the mesh, not just the + # exterior/interior ones. We have to filter it. Note that these + # dats are ragged because support sizes are not consistent. + all_facets_support_dat = self._support_dats[self.facet_label][self.cell_label] + + if facet_type == "exterior": + facet_axes = self.exterior_facets + selected_facets_is = dmcommon.section_offsets( + self._old_to_new_facet_numbering, self._exterior_facet_plex_indices, sort=True + ) + arity = 1 + else: + facet_axes = self.interior_facets + selected_facets_is = dmcommon.section_offsets( + self._old_to_new_facet_numbering, self._interior_facet_plex_indices, sort=True + ) + if only_owned: + arity = 2 + else: + # ragged + arity = None + + if only_owned: + facet_axes = facet_axes.owned + + # Remove ghost indices + selected_facets_is = dmcommon.filter_is( + selected_facets_is, 0, self.facets.owned.local_size + ) + assert selected_facets_is.size == facet_axes.local_size + + facet_axis = facet_axes.as_axis() + facet_selector = op3.Slice( + all_facets_support_dat.axes.root.label, + [ + op3.Subset( + all_facets_support_dat.axes.root.component.label, + op3.Dat.from_array(selected_facets_is.indices), + label=facet_axis.component.label, + ) + ], + label=facet_axis.label, ) - return op3.Map(supports, name="support") + arity_axis = all_facets_support_dat.axes.leaf_axis + arity_slice = op3.Slice( + arity_axis.label, + [op3.AffineSliceComponent(arity_axis.component.label, stop=arity)], + label="support", + ) + return all_facets_support_dat[facet_selector, arity_slice] - # TODO: Redesign all this, this sucks for extruded meshes @cached_property def _support_dats(self): def support_func(pt): @@ -1228,62 +1289,15 @@ def support_func(pt): iterset_axis, op3.Axis(size_dat) ]) support_dat = op3.Dat(support_axes.regionless().materialize(), data=data, prefix="support") - owned_support_dat = op3.Dat( - support_axes.owned.regionless().materialize(), data=support_dat.data_ro, prefix="support" - ) + # owned_support_dat = op3.Dat( + # support_axes.owned.regionless().materialize(), data=support_dat.data_ro, prefix="support" + # ) - supports.append({map_dim: (support_dat, owned_support_dat)}) + # supports.append({map_dim: (support_dat, owned_support_dat)}) + supports.append({map_dim: support_dat}) return tuple(supports) - # this is almost completely pointless - def _facet_support_dat(self, facet_type: Literal["exterior"] | Literal["interior"]) -> op3.Dat: - assert facet_type in {"exterior", "interior"} - - # Get the support map for *all* facets in the mesh, not just the - # exterior/interior ones. We have to filter it. Note that these - # dats are ragged because support sizes are not consistent. - _, facet_support_dat = self._support_dats[self.facet_label][self.cell_label] - - if facet_type == "exterior": - facet_axis = self.exterior_facets.owned.as_axis() - selected_facets_is = dmcommon.section_offsets( - self._old_to_new_facet_numbering, self._exterior_facet_plex_indices, sort=True - ) - arity = 1 - else: - facet_axis = self.interior_facets.owned.as_axis() - selected_facets_is = dmcommon.section_offsets( - self._old_to_new_facet_numbering, self._interior_facet_plex_indices, sort=True - ) - arity = 2 - - # Remove ghost indices - new_selected_facets_is = dmcommon.filter_is(selected_facets_is, 0, self.facets.owned.local_size) - selected_facets = new_selected_facets_is.indices - assert selected_facets.size == facet_axis.local_size - - mysubset = op3.Slice( - facet_support_dat.axes.root.label, - [ - op3.Subset( - facet_support_dat.axes.root.component.label, - op3.Dat.from_array(selected_facets, comm=self.comm), - label=facet_axis.component.label, - ) - ], - label=facet_axis.label, - ) - - *others, (leaf_axis_label, leaf_component_label) = facet_support_dat.axes.leaf_path.items() - myslice = op3.Slice(leaf_axis_label, [op3.AffineSliceComponent(leaf_component_label, stop=arity)], label="support") - - # TODO: This should ideally work - # return facet_support_dat[mysubset, slice(arity)] - specialized_by_type_facet_support_dat = facet_support_dat[mysubset, myslice] - assert specialized_by_type_facet_support_dat.axes.local_size == facet_axis.local_size * arity - return specialized_by_type_facet_support_dat - # delete? def create_section(self, nodes_per_entity, real_tensorproduct=False, block_size=1): diff --git a/firedrake/mg/opencascade_mh.py b/firedrake/mg/opencascade_mh.py index 492dd1bd42..98671794ae 100644 --- a/firedrake/mg/opencascade_mh.py +++ b/firedrake/mg/opencascade_mh.py @@ -136,7 +136,7 @@ def project_mesh_to_cad_3d(mesh, cad): from OCC.Core.GeomAPI import GeomAPI_ProjectPointOnSurf, GeomAPI_ProjectPointOnCurve coorddata = mesh.coordinates.dat.data - ids = mesh.exterior_facets.unique_markers + ids = mesh.facet_markers filt = lambda arr: arr[numpy.where(arr < mesh.coordinates.function_space().axes.local_size)[0]] boundary_nodes = {id: filt(mesh.coordinates.function_space().boundary_nodes(int(id))) for id in ids} @@ -211,7 +211,7 @@ def project_mesh_to_cad_2d(mesh, cad): from OCC.Core.GeomAPI import GeomAPI_ProjectPointOnCurve coorddata = mesh.coordinates.dat.data - ids = mesh.exterior_facets.unique_markers + ids = mesh.facet_markers filt = lambda arr: arr[numpy.where(arr < mesh.coordinates.function_space().axes.owned.local_size)[0]] boundary_nodes = {id: filt(mesh.coordinates.function_space().boundary_nodes(int(id))) for id in ids} diff --git a/firedrake/parloops.py b/firedrake/parloops.py index 3b01880760..db200a4784 100644 --- a/firedrake/parloops.py +++ b/firedrake/parloops.py @@ -141,20 +141,9 @@ def par_loop(kernel, measure, args, kernel_kwargs=None, **kwargs): :class:`.Function`\s or components of mixed :class:`.Function`\s and indicates how these :class:`.Function`\s are to be accessed. :arg kernel_kwargs: keyword arguments to be passed to the - ``pyop2.Kernel`` constructor + ``pyop3.Function`` constructor :arg kwargs: additional keyword arguments are passed to the underlying - ``pyop2.par_loop`` - - :kwarg iterate: Optionally specify which region of an - :class:`pyop2.types.set.ExtrudedSet` to iterate over. - Valid values are the following objects from pyop2: - - - ``ON_BOTTOM``: iterate over the bottom layer of cells. - - ``ON_TOP`` iterate over the top layer of cells. - - ``ALL`` iterate over all cells (the default if unspecified) - - ``ON_INTERIOR_FACETS`` iterate over all the layers - except the top layer, accessing data two adjacent (in - the extruded direction) cells at a time. + ``pyop3.loop`` **Example** @@ -302,7 +291,7 @@ def par_loop(kernel, measure, args, kernel_kwargs=None, **kwargs): domain, = domains mesh = domain - with heavy_caches({mesh.topology}): + with heavy_caches([mesh.topology]): kernel_domains, instructions = kernel function = _form_loopy_kernel(kernel_domains, instructions, measure, args, **kernel_kwargs) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 85060eb1ac..5281a9ac03 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -227,9 +227,9 @@ def local_bc(bc, cellwise): sub_domain = list(bc.sub_domain) if "on_boundary" in sub_domain: sub_domain.remove("on_boundary") - sub_domain.extend(V.mesh().unique().exterior_facets.unique_markers) + sub_domain.extend(V.mesh().unique().facet_markers) - valid_markers = Vsub.mesh().unique().exterior_facets.unique_markers + valid_markers = Vsub.mesh().unique().facet_markers sub_domain = list(set(sub_domain) & set(valid_markers)) bc = bc.reconstruct(V=Vsub, g=0, sub_domain=sub_domain) if cellwise: diff --git a/pyop3/axis_tree/tree.py b/pyop3/axis_tree/tree.py index 9df344a454..111342f4cd 100644 --- a/pyop3/axis_tree/tree.py +++ b/pyop3/axis_tree/tree.py @@ -28,6 +28,7 @@ from petsc4py import PETSc import pyop3.cache +import pyop3.labeled_tree import pyop3.record from pyop3.cache import cached_on, memory_cache, cached_method from pyop3.collections import StrictlyUniqueDict, OrderedSet, OrderedFrozenSet @@ -40,7 +41,7 @@ from pyop3.labeled_tree import ( as_node_map, LabelledNodeComponent, - LabelledTree, + LabeledTree, MultiComponentLabelledNode, MutableLabelledTreeMixin, accumulate_path, @@ -775,7 +776,7 @@ def _getitem_cache_key(indices, *, strict=False) -> Hashable: return (indices, strict) -class AbstractAxisTreeLike(pyop3.obj.Pyop3Object): +class AbstractAxisTreeLike(pyop3.labeled_tree.AbstractLabeledTreeLike): """Base class for things that look like axis trees or forests.""" # {{{ abstract methods @@ -827,6 +828,13 @@ def buffer_size(self, *, include_ghosts: bool) -> int: def block_shape(self) -> tuple[int, ...]: pass + # TODO: This is actually a property of a LabeledTree, we should have LabeledTreeLike + # for axis forests + @property + @abc.abstractmethod + def is_linear(self) -> bool: + pass + @property @abc.abstractmethod def sf(self) -> StarForest: @@ -891,7 +899,7 @@ def __contains__(self, obj: Any, /) -> bool: node_map = idict({idict(): None}) -class AbstractNonUnitAxisTree(AbstractAxisTree, ContextFreeLoopIterable, LabelledTree): +class AbstractNonUnitAxisTree(LabeledTree, AbstractAxisTree, ContextFreeLoopIterable): """Base class for non-unit axis trees.""" # {{{ abstract methods @@ -1434,7 +1442,7 @@ def index(self) -> LoopIndex: -@pyop3.record.frozenrecord() +@pyop3.record.frozenrecord(repr=False) class AxisTree(MutableLabelledTreeMixin, AbstractNonUnitAxisTree, AbstractUnindexedAxisTree): # {{{ instance attrs @@ -1675,7 +1683,7 @@ def global_numbering(self) -> Dat[IntType]: return Dat(self, data=numbering, constant=True) -@pyop3.record.frozenrecord() +@pyop3.record.frozenrecord(repr=False) class IndexedAxisTree(AbstractNonUnitAxisTree, AbstractIndexedAxisTree): # {{{ instance attrs @@ -2371,6 +2379,10 @@ def block_shape(self) -> tuple[int, ...]: tree.block_shape[-min_block_shape_size:] for tree in self.trees )) + @property + def is_linear(self) -> bool: + return utils.single_valued(t.is_linear for t in self.trees) + # }}} def __str__(self, /) -> str: diff --git a/pyop3/buffer.py b/pyop3/buffer.py index a13a234f37..8d8cb63f37 100644 --- a/pyop3/buffer.py +++ b/pyop3/buffer.py @@ -243,7 +243,7 @@ def handle(self, *, nest_indices: tuple[tuple[int, ...], ...] = ()) -> Any: """The underlying data structure.""" -@pyop3.record.record() +@pyop3.record.record(repr=False) class ArrayBuffer(AbstractArrayBuffer, ConcreteBuffer): """A buffer whose underlying data structure is a lazily-evaluated NumPy/CuPy array.""" diff --git a/pyop3/exceptions.py b/pyop3/exceptions.py index 143cd60336..1494c7e082 100644 --- a/pyop3/exceptions.py +++ b/pyop3/exceptions.py @@ -35,6 +35,13 @@ class EmptyIterableException(Pyop3Exception): class NonUnitIterableException(Pyop3Exception): pass +# {{{ expressions + +class ExpressionUnchangedException(Pyop3Exception): + pass + +# }}} + # {{{ axis trees class IncompatibleAxisTargetException(Pyop3Exception): diff --git a/pyop3/expr/buffer.py b/pyop3/expr/buffer.py index 66090d97db..9920dc3f01 100644 --- a/pyop3/expr/buffer.py +++ b/pyop3/expr/buffer.py @@ -7,6 +7,7 @@ from immutabledict import immutabledict as idict from typing import ClassVar +import pyop3.axis_tree import pyop3.record from pyop3 import utils from pyop3.node import NodeVisitor @@ -220,8 +221,6 @@ class NonlinearDatBufferExpression(DatBufferExpression, NonlinearBufferExpressio This class is useful for describing dats whose layouts have been optimised. - Unlike `_ExpressionDat` a `_ConcretizedDat` is permitted to be multi-component. - """ # {{{ instance attrs @@ -239,12 +238,7 @@ def get_disk_cache_key(self, visitor) -> Hashable: return (type(self), visitor(self._buffer), layouts_key) def __post_init__(self) -> None: - from pyop3.expr.visitors import check_valid_layout - - assert isinstance(self._buffer, AbstractBuffer) - assert isinstance(self.layouts, idict) - for l in self.layouts.values(): - check_valid_layout(l) + pass # }}} @@ -432,9 +426,17 @@ def _(expr: LinearDatBufferExpression) -> LinearDatBufferExpression: def _(dat: Dat) -> LinearDatBufferExpression: assert dat.transform is None if not dat.axes.is_linear: - raise ValueError("The provided Dat must be linear") + raise ValueError("The provided dat must be linear") axes = dat.axes.regionless() + # We assume that if we hit an axis forest at this point then any layout + # expression is valid. + # This can happen if we use maps with multiple possible matches (e.g. mapping + # from cells or owned cells). + if isinstance(axes, pyop3.axis_tree.AxisForest): + # FIXME, merge? + axes = axes.trees[-1] + layout = utils.just_one(axes.leaf_subst_layouts.values()) return LinearDatBufferExpression(dat.buffer, layout) diff --git a/pyop3/expr/tensor/dat.py b/pyop3/expr/tensor/dat.py index 98c76865e6..9bfb6dcb72 100644 --- a/pyop3/expr/tensor/dat.py +++ b/pyop3/expr/tensor/dat.py @@ -705,6 +705,13 @@ def __init__(self, axis_tree, exprs) -> None: exprs = idict(exprs) object.__setattr__(self, "axis_tree", axis_tree) object.__setattr__(self, "exprs", exprs) + self.__post_init__() + + def __post_init__(self): + pass + # expected = "(dat_222_buffer[array_404[dat_307_buffer[((support_5_buffer[(array_1077[dat_625_buffer[i_{_label_Slice_540_owned}]] + i_{support})] * 4) + i_{closure})]]] + (array_260[(orientations_1_buffer[((support_5_buffer[(array_1077[dat_625_buffer[i_{_label_Slice_540_owned}]] + i_{support})] * 15) + i_{closure})] + i_{dof0})] * 3))" + # if expected in str(list(self.exprs.values())[0]): + # breakpoint() # }}} diff --git a/pyop3/expr/visitors/__init__.py b/pyop3/expr/visitors/__init__.py index bce71ec7e6..021e616778 100644 --- a/pyop3/expr/visitors/__init__.py +++ b/pyop3/expr/visitors/__init__.py @@ -262,8 +262,8 @@ def _(array: pyop3.expr.Tensor, /, loop_context): def replace_terminals(obj: Any, /, replace_map, *, assert_modified: bool = False) -> ExpressionT: new_obj = _replace_terminals(obj, replace_map) - if assert_modified: - assert new_obj != obj + if assert_modified and new_obj == obj: + raise pyop3.exceptions.ExpressionUnchangedException return new_obj diff --git a/pyop3/index_tree/tree.py b/pyop3/index_tree/tree.py index 2fce6cc9a1..1fb14786e9 100644 --- a/pyop3/index_tree/tree.py +++ b/pyop3/index_tree/tree.py @@ -50,7 +50,7 @@ from pyop3.labeled_tree import ( as_node_map, LabelledNodeComponent, - LabelledTree, + LabeledTree, MultiComponentLabelledNode, MutableLabelledTreeMixin, accumulate_path, @@ -79,7 +79,7 @@ class Index(MultiComponentLabelledNode): # nonsense. Instead I think they should just advertise a degree and then attach # to matching index (instead of label). @pyop3.record.frozenrecord() -class IndexTree(MutableLabelledTreeMixin, LabelledTree): +class IndexTree(MutableLabelledTreeMixin, LabeledTree): # {{{ instance attrs @@ -859,8 +859,6 @@ def axes(self) -> IndexedAxisTree: input_target_path = merge_dicts(t.path for t in input_target) if input_target_path in self.connectivity: - found = True - if len(self.connectivity[input_target_path]) > 1: raise UnspecialisedCalledMapException( "Multiple (equivalent) output paths are generated by the map. " @@ -871,13 +869,14 @@ def axes(self) -> IndexedAxisTree: # make a method subaxis, subtargets = _make_leaf_axis_from_called_map_new( - self, self.name, output_spec, input_target, + self, self.name, output_spec, input_target ) axes_ = axes_.add_axis(input_leaf_path, subaxis) for subtarget_key, subtarget_value in subtargets.items(): targets[input_leaf_path | subtarget_key] = subtarget_value + found = True break assert found diff --git a/pyop3/labeled_tree.py b/pyop3/labeled_tree.py index 1097900bc8..177f8bac5c 100644 --- a/pyop3/labeled_tree.py +++ b/pyop3/labeled_tree.py @@ -97,7 +97,15 @@ def component_label(self): return just_one(self.component_labels) -class LabelledTree(pyop3.obj.Pyop3Object): +class AbstractLabeledTreeLike(pyop3.obj.Pyop3Object): + + @property + @abc.abstractmethod + def is_linear(self) -> bool: + pass + + +class LabeledTree(AbstractLabeledTreeLike): # {{{ abstract methods @@ -116,7 +124,7 @@ def as_node(self, obj: Any) -> Node: # {{{ constructors @classmethod - def from_iterable(cls, iterable: Iterable) -> LabelledTree: + def from_iterable(cls, iterable: Iterable) -> LabeledTree: if not iterable: return cls() @@ -129,7 +137,7 @@ def from_iterable(cls, iterable: Iterable) -> LabelledTree: return cls(node_map) @classmethod - def from_nest(cls, nest: Mapping[Node, Sequence[Mapping | Node]] | Node) -> LabelledTree: + def from_nest(cls, nest: Mapping[Node, Sequence[Mapping | Node]] | Node) -> LabeledTree: if isinstance(nest, Node): return cls(nest) else: @@ -637,7 +645,7 @@ def add_node(self, path: PathT | None, node: Node) -> MutableLabelledTreeMixin: return type(self)(self.node_map | {path: node}) - def add_subtree(self, path: PathT | None, subtree: LabelledTree) -> MutableLabelledTreeMixin: + def add_subtree(self, path: PathT | None, subtree: LabeledTree) -> MutableLabelledTreeMixin: """Attach another tree to a leaf of the current tree.""" if path is None: path = self.leaf_path diff --git a/pyop3/record.py b/pyop3/record.py index dd64d5b5c1..946fbb75c8 100644 --- a/pyop3/record.py +++ b/pyop3/record.py @@ -28,12 +28,14 @@ from pyop3.exceptions import UnhashableObjectException -def record(): - return _make_record_class(eq=False) +def record(**kwargs): + assert "eq" not in kwargs + return _make_record_class(eq=False, **kwargs) -def frozenrecord(): - return _make_record_class(frozen=True) +def frozenrecord(**kwargs): + assert "frozen" not in kwargs + return _make_record_class(frozen=True, **kwargs) def _make_record_class(**kwargs): diff --git a/pyop3/utils.py b/pyop3/utils.py index 0a056d78fb..4d6ce4a0ca 100644 --- a/pyop3/utils.py +++ b/pyop3/utils.py @@ -302,6 +302,7 @@ def just_one(iterable: collections.abc.Iterable) -> Any: try: first = next(iterator) except StopIteration: + breakpoint() raise pyop3.exceptions.EmptyIterableException("Iterable is empty") try: diff --git a/tests/firedrake/regression/test_exodus_mesh.py b/tests/firedrake/regression/test_exodus_mesh.py index bb2cac177b..eee8bf3668 100644 --- a/tests/firedrake/regression/test_exodus_mesh.py +++ b/tests/firedrake/regression/test_exodus_mesh.py @@ -27,4 +27,4 @@ def test_sidesets(exodus_mesh): if exodus_mesh is None: pytest.skip("PETSc not configured with exodusII") else: - assert (exodus_mesh.exterior_facets.unique_markers == [200, 201]).all() + assert (exodus_mesh.facet_markers == [200, 201]).all() From 71da8ac82a497859fb787c59611cd269a41c3fe0 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Tue, 23 Jun 2026 15:17:03 +0100 Subject: [PATCH 07/24] fix use_fuse issue --- firedrake/utility_meshes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index be87907bc0..aa342a636c 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -1618,6 +1618,7 @@ def TensorBoxMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a mesh of a 3D box. From 4fd6e9ac0f2c39738bbc5364ad7f5c85614cd587 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Tue, 23 Jun 2026 16:38:36 +0100 Subject: [PATCH 08/24] missed a propogation of use_fuse --- firedrake/utility_meshes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index aa342a636c..cf9a5e16c6 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -331,6 +331,7 @@ def UnitIntervalMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) From 8d13a629d01c19d5a52542fc1f96ff5b103b61cd Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Tue, 23 Jun 2026 17:12:22 +0100 Subject: [PATCH 09/24] Parallel submesh fixes --- firedrake/mesh.py | 17 +++++++++-------- pyop3/expr/visitors/__init__.py | 20 ++++++++++---------- pyop3/insn/visitors.py | 1 - 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 48bc467091..0e131a11d9 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -633,16 +633,17 @@ def _local_facet_numbers_dat(self, facet_type: Literal["exterior"] | Literal["in arity = 2 local_facet_numbers = dmcommon.local_facet_number(self, facet_type) - owned_local_facet_numbers = local_facet_numbers[:facet_axes.owned.local_size] + # owned_local_facet_numbers = local_facet_numbers[:facet_axes.owned.local_size] # only ghost facets can have negative entries - utils.debug_assert(lambda: (owned_local_facet_numbers >= 0).all()) + # utils.debug_assert(lambda: (owned_local_facet_numbers >= 0).all()) # FIXME: cast dtype, should be avoidable - owned_local_facet_numbers = owned_local_facet_numbers.astype(np.uint32) + # owned_local_facet_numbers = owned_local_facet_numbers.astype(np.uint32) + local_facet_numbers = local_facet_numbers.astype(np.uint32) - axes = op3.AxisTree.from_iterable([facet_axes.owned.as_axis(), arity]) - return op3.Dat(axes, data=owned_local_facet_numbers.flatten()) + axes = op3.AxisTree.from_iterable([facet_axes.as_axis(), arity]) + return op3.Dat(axes, data=local_facet_numbers.flatten()) @cached_property def _exterior_facet_local_orientation_dat(self) -> op3.Dat: @@ -657,12 +658,12 @@ def _local_facet_orientation_dat(self, facet_type: Literal["exterior", "interior if facet_type == "exterior": local_facet_numbers_dat = self._exterior_facet_local_numbers_dat arity = 1 - facet_to_cell_map = self._facet_support_dats("exterior")[1].data_ro + facet_to_cell_map = self._facet_support_dat("exterior", only_owned=False).data_ro else: assert facet_type == "interior" local_facet_numbers_dat = self._interior_facet_local_numbers_dat arity = 2 - facet_to_cell_map = self._facet_support_dats("interior")[1].data_ro + facet_to_cell_map = self._facet_support_dat("interior", only_owned=False).data_ro facet_to_cell_map = facet_to_cell_map.reshape((-1, arity)) @@ -700,7 +701,7 @@ def _local_facet_orientation_dat(self, facet_type: Literal["exterior", "interior # form = FacetNormal(meshA)[0] * ds(meshB, interface) # # Reshape local_facets as (-1, self._rank) to uniformly handle exterior and interior facets. - local_facets = local_facet_numbers_dat.data_ro.reshape((-1, arity)) + local_facets = local_facet_numbers_dat.data_ro_with_halos.reshape((-1, arity)) # Make slice for masking out rows for which orientations are not needed. slice_ = (facet_to_cell_map != -1).all(axis=1) data = np.full_like(local_facets, np.iinfo(dtype).max) diff --git a/pyop3/expr/visitors/__init__.py b/pyop3/expr/visitors/__init__.py index 021e616778..fa4e83423a 100644 --- a/pyop3/expr/visitors/__init__.py +++ b/pyop3/expr/visitors/__init__.py @@ -575,9 +575,9 @@ def _(self, dat_expr: pyop3.expr.NonlinearDatBufferExpression, index, /, *, axis axis_tree = utils.just_one(axis_trees) candidates = {} - for path, layout in dat_expr.layouts.items(): - selector_ = selector[index, path] if selector is not None else None - candidates[index, path] = collect_candidate_indirections( + for i, (path, layout) in enumerate(dat_expr.layouts.items()): + selector_ = selector[index, i] if selector is not None else None + candidates[index, i] = collect_candidate_indirections( layout, axis_tree.linearize(path), selector=selector_, **kwargs ) return idict(candidates) @@ -607,9 +607,9 @@ def _(self, mat_expr: pyop3.expr.MatArrayBufferExpression, index, /, *, axis_tre candidates = {} layoutss = [mat_expr.row_layouts, mat_expr.column_layouts] for i, (axis_tree, layouts) in enumerate(zip(axis_trees, layoutss, strict=True)): - for path, layout in layouts.items(): - selector_ = selector[index, i, path] if selector is not None else None - candidates[index, i, path] = collect_candidate_indirections( + for j, (path, layout) in enumerate(layouts.items()): + selector_ = selector[index, i, j] if selector is not None else None + candidates[index, i, j] = collect_candidate_indirections( layout, axis_tree.linearize(path), loop_indices, compress=compress, selector=selector_ ) return idict(candidates) @@ -782,8 +782,8 @@ def _(self, buffer_expr: pyop3.expr.LinearDatBufferExpression, index, layouts, k @process.register(pyop3.expr.NonlinearDatBufferExpression) def _(self, buffer_expr: pyop3.expr.NonlinearDatBufferExpression, index, layouts, key): new_layouts = {} - for leaf_path in buffer_expr.layouts.keys(): - layout = layouts[key + ((index, leaf_path),)] + for i, leaf_path in enumerate(buffer_expr.layouts.keys()): + layout = layouts[key + ((index, i),)] new_layouts[leaf_path] = linearize_expr(layout, path=leaf_path) new_layouts = idict(new_layouts) return buffer_expr.__record_init__(layouts=new_layouts) @@ -803,8 +803,8 @@ def _(self, buffer_expr: pyop3.expr.MatArrayBufferExpression, index, /, layouts, buffer_layoutss = [buffer_expr.row_layouts, buffer_expr.column_layouts] for i, buffer_layouts in enumerate(buffer_layoutss): new_layouts = {} - for leaf_path in buffer_layouts.keys(): - layout = layouts[key + ((index, i, leaf_path),)] + for j, leaf_path in enumerate(buffer_layouts.keys()): + layout = layouts[key + ((index, i, j),)] new_layouts[leaf_path] = linearize_expr(layout, path=leaf_path) new_buffer_layoutss.append(utils.freeze(new_layouts)) return buffer_expr.__record_init__(row_layouts=new_buffer_layoutss[0], column_layouts=new_buffer_layoutss[1]) diff --git a/pyop3/insn/visitors.py b/pyop3/insn/visitors.py index a1632dc4b7..d7471a4405 100644 --- a/pyop3/insn/visitors.py +++ b/pyop3/insn/visitors.py @@ -617,7 +617,6 @@ def materialize_indirections(insn: pyop3.insn.Instruction, *, compress: bool = F # Drop cost information from 'best_candidate' best_candidate = {key: expr for key, (expr, _, _) in best_candidate.items()} - else: materialize_idxss = insn.comm.bcast(None) From e27fd4d9fb019de29ae03866ff791dc569da127e Mon Sep 17 00:00:00 2001 From: India Marsden Date: Tue, 23 Jun 2026 17:21:20 +0100 Subject: [PATCH 10/24] more additions to utility, fix for naming --- firedrake/mesh.py | 7 ++----- firedrake/pack.py | 3 --- firedrake/utility_meshes.py | 20 ++++++++++++++++++-- 3 files changed, 20 insertions(+), 10 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index d001f63f08..5bc4fefae2 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -1458,13 +1458,11 @@ def _fiat_cell_closures(self) -> np.ndarray: self.submesh_parent._fiat_cell_closures, entity_per_cell, ) - elif hasattr(self.ufl_cell(), "to_fiat") and self.ufl_cell().cellname == "tetrahedron": - # TODO better way to identify fuse use - use env var + elif self._use_fuse and self.ufl_cell().cellname.split("_")[-1] == "tetrahedron": return self._reorder_closure_fuse_tet(plex_closures) elif self.ufl_cell().is_simplex: return self._reorder_closure_fiat_simplex(plex_closures) - - elif self.ufl_cell().cellname == "quadrilateral": + elif self.ufl_cell().cellname.split("_")[-1] == "quadrilateral": return self._reorder_closure_fiat_quad(plex_closures) else: @@ -3508,7 +3506,6 @@ def _renumber_entities(self, reorder): swarm_parent_cell_nums = swarm.getField(cell_id_name).ravel() parent_renum = self._parent_mesh._new_to_old_point_renumbering.getIndices() pStart, _ = parent.getChart() - parent_renum_inv = np.empty_like(parent_renum) parent_renum_inv[parent_renum - pStart] = np.arange(len(parent_renum)) # Use kind = 'stable' to make the ordering deterministic. perm = np.argsort(parent_renum_inv[swarm_parent_cell_nums - pStart], kind='stable').astype(IntType) diff --git a/firedrake/pack.py b/firedrake/pack.py index 72252f6079..82209fbf0a 100644 --- a/firedrake/pack.py +++ b/firedrake/pack.py @@ -641,9 +641,6 @@ def construct_switch_statement(space, mats: dict, n: int, idx: int, args: list, def get_utility_kernels(ns: tuple[int]) -> tuple: strns = "".join([str(n) for n in ns]) - print(strns) - if "," in strns: - breakpoint() if len(ns) == 1: row_idx = "j" col_idx = "" diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index cf9a5e16c6..e6abd8a642 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -1773,6 +1773,7 @@ def BoxMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a mesh of a 3D box. @@ -1839,6 +1840,7 @@ def BoxMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) @@ -1901,6 +1903,7 @@ def CubeMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) @@ -1959,6 +1962,7 @@ def UnitCubeMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse = use_fuse ) @@ -2225,6 +2229,7 @@ def PeriodicUnitCubeMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) @@ -2403,6 +2408,7 @@ def UnitIcosahedralSphereMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) @@ -2619,6 +2625,7 @@ def UnitOctahedralSphereMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) @@ -2869,6 +2876,7 @@ def UnitCubedSphereMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) @@ -2885,6 +2893,7 @@ def TorusMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a toroidal mesh @@ -2979,6 +2988,7 @@ def AnnulusMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate an annulus mesh periodically extruding an interval mesh @@ -3009,7 +3019,8 @@ def AnnulusMesh( comm=comm, name=base_name, distribution_name=distribution_name, - permutation_name=permutation_name) + permutation_name=permutation_name, + use_fuse=use_fuse) bar = ExtrudedMesh(base, layers=nt, layer_height=2 * np.pi / nt, extrusion_type="uniform", periodic=True) x, y = ufl.SpatialCoordinate(bar) V = bar.coordinates.function_space() @@ -3032,6 +3043,7 @@ def SolidTorusMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False, ): """Generate a solid toroidal mesh (with axis z) periodically extruding a disk mesh @@ -3061,7 +3073,8 @@ def SolidTorusMesh( distribution_parameters=distribution_parameters, comm=comm, distribution_name=distribution_name, - permutation_name=permutation_name) + permutation_name=permutation_name, + use_fuse=use_fuse) x, y = ufl.SpatialCoordinate(unit) V = unit.coordinates.function_space() coord = Function(V).interpolate(ufl.as_vector([r * x + R, r * y])) @@ -3093,6 +3106,7 @@ def CylinderMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False ): """Generates a cylinder mesh. @@ -3264,6 +3278,7 @@ def PartiallyPeriodicRectangleMesh( name=DEFAULT_MESH_NAME, distribution_name=None, permutation_name=None, + use_fuse=False ): """Generate a RectangleMesh that is periodic in the x or y direction. @@ -3321,6 +3336,7 @@ def PartiallyPeriodicRectangleMesh( name=name, distribution_name=distribution_name, permutation_name=permutation_name, + use_fuse=use_fuse ) From 63ba7a798b8395ee9aa97e2dceddf4729d88401d Mon Sep 17 00:00:00 2001 From: India Marsden Date: Wed, 24 Jun 2026 12:25:21 +0100 Subject: [PATCH 11/24] add use_fuse to create element --- firedrake/functionspaceimpl.py | 6 +++--- firedrake/mesh.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index a5ed21f6d3..85f31b5649 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -102,8 +102,8 @@ def check_element(element, top=True): check_element(e, top=False) -def create_element(ufl_element): - finat_element = _create_element(ufl_element) +def create_element(ufl_element, use_fuse=False): + finat_element = _create_element(ufl_element, use_fuse=use_fuse) if isinstance(finat_element, finat.TensorFiniteElement): # Retrieve scalar element finat_element = finat_element.base_element @@ -1078,7 +1078,7 @@ def __init__(self, mesh, element, name=None, *, layout=None): self.comm = mesh.comm self.element = element - self.finat_element = create_element(element) + self.finat_element = create_element(element, use_fuse=mesh._use_fuse) entity_dofs = self.finat_element.entity_dofs() nodes_per_entity = tuple(len(entity_dofs[d][0]) for d in sorted(entity_dofs)) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 5bc4fefae2..959b22f7f8 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -3506,6 +3506,7 @@ def _renumber_entities(self, reorder): swarm_parent_cell_nums = swarm.getField(cell_id_name).ravel() parent_renum = self._parent_mesh._new_to_old_point_renumbering.getIndices() pStart, _ = parent.getChart() + parent_renum_inv = np.empty_like(parent_renum) parent_renum_inv[parent_renum - pStart] = np.arange(len(parent_renum)) # Use kind = 'stable' to make the ordering deterministic. perm = np.argsort(parent_renum_inv[swarm_parent_cell_nums - pStart], kind='stable').astype(IntType) From b111881b0dd12a577e546624504e55a0622a8716 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Wed, 24 Jun 2026 15:35:07 +0100 Subject: [PATCH 12/24] Fix for interpolation with GCC, awful compiler bug --- firedrake/checkpointing.py | 2 + firedrake/cython/dmcommon.pyx | 59 ------------------- firedrake/interpolation.py | 29 +++++---- firedrake/mesh.py | 54 +++++++++-------- firedrake/mg/netgen.py | 8 +-- pyop3/insn/exec.py | 2 +- ...test_io_freeze_distribution_permutation.py | 3 +- tests/firedrake/output/test_io_function.py | 6 +- tests/firedrake/submesh/test_submesh_base.py | 18 +++--- tsfc/loopy.py | 27 ++++++++- 10 files changed, 92 insertions(+), 116 deletions(-) diff --git a/firedrake/checkpointing.py b/firedrake/checkpointing.py index dee90b2b4d..5b4891e605 100644 --- a/firedrake/checkpointing.py +++ b/firedrake/checkpointing.py @@ -1520,6 +1520,8 @@ def _load_function_topology(self, tmesh, element, tf_name, idx=None): self.viewer.pushTimestepping() self.viewer.setTimestep(idx) if element.family() == "Real": + if tV.comm.size > 1: + raise NotImplementedError assert not isinstance(element, (finat.ufl.VectorElement, finat.ufl.TensorElement)) value = self.get_attr(path, "_".join([PREFIX, "value" if idx is None else "value_" + str(idx)])) tf.dat.data[...] = value diff --git a/firedrake/cython/dmcommon.pyx b/firedrake/cython/dmcommon.pyx index 706a3cbae2..8f0e19a4af 100644 --- a/firedrake/cython/dmcommon.pyx +++ b/firedrake/cython/dmcommon.pyx @@ -2465,65 +2465,6 @@ def mark_entity_classes_using_cell_dm(PETSc.DM swarm): CHKERR(PetscFree(plex_cell_classes)) -@cython.boundscheck(False) -@cython.wraparound(False) -def get_cell_markers(PETSc.DM dm, np.ndarray cell_numbering, - subdomain_id): - """Get the cells marked by a given subdomain_id. - - :arg dm: The DM for the mesh topology - :arg cell_numbering: Section mapping dm cell points to firedrake cell indices. - :arg subdomain_id: The subdomain_id to look for. - - :raises ValueError: if the subdomain_id is not valid. - :returns: A numpy array (possibly empty) of the cell ids. - """ - cdef: - PetscInt i, j, n, offset, c, cStart, cEnd, ncells - np.ndarray cells - np.ndarray indices - - if not dm.hasLabel(CELL_SETS_LABEL): - return np.empty(0, dtype=IntType) - vals = dm.getLabelIdIS(CELL_SETS_LABEL).indices - comm = dm.comm.tompi4py() - - def merge_ids(x, y, datatype): - return x.union(y) - - op = MPI.Op.Create(merge_ids, commute=True) - - all_ids = np.asarray(sorted(comm.allreduce(set(vals), op=op)), - dtype=IntType) - op.Free() - if subdomain_id not in all_ids: - raise ValueError("Invalid subdomain_id %d not in %s" % (subdomain_id, vals)) - - if subdomain_id not in vals: - return np.empty(0, dtype=IntType) - - n = dm.getStratumSize(CELL_SETS_LABEL, subdomain_id) - if n == 0: - return np.empty(0, dtype=IntType) - indices = dm.getStratumIS(CELL_SETS_LABEL, subdomain_id).indices - ncells = 0 - get_height_stratum(dm.dm, 0, &cStart, &cEnd) - for i in range(n): - c = indices[i] - if cStart <= c < cEnd: - ncells += 1 - if ncells == 0: - return np.empty(0, dtype=IntType) - cells = np.empty(ncells, dtype=IntType) - j = 0 - for i in range(n): - c = indices[i] - if cStart <= c < cEnd: - cells[j] = cell_numbering[c] - j += 1 - return cells - - @cython.boundscheck(False) @cython.wraparound(False) def get_facet_ordering(PETSc.DM plex, PETSc.Section facet_numbering): diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index a35b798764..da746f4b08 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -7,9 +7,10 @@ import tempfile import abc +import dataclasses from functools import cached_property, partial from typing import Hashable, Literal, Callable, Iterable -from dataclasses import asdict, dataclass +from dataclasses import dataclass from numbers import Number from ufl.algorithms import extract_arguments, replace @@ -151,7 +152,13 @@ def __init__(self, expr: Expr, V: WithGeometry | BaseForm, **kwargs): def _ufl_expr_reconstruct_( self, expr: Expr, v: WithGeometry | BaseForm | None = None, **interp_data ): - interp_data = interp_data or asdict(self.options) + # Note that we can't use dataclasses.asdict here because we can't deepcopy + # PETSc objects, this is the recommended workaround. + options = { + field.name: getattr(self.options, field.name) + for field in dataclasses.fields(self.options) + } + interp_data = options | interp_data return UFLInterpolate._ufl_expr_reconstruct_(self, expr, v=v, **interp_data) @property @@ -1073,12 +1080,10 @@ def _build_interpolation_callables( target_element = runtime_quadrature_element(source_mesh, target_element, rt_var_name=rt_var_name) - iter_spec = get_iteration_spec(target_mesh, "cell") - - if not (subset is None or subset is Ellipsis): - raise NotImplementedError - assert subset.superset == cell_set - cell_set = subset + if subset is not None: + iter_spec = get_iteration_spec(target_mesh, "cell", subdomain_id=subset) + else: + iter_spec = get_iteration_spec(target_mesh, "cell") parameters = {} parameters['scalar_type'] = ScalarType @@ -1667,8 +1672,12 @@ def _build_matnest( """Return a PETSc nested matrix built from sub-interpolator matrices.""" shape = tuple(len(a.function_space()) for a in self.interpolate_args) blocks = numpy.full(shape, PETSc.Mat(), dtype=object) - for indices, (interp, sub_bcs) in Isub.items(): - blocks[indices] = interp._get_callable(bcs=sub_bcs, mat_type=sub_mat_type)() + for (ridx, cidx), (interp, sub_bcs) in Isub.items(): + if ridx is None: + ridx = 0 + if cidx is None: + cidx = 0 + blocks[ridx, cidx] = interp._get_callable(bcs=sub_bcs, mat_type=sub_mat_type)() return PETSc.Mat().createNest(blocks) def _build_aij( diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 0e131a11d9..2a62f9276b 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -358,12 +358,6 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, Submesh parent. """ - if comm.size == 1: - # in serial the point sf isn't initialised - p_start, p_end = topology_dm.getChart() - serial_sf = op3.sf.local_sf(p_end-p_start, comm) - topology_dm.setPointSF(serial_sf.sf) - dmcommon.validate_mesh(topology_dm) topology_dm.setFromOptions() self.topology_dm = topology_dm @@ -380,16 +374,12 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, self._distribute() self._grown_halos = False - self.name = name - self._did_reordering = bool(reorder) - if self.comm.size > 1: self._add_overlap() if self.sfXB is not None: self.sfXC = sfXB.compose(self.sfBC) if self.sfBC else self.sfXB dmcommon.label_facets(self.topology_dm) # this is there twice, why? dmcommon.complete_facet_labels(self.topology_dm) - # TODO: Allow users to set distribution name if they want to save # conceptually the same mesh but with different distributions, # e.g., those generated by different partitioners. @@ -467,6 +457,8 @@ def __init__(self, topology_dm, name, reorder, sfXB, perm_is, distribution_name, self._new_to_old_point_renumbering = new_to_old_point_numbering self._old_to_new_point_renumbering = new_to_old_point_numbering.invertPermutation() + self.name = name + self._did_reordering = bool(reorder) # Set/Generate names to be used when checkpointing. self._distribution_name = distribution_name or _generate_default_mesh_topology_distribution_name(self.topology_dm.comm.size, self._distribution_parameters) self._permutation_name = permutation_name or _generate_default_mesh_topology_permutation_name(reorder) @@ -800,7 +792,8 @@ def _plex_to_entity_numbering(self, dim): @cached_property def _global_old_to_new_vertex_numbering(self) -> PETSc.Section: # NOTE: This will return negative entries for ghosts - return self._old_to_new_vertex_numbering.createGlobalSection(self.topology_dm.getPointSF()) + + return self._old_to_new_vertex_numbering.createGlobalSection(self.point_sf.sf) @property def comm(self): @@ -812,8 +805,10 @@ def mpi_comm(self): @cached_property def point_sf(self) -> op3.StarForest: - petsc_sf = self.topology_dm.getPointSF() - return op3.StarForest(petsc_sf, self.num_points) + if self.comm.size == 1: + return op3.sf.local_sf(self.num_points, self.comm) + else: + return op3.StarForest(self.topology_dm.getPointSF(), self.num_points) @property def topology(self): @@ -2132,6 +2127,27 @@ def vertices(self): def cell_set(self): return self.cells.owned + @PETSc.Log.EventDecorator() + def cell_subset(self, subdomain_id, all_integer_subdomain_ids=None) -> op3.Slice: + """Return a subset over cells with the given subdomain_id. + + :arg subdomain_id: The subdomain of the mesh to iterate over. + Either an integer, an iterable of integers or the special + subdomains ``"everywhere"`` or ``"otherwise"``. + :arg all_integer_subdomain_ids: Information to interpret the + ``"otherwise"`` subdomain. ``"otherwise"`` means all + entities not explicitly enumerated by the integer + subdomains provided here. For example, if + all_integer_subdomain_ids is empty, then ``"otherwise" == + "everywhere"``. If it contains ``(1, 2)``, then + ``"otherwise"`` is all entities except those marked by + subdomains 1 and 2. + + :returns: A :class:`pyop2.types.set.Subset` for iteration. + """ + # FIXME: This method is now very pointless. interpolate should instead take a subdomain_id argument + return subdomain_id + @PETSc.Log.EventDecorator() def _set_partitioner(self, plex, distribute, partitioner_type=None): """Set partitioner for (re)distributing underlying plex over comm. @@ -2195,9 +2211,7 @@ def mark_entities(self, tf, label_value, label_name=None): "ghost", "exterior_facets", "interior_facets", - "pyop2_core", - "pyop2_owned", - "pyop2_ghost"): + "firedrake_is_ghost"): raise ValueError(f"Label name {label_name} is reserved") if not isinstance(tf, function.CoordinatelessFunction): raise TypeError(f"tf must be an instance of CoordinatelessFunction: {type(tf)} is not CoordinatelessFunction") @@ -2661,13 +2675,7 @@ def flat_points(self): base_mesh_axis = self._base_mesh.flat_points npoints = base_mesh_axis.component.local_size * column_height - # NOTE: In serial the point SF isn't set up in a valid state so we do this. It - # would be nice to avoid this branch. - if self.comm.size > 1: - point_sf = self.topology_dm.getPointSF() - else: - point_sf = op3.local_sf(self.num_points, self.comm).sf - + point_sf = self.point_sf.sf point_sf_renum = op3.sf.renumber_petsc_sf(point_sf, self._new_to_old_point_renumbering) point_sf_renum = op3.StarForest(point_sf_renum, self.comm) diff --git a/firedrake/mg/netgen.py b/firedrake/mg/netgen.py index 83f98bbbef..7376bfe51c 100644 --- a/firedrake/mg/netgen.py +++ b/firedrake/mg/netgen.py @@ -155,9 +155,7 @@ def uniformRefinementRoutine(ngmesh, cdm): logger.info(f"\t\t\t[{time.time()}]Refining the plex") cdm.setRefinementUniform(True) rdm = cdm.refine() - rdm.removeLabel("pyop2_core") - rdm.removeLabel("pyop2_owned") - rdm.removeLabel("pyop2_ghost") + rdm.removeLabel("firedrake_is_ghost") logger.info(f"\t\t\t[{time.time()}]Mapping the mesh to Netgen mesh") tic = time.time() mapping = MeshMapping(rdm, geo=ngmesh) @@ -266,9 +264,7 @@ def NetgenHierarchy(mesh, levs, flags, distribution_parameters=None): mesh = reconstruct_mesh(mesh, coordinates) # Make a plex (cdm) without overlap. cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "depth", tdim, True) - cdm.removeLabel("pyop2_core") - cdm.removeLabel("pyop2_owned") - cdm.removeLabel("pyop2_ghost") + cdm.removeLabel("firedrake_is_ghost") no = impl.create_lgmap(cdm) o = impl.create_lgmap(mesh.topology_dm) lgmaps.append((no, o)) diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index e8cb07f54b..cf9b1c0d3d 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -486,7 +486,7 @@ def __call__(self, **kwargs) -> None: This code is performance critical. """ - # if "form" in str(self): + # if "expression" in str(self): # breakpoint() if not kwargs: # shortcut for the most common case diff --git a/tests/firedrake/output/test_io_freeze_distribution_permutation.py b/tests/firedrake/output/test_io_freeze_distribution_permutation.py index bc81af612d..564f940847 100644 --- a/tests/firedrake/output/test_io_freeze_distribution_permutation.py +++ b/tests/firedrake/output/test_io_freeze_distribution_permutation.py @@ -1,3 +1,4 @@ +"""Test that writing then reading functions in parallel will preserve the distribution.""" import pytest from firedrake import * from pyop3.mpi import COMM_WORLD @@ -10,7 +11,7 @@ func_name = "f" -@pytest.mark.parallel(nprocs=7) +@pytest.mark.parallel(7) @pytest.mark.parametrize('case', ["interval", "interval_small", "interval_periodic", diff --git a/tests/firedrake/output/test_io_function.py b/tests/firedrake/output/test_io_function.py index 90dbb6c179..b81b3b2b96 100644 --- a/tests/firedrake/output/test_io_function.py +++ b/tests/firedrake/output/test_io_function.py @@ -147,7 +147,7 @@ def _load_check_save_functions(filename, func_name, comm, method, mesh_name): @pytest.mark.parallel(2) -@pytest.mark.parametrize('cell_family_degree', [ +@pytest.mark.parametrize('cell_type,family,degree', [ ("triangle_small", "P", 1), ("triangle_small", "P", 6), ("triangle_small", "DP", 0), @@ -185,9 +185,7 @@ def _load_check_save_functions(filename, func_name, comm, method, mesh_name): ("triangle_3d", "BDMF", 4), ("quad_3d", "RTCF", 4) ]) -def test_io_function_base(cell_family_degree, tmpdir): - # Parameters - cell_type, family, degree = cell_family_degree +def test_io_function_base(cell_type, family, degree, tmpdir): filename = join(str(tmpdir), "test_io_function_base_dump.h5") filename = COMM_WORLD.bcast(filename, root=0) meshA = _get_mesh(cell_type, COMM_WORLD) diff --git a/tests/firedrake/submesh/test_submesh_base.py b/tests/firedrake/submesh/test_submesh_base.py index 58d379bf5a..1a461435d0 100644 --- a/tests/firedrake/submesh/test_submesh_base.py +++ b/tests/firedrake/submesh/test_submesh_base.py @@ -189,21 +189,21 @@ def test_submesh_base_facet_integral_hex_1_process(family_degree, nelem): _test_submesh_base_facet_integral_hex(family_degree, nelem) -@pytest.mark.parallel(nprocs=2) +@pytest.mark.parallel(2) @pytest.mark.parametrize('family_degree', [("Q", 3), ]) @pytest.mark.parametrize('nelem', [2, 4, 8]) def test_submesh_base_facet_integral_hex_2_processes(family_degree, nelem): _test_submesh_base_facet_integral_hex(family_degree, nelem) -@pytest.mark.parallel(nprocs=4) +@pytest.mark.parallel(4) @pytest.mark.parametrize('family_degree', [("Q", 3), ]) @pytest.mark.parametrize('nelem', [2, 4, 8]) def test_submesh_base_facet_integral_hex_4_processes(family_degree, nelem): _test_submesh_base_facet_integral_hex(family_degree, nelem) -@pytest.mark.parallel(nprocs=2) +@pytest.mark.parallel(2) def test_submesh_base_entity_maps(): # 3---9--(5)-(12)(7) (7)-(13)-3---9---5 @@ -232,20 +232,16 @@ def test_submesh_base_entity_maps(): submesh = Submesh(mesh, dim, label_value) submesh.topology_dm.viewFromOptions("-dm_view") subdm = submesh.topology.topology_dm + raise NotImplementedError("interior_facets.facets etc") if rank == 0: - assert subdm.getLabel("pyop2_core").getStratumSize(1) == 0 - assert subdm.getLabel("pyop2_owned").getStratumSize(1) == 9 - assert subdm.getLabel("pyop2_ghost").getStratumSize(1) == 0 - assert (subdm.getLabel("pyop2_owned").getStratumIS(1).getIndices() == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])).all() + assert subdm.getLabel("firedrake_is_ghost").getStratumSize(1) == 0 assert (mesh.interior_facets.facets == np.array([11])).all assert (mesh.exterior_facets.facets == np.array([8, 9, 10, 12, 13, 14])).all assert (submesh.interior_facets.facets == np.array([])).all assert (submesh.exterior_facets.facets == np.array([5, 8, 6, 7])).all() else: - assert subdm.getLabel("pyop2_core").getStratumSize(1) == 0 - assert subdm.getLabel("pyop2_owned").getStratumSize(1) == 0 - assert subdm.getLabel("pyop2_ghost").getStratumSize(1) == 9 - assert (subdm.getLabel("pyop2_ghost").getStratumIS(1).getIndices() == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])).all() + assert subdm.getLabel("firedrake_is_ghost").getStratumSize(1) == 9 + assert (subdm.getLabel("firedrake_is_ghost").getStratumIS(1).getIndices() == np.array([0, 1, 2, 3, 4, 5, 6, 7, 8])).all() assert (mesh.interior_facets.facets == np.array([8])).all assert (mesh.exterior_facets.facets == np.array([9, 10, 11, 12, 13, 14])).all assert (submesh.interior_facets.facets == np.array([])).all diff --git a/tsfc/loopy.py b/tsfc/loopy.py index 6826f0b672..d832b180d3 100644 --- a/tsfc/loopy.py +++ b/tsfc/loopy.py @@ -334,11 +334,36 @@ def statement_accumulate(leaf, ctx): @statement.register(imp.Return) def statement_return(leaf, ctx): + import pyop3.compile + lhs = expression(leaf.variable, ctx) rhs = expression(leaf.expression, ctx) if ctx.return_increments: rhs = lhs + rhs - return [lp.Assignment(lhs, rhs, within_inames=ctx.active_inames())] + + insns = [lp.Assignment(lhs, rhs, within_inames=ctx.active_inames())] + + # FIXME: This is really hacky + if pyop3.compile._compiler is None: + pyop3.compile.set_default_compiler("mpicc") + assert pyop3.compile._compiler is not None + + # GCC has a race condition bug for non-increment returns which result + # in numerical nonsense. We add a '__sync_synchronize()' call to + # prevent it. + # + # The bug has been observed: + # * On an x86 machine with both GCC 15.2 and 16.1 + # * With -O0 + if ( + not ctx.return_increments + and pyop3.compile._compiler.func is pyop3.compile.LinuxGnuCompiler + ): + insns.append( + lp.CInstruction((), "__sync_synchronize();", within_inames=ctx.active_inames()) + ) + + return insns @statement.register(imp.ReturnAccumulate) From 36407b6e09d74a0789c0b02aef0539cc18d13f00 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Thu, 25 Jun 2026 13:54:25 +0100 Subject: [PATCH 13/24] try to fix cython issue --- requirements-build.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements-build.txt b/requirements-build.txt index f202ca3776..ba05685e6f 100644 --- a/requirements-build.txt +++ b/requirements-build.txt @@ -1,5 +1,6 @@ # Core build dependencies (adapted from pyproject.toml) Cython>=3.0 +Cython>=3.2.5 firedrake-rtree libsupermesh>=2026.0 mpi4py>3; python_version >= '3.13' From dc50bffee08f2c0aee7caebbd6afef43cf9afd3e Mon Sep 17 00:00:00 2001 From: India Marsden Date: Thu, 25 Jun 2026 14:38:37 +0100 Subject: [PATCH 14/24] change requirements --- requirements-build.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/requirements-build.txt b/requirements-build.txt index ba05685e6f..5b59914dfa 100644 --- a/requirements-build.txt +++ b/requirements-build.txt @@ -1,6 +1,5 @@ # Core build dependencies (adapted from pyproject.toml) -Cython>=3.0 -Cython>=3.2.5 +Cython>=3.0, <=3.2.5 firedrake-rtree libsupermesh>=2026.0 mpi4py>3; python_version >= '3.13' From 279e45ecbca7985da6d16a190a0cdf357cb74a39 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Thu, 25 Jun 2026 17:10:01 +0100 Subject: [PATCH 15/24] fix use fuse issue for extruded meshes --- firedrake/mesh.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 60cb3060e1..eaef9bc607 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -5018,9 +5018,9 @@ def ExtrudedMesh(mesh, layers, layer_height=None, extrusion_type='uniform', peri if extrusion_type == 'radial_hedgehog': helement = helement.reconstruct(family="DG", variant="equispaced") if periodic: - velement = finat.ufl.FiniteElement("DP", as_cell("interval", self._use_fuse), 1, variant="equispaced") + velement = finat.ufl.FiniteElement("DP", as_cell("interval", mesh._use_fuse), 1, variant="equispaced") else: - velement = finat.ufl.FiniteElement("Lagrange", as_cell("interval", self._use_fuse), 1) + velement = finat.ufl.FiniteElement("Lagrange", as_cell("interval", mesh._use_fuse), 1) element = finat.ufl.TensorProductElement(helement, velement) if gdim is None: From b1cc331d7bf5967e9aab35ed9cacc9da9c745761 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Wed, 1 Jul 2026 09:34:39 +0100 Subject: [PATCH 16/24] small edits to get tensors further --- firedrake/functionspaceimpl.py | 2 +- firedrake/mesh.py | 3 ++- firedrake/pack.py | 4 ++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 85f31b5649..c8d8d24119 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -82,7 +82,7 @@ def check_element(element, top=True): type(element) is not finat.ufl.MixedElement: raise ValueError("MixedElement modifier must be outermost") if element.cell.cellname == "hexahedron" and \ - element.family() not in ["Q", "DQ", "Real"]: + element.family() not in ["Q", "DQ", "Real", "IT"]: raise NotImplementedError("Currently can only use 'Q', 'DQ', and/or 'Real' elements on hexahedral meshes, not", element.family()) if type(element) in (finat.ufl.BrokenElement, finat.ufl.RestrictedElement, finat.ufl.HDivElement, finat.ufl.HCurlElement): diff --git a/firedrake/mesh.py b/firedrake/mesh.py index eaef9bc607..a786114c59 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -1476,7 +1476,8 @@ def _fiat_cell_closures(self) -> np.ndarray: return self._reorder_closure_fiat_quad(plex_closures) else: - assert self.ufl_cell() == ufl.hexahedron + # assert self.ufl_cell() == ufl.hexahedron + assert self.ufl_cell().cellname.split("_")[-1] == "hexahedron" return self._reorder_closure_fiat_hex(plex_closures) @cached_property diff --git a/firedrake/pack.py b/firedrake/pack.py index 82209fbf0a..8996335a02 100644 --- a/firedrake/pack.py +++ b/firedrake/pack.py @@ -809,8 +809,8 @@ def fuse_orientations(spaces: list[WithGeometry]): reversed_mats += [mat_list[i]] t_dim = space._mesh.cell_dimension() # t_dim = space.ufl_element().cell._tdim - if space._mesh.cell_dimension() != space.ufl_element().cell._tdim: - raise NotImplementedError("FUSE: cell dimension of mesh doesn't match topological dimension of element cell") + # if space._mesh.cell_dimension() != space.ufl_element().cell._tdim: + # raise NotImplementedError("FUSE: cell dimension of mesh doesn't match topological dimension of element cell") os = mats[-1][t_dim][0] ns += (os[next(iter(os.keys()))].shape[0],) strns = "".join([str(n) for n in ns]) From 97921de3f9a81f3a3b7491098221c7e288f323b3 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Thu, 2 Jul 2026 11:20:14 +0100 Subject: [PATCH 17/24] test out skipping reordering --- firedrake/mesh.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index a786114c59..fdde3d83fd 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -1916,6 +1916,7 @@ def entity_orientations(self): @cached_property def entity_orientations_dat_fuse(self): # Needed in this order for FUSE orientations + return self.entity_orientations_dat cell_numbering = self._old_to_new_cell_numbering_is.getIndices() entity_orientations_original = dmcommon.entity_orientations(self, self._fiat_cell_closures) entity_orientations = [[] for i in range(len(cell_numbering))] @@ -1941,8 +1942,8 @@ def entity_orientations_dat(self): cell_axis = self.cells.root # # so instead we do # cell_axis = op3.Axis([self.points.root.components[0]], self.points.root.label) - if self._use_fuse: - return self.entity_orientations_dat_fuse + # if self._use_fuse: + # return self.entity_orientations_dat_fuse # TODO: This is quite a funky way of getting this. We should be able to get # it without calling the map. From 75c38b7c7678fb21fba32d1cacdc7f88cbf19fe1 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Thu, 2 Jul 2026 14:28:43 +0100 Subject: [PATCH 18/24] refactor --- firedrake/mesh.py | 25 ------------------------- firedrake/pack.py | 12 ++++++++---- 2 files changed, 8 insertions(+), 29 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index fdde3d83fd..afc9c967e3 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -1913,37 +1913,12 @@ def entity_orientations(self): # return np.zeros_like(self._fiat_cell_closures) return dmcommon.entity_orientations(self, self._fiat_cell_closures)[self._new_to_old_cell_numbering] - @cached_property - def entity_orientations_dat_fuse(self): - # Needed in this order for FUSE orientations - return self.entity_orientations_dat - cell_numbering = self._old_to_new_cell_numbering_is.getIndices() - entity_orientations_original = dmcommon.entity_orientations(self, self._fiat_cell_closures) - entity_orientations = [[] for i in range(len(cell_numbering))] - for row,i in zip(entity_orientations_original, cell_numbering): - entity_orientations[i] = row - entity_orientations = np.array(entity_orientations) - - # FIXME: the following does not work because the labels change - cell_axis = self.cells.root - # # so instead we do - # cell_axis = op3.Axis([self.points.root.components[0]], self.points.root.label) - - # TODO: This is quite a funky way of getting this. We should be able to get - # it without calling the map. - closure_axis = self.closure(self.cells.iter()).axes.root - axis_tree = op3.AxisTree.from_nest({cell_axis: [closure_axis]}) - assert axis_tree.local_size == entity_orientations.size - return op3.Dat(axis_tree, data=entity_orientations.flatten(), prefix="orientations") - @cached_property def entity_orientations_dat(self): # FIXME: the following does not work because the labels change cell_axis = self.cells.root # # so instead we do # cell_axis = op3.Axis([self.points.root.components[0]], self.points.root.label) - # if self._use_fuse: - # return self.entity_orientations_dat_fuse # TODO: This is quite a funky way of getting this. We should be able to get # it without calling the map. diff --git a/firedrake/pack.py b/firedrake/pack.py index 8996335a02..ea24c76601 100644 --- a/firedrake/pack.py +++ b/firedrake/pack.py @@ -174,7 +174,7 @@ def transform_packed_cell_closure_dat( warnings.warn("Int Type dats cannot be transformed using fuse transforms, using old rules") packed_dat = _orient_dofs(packed_dat, space, cell_index, depth=depth) elif transform_in_kernel and transform_out_kernel: - orientations = space.mesh().entity_orientations_dat_fuse + orientations = space.mesh().entity_orientations_dat mat_work_array = op3.Dat.null(op3.AxisTree.from_iterable([packed_dat.size, packed_dat.size]), dtype=utils.ScalarType, prefix="trans") @@ -242,8 +242,8 @@ def transform_packed_cell_closure_mat( column_depth=column_depth, ) elif transform_in_kernel and transform_out_kernel: - orientations = row_space.mesh().entity_orientations_dat_fuse - orientations_c = column_space.mesh().entity_orientations_dat_fuse + orientations = row_space.mesh().entity_orientations_dat + orientations_c = column_space.mesh().entity_orientations_dat mat_work_array_row = op3.Dat.null(op3.AxisTree.from_iterable([packed_mat.nrows, packed_mat.nrows]), dtype=utils.ScalarType, prefix="trans") mat_work_array_col = op3.Dat.null(op3.AxisTree.from_iterable([packed_mat.ncols, packed_mat.ncols]), dtype=utils.ScalarType, prefix="trans") @@ -585,6 +585,10 @@ def construct_switch_statement(space, mats: dict, n: int, idx: int, args: list, indent += 1 string += indent*"\t" + f"o_val = o{idx}[i + closure_size_acc]; \n " string += [indent*"\t" + "switch (i) { \n"] + if isinstance(dim, tuple): + dim_str = "_".join([str(d) for d in dim]) + else: + dim_str = str(dim) indent += 1 for i in range(closure_sizes[dim]): string += indent*"\t" + f"case {i}:\n " @@ -594,7 +598,7 @@ def construct_switch_statement(space, mats: dict, n: int, idx: int, args: list, for val in sorted(mats[dim][i].keys()): string += indent*"\t" + f"case {val}:\n " indent += 1 - matname = f"mat{dim}_{i}_{val}" + matname = f"mat{dim_str}_{i}_{val}" #if dim > 1: # string += indent*"\t" + f"printf(\"{matname}\\n\"); \n" string += indent*"\t" + f"a{idx} = {matname};\n" From 98f0da7b1d979bae86b2bc68b66a1e369fc6fc4d Mon Sep 17 00:00:00 2001 From: India Marsden Date: Thu, 2 Jul 2026 15:31:30 +0100 Subject: [PATCH 19/24] generalise for number of factors in tp --- firedrake/extrusion_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/extrusion_utils.py b/firedrake/extrusion_utils.py index 31ec3978da..ae384b79a6 100644 --- a/firedrake/extrusion_utils.py +++ b/firedrake/extrusion_utils.py @@ -320,7 +320,7 @@ def is_real_tensor_product_element(element): assert not isinstance(element, finat.TensorFiniteElement) if isinstance(element, finat.TensorProductElement): - _, factor = element.factors + factor = element.factors[-1] return isinstance(factor, finat.Real) else: return False From b3f5d4720329d9ed4c30850780d881e5c3b50bc0 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Thu, 2 Jul 2026 16:59:51 +0100 Subject: [PATCH 20/24] add warning --- firedrake/functionspaceimpl.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index c8d8d24119..a8e4699efa 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -482,6 +482,8 @@ def collapse(self): def make_function_space(cls, mesh, element, name=None, _labels=None, **kwargs): r"""Factory method for :class:`.WithGeometryBase`.""" topology = mesh.topology + if hasattr(element, 'triple') and not mesh._use_fuse: + raise NotImplementedError("FUSE defined element not on fuse enabled mesh.") # Create a new abstract (Mixed/Real)FunctionSpace, these are neither primal nor dual. if type(element) is finat.ufl.MixedElement: if isinstance(mesh, MeshGeometry): From 4f362a27cc3385ac1bc04db8d70a91effd815fc8 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Tue, 7 Jul 2026 14:48:51 +0100 Subject: [PATCH 21/24] fix some issues where use fuse isn't passed through correctly --- firedrake/utility_meshes.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index e6abd8a642..08045ef04d 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -683,7 +683,7 @@ def UnitTriangleMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, - use_fuse=False, + use_fuse=use_fuse, ) @@ -869,7 +869,7 @@ def TensorRectangleMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, - use_fuse=False + use_fuse=use_fuse ) @@ -1143,7 +1143,7 @@ def PeriodicRectangleMesh( distribution_name=distribution_name, permutation_name=permutation_name, comm=comm, - use_fuse=False + use_fuse=use_fuse ) From e5aa5a92fa43990264ce60ea05f8dca3870c39c8 Mon Sep 17 00:00:00 2001 From: India Marsden <37078108+indiamai@users.noreply.github.com> Date: Mon, 13 Jul 2026 13:48:25 +0100 Subject: [PATCH 22/24] Refactor pack (#5241) * refactor building of packing loopy code * fix missing self --- firedrake/pack.py | 439 +++++++++++++++++++++------------------------- 1 file changed, 201 insertions(+), 238 deletions(-) diff --git a/firedrake/pack.py b/firedrake/pack.py index ea24c76601..c4168261f6 100644 --- a/firedrake/pack.py +++ b/firedrake/pack.py @@ -569,141 +569,208 @@ def modified_lgmaps(mat: op3.Mat, indices, lgmaps): petscmat.setLGMap(*orig_lgmaps) -def construct_switch_statement(space, mats: dict, n: int, idx: int, args: list, var_list: list[str]) -> str: - string = [] - string += f"a{idx} = iden; \n " - string += "\nswitch (dim) { \n" - - var_list += ["iden"] - args += [lp.TemporaryVariable("iden", initializer=np.identity(n), dtype=utils.ScalarType, read_only=True, address_space=lp.AddressSpace(1))] - - closure_sizes = space._mesh._closure_sizes[space._mesh.cell_dimension()] - closure_size_acc = 0 - indent = 0 - for dim_i, dim in enumerate(list(closure_sizes.keys())[:-1]): - string += f"case {dim_i}:\n " - indent += 1 - string += indent*"\t" + f"o_val = o{idx}[i + closure_size_acc]; \n " - string += [indent*"\t" + "switch (i) { \n"] - if isinstance(dim, tuple): - dim_str = "_".join([str(d) for d in dim]) + +class FuseMatrixApplyBuilder(object): + + build_id = itertools.count() + + def __init__(self, spaces, ns: tuple[int], closures: np.ndarray): + self.ns = ns + self.spaces = spaces + self.id = next(self.build_id) + self.closures = closures + self.args = [lp.ValueArg("d", dtype=utils.IntType), + lp.ValueArg("closure_size_acc", dtype=utils.IntType), + lp.ValueArg("o_val", dtype=utils.IntType)] + self.args += [lp.GlobalArg(f"o{i}", dtype=utils.IntType, shape=(sum(self.closures)), is_input=True) for i in range(len(self.ns))] + self.args += [lp.GlobalArg(f"a{i}", dtype=utils.ScalarType, shape=(self.ns[i], self.ns[i]), is_input=True, is_output=False) for i in range(len(self.ns))] + self.args += [lp.GlobalArg("b", dtype=utils.ScalarType, shape=self.ns, is_input=True, is_output=True), + lp.GlobalArg("res", dtype=utils.ScalarType, shape=self.ns, is_input=True, is_output=True)] + + self.a_list = ",".join([f"a{i}[:,:]" for i in range(len(self.ns))]) + self.o_list = ",".join([f"o{i}[:]" for i in range(len(self.ns))]) + self.var_list = [f"o{i}" for i in range(len(self.ns))] + ["d", "i", "o_val", "dim"] + + def get_utility_kernels(self) -> tuple: + if len(self.ns) == 1: + row_idx, col_idx, iter_idx, all_elems = "j", "", "i", ":" + all_idxs = f"{{[{row_idx}]:0 <= j < {self.ns[0]}}}", + elif len(self.ns) == 2: + row_idx, col_idx, iter_idx, all_elems = "i", "j", "k", ":,:" + all_idxs = f"{{[i,j]:0 <= i < {self.ns[0]} and 0 <= j < {self.ns[1]}}}", + else: + raise NotImplementedError("Fuse orientations cannot handle tensors") + a_idx = ",".join([i for i in [row_idx, iter_idx] if i != ""]) + res_idx = ",".join([i for i in [row_idx, col_idx] if i != ""]) + b_idx = ",".join([i for i in [iter_idx, col_idx] if i != ""]) + all_idx = ",".join([i for i in [row_idx, iter_idx, col_idx] if i != ""]) + matmuls = [] + if len(self.ns) == 1: + # computes res = Ab + matmuls += [lp.make_function( + f"{{[{all_idx}]:0 <= {all_idx} < {self.ns[0]}}}", + f""" + res[{res_idx}] = res[{res_idx}] + a[{a_idx}]*b[{b_idx}] + """, name=f"matmul{self.id}n0", lang_version=op3.LOOPY_LANG_VERSION, target=lp.CWithGNULibcTarget())] else: - dim_str = str(dim) - indent += 1 - for i in range(closure_sizes[dim]): - string += indent*"\t" + f"case {i}:\n " + # computes res = A B + matmuls += [lp.make_function( + f"{{[i,j,k]:0 <= i,k < {self.ns[0]} and 0 <= j < {self.ns[1]}}}", + f""" + res[i,j] = res[i,j] + a[i, k]*b[k,j] + """, name=f"matmul{self.id}n0", lang_version=op3.LOOPY_LANG_VERSION, target=lp.CWithGNULibcTarget())] + # computes res = B revA + matmuls += [lp.make_function( + f"{{[i,j,k]:0 <= i < {self.ns[0]} and 0 <= j,k < {self.ns[1]}}}", + f""" + res[i,j] = res[i,j] + b[i, k]*a[k, j] + """, name=f"matmul{self.id}n1", lang_version=op3.LOOPY_LANG_VERSION, target=lp.CWithGNULibcTarget())] + + set_args = [lp.GlobalArg("b", dtype=utils.ScalarType, shape=self.ns, is_input=True, is_output=True), + lp.GlobalArg("res", dtype=utils.ScalarType, shape=self.ns, is_input=True)] + set_knl = lp.make_function( + all_idxs, + [f"b[{res_idx}] = res[{res_idx}]"], + kernel_data=set_args, + name=f"set{self.id}", + lang_version=op3.LOOPY_LANG_VERSION, + target=lp.CWithGNULibcTarget() + ) + zero_knl = lp.make_function( + all_idxs, + [f"res[{res_idx}] = 0"], + [lp.GlobalArg("res", shape=self.ns, dtype=int, is_input=True, is_output=True)], + lang_version=op3.LOOPY_LANG_VERSION, + target=lp.CWithGNULibcTarget(), + name=f"zero{self.id}", + ) + self.all_elems = all_elems + return matmuls + [set_knl, zero_knl] + + + def construct_switch_statement(self, mats: dict, idx: int) -> str: + var_list = [] + args = [] + string = [] + string += f"a{idx} = iden; \n " + string += "\nswitch (dim) { \n" + + var_list += ["iden"] + args += [lp.TemporaryVariable("iden", initializer=np.identity(self.ns[idx]), dtype=utils.ScalarType, read_only=True, address_space=lp.AddressSpace(1))] + + closure_sizes = self.spaces[idx]._mesh._closure_sizes[self.spaces[idx]._mesh.cell_dimension()] + closure_size_acc = 0 + indent = 0 + for dim_i, dim in enumerate(list(closure_sizes.keys())[:-1]): + string += f"case {dim_i}:\n " indent += 1 - string += indent*"\t" + "switch (o_val) { \n" + string += indent*"\t" + f"o_val = o{idx}[i + closure_size_acc]; \n " + string += [indent*"\t" + "switch (i) { \n"] + if isinstance(dim, tuple): + dim_str = "_".join([str(d) for d in dim]) + else: + dim_str = str(dim) indent += 1 - for val in sorted(mats[dim][i].keys()): - string += indent*"\t" + f"case {val}:\n " + for i in range(closure_sizes[dim]): + string += indent*"\t" + f"case {i}:\n " indent += 1 - matname = f"mat{dim_str}_{i}_{val}" - #if dim > 1: - # string += indent*"\t" + f"printf(\"{matname}\\n\"); \n" - string += indent*"\t" + f"a{idx} = {matname};\n" - string += indent*"\t" + "break;\n" - var_list += [matname] - if dim == 2: - #print(mats[dim][i][val][np.ix_(list(range(22, 34)), list(range(22,34)))]) - #mat = np.eye(mats[dim][i][val].shape[0], dtype=utils.ScalarType) - mat = np.array(mats[dim][i][val], dtype=utils.ScalarType) - else: + string += indent*"\t" + "switch (o_val) { \n" + indent += 1 + for val in sorted(mats[dim][i].keys()): + string += indent*"\t" + f"case {val}:\n " + indent += 1 + matname = f"mat{dim_str}_{i}_{val}" + string += indent*"\t" + f"a{idx} = {matname};\n" + string += indent*"\t" + "break;\n" + var_list += [matname] mat = np.array(mats[dim][i][val], dtype=utils.ScalarType) - args += [lp.TemporaryVariable(matname, initializer=mat, dtype=utils.ScalarType, read_only=True, address_space=lp.AddressSpace(1))] + args += [lp.TemporaryVariable(matname, initializer=mat, dtype=utils.ScalarType, read_only=True, address_space=lp.AddressSpace(1))] + indent -= 1 indent -= 1 + string += indent*"\t" + "default: break;}break;\n" + indent -= 1 + string += indent*"\t" + "default: break; }break;\n" + closure_size_acc += closure_sizes[dim] indent -= 1 - string += indent*"\t" + "default: break;}break;\n" - indent -= 1 - string += indent*"\t" + "default: break; }break;\n" - closure_size_acc += closure_sizes[dim] - indent -= 1 - string += "default: break; }\n" + string += "default: break; }\n" + return string, args, var_list + + def debug(self) -> tuple[lp.CInstruction]: + if (self.ns[0] == 30) and len(self.ns) == 1: + if self.ns[0] == 84: + print_range = range(37,43) + elif self.ns[0] == 35: + print_range = range(22,25) + elif self.ns[0] == 30: + print_range = range(18, self.ns[0]) + elif self.ns[0] == 45: + print_range = range(18,32) + elif self.ns[0] == 20: + print_range = [2,3,8,9,16] + elif self.ns[0] == 60: + print_range = list(np.concat([[i*3, i*3 + 1, i*3 +2] for i in [2,3,8,9,16]])) + else: + print_range = range(0, self.ns[0]) + print_insn = lp.CInstruction(tuple(), + f"""printf(\"initial b: {" ".join('%f' for i in print_range)}\\n\", {', '.join(f"b[{j}]" for j in print_range)}); + printf(\"o: {" ".join('%d' for i in range(sum(self.closures)))}\\n\", {', '.join(f"o0[{j}]" for j in range(sum(self.closures)))}); + """, assignees=(), read_variables=frozenset([]), id="print") + print_insn1 = lp.CInstruction(tuple(), + f"""printf(\"final res: {" ".join('%f' for i in print_range)}\\n\", {', '.join(f"res[{j}]" for j in print_range)}); + """, assignees=(), read_variables=frozenset(["res"]), depends_on="replace") + else: + print_insn = lp.CInstruction(tuple(), "", assignees=(), read_variables=frozenset([]), id="print") + print_insn1 = lp.CInstruction(tuple(),"", assignees=(), read_variables=frozenset(["res"]), depends_on="replace") + return print_insn, print_insn1 + + def switch(self, mats, i, name): + dim_arg = [lp.ValueArg("dim", dtype=utils.IntType)] + switch_string, extra_args, extra_vars = self.construct_switch_statement(mats, i) + transform_insn = lp.CInstruction(tuple(), "".join(switch_string), assignees=(f"a{i}", "o_val"), read_variables=frozenset(self.var_list + extra_vars), within_inames=frozenset(["i"]), id="assign", depends_on="zero") + matmul_insn = f"res[{self.all_elems}] = matmul{self.id}n{i}(a{i}, b, res) {{id=matmul, dep=*, dep=assign, inames=i}}" + print_insn1 = lp.CInstruction(tuple(), + f"""""", assignees=(), read_variables=frozenset(["res"]), within_inames=frozenset(["i"]), depends_on="matmul", id="print") + return lp.make_function( + "{[i]:0<= i < d }", + [f"res[{self.all_elems}] = zero{self.id}(res) {{id=zero, inames=i}}", + transform_insn, matmul_insn, print_insn1, + f"b[{self.all_elems}] = set{self.id}(b[{self.all_elems}], res[{self.all_elems}]) {{id=set, dep=print, inames=i}}" + ], + name=f"{name}{i}_switch_on_o_{self.id}", + kernel_data=dim_arg + self.args + extra_args, + lang_version=op3.LOOPY_LANG_VERSION, + target=lp.CWithGNULibcTarget()) + + def loop_dims(self, direction): + closure_arg = [lp.TemporaryVariable("closure_sizes", initializer=np.array(self.closures, dtype=np.int32), dtype=utils.IntType, read_only=True, address_space=lp.AddressSpace(1))] + num_switch = len(self.all_elems.split(",")) + labelling = [ f"{{id=switch{chr(i+65)} " + (",dep=*}" if i == 0 else f",dep=switch{chr(i+64)},dep=*}}") for i in range(num_switch)] + switches = [f""" + b[{self.all_elems}], res[{self.all_elems}] = {direction + str(i)}_switch_on_o_{self.id}(dim, d, closure_size_acc, o_val, {self.o_list}, {self.a_list}, b[{self.all_elems}], res[{self.all_elems}]) {labelling[i]}""" + for i in range(num_switch)] + return lp.make_function( + f"{{[dim]:{0} <= dim <= {len(self.closures) - 1}}}", + ["d = closure_sizes[dim] {id=closure}"] + switches + + [f"closure_size_acc = closure_size_acc + d {{id=replace, dep=switch{chr(65 + num_switch-1)}, inames=dim}}"], + name=f"{direction}_loop_over_dims_{self.id}", + kernel_data=closure_arg + self.args, + lang_version=op3.LOOPY_LANG_VERSION, + target=lp.CWithGNULibcTarget()) - # string += indent*"\t" + "if ((i == 2) && dim == 1) {\n" - # indent += 1 - # string += indent*"\t" + f"printf(\"o : '%d', d: '%d', i : '%d'\\n\", o_val, d, i); \n" - # dof_range = [] - # if n == 20: - # dof_range = [2,3,8,9,16] - # if n == 60: - # #,16 - # dof_range = list(np.concat([[i*3, i*3 + 1, i*3 +2] for i in [2,3,8,9]])) - # # dof_range = list(range(12, 14)) + list(range(18, 20)) - # for i in dof_range: - # string += indent*"\t" + f"printf(\"a{idx} row {i}: {" ".join('%f' for i in dof_range)}\\n\", {", ".join(f"a{idx}[{i*n + j}]" for j in dof_range)}); \n" - #string += indent*"\t" + f"printf(\"should be \\n\"); \n" - #for i in range(4,6): - # string += indent*"\t" + f"printf(\"mat1_0_1 row {i}: {" ".join('%f' for i in range(n))}\\n\", {", ".join(f"mat1_0_1[{i*n + j}]" for j in range(n))}); \n" - #string += indent*"\t" + f"printf(\"a row 1: {" ".join('%f' for i in range(n))}\\n\", {", ".join(f"a[{n + i}]" for i in range(n))}); \n" - #string += indent*"\t" + f"printf(\"a row 2: {" ".join('%f' for i in range(n))}\\n\", {", ".join(f"a[{2*n + i}]" for i in range(n))}); \n" - #string += indent*"\t" + "printf(\"a rows ...\\n\"); \n" - #string += indent*"\t" + "printf(\"\\n\");\n" - # string += indent*"\t" + "}\n" - return string, args, var_list - -def get_utility_kernels(ns: tuple[int]) -> tuple: - strns = "".join([str(n) for n in ns]) - if len(ns) == 1: - row_idx = "j" - col_idx = "" - iter_idx = "i" - all_elems = ":" - all_idxs = f"{{[{row_idx}]:0 <= j < {ns[0]}}}", - elif len(ns) == 2: - row_idx = "i" - col_idx = "j" - iter_idx = "k" - all_elems = ":,:" - all_idxs = f"{{[i,j]:0 <= i < {ns[0]} and 0 <= j < {ns[1]}}}", - else: - raise NotImplementedError("Fuse orientations cannot handle tensors") - a_idx = ",".join([i for i in [row_idx, iter_idx] if i != ""]) - res_idx = ",".join([i for i in [row_idx, col_idx] if i != ""]) - b_idx = ",".join([i for i in [iter_idx, col_idx] if i != ""]) - all_idx = ",".join([i for i in [row_idx, iter_idx, col_idx] if i != ""]) - matmuls = [] - if len(ns) == 1: - # computes res = Ab - matmuls += [lp.make_function( - f"{{[{all_idx}]:0 <= {all_idx} < {ns[0]}}}", - f""" - res[{res_idx}] = res[{res_idx}] + a[{a_idx}]*b[{b_idx}] - """, name=f"matmul{strns}n0", lang_version=op3.LOOPY_LANG_VERSION, target=lp.CWithGNULibcTarget())] - else: - # computes res = A B - matmuls += [lp.make_function( - f"{{[i,j,k]:0 <= i,k < {ns[0]} and 0 <= j < {ns[1]}}}", - f""" - res[i,j] = res[i,j] + a[i, k]*b[k,j] - """, name=f"matmul{strns}n0", lang_version=op3.LOOPY_LANG_VERSION, target=lp.CWithGNULibcTarget())] - # computes res = B revA - matmuls += [lp.make_function( - f"{{[i,j,k]:0 <= i < {ns[0]} and 0 <= j,k < {ns[1]}}}", - f""" - res[i,j] = res[i,j] + b[i, k]*a[k, j] - """, name=f"matmul{strns}n1", lang_version=op3.LOOPY_LANG_VERSION, target=lp.CWithGNULibcTarget())] - - set_args = [lp.GlobalArg("b", dtype=utils.ScalarType, shape=ns, is_input=True, is_output=True), - lp.GlobalArg("res", dtype=utils.ScalarType, shape=ns, is_input=True)] - set_knl = lp.make_function( - all_idxs, - [f"b[{res_idx}] = res[{res_idx}]"], - kernel_data=set_args, - name=f"set{strns}", + def overall(self, direction): + print_insn, print_insn1 = self.debug() + return lp.make_kernel( + "{:}", + [print_insn, + f"b[{self.all_elems}], res[{self.all_elems}] = {direction}_loop_over_dims_{self.id}(0,0,0, {self.o_list}, {self.a_list}, b[{self.all_elems}], res[{self.all_elems}]) {{dep=print,id=loop}}", + f"res[{self.all_elems}] = set{self.id}(res[{self.all_elems}], b[{self.all_elems}]) {{id=replace, dep=loop}}", + f"b[{self.all_elems}] = zero{self.id}(b[{self.all_elems}]) {{dep=replace, id=zerob}}", + print_insn1], + name=f"{direction}_transform_{self.id}", + kernel_data=self.args[3:], lang_version=op3.LOOPY_LANG_VERSION, - target=lp.CWithGNULibcTarget() - ) - zero_knl = lp.make_function( - all_idxs, - [f"res[{res_idx}] = 0"], - [lp.GlobalArg("res", shape=ns, dtype=int, is_input=True, is_output=True)], - lang_version=op3.LOOPY_LANG_VERSION, - target=lp.CWithGNULibcTarget(), - name=f"zero{strns}", - ) - return matmuls + [set_knl, zero_knl], all_elems + target=lp.CWithGNULibcTarget()) def combine_matrices(matrices): """ Combine matrix dictionaries for Vector Matrices @@ -789,8 +856,6 @@ def check_fuse_enriched(element): return any(fuse), any(matrix), any(apply), None, None - - def fuse_orientations(spaces: list[WithGeometry]): fuse_defined_spaces, fuse_matrix_spaces, fuse_needs_matrices, mat_list, reversed_mat_list = list(zip(*[check_fuse(space.ufl_element()) for space in spaces])) @@ -812,129 +877,27 @@ def fuse_orientations(spaces: list[WithGeometry]): mats += [reversed_mat_list[i]] reversed_mats += [mat_list[i]] t_dim = space._mesh.cell_dimension() - # t_dim = space.ufl_element().cell._tdim - # if space._mesh.cell_dimension() != space.ufl_element().cell._tdim: - # raise NotImplementedError("FUSE: cell dimension of mesh doesn't match topological dimension of element cell") os = mats[-1][t_dim][0] ns += (os[next(iter(os.keys()))].shape[0],) - strns = "".join([str(n) for n in ns]) + closures_dict = mesh._closure_sizes[fs._mesh.cell_dimension()] closures = [closures_dict[c] for c in sorted(closures_dict.keys())] - utilities, all_elems = get_utility_kernels(ns) - args = [lp.ValueArg("d", dtype=utils.IntType), - lp.ValueArg("closure_size_acc", dtype=utils.IntType), - lp.ValueArg("o_val", dtype=utils.IntType)] + [lp.GlobalArg(f"o{i}", dtype=utils.IntType, shape=(sum(closures)), is_input=True) for i in range(len(ns))] + [lp.GlobalArg(f"a{i}", dtype=utils.ScalarType, shape=(ns[i], ns[i]), is_input=True, is_output=False) for i in range(len(ns))] + [lp.GlobalArg("b", dtype=utils.ScalarType, shape=ns, is_input=True, is_output=True), - lp.GlobalArg("res", dtype=utils.ScalarType, shape=ns, is_input=True, is_output=True)] - - a_list = ",".join([f"a{i}[:,:]" for i in range(len(ns))]) - o_list = ",".join([f"o{i}[:]" for i in range(len(ns))]) - var_list = [f"o{i}" for i in range(len(ns))] + ["d", "i", "o_val", "dim"] - - def switch(space, mats, n, i, args, var_list, all_elems, name, reverse=False): - dim_arg = [lp.ValueArg("dim", dtype=utils.IntType)] - switch_string, args, var_list = construct_switch_statement(fs, mats, n, i, args, var_list) - transform_insn = lp.CInstruction(tuple(), "".join(switch_string), assignees=(f"a{i}", "o_val"), read_variables=frozenset(var_list), within_inames=frozenset(["i"]), id="assign", depends_on="zero") - matmul_insn = f"res[{all_elems}] = matmul{strns}n{i}(a{i}, b, res) {{id=matmul, dep=*, dep=assign, inames=i}}" - print_insn1 = lp.CInstruction(tuple(), - f"""""", assignees=(), read_variables=frozenset(["res"]), within_inames=frozenset(["i"]), depends_on="matmul", id="print") - #print_insn1 = lp.CInstruction(tuple(), - # f"""printf(\"mid res: {" ".join('%f' for i in range(22, ns[0]))}\\n\", {', '.join(f"res[{j}]" for j in range(22, ns[0]))}); - # """, assignees=(), read_variables=frozenset(["res"]), within_inames=frozenset(["i"]), depends_on="matmul", id="print") - return lp.make_function( - "{[i]:0<= i < d }", - [f"res[{all_elems}] = zero{strns}(res) {{id=zero, inames=i}}", - transform_insn, matmul_insn, print_insn1, - f"b[{all_elems}] = set{strns}(b[{all_elems}], res[{all_elems}]) {{id=set, dep=print, inames=i}}" - ], - name=name + "_switch_on_o", - kernel_data=dim_arg + args, - lang_version=op3.LOOPY_LANG_VERSION, - target=lp.CWithGNULibcTarget()) - in_switches = [switch(spaces[i], mats[i], ns[i], i, args, var_list, all_elems, name="in"+str(i)) for i in range(len(spaces))] - out_switches = [switch(spaces[i], reversed_mats[i], ns[i], i, args, var_list, all_elems, name="out"+str(i)) for i in range(len(spaces))] - - closure_arg = [lp.TemporaryVariable("closure_sizes", initializer=np.array(closures, dtype=np.int32), dtype=utils.IntType, read_only=True, address_space=lp.AddressSpace(1))] - - def loop_dims(direction, all_elems): - num_switch = len(all_elems.split(",")) - labelling = [ f"{{id=switch{chr(i+65)} " + (",dep=*}" if i == 0 else f",dep=switch{chr(i+64)},dep=*}}") for i in range(num_switch)] - switches = [f""" - b[{all_elems}], res[{all_elems}] = {direction + str(i)}_switch_on_o(dim, d, closure_size_acc, o_val, {o_list}, {a_list}, b[{all_elems}], res[{all_elems}]) {labelling[i]}""" - for i in range(num_switch)] - return lp.make_function( - f"{{[dim]:{0} <= dim <= {len(closures) - 1}}}", - ["d = closure_sizes[dim] {id=closure}"] + switches + - [f"closure_size_acc = closure_size_acc + d {{id=replace, dep=switch{chr(65 + num_switch-1)}, inames=dim}}"], - name=f"{direction}_loop_over_dims", - kernel_data=closure_arg + args, - lang_version=op3.LOOPY_LANG_VERSION, - target=lp.CWithGNULibcTarget()) - # printf("o: {" ".join('%d' for i in range(sum(closures)))}\\n\", {', '.join(f"o0[{j}]" for j in range(sum(closures)))});") - if (ns[0] == 30) and len(ns) == 1: - if ns[0] == 84: - print_range = range(37,43) - elif ns[0] == 35: - print_range = range(22,25) - elif ns[0] == 30: - print_range = range(18,ns[0]) - elif ns[0] == 45: - print_range = range(18,32) - elif ns[0] == 20: - print_range = [2,3,8,9,16] - elif ns[0] == 60: - # print_range = [2,3,8,9,16] + [22, 23, 28, 29, 36] - print_range = list(np.concat([[i*3, i*3 + 1, i*3 +2] for i in [2,3,8,9,16]])) - else: - print_range = range(0, ns[0]) - if len(ns) > 1: - print_insn = lp.CInstruction(tuple(), "", assignees=(), read_variables=frozenset([]), id="print") - # print_insn = lp.CInstruction(tuple(), - # f"""printf(\"initial b: {" ".join('%f' for i in range(34, 41))}\\n\", {', '.join(f"b[{i*ns[0] + j}]" for j in range(12, 16) for i in range(12,16))}); - # """, assignees=(), read_variables=frozenset([]), id="print") - else: - print_insn = lp.CInstruction(tuple(), - f"""printf(\"initial b: {" ".join('%f' for i in print_range)}\\n\", {', '.join(f"b[{j}]" for j in print_range)}); - printf(\"o: {" ".join('%d' for i in range(sum(closures)))}\\n\", {', '.join(f"o0[{j}]" for j in range(sum(closures)))}); - """, assignees=(), read_variables=frozenset([]), id="print") - - if len(ns) > 1: - print_insn1 = lp.CInstruction(tuple(),"", assignees=(), read_variables=frozenset(["res"]), depends_on="replace") - # print_insn1 = lp.CInstruction(tuple(), - # f"""printf(\"final res: {" ".join('%f' for i in range(0, 16))}\\n\", {', '.join(f"res[{i*ns[0] + j}]" for j in range(12, 16) for i in range(12, 16))}); - # """, assignees=(), read_variables=frozenset(["res"]), depends_on="replace") - else: - # range(34, ns[0]) - print_insn1 = lp.CInstruction(tuple(), - f"""printf(\"final res: {" ".join('%f' for i in print_range)}\\n\", {', '.join(f"res[{j}]" for j in print_range)}); - """, assignees=(), read_variables=frozenset(["res"]), depends_on="replace") - else: - print_insn = lp.CInstruction(tuple(), "", assignees=(), read_variables=frozenset([]), id="print") - - print_insn1 = lp.CInstruction(tuple(),"", assignees=(), read_variables=frozenset(["res"]), depends_on="replace") + builder = FuseMatrixApplyBuilder(spaces, ns, closures) + utilities = builder.get_utility_kernels() - def overall(direction, all_elems): - return lp.make_kernel( - "{:}", - [print_insn, - f"b[{all_elems}], res[{all_elems}] = {direction}_loop_over_dims(0,0,0, {o_list}, {a_list}, b[{all_elems}], res[{all_elems}]) {{dep=print,id=loop}}", - f"res[{all_elems}] = set{strns}(res[{all_elems}], b[{all_elems}]) {{id=replace, dep=loop}}", - f"b[{all_elems}] = zero{strns}(b[{all_elems}]) {{dep=replace, id=zerob}}", - print_insn1], - name=f"{direction}_transform", - kernel_data=args[3:], - lang_version=op3.LOOPY_LANG_VERSION, - target=lp.CWithGNULibcTarget()) - - in_knl = lp.merge([overall("in", all_elems), loop_dims("in", all_elems)] + in_switches + utilities) - out_knl = lp.merge([overall("out", all_elems), loop_dims("out", all_elems)] + out_switches + utilities) + in_switches = [builder.switch(mats[i], i, name="in") for i in range(len(spaces))] + out_switches = [builder.switch(reversed_mats[i], i, name="out") for i in range(len(spaces))] + + in_knl = lp.merge([builder.overall("in"), builder.loop_dims("in")] + in_switches + utilities) + out_knl = lp.merge([builder.overall("out"), builder.loop_dims("out")] + out_switches + utilities) + # b is modified in the transform functions but the result is written to res and therefore is not needed further. transform_in = op3.Function(in_knl, [op3.READ for n in ns] + [op3.WRITE for n in ns] + [op3.READ, op3.RW]) transform_out = op3.Function(out_knl, [op3.READ for n in ns] + [op3.WRITE for n in ns] + [op3.READ, op3.RW]) return transform_in, transform_out elif fuse_defined_spaces and sum(fuse_matrix_spaces) == 0: - print("not matrix space") return None, None elif fuse_defined_spaces and sum(fuse_defined_spaces) != sum(fuse_matrix_spaces): raise ValueError("If a fuse space is used, all spaces must be fuse spaces") From 03dc3316c2c4b7ef2d1216964f2417fa436bb075 Mon Sep 17 00:00:00 2001 From: India Marsden Date: Wed, 15 Jul 2026 14:53:31 +0100 Subject: [PATCH 23/24] allow any element on hex if fuse enabled --- firedrake/functionspaceimpl.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index a8e4699efa..334bdf4fb6 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -49,7 +49,7 @@ from firedrake.utils import IntType, deprecated -def check_element(element, top=True): +def check_element(element, top=True, use_fuse=False): """Run some checks on the provided element. The :class:`finat.ufl.mixedelement.VectorElement` and @@ -82,7 +82,7 @@ def check_element(element, top=True): type(element) is not finat.ufl.MixedElement: raise ValueError("MixedElement modifier must be outermost") if element.cell.cellname == "hexahedron" and \ - element.family() not in ["Q", "DQ", "Real", "IT"]: + element.family() not in ["Q", "DQ", "Real", "IT"] and not use_fuse: raise NotImplementedError("Currently can only use 'Q', 'DQ', and/or 'Real' elements on hexahedral meshes, not", element.family()) if type(element) in (finat.ufl.BrokenElement, finat.ufl.RestrictedElement, finat.ufl.HDivElement, finat.ufl.HCurlElement): @@ -99,7 +99,7 @@ def check_element(element, top=True): else: inner = () for e in inner: - check_element(e, top=False) + check_element(e, top=False, use_fuse=use_fuse) def create_element(ufl_element, use_fuse=False): @@ -499,7 +499,7 @@ def make_function_space(cls, mesh, element, name=None, _labels=None, **kwargs): if isinstance(mesh, MeshSequenceGeometry): raise TypeError(f"mesh must not be MeshSequenceGeometry: got {mesh}") # Check that any Vector/Tensor/Mixed modifiers are outermost. - check_element(element) + check_element(element, use_fuse=mesh._use_fuse) if element.family() == "Real": new = RealFunctionSpace(topology, element, name=name, **kwargs) else: From 07d57b5188470c8eefb94ff0950943c3094c2d1a Mon Sep 17 00:00:00 2001 From: India Marsden Date: Mon, 20 Jul 2026 13:58:44 +0100 Subject: [PATCH 24/24] add two hex mesh --- firedrake/utility_meshes.py | 79 +++++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/firedrake/utility_meshes.py b/firedrake/utility_meshes.py index 08045ef04d..a9fb3621f4 100644 --- a/firedrake/utility_meshes.py +++ b/firedrake/utility_meshes.py @@ -216,6 +216,85 @@ def TwoTetMesh( return m +@PETSc.Log.EventDecorator() +def TwoHexMesh( + perm=None, + right=None, + distribution_parameters=None, + reorder=False, + comm=COMM_WORLD, + name=DEFAULT_MESH_NAME, + distribution_name=None, + permutation_name=None, + use_fuse=False, +): + """ + Mesh with only two hexahedra, sharing a single quadrilateral face. + + Analogous to :func:`TwoTetMesh`: the two cells share a unit-square + face in the ``z == 0`` plane (cell A occupies ``z in [-1, 0]``, cell + B occupies ``z in [0, 1]``). Cell B's local vertex ordering is fixed; + ``perm`` (if given) is applied to cell A's local vertex list, letting + the caller sweep the relative orientation of the shared face through + the quadrilateral's full 8-element dihedral symmetry group -- useful + for exhaustively testing H(div)/H(curl) orientation handling on + hexahedral cells, the way :func:`TwoTetMesh` does for tetrahedra + (whose shared triangular face only has a 6-element symmetry group). + + :arg perm: (optional) a callable (e.g. a `sympy.combinatorics.Permutation` + of length 8) applied to cell A's local vertex list + ``[4, 5, 6, 7, 0, 3, 2, 1]`` (positions 0-3 hold cell A's own + private vertices, positions 4-7 the shared-face vertices). + ``perm`` must be a symmetry of the cube's vertex-position graph, + i.e. a D4 element acting simultaneously on BOTH blocks (the two + blocks are matched by the vertical vertex pairing of the DMPlex + hexahedron cone convention: bottom position ``i`` sits below top + position ``[0, 3, 2, 1][i]``). Permuting only positions 4-7 + while fixing 0-3 is NOT a cube symmetry: it produces a twisted + trilinear cell (cell A's volume comes out != 1 -- 2/3 or 1/3 + for face rotations, 1/sqrt(3) for face reflections), silently + invalidating anything computed on the mesh. Sweeping the paired + D4 elements realises all 8 relative orientations of the shared + face with valid geometry. + :kwarg distribution_parameters: options controlling mesh + distribution, see :func:`.Mesh` for details. + :kwarg reorder: (optional), should the mesh be reordered? + :kwarg comm: Optional communicator to build the mesh on. + :kwarg name: Optional name of the mesh. + :kwarg distribution_name: the name of parallel distribution used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg permutation_name: the name of entity permutation (reordering) used + when checkpointing; if `None`, the name is automatically + generated. + :kwarg use_fuse: Flag indicting if mesh should be created using a fuse cell + Default False, mesh is based on the ufc cell. + """ + coords = np.array([ + [0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0], # 0-3: shared face + [0, 0, -1], [0, 1, -1], [1, 1, -1], [1, 0, -1], # 4-7: cell A private + [0, 0, 1], [1, 0, 1], [1, 1, 1], [0, 1, 1], # 8-11: cell B private + ]) + cellA = [4, 5, 6, 7, 0, 3, 2, 1] + cellB = [0, 1, 2, 3, 8, 9, 10, 11] + if perm is not None: + cellA = perm(cellA) + cells = np.array([cellA, cellB]) + plex = plex_from_cell_list( + 3, cells, coords, comm, _generate_default_mesh_topology_name(name) + ) + m = Mesh( + plex, + reorder=reorder, + distribution_parameters=distribution_parameters, + name=name, + distribution_name=distribution_name, + permutation_name=permutation_name, + comm=comm, + use_fuse=use_fuse + ) + return m + @PETSc.Log.EventDecorator() def IntervalMesh(