Scalable point location

firedrake/cython/rtree.pyx

@@ -1,11 +1,35 @@
 # cython: language_level=3
 
+from firedrake.utils import IntType
 cimport numpy as np
 import numpy as np
 import ctypes
 import cython
 from libc.stddef cimport size_t
 from libc.stdint cimport uintptr_t, uint32_t, int64_t
+from libc.stdlib cimport free, malloc
+
+cimport mpi4py.MPI as MPI
+from mpi4py.libmpi cimport (
+    MPI_ANY_SOURCE,
+    MPI_REQUEST_NULL,
+    MPI_STATUS_IGNORE,
+    MPI_STATUSES_IGNORE,
+    MPI_TYPECLASS_INTEGER,
+    MPI_Datatype,
+    MPI_Get_count,
+    MPI_Ibarrier,
+    MPI_Iprobe,
+    MPI_Issend,
+    MPI_Recv,
+    MPI_Request,
+    MPI_Status,
+    MPI_Test,
+    MPI_Testall,
+    MPI_Type_match_size,
+)
+
+include "petschdr.pxi"
 
 cdef extern from "rtree-capi.h":
     ctypedef enum RTreeError:
@@ -27,6 +51,38 @@ cdef extern from "rtree-capi.h":
 
     RTreeError rtree_free(RTreeH *tree)
 
+    RTreeError rtree_free_ids(int64_t *ids, size_t n)
+
+    RTreeError rtree_free_offsets(size_t *offsets, size_t n)
+
+    RTreeError rtree_free_point_indices(size_t *point_indices, size_t n)
+
+    RTreeError rtree_locate_points_grouped_by_id_unique(
+        const RTreeH *tree,
+        const double *points,
+        size_t n_points,
+        int64_t **ids_out,
+        size_t **offsets_out,
+        size_t **point_indices_out,
+        size_t *n_ids_out
+    )
+
+    RTreeError rtree_depth(const RTreeH *tree, size_t *depth_out)
+
+    RTreeError rtree_collect_bounding_boxes(
+        const RTreeH *tree,
+        size_t level,
+        double **mins_out,
+        double **maxs_out,
+        size_t *nboxes_out
+    )
+
+    RTreeError rtree_free_bounding_boxes(
+        double *mins,
+        double *maxs,
+        size_t nboxes,
+        uint32_t dim
+    )
 
 cdef class RTree(object):
     """Python class for holding an Rtree."""
@@ -101,3 +157,175 @@ def build_from_aabb(np.ndarray[np.float64_t, ndim=2, mode="c"] coords_min,
         raise RuntimeError("rtree_bulk_load failed")
 
     return RTree(<uintptr_t>rtree)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def discover_remote_roots(
+        RTree rtree,
+        np.ndarray[np.float64_t, ndim=2, mode="c"] points,
+        MPI.Comm comm):
+    """Build the remote array for a candidate star forest. Collective.
+
+    This implements the non-blocking consensus algorithm from Hoefler et al.
+    'Scalable communication protocols for dynamic sparse data exchange'.
+
+    This routine duplicates the given comm since we send tag-0 messages
+    and listen with MPI_IProbe(MPI_ANY_SOURCE).
+
+    Parameters
+    ----------
+    rtree : RTree
+        The distributed Rtree built by :func:`build_from_aabb` with rank
+        numbers as leaf ids.
+    points : (n_points, gdim) float64 array
+        Local point coordinates.
+    comm : mpi4py.MPI.Comm
+        The MPI communicator.
+
+    Returns
+    -------
+    remote : (nleaves, 2) IntType array
+        For every local candidate leaf, the MPI rank and local index of its
+        remote root point.
+    """
+    cdef:
+        MPI.Comm dup_comm = comm.Dup()
+        MPI.MPI_Comm mpi_comm = dup_comm.ob_mpi
+        MPI_Request *send_requests = NULL
+        MPI_Request barrier_request = MPI_REQUEST_NULL
+        MPI_Status status
+        MPI_Datatype point_index_type
+        int count, source_rank, message_ready, sends_complete = 0
+        int barrier_complete = 0
+        Py_ssize_t nleaves = 0, recv_offset = 0
+        int64_t *send_ranks = NULL
+        size_t *point_indices = NULL
+        size_t *send_offsets = NULL
+        size_t n_points = points.shape[0], nranks_to = 0
+        np.ndarray[np.uintp_t, ndim=1, mode="c"] received
+        np.ndarray[PetscInt, ndim=2, mode="c"] remote
+        list recv_messages = []
+
+    try:
+        # MPI does not have a builtin type for size_t
+        CHKERRMPI(MPI_Type_match_size(MPI_TYPECLASS_INTEGER, sizeof(size_t), &point_index_type))
+
+        if rtree_locate_points_grouped_by_id_unique(
+            rtree.tree,
+            <const double *>points.data,
+            n_points,
+            &send_ranks,  # unique ranks I am sending to
+            &send_offsets,  # offsets into `point_indices` of points I need to send
+            &point_indices,  # the points I need to send
+            &nranks_to,  # number of ranks I am sending to (len(send_ranks))
+        ) != Success:
+            raise RuntimeError("rtree_locate_points_grouped_by_id_unique failed")
+
+        # Each rank knows where to send to, but not where it is going to receive from.
+        # The following algorithm is the 'non-blocking consensus' algorithm (Hoefler et al.)
+        # which is a sparse alternative to MPI_Alltoallv
+        if nranks_to != 0:
+            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.
+        # This is so that a rank cannot hit the barrier while any of its outgoing messages remain unmatched.
+        for k in range(nranks_to):
+            count = send_offsets[k + 1] - send_offsets[k]
+            CHKERRMPI(MPI_Issend(
+                <const void *>&point_indices[send_offsets[k]],
+                count, point_index_type,
+                <int>send_ranks[k],
+                0,
+                mpi_comm,
+                &send_requests[k],
+            ))
+
+        # MPI_Ibarrier is a non-blocking barrier. It posts into `barrier_request`. This request is
+        # not considered complete until all ranks on the comm post the barrier request.
+        # This is used to know when a rank has sent all of its messages.
+        # In the case we are sending no messages, we do this now.
+        if nranks_to == 0:
+            CHKERRMPI(MPI_Ibarrier(mpi_comm, &barrier_request))
+
+        # Enter a loop where we continuously probe for incoming messages.
+        # `barrier_complete` is true when all ranks have hit the barrier.
+        while not barrier_complete:
+            # We don't know where we are going to receive from, so we use MPI_ANY_SOURCE. If we're
+            # about to receive a message, then `message_ready` is true.
+            CHKERRMPI(MPI_Iprobe(MPI_ANY_SOURCE, 0, mpi_comm, &message_ready, &status))
+
+            # We're receiving a message: determine the number of entries, and do the blocking receive.
+            if message_ready:
+                CHKERRMPI(MPI_Get_count(&status, point_index_type, &count))
+                source_rank = status.MPI_SOURCE
+                received = np.empty(count, dtype=np.uintp)
+                CHKERRMPI(MPI_Recv(<void *>received.data, count, point_index_type, source_rank, 0, mpi_comm, MPI_STATUS_IGNORE))
+                recv_messages.append((source_rank, received))
+                nleaves += count
+
+            if barrier_request == MPI_REQUEST_NULL:
+                # This rank has not hit the barrier. Test if all the sends have been complete. If they
+                # have, then this ranks hits the barrier to mark that its sends have completed.
+                CHKERRMPI(MPI_Testall(nranks_to, send_requests, &sends_complete, MPI_STATUSES_IGNORE))
+                if sends_complete:
+                    CHKERRMPI(MPI_Ibarrier(mpi_comm, &barrier_request))
+            else:
+                # This rank has hit the barrier and has posted `barrier_request`. This request completes
+                # when _all_ other ranks send barrier_request. Then MPI_Test will set `barrier_complete`
+                # to true which ends the loop.
+                CHKERRMPI(MPI_Test(&barrier_request, &barrier_complete, MPI_STATUS_IGNORE))
+    finally:
+        dup_comm.Free()
+        if send_requests != NULL:
+            free(send_requests)
+        if rtree_free_ids(send_ranks, nranks_to) != Success:
+            raise RuntimeError("rtree_free_ids failed")
+        if rtree_free_point_indices(point_indices, send_offsets[nranks_to]) != Success:
+            raise RuntimeError("rtree_free_point_indices failed")
+        if rtree_free_offsets(send_offsets, nranks_to + 1) != Success:
+            raise RuntimeError("rtree_free_offsets failed")
+
+    # create remote array for candidate star forest
+    remote = np.empty((nleaves, 2), dtype=IntType)
+    for source_rank, received in recv_messages:
+        count = received.shape[0]
+        remote[recv_offset:recv_offset + count, 0] = source_rank
+        remote[recv_offset:recv_offset + count, 1] = received
+        recv_offset += count
+
+    return remote
+
+
+def bounding_boxes_at_level(RTree rtree, size_t level, uint32_t dim):
+    """Return all bounding boxes at the specified level of the Rtree."""
+    cdef:
+        double *mins = NULL
+        double *maxs = NULL
+        size_t n_boxes = 0
+        np.ndarray[np.float64_t, ndim=3, mode="c"] boxes
+
+    if rtree_collect_bounding_boxes(rtree.tree, level, &mins, &maxs, &n_boxes) != Success:
+        raise RuntimeError("rtree_bounding_boxes failed")
+
+    boxes = np.empty((n_boxes, 2, dim), dtype=np.float64)
+
+    for i in range(n_boxes):
+        for j in range(dim):
+            boxes[i, 0, j] = mins[i * dim + j]
+            boxes[i, 1, j] = maxs[i * dim + j]
+
+    rtree_free_bounding_boxes(mins, maxs, n_boxes, dim)
+
+    return boxes
+
+
+def tree_depth(RTree rtree):
+    """Return the depth of the Rtree."""
+    cdef:
+        size_t depth = 0
+
+    if rtree_depth(rtree.tree, &depth) != Success:
+        raise RuntimeError("rtree_depth failed")
+    return depth

firedrake/evaluate.h

@@ -52,17 +52,19 @@ typedef PetscReal (*ref_cell_l1_dist_xtr)(void *data_,
 				double *x);
 
 extern PetscErrorCode locate_cell_from_candidates(struct Function *f,
-		       double *x,
-		       ref_cell_l1_dist try_candidate,
-		       ref_cell_l1_dist_xtr try_candidate_xtr,
-		       void *temp_ref_coords,
-		       void *found_ref_coords,
-		       PetscReal *found_ref_cell_dist_l1,
-			   size_t nids,
-			   const int64_t *ids,
-			   size_t ncells_ignore,
-			   const PetscInt *cells_ignore,
-			   PetscInt *cell_out);
+                                                  double *x,
+                                                  ref_cell_l1_dist try_candidate,
+                                                  ref_cell_l1_dist_xtr try_candidate_xtr,
+                                                  void *temp_ref_coords,
+                                                  void *found_ref_coords,
+                                                  PetscReal *found_ref_cell_dist_l1,
+                                                  size_t nids,
+                                                  const int64_t *ids,
+                                                  size_t ncells_ignore,
+                                                  const PetscInt *cells_ignore,
+                                                  const PetscInt *cell_owner_ranks,
+                                                  PetscInt *cell_out,
+                                                  PetscInt *owner_out);
 
 extern int evaluate(struct Function *f,
 		    double *x,

firedrake/locate.c

@@ -11,7 +11,9 @@ PetscErrorCode locate_cell_from_candidates(struct Function *f,
         const int64_t *ids,
         size_t ncells_ignore,
         const PetscInt *cells_ignore,
-        PetscInt *cell_out)
+        const PetscInt *cell_owner_ranks,
+        PetscInt *cell_out,
+        PetscInt *owner_out)
 {
     bool cell_ignore_found = false;
     /* NOTE: temp_ref_coords and found_ref_coords are actually of type
@@ -21,13 +23,20 @@ PetscErrorCode locate_cell_from_candidates(struct Function *f,
     surrounds this is declared in pointquery_utils.py. We cast when we use the
     ref_coords_copy function and trust that the underlying memory which the
     pointers refer to is updated as necessary. */
-    PetscReal ref_cell_dist_l1 = PETSC_MAX_REAL;
-    PetscReal current_ref_cell_dist_l1 =  -0.5;
+    PetscReal best_distance = PETSC_MAX_REAL;
+    PetscInt best_owner = -1;
+    PetscInt best_cell = -1;
     /* NOTE: `tolerance`, which is used throughout this function, is a static
        variable defined outside this function when putting together all the C
        code that needs to be compiled - see pointquery_utils.py */
 
+    /* NOTE: `cell_owner_ranks` and `owner_out` may be NULL. In this case, every
+       candidate is treated as having owner rank 0. This is for backwards compatibility
+       with the `Function.at` code path which is being deprecated soon. */
     *cell_out = -1;
+    if (owner_out) {
+        *owner_out = -1;
+    }
     for (size_t i = 0; i < nids; ++i) {
         /* Check that casting the ids from int64 to PetscInt is safe (for 32 bit petsc builds). 
         Since the ids are mesh cell ids this *should* always be safe, but better to check
@@ -47,33 +56,34 @@ PetscErrorCode locate_cell_from_candidates(struct Function *f,
             continue;
         }
 
+        PetscReal distance;
+        PetscInt owner;
         if (f->extruded) {
             PetscInt nlayers = f->n_layers;
             PetscInt c = candidate / nlayers;
             PetscInt l = candidate % nlayers;
-            current_ref_cell_dist_l1 = (*try_candidate_xtr)(temp_ref_coords, f, c, l, x);
+            distance = (*try_candidate_xtr)(temp_ref_coords, f, c, l, x);
+            owner = cell_owner_ranks ? cell_owner_ranks[c] : 0;
         }
         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;
         }
-
-        if (current_ref_cell_dist_l1 <= 0.0) {
-            /* 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)) {
+            best_distance = distance;
+            best_owner = owner;
+            best_cell = candidate;
             memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords));
-            found_ref_cell_dist_l1[0] = current_ref_cell_dist_l1;
-            break;
         }
-        else if (current_ref_cell_dist_l1 < ref_cell_dist_l1) {
-            /* getting closer... */
-            ref_cell_dist_l1 = current_ref_cell_dist_l1;
-            if (ref_cell_dist_l1 < tolerance) {
-                /* Close to cell within tolerance so could be this cell */
-                *cell_out = candidate;
-                memcpy(found_ref_coords, temp_ref_coords, sizeof(struct ReferenceCoords));
-                found_ref_cell_dist_l1[0] = ref_cell_dist_l1;
-            }
+    }
+
+    if (best_cell != -1 && (best_distance <= 0.0 || best_distance < tolerance)) {
+        *cell_out = best_cell;
+        if (owner_out) {
+            *owner_out = best_owner;
         }
+        *found_ref_cell_dist_l1 = best_distance;
     }
     return PETSC_SUCCESS;
 }