Scalable point location by leo-collins · Pull Request #5396 · firedrakeproject/firedrake

leo-collins · 2026-08-27T12:06:11Z

Introduces the partition_rtree, which is an R-tree of bounding boxes covering each partition of the mesh. This R-tree is replicated on all ranks by allgathering the bounding boxes. The bounding boxes on each rank are obtained by traversing down the local R-tree. Each rank may submit more than one bounding box to the partition_rtree (we want to do this to minimise candidates since the mesh partitions may not be nice and convex, and might even be disconnected).

We choose the number of bounding boxes via a simple heuristic (we stop traversing when the total volume of the bounding boxes stop decreasing - see _box_ratio_heuristic). This should probably be improved to take into account the total number of bounding boxes so we don't blow up the allgather.

Constructing the partition rtree uses an allgather, although each rank normally submits O(10s) of boxes so I have not seen this kill us yet when testing up to 2048 ranks. If doing huge number of cores we might want to consider implementing some sort of hierarchy of partition rtrees.

Each rank queries this R-tree with the points local to that rank to determine the 'candidate' ranks for each point. We then build the candidate_sf which is a PETScSF mapping input points to their candidates. The local cell location process proceeds as usual, solving parallel conflicts with broadcasts/reductions via the candidate_sf.

This current implementation gets around 80% weak scaling efficiency in my tests up to 2048 cores.

more logging

one more

doing this fixes test errors...

connorjward · 2026-09-01T15:45:45Z

+                nleaves += count
+
+            if not sends_complete:
+                CHKERRMPI(MPI_Testall(nranks_to, requests, &sends_complete, MPI_STATUSES_IGNORE))


Suggested change

CHKERRMPI(MPI_Testall(nranks_to, requests, &sends_complete, MPI_STATUSES_IGNORE))

CHKERRMPI(MPI_Testall(nranks_to, send_requests, &sends_complete, MPI_STATUSES_IGNORE))

connorjward

I am still quite confused by this PR. There is a lot going on.

connorjward · 2026-09-02T17:03:13Z

        self._parent_mesh = parentmesh

-        super().__init__(swarm, name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm)
+        super().__init__(swarm.dm, name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm)


This now means that we have to keep track of the Firedrake swarm object as well as the PETSc one. I would prefer it if they were the same, either by keeping the current inheritance pattern or by adding PETSc.DMSwarm methods to FiredrakeDMSwarm. I think arguably the latter makes more sense because then (e.g.) swarm.clone() will give you back something with the right type.

Therefore I propose:

Suggested change

super().__init__(swarm.dm, name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm)

super().__init__(swarm, name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm)

I think inheritance will be easier. AbstractMeshTopology calls a bunch of cython routines which expect a PETSc object.

Alternatively make it a DMShell?

That seems like an awful lot of effort for little gain

I don't mind what approach you take. I just really think that swarm and topology_dm should be the same object. I would eventually like to drop VertexOnlyMesh and fold everything into a single Mesh class. We shouldn't be adding additional differences between the classes.

Made FiredrakeDMSwarm inherit from PETSc.DMSwarm.

connorjward · 2026-09-02T17:03:46Z

        mesh cell order.
        """
-        with self.topology_dm.field("parentcellnum") as parentcellnum_field:
+        with self.swarm.field("parentcellnum") as parentcellnum_field:


And if you did my suggestion then all of these changes would not be needed

Undone these now

connorjward · 2026-09-02T17:08:40Z

+        """Initialize a Firedrake DMSwarm.
+
+        Parameters
+        ----------
+        dm : PETSc.DMSwarm
+            The underlying PETSc DMSwarm.
+        extruded : bool
+            Whether the swarm is embedded in an extruded mesh.
+        """


Put this in the main class docstring (that's what numpydoc suggests). It also means you don't have to write "Initialize a Firedrake DMSwarm.", which is trivially obvious.

connorjward · 2026-09-02T17:10:02Z

+        if not isinstance(dm, PETSc.DMSwarm):
+            raise TypeError(f"`dm` must be a `PETSc.DMSwarm`, not a {type(dm).__name__}")
+        self.dm = dm
+        self.extruded = extruded


should this be inferred from whether parentcellbasenum exists?

But we create parentcellbasenum based on whether extruded is true or false...

Is that true? In https://github.com/firedrakeproject/firedrake/pull/5396/changes#diff-7b656fa2d868bdf9f07513bd2565b7b30f1896c30352e3d8b92ae435ec3befd2R4115 it looks like you create the field independently.

As an alternative phrasing, do you ever actually inspect FiredrakeDMSwarm.extruded? A quick search makes me think that you don't.

Oh right that is extruded and not self.extruded...

connorjward · 2026-09-02T17:12:35Z

+            return rtree.build_from_aabb(coords_min, coords_max)
+
+    @PETSc.Log.EventDecorator()
+    def bounding_boxes_total_volume(self, bounding_boxes: np.ndarray):


make a private attribute? can't imagine users wanting this

connorjward · 2026-09-02T17:13:35Z

-        return self._rtree
+            return rtree.build_from_aabb(coords_min, coords_max)
+
+    @PETSc.Log.EventDecorator()


Suggested change

@PETSc.Log.EventDecorator()

looks trivial, doesn't need to be profiled

connorjward · 2026-09-02T17:14:41Z

+                owner_ranks_data,
+                npoints,
+                cells_ignore.shape[1],
+                cells_ignore,


I thought this was going?

The locate_cells_.. methods are all public API that need this functionality

I can't imagine anyone using cells_ignore in their own code, but fine. Maybe raise a deprecation warning?

I will open another PR to do this.

OK. Once you open that PR (or create an issue) then this can be resolved

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

connorjward

Right, read the whole thing now. This is really heroic.

connorjward · 2026-09-03T14:39:20Z

+            send_requests = <MPI_Request *>malloc(nranks_to * sizeof(MPI_Request))
+
+        # Post non-blocking synchronous sends. 'Synchronous' means that the send is not
+        # considered as complete by MPI_Test until the destination starts the matching receive.


Why is this important?

Added a comment

connorjward · 2026-09-03T14:44:08Z

        else {
-            current_ref_cell_dist_l1 = (*try_candidate)(temp_ref_coords, f, candidate, x);
+            distance = (*try_candidate)(temp_ref_coords, f, candidate, x);
+            owner = cell_owner_ranks ? cell_owner_ranks[candidate] : 0;


Is this a backwards compatibility thing? (cell_owner_ranks may be undefined?) If so a comment wouldn't go amiss.

Yes. Added a comment

connorjward · 2026-09-03T14:50:05Z

-            /* Found cell! */
-            *cell_out = candidate;
+        /* Select owning cell by minimum L1 distance, with ties broken by highest owning rank. */
+        if (distance < best_distance || (distance == best_distance && owner > best_owner)) {


I still worry about this. Consider an interval mesh like

global: x-----x-----x o rank 0: x-----x o rank 1: x-----x o

(which is allowed since we don't always need a full cell of overlap)

If o is the point we want to locate then rank 0 and rank 1 both have one entry in cell_owner_ranks - they will each think that they own the point.

A rank doesn't decide at this point that they own the point. locate.c just determines whether the point lies in one of their owned cells (and L^1 distances and ref coords, etc.).

The 'global' selection in _parent_mesh_embedding is what determines the owning rank. In this case the points both have a zero L^1 distance, so the owner is going to be rank 1 since we break ties by choosing the highest owning rank.

Will it work if both ranks think that the other rank is the owner?

If I'm reading your example correctly, rank 0 can see no cells of rank 1 and vice-versa. So I can't see how a rank could think the point lived in a cell owned by the other rank in this case.

If we do have halos then this can happen, but by selecting the single highest owning rank we can never end up that situation.

It doesn't apply to the example I gave earlier. Only if there are halos.

I guess my concern comes down to my mental model of floating point in parallel. Can we be sure that two ranks will compute the distance of a point on the cell boundary as exactly zero under any circumstance?

connorjward · 2026-09-03T14:59:22Z

+            ):
+                parent_order = parent_renum_inv[swarm_parent_cell_nums.ravel() - pStart]
+                # sort by parent cell order, with ties broken by point global index
+                perm = np.lexsort((swarm_global_indices.ravel(), parent_order)).astype(IntType)


I'm surprised that this needs any global information at all. This numbering is normally done with only local information.

connorjward · 2026-09-03T15:02:22Z

+                owner_ranks_data,
+                npoints,
+                cells_ignore.shape[1],
+                cells_ignore,


OK. Once you open that PR (or create an issue) then this can be resolved

connorjward · 2026-09-03T15:04:37Z

+        root_values: np.ndarray,
+        leaf_values: np.ndarray,
+    ) -> None:
+        # TODO: make these collective


I worry that that would be quite expensive

Maybe we could add a collective: bool argument which is False by default? I found these checks quite useful for debugging purposes.

connorjward · 2026-09-03T15:09:37Z

+    swarm.set_field("inputindex", embedded_sf.input_indices[swarm_indices].astype(IntType))
+    if parent_mesh.extruded:
+        swarm.set_field("parentcellbasenum", swarm_base_cells)
+        swarm.set_field("parentcellextrusionheight", swarm_extrusion_heights)


Why are some fields set in create_with_fields but not all?

create_with_fields registers our fields but doesn't set them.

Is that the best API?

connorjward · 2026-09-03T15:12:18Z

+    keep &= ref_cell_dists == candidate_sf.broadcast(root_distance_min)
+
+    # multiple ranks may claim the minimum L1 distance. Break ties
+    # by choosing the highest numbered rank.


If you're doing this then I wonder if the parallel issue I flagged above is applicable. Could owned_cell_ranks just go? The C code could just see if a rank claims to own the point or not.

I think you only need rank information here, not there as well.

I think we need to have rank information inside locate.c as long as we support exclude_halos=False (which is what VertexOnlyMesh uses).

In the future I want to get rid of this behaviour (i.e. always exclude halos from the final VOM) because a) the idea of a halo cell in a vom doesn't really make sense to me and we don't use this anywhere, and b) it substantially simplifies the vom construction code.

Yeah I have always felt that halo points in a VoM was the wrong thing to do.

You need halos in order to build the correct stencils, but VoMs have no connectivity and so it's really pointless.

Happy to punt on this for now, I just always think it's worth discussing the design of these things.

fix typo Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

leo-collins added 30 commits July 7, 2026 13:40

use rstar capi for cell location

1accc33

add logging

6b5eea5

remove allgather logs

06da607

update comments

4709136

fix cached property

0ce05f9

merge logging

793a92d

more logging

more logs

ff062ff

more logging

1dd92d4

adapt to new API

2fc20a6

rebase

ea48fb6

partial fix

f413029

tidy up cython

d792583

fixes

923db26

renaming and use firedrake_rtree package

cb25eb6

firedrake-rtree fixes

7418224

remove unused stuff from rtree cython

a66d16d

remove unnecessary stuff from petschdr

ac08113

misc tidying

44863a6

one more

fixes

5a60847

fixup mesh.py

df6df0b

add rtree libspatialindex back in

789b486

add spatialindex to supermesh extension

088e8c7

add extra runtime path

a0be450

lowercase r

431daec

fix

aae03f4

libsupermesh fix

6719b40

remove library

4cfd969

remove rtree import from setup.py

17cdaa4

use libsupermesh branch for docs

db7972b

installation and workflow

bf7037e

do exact floating point equality

a6fc0a7

doing this fixes test errors...

leo-collins added ci:complex Run the test suite in complex mode ci:macos Run the test suite on the Mac runners ci:int64 Do a 64 bit integer build labels Aug 27, 2026

leo-collins added 3 commits August 27, 2026 17:40

rename test

425390b

Merge branch 'main' into leo/distributed-rtree-refactor

39b7b1b

int64 fix

3d7c6f6

leo-collins marked this pull request as ready for review August 29, 2026 17:55

Merge branch 'main' into leo/distributed-rtree-refactor

c321012

connorjward requested changes Sep 1, 2026

View reviewed changes

leo-collins added 3 commits September 2, 2026 10:13

add comments to parallel algorithm

e7cac8c

duplicate comm inside of collective function

d99abf5

lint

e209f8f

leo-collins requested a review from connorjward September 2, 2026 11:23

connorjward requested changes Sep 2, 2026

View reviewed changes

leo-collins and others added 7 commits September 3, 2026 12:05

move docstring

b1de9fd

inherit

c211427

add docstring

5570933

tidy up VertexOnlyMeshSF

875f4c3

final review suggestions

9e3407e

tidy

d39ac1e

Apply suggestion from @connorjward

2ad2709

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

leo-collins requested review from JHopeCollins and connorjward September 3, 2026 12:55

remove init

4ed302b

connorjward requested changes Sep 3, 2026

View reviewed changes

leo-collins and others added 4 commits September 7, 2026 16:04

add comment

9f10b94

add comment about NULL values in locate.c

fde4af2

Update firedrake/cython/rtree.pyx

62253a1

fix typo Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

Merge branch 'main' into leo/distributed-rtree-refactor

db515cd

	CHKERRMPI(MPI_Testall(nranks_to, requests, &sends_complete, MPI_STATUSES_IGNORE))
	CHKERRMPI(MPI_Testall(nranks_to, send_requests, &sends_complete, MPI_STATUSES_IGNORE))

	super().__init__(swarm.dm, name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm)
	super().__init__(swarm, name, reorder, None, perm_is, distribution_name, permutation_name, parentmesh.comm)

Uh oh!

Conversation

leo-collins commented Aug 27, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

connorjward left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

connorjward left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment