compiler: Revamp MPI hoisting and merging #2629

FabioLuporini · 2025-06-11T15:30:05Z

besides making things generally more efficient, by generating code that sometimes avoids redundant halo exchanges, we are here achieving:

fixes #2622

codecov · 2025-06-11T15:47:38Z

Codecov Report

Attention: Patch coverage is 95.17241% with 14 lines in your changes missing coverage. Please review.

Project coverage is 92.01%. Comparing base (2a0076f) to head (b3a3e00).
Report is 9 commits behind head on main.

Files with missing lines	Patch %	Lines
devito/mpi/halo_scheme.py	69.56%	4 Missing and 3 partials ⚠️
devito/passes/iet/mpi.py	96.22%	2 Missing and 2 partials ⚠️
tests/test_mpi.py	97.47%	1 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2629      +/-   ##
==========================================
+ Coverage   91.99%   92.01%   +0.01%     
==========================================
  Files         245      245              
  Lines       48495    48710     +215     
  Branches     4261     4292      +31     
==========================================
+ Hits        44612    44819     +207     
- Misses       3202     3205       +3     
- Partials      681      686       +5

Flag	Coverage Δ
pytest-gpu-aomp-amdgpuX	`72.59% <82.24%> (+0.04%)`	⬆️
pytest-gpu-nvc-nvidiaX	`73.66% <83.43%> (+0.04%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

mloubout

Minor comments but it looks all sane, will test it as much as possible

mloubout · 2025-06-11T17:23:29Z

devito/passes/iet/mpi.py

+    otherwise.
+    """
+
+    def rule0(dep):


These two don't need to be inside _is_iter_carried

sure, but they are so local...

I would keep them there purely for neatness

mloubout · 2025-06-11T17:30:56Z

devito/ir/iet/visitors.py

+        if flag and self.rule(self.match, o):
+            found.append(o)
+        for i in o.children:
+            found, flag = self._visit(i, ret=(found, flag))


Should probably exit if flag switch from true to false. Or it's gonnavisit everything afterward for nothing.

but then I'd have to add logic to keep track that a switch from False to True had already happened (that is, that a start-stop has been encountered as opposed to not yet), which I prefer to avoid (less complexity)

Somethign like

for i in o.children: found, newflag = self._visit(i, ret=(found, flag)) if flag and not newflag: return found, newflag else: flag = newflag

Should do it and avoid quite a lot of visit

mloubout · 2025-06-11T17:33:10Z

devito/passes/iet/mpi.py

+
+                # Find candidate for subsequent merging
+                for hs1 in halo_spots:
+                    if hs0 is hs1 or cond_mapper[hs1]:


Could check if cond_mapper[hs1] is a "subset" of "cond_mapper[hs0]" , i.e less conditions in which case it could be hoisted

I could, but that really is an impossible situation

mloubout · 2025-06-11T19:44:57Z

devito/passes/iet/mpi.py

+        return False

+    # Then, check the data dependences would be satisfied
+    if not _is_iter_carried(hsf1, scope):


just return _is_iter_carried(hsf1, scope)

mloubout · 2025-06-11T19:45:57Z

devito/passes/iet/mpi.py

+    and ends at the HaloSpot `hs1`.
+    """
+    expressions = FindWithin(Expression, hs0, stop=hs1).visit(it)
+    assert len(expressions) > 0, \


Is it even possible? hs0 will always be in it so not sure worth the assert

yah I'll drop it

EdCaunt · 2025-06-12T09:27:43Z

devito/ir/iet/visitors.py

        return ret


+class FindWithin(FindNodes):


This will want rebasing on top of #2621 in due course

it's more likely the other way round

EdCaunt · 2025-06-12T09:30:10Z

devito/ir/support/basic.py

            if self.function._mem_shared:
                # Special case: the distance between two regular, thread-shared
-                # objects fallbacks to zero, as any other value would be nonsensical
+                # objects fallbacks to zero, as any other value would be


Nitpick: "falls back" or "is set to zero as a fallback"

EdCaunt · 2025-06-12T15:41:52Z

devito/ir/support/basic.py

+    same value, False otherwise.
+    """
+    # Case 1: ModuloDimensions of size 1
+    try:


Could this just be:

return d0.is_Modulo and d1.is_Modulo and d0.modulo == d1.modulo == 1

in what case would the attribute error be triggered? Regardless the if can just be a return

you're right

the try is necessary because not all Dimensions may be ModuloDimensions

anyway, dropping the if

ah no, now I remember why I made it that way...

as you can see, the function starts with

# Case 1: ModuloDimensions of size 1

as I'm retaining the possibility that at some point in the future this function may be extended. In which case, the if makes more sense. So I'm keeping it

Seems reasonable. Possibly leave a comment mentioning this so it doesn't get factored out in the future?

EdCaunt · 2025-06-12T15:42:21Z

devito/mpi/halo_scheme.py

        return hash((self.loc_indices, self.loc_dirs, self.halos, self.dims,
                     self.bundle))

+    @property


Cached property?

EdCaunt · 2025-06-12T15:47:56Z

devito/passes/iet/mpi.py

-    iet = _hoist_invariant(iet)
+    iet = _hoist_redundant_from_conditionals(iet)
    iet = _merge_halospots(iet)
+    iet = _hoist_invariant(iet)


Does this switch also yield a bit of compiler speedup, since halospots are merged before hoisting, so there are fewer to hoist?

maybe but not sure

EdCaunt · 2025-06-12T16:16:04Z

devito/passes/iet/mpi.py

-                        not isinstance(i.condition, GuardFactorEq)}
-        cond_mapper[hs] = conditionals
+    return {hs: tuple(i for i in v if i.is_Conditional)
+            for hs, v in MapHaloSpots().visit(iet).items()}


MapHaloSpots().visit(iet).items() <- technically fine, but upsetting

OK, assigning mapper = MapHaloSpots().visit(iet) first

EdCaunt · 2025-06-12T16:16:24Z

devito/passes/iet/mpi.py

+    and ends at the HaloSpot `hs1`.
+    """
+    expressions = FindWithin(Expression, hs0, stop=hs1).visit(it)
+    return Scope([e.expr for e in expressions])


Can this just use a generator without the list?

EdCaunt · 2025-06-12T16:17:36Z

devito/passes/iet/mpi.py

+    otherwise.
+    """
+
+    def rule0(dep):


I would keep them there purely for neatness

EdCaunt · 2025-06-12T16:18:44Z

devito/passes/iet/mpi.py

+
+    for f, v in hsf.fmapper.items():
+        for dep in scope.d_flow.project(f):
+            if not rule0(dep) and not rule1(dep, v.loc_indices):


rule0 has a double not here. Why not fold these nots into the rules themselves?

yes I could, but the semantics of rule0 and rule1 hasn't been changed in years, is it really worth it?

Probably not, just figured it was worth mentioning

EdCaunt · 2025-06-12T16:20:31Z

tests/test_mpi.py

            assert np.allclose(h.data_ro_domain[0, 5:], [4.4, 4.4, 4.4, 3.4, 3.1], rtol=R)

+    @pytest.mark.parallel(mode=1)
+    def test_merge_and_hoist_haloupdate_if_diff_locindices_v2(self, mode):


Snappy 😂

who needs docstrings ^^

EdCaunt

Looks ready to go

FabioLuporini requested review from EdCaunt, JDBetteridge, georgebisbas and mloubout June 11, 2025 15:30

FabioLuporini added the bug-C bug in the generated code label Jun 11, 2025

mloubout reviewed Jun 11, 2025

View reviewed changes

FabioLuporini force-pushed the fix-mpi-hoisting-2 branch from 3b80b02 to 8fdbe7f Compare June 12, 2025 08:21

EdCaunt reviewed Jun 12, 2025

View reviewed changes

EdCaunt approved these changes Jun 13, 2025

View reviewed changes

mloubout approved these changes Jun 13, 2025

View reviewed changes

FabioLuporini added 8 commits June 13, 2025 08:46

compiler: Fix DDA with degenerating ModuloDimensions

0998923

compiler: Improve halo-exchange hoisting

00813ea

compiler: Rework MPI hoist and merge

4491316

compiler: Revamp HaloUpdate hoisting and merging

09bbb1f

compiler: Cleanup MPI passes

fcdd6a4

compiler: Polish MPI passes

32c6b61

compiler: Speed up FindWithin

82fa9e6

compiler: Use more property caching and generators

b3a3e00

mloubout force-pushed the fix-mpi-hoisting-2 branch from 8ca5173 to b3a3e00 Compare June 13, 2025 12:46

mloubout merged commit b3408d8 into main Jun 13, 2025
35 checks passed

mloubout deleted the fix-mpi-hoisting-2 branch June 13, 2025 16:59

compiler: Revamp MPI hoisting and merging #2629

compiler: Revamp MPI hoisting and merging #2629

Uh oh!

Conversation

FabioLuporini commented Jun 11, 2025

Uh oh!

codecov bot commented Jun 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

mloubout left a comment

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!

mloubout Jun 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

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

codecov bot commented Jun 11, 2025 •

edited

Loading

mloubout Jun 12, 2025 •

edited

Loading