From 781168644c71aa5bbd29787bae19952d6addad62 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 19 Aug 2026 15:22:34 +1000 Subject: [PATCH 1/2] Swarm.migrate: shrink the working set instead of reclassifying every round MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The round loop re-offered the whole local array to points_in_domain on every round, so the classification work per rank grew with the round count. Measured on a fixed 47k-point global set through global_evaluate, points offered to points_in_domain: 47k at np=1, 142k at np=2, 238k at np=4. A point this rank has found to be in its domain stays in its domain — neither the coordinates nor the mesh change during migration — so it does not need retesting. The loop now remembers what it has claimed and classifies only the rest. The claimed set is keyed by COORDINATE rather than 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. Two particles sharing a coordinate share the answer, so a key collision is harmless. Exact bit-pattern comparison, vectorised through a void row view (15 ms against points_in_domain's 58 ms at 47k, so it pays whenever more than about a quarter of the local array is already claimed). Measured after, same probe: np global_evaluate points offered to points_in_domain 1 0.208 -> 0.213 s 46 901 -> 46 901 2 0.326 -> 0.232 s 141 516 -> 47 558 4 0.443 -> 0.257 s 237 826 -> 50 025 42% off global_evaluate at np=4, and the growth with rank count is gone: the classification now sees roughly the local set once whatever np is. The nearest-centroid walk shrinks with it, since it is called from inside the classification. The call to points_in_domain is UNCONDITIONAL. It is collective — it reaches get_max_radius() before any short-circuit precisely so a rank with nothing to classify still joins the reduction (#405) — and an earlier draft of this skipped it when a rank had no undecided points, which deadlocked at np=4 with every rank inside the call. Verified behaviour-preserving rather than merely green: the per-rank partition fingerprint after migration is byte-identical to development at np=2 and np=4. Full ./uw test 1556 passed, matching development. test_0776 covers the loop, with a premise test that the fixture actually migrates anything, and its own negative control: reintroducing the conditional collective makes it time out at np=4 (rc=241), and the fix makes it pass. Underworld development team with AI support from Claude Code --- src/underworld3/swarm.py | 62 +++++++- .../test_0776_migrate_working_set_mpi.py | 136 ++++++++++++++++++ 2 files changed, 197 insertions(+), 1 deletion(-) create mode 100644 tests/parallel/test_0776_migrate_working_set_mpi.py 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_0776_migrate_working_set_mpi.py b/tests/parallel/test_0776_migrate_working_set_mpi.py new file mode 100644 index 000000000..3ebf45fa0 --- /dev/null +++ b/tests/parallel/test_0776_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_0776_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") From e29d6eac8ab3ff176d02792716ef170e7586f6e2 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 19 Aug 2026 16:32:35 +1000 Subject: [PATCH 2/2] Renumber to 0777: 0776 was already taken MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tests/parallel/test_0776_linear_rbf_proxy_parallel.py already exists, so the new file collided with it — the same defect as #600, where two subjects ended up sharing 1029. Caught while auditing which parallel files scripts/test.sh actually runs. Underworld development team with AI support from Claude Code --- ..._working_set_mpi.py => test_0777_migrate_working_set_mpi.py} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename tests/parallel/{test_0776_migrate_working_set_mpi.py => test_0777_migrate_working_set_mpi.py} (98%) diff --git a/tests/parallel/test_0776_migrate_working_set_mpi.py b/tests/parallel/test_0777_migrate_working_set_mpi.py similarity index 98% rename from tests/parallel/test_0776_migrate_working_set_mpi.py rename to tests/parallel/test_0777_migrate_working_set_mpi.py index 3ebf45fa0..a6992345a 100644 --- a/tests/parallel/test_0776_migrate_working_set_mpi.py +++ b/tests/parallel/test_0777_migrate_working_set_mpi.py @@ -21,7 +21,7 @@ Run under MPI:: mpirun -np 4 python -m pytest --with-mpi \ - tests/parallel/test_0776_migrate_working_set_mpi.py + tests/parallel/test_0777_migrate_working_set_mpi.py """ import numpy as np