diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 86fb29aae..1c55d166c 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -76,6 +76,25 @@ def __setitem__(self, key, value): raise ValueError(self._GUIDANCE) +def _coordinate_row_keys(coords): + """One opaque key per coordinate row, for exact set membership. + + Used by :meth:`Swarm.migrate` to remember which points a rank has already + claimed, so the round loop does not reclassify them. The identity has to be + the coordinate rather than the index: ``dm.migrate`` does not preserve the + local ordering, so an index from the previous round refers to a different + particle after the move. + + The comparison is on the exact bit pattern, so this matches only points + whose coordinates are identical — which is what migration produces, since + it moves values without touching them. + """ + + a = np.ascontiguousarray(coords, dtype=np.float64) + + return a.view(np.dtype((np.void, a.dtype.itemsize * a.shape[1]))).reshape(-1) + + class SwarmType(Enum): """ PETSc swarm type specification. @@ -3728,6 +3747,11 @@ def migrate( swarm_coord_array, ) + # The working set for the round loop below. Seeded with what this first + # pass claimed; it only ever grows, because a point in this rank's + # domain stays in it. + claimed_keys = _coordinate_row_keys(swarm_coord_array[in_or_not]) + num_points_in_domain = np.count_nonzero(in_or_not == True) num_points_not_in_domain = np.count_nonzero(in_or_not == False) not_my_points = np.where(in_or_not == False)[0] @@ -3787,7 +3811,43 @@ def migrate( uw.mpi.barrier() swarm_coord_array = self.dm.getField("DMSwarmPIC_coor").reshape(-1, self.cdim) - in_or_not = self.mesh.points_in_domain(swarm_coord_array) + + # Only classify what we have not already claimed. A point this + # rank has found to be in its domain stays in its domain: the + # coordinates do not change during migration and neither does + # the mesh. Re-testing it on every round is what made the + # per-rank classification work grow with the round count (a + # fixed 47k-point global set was offered to points_in_domain + # 238k times at np=4). + # + # The bookkeeping is by COORDINATE rather than by index because + # `dm.migrate` does not preserve the local ordering — measured, + # retained points are not left at the front — so an index from + # the previous round means nothing after the move. Two + # particles sharing a coordinate share the answer, so a + # collision is harmless. + keys_now = _coordinate_row_keys(swarm_coord_array) + already = np.isin(keys_now, claimed_keys) + + # UNCONDITIONAL: points_in_domain is COLLECTIVE — it reaches + # get_max_radius() before any short-circuit, precisely so that a + # rank with nothing to classify still joins the reduction (the + # #405 treatment, stated in its own source). Calling it only + # when this rank has undecided points deadlocks as soon as one + # rank runs out of them, which is what an earlier draft of this + # did at np=4. + in_or_not = already.copy() + undecided = np.where(~already)[0] + in_or_not[undecided] = self.mesh.points_in_domain( + swarm_coord_array[undecided] + ) + + newly_claimed = np.where(in_or_not & ~already)[0] + if newly_claimed.size: + claimed_keys = np.concatenate( + [claimed_keys, keys_now[newly_claimed]] + ) + self.dm.restoreField("DMSwarmPIC_coor") num_points_in_domain = np.count_nonzero(in_or_not == True) diff --git a/tests/parallel/test_0777_migrate_working_set_mpi.py b/tests/parallel/test_0777_migrate_working_set_mpi.py new file mode 100644 index 000000000..a6992345a --- /dev/null +++ b/tests/parallel/test_0777_migrate_working_set_mpi.py @@ -0,0 +1,136 @@ +"""Swarm.migrate's round loop: correctness, and the collective it must not skip. + +The loop classifies points once per round. It used to re-offer the whole local +array to `points_in_domain` every round, so the classification work per rank grew +with the round count — a fixed 47k-point set was offered 238k times at np=4. It +now remembers what it has already claimed and only classifies the rest. + +The claimed set is keyed by COORDINATE, not by index: `dm.migrate` does not +preserve the local ordering (measured — retained points are not left at the +front), so an index from the previous round names a different particle after the +move. + +The second test is here because the first draft of that change deadlocked. It +made the `points_in_domain` call conditional on this rank having undecided +points, and `points_in_domain` is collective — it reaches `get_max_radius()` +before any short-circuit precisely so a rank with nothing to classify still +joins the reduction (the #405 treatment, stated in its own source). At np=4 a +rank ran out of undecided points, skipped the reduction, and the job hung with +every rank inside the call. + +Run under MPI:: + + mpirun -np 4 python -m pytest --with-mpi \ + tests/parallel/test_0777_migrate_working_set_mpi.py +""" + +import numpy as np +import pytest + +import underworld3 as uw + +pytestmark = [ + pytest.mark.mpi(min_size=2), + pytest.mark.timeout(180), + pytest.mark.level_1, + pytest.mark.tier_a, +] + +SEED = 20260819 + + +def _scattered_swarm(cell_size=0.15, n_points=2000): + """A swarm whose points are spread over the WHOLE domain. + + Populate alone leaves every particle already owned by its own rank, so + nothing migrates and the round loop never runs — the case that makes this + file look like it is testing something when it is not. + """ + + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=cell_size, + regular=False, qdegree=2) + swarm = uw.swarm.Swarm(mesh=mesh) + swarm.populate(fill_param=3) + + rng = np.random.default_rng(SEED) + everywhere = rng.uniform(0.01, 0.99, size=(n_points, 2)) + mine = np.array_split(everywhere, uw.mpi.size)[uw.mpi.rank] + + coords = swarm.dm.getField("DMSwarmPIC_coor").reshape(-1, mesh.dim) + n = min(coords.shape[0], mine.shape[0]) + coords[:n] = mine[:n] + swarm.dm.restoreField("DMSwarmPIC_coor") + + return mesh, swarm + + +def _local_coords(swarm, dim): + out = swarm.dm.getField("DMSwarmPIC_coor").reshape(-1, dim).copy() + swarm.dm.restoreField("DMSwarmPIC_coor") + + return out + + +def test_premise_the_fixture_actually_migrates(): + """Without points crossing ranks the round loop never runs.""" + + mesh, swarm = _scattered_swarm() + before = _local_coords(swarm, mesh.dim) + owned_before = int(np.count_nonzero(mesh.points_in_domain(before))) + strangers = uw.mpi.comm.allreduce(before.shape[0] - owned_before) + + assert strangers > 0, ( + "every point is already owned by the rank holding it, so migration has " + "nothing to do and neither of the tests below exercises the loop") + + +def test_migrate_returns_and_every_point_lands_on_a_rank_that_holds_it(): + """The loop terminates, and its answer is right. + + The deadlock this guards against leaves every rank inside `points_in_domain`, + so the failure is a timeout rather than an assertion. + """ + + mesh, swarm = _scattered_swarm() + before = _local_coords(swarm, mesh.dim) + n_before = uw.mpi.comm.allreduce(before.shape[0]) + + swarm.migrate() + + after = _local_coords(swarm, mesh.dim) + n_after = uw.mpi.comm.allreduce(after.shape[0]) + + assert n_after == n_before, ( + f"migration changed the global particle count: {n_before} -> {n_after}") + + if after.shape[0]: + owned = mesh.points_in_domain(after) + assert bool(owned.all()), ( + f"rank {uw.mpi.rank} holds {int(np.count_nonzero(~owned))} points " + "its own mesh does not contain") + else: + mesh.points_in_domain(after) # collective: join the reduction + + +def test_the_partition_is_the_same_at_every_rank_count(): + """The global set of owned coordinates is the input set, whatever np is. + + Keyed on the global multiset rather than per-rank counts, which are a + partition detail. This is what a change to the claimed-set bookkeeping would + break: dropping a point, or claiming one twice. + """ + + mesh, swarm = _scattered_swarm() + before = _local_coords(swarm, mesh.dim) + swarm.migrate() + after = _local_coords(swarm, mesh.dim) + + def gathered_rows(a): + rows = uw.mpi.comm.allgather(np.ascontiguousarray(a)) + stacked = np.concatenate([r for r in rows if r.size], axis=0) + return stacked[np.lexsort((stacked[:, 1], stacked[:, 0]))] + + assert np.array_equal(gathered_rows(before), gathered_rows(after)), ( + "the multiset of particle coordinates changed across migration — a " + "point was lost or duplicated")