compiler: Rework HaloSpot optimization #2477

georgebisbas · 2024-11-04T11:16:47Z

Fix for issue #2448, supplemented with some code edits

devito/passes/iet/mpi.py

tests/test_mpi.py

codecov · 2024-11-04T11:36:02Z

Codecov Report

Attention: Patch coverage is 98.21429% with 4 lines in your changes missing coverage. Please review.

Project coverage is 87.34%. Comparing base (279f074) to head (eb1435a).

Files with missing lines	Patch %	Lines
devito/mpi/halo_scheme.py	90.47%	1 Missing and 1 partial ⚠️
devito/passes/iet/mpi.py	96.72%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2477      +/-   ##
==========================================
+ Coverage   87.29%   87.34%   +0.05%     
==========================================
  Files         238      238              
  Lines       45366    45522     +156     
  Branches     4029     4030       +1     
==========================================
+ Hits        39602    39763     +161     
+ Misses       5083     5079       -4     
+ Partials      681      680       -1

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

review-notebook-app · 2024-11-04T12:58:01Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

devito/ir/iet/nodes.py

EdCaunt · 2024-11-04T13:21:33Z

devito/mpi/halo_scheme.py

+                     frozenset(self.halos),
+                     frozenset(self.dims)))
+
+    def rebuild(self, **kwargs):


Subclass Reconstructable and add rargs?

EdCaunt · 2024-11-04T13:37:49Z

devito/mpi/halo_scheme.py

@@ -95,8 +127,16 @@ def __init__(self, exprs, ispace):
        self._honored = frozendict(self._honored)

    def __repr__(self):


Could this be attached to a mixin class for both HaloSpot and HaloSchemeEntry?

yes but where would you add it?

I would just drop all these repr it's a lot of code for what at the end of the day?

When you print the IET, you see more useful details, it makes dev/debug life easier!

anyway, Ed has a point, this is redundant code which should somehow be factorized (a private method of HaloScheme, also called by HaloSpot.__repr__? not sure)

I was also thinking of this, but fmapper is the problem, we need to drop it from the entry? and thus be a class itself?
needs some thought, I may restructure more

tests/test_mpi.py

devito/ir/iet/nodes.py

FabioLuporini · 2024-11-05T08:04:13Z

devito/mpi/halo_scheme.py

@@ -28,7 +28,39 @@ class HaloLabel(Tag):
 STENCIL = HaloLabel('stencil')


-HaloSchemeEntry = namedtuple('HaloSchemeEntry', 'loc_indices loc_dirs halos dims')
+class HaloSchemeEntry:


as Ed wrote below, this should inherit from Reconstructable

FabioLuporini · 2024-11-05T08:04:59Z

devito/mpi/halo_scheme.py

-HaloSchemeEntry = namedtuple('HaloSchemeEntry', 'loc_indices loc_dirs halos dims')
+class HaloSchemeEntry:
+
+    def __init__(self, loc_indices, loc_dirs, halos, dims):


can it not inherit, alternatively, from EnrichedTuple?

THere are some 'getters', that look redundant ?
Reconstrcutable seems to work fine

Is that ok now?

that look redundant

with EnrichedTuple, once you properly override __rargs__ and __rkwargs__, you probably don't need to override __init__ and __eq__

I fail to do so, so far,, the already existing code of DimensionTuple does not seem to help me....

FabioLuporini · 2024-11-05T08:05:57Z

devito/mpi/halo_scheme.py

@@ -95,8 +127,16 @@ def __init__(self, exprs, ispace):
        self._honored = frozendict(self._honored)

    def __repr__(self):


yes but where would you add it?

I would just drop all these repr it's a lot of code for what at the end of the day?

tests/test_mpi.py

FabioLuporini · 2024-11-05T08:28:24Z

devito/passes/iet/mpi.py

-    # a stopper to halo merging
+    # Loop over the functions in the HaloSpots
+    for f, v in hs1.fmapper.items():
+        # If no time accesses, skip


"E.g., if no time accesses, skip"

you have to say it's an example...

wdym by 'say it's an example' ?

FabioLuporini · 2024-11-05T08:28:57Z

devito/passes/iet/mpi.py

-        return not any(d in hs.dimensions or dep.distance_mapper[d] is S.Infinity
-                       for d in dep.cause)
+        # If the function is not in both HaloSpots, skip
+        if (*hs0.functions, *hs1.functions).count(f) < 2:


.intersection(...), again as mentioned above?

no needed anymore

devito/passes/iet/mpi.py

FabioLuporini · 2024-11-05T08:31:09Z

devito/passes/iet/mpi.py


    # Precompute scopes to save time
    scopes = {i: Scope([e.expr for e in v]) for i, v in MapNodes().visit(iet).items()}

+    cond_mapper = _create_cond_mapper(iet)


typically we use the word "make", so _make_cond_mapper(iet), but obviously I'm nitpicking here

FabioLuporini · 2024-11-05T08:36:01Z

devito/passes/iet/mpi.py


-                    candidates = [i.dim._defines for i in iters[n:]]
+            for hs1 in halo_spots[i+1:]:


maybe you could avoid an indentation level but rather considering all pairs

from itertools import combinations ... for hs0, hs1 in combinations(halo_schemes, 2): ...

btw, are we sure that the current code works if there are 2 hoisting opportunities, instead of 1 ?

could we add a test?

I am working on this now

georgebisbas · 2024-11-08T11:15:29Z

devito/mpi/halo_scheme.py

-HaloSchemeEntry = namedtuple('HaloSchemeEntry', 'loc_indices loc_dirs halos dims')
+class HaloSchemeEntry:
+
+    def __init__(self, loc_indices, loc_dirs, halos, dims):


Is that ok now?

georgebisbas · 2024-11-08T11:15:50Z

devito/passes/iet/mpi.py

    """
-    iet = _drop_halospots(iet)
-    iet = _hoist_halospots(iet)
+    iet = _drop_reduction_halospots(iet)


this is just renamed

georgebisbas · 2024-11-08T11:16:26Z

devito/passes/iet/mpi.py

+
+    for it, halo_spots in iter_mapper.items():
+
+        for hs0, hs1 in combinations(halo_spots, r=2):


two-level combo as proposed

georgebisbas · 2024-11-08T11:17:31Z

devito/passes/iet/mpi.py

    iet = Transformer(mapper, nested=True).visit(iet)
-
-    # Clean up: de-nest HaloSpots if necessary


there is no need for that anymore

tests/test_mpi.py

EdCaunt · 2024-11-08T10:18:21Z

devito/mpi/halo_scheme.py

-HaloSchemeEntry = namedtuple('HaloSchemeEntry', 'loc_indices loc_dirs halos dims')
+class HaloSchemeEntry(Reconstructable):
+
+    __rkwargs__ = ('loc_indices', 'loc_dirs', 'halos', 'dims')


Should this be __rargs__?

EdCaunt · 2024-11-11T14:24:06Z

devito/passes/iet/mpi.py

+    return cond0 != cond1
+
+
+def motion_rules():


Perhaps define a motion_rules = _motion_rules()?

this should not be a function

And I'd call this simply "rules"

def _rule0(...): """....""" <impl> def _rule1(...): """....""" <impl> rules = (_rule0, _rule1)

FabioLuporini · 2024-11-12T08:18:40Z

devito/mpi/halo_scheme.py

-HaloSchemeEntry = namedtuple('HaloSchemeEntry', 'loc_indices loc_dirs halos dims')
+class HaloSchemeEntry(Reconstructable):
+
+    __rkwargs__ = ('loc_indices', 'loc_dirs', 'halos', 'dims')


FabioLuporini · 2024-11-12T08:20:52Z

devito/mpi/halo_scheme.py

-HaloSchemeEntry = namedtuple('HaloSchemeEntry', 'loc_indices loc_dirs halos dims')
+class HaloSchemeEntry:
+
+    def __init__(self, loc_indices, loc_dirs, halos, dims):


that look redundant

with EnrichedTuple, once you properly override __rargs__ and __rkwargs__, you probably don't need to override __init__ and __eq__

FabioLuporini · 2024-11-12T08:22:40Z

devito/mpi/halo_scheme.py

@@ -95,8 +127,16 @@ def __init__(self, exprs, ispace):
        self._honored = frozendict(self._honored)

    def __repr__(self):


anyway, Ed has a point, this is redundant code which should somehow be factorized (a private method of HaloScheme, also called by HaloSpot.__repr__? not sure)

FabioLuporini · 2024-11-12T08:23:14Z

devito/mpi/halo_scheme.py

-                hse = HaloSchemeEntry(frozendict(loc_indices),
-                                      frozendict(loc_dirs),
-                                      hse0.halos, hse0.dims)
+                hse = hse0._rebuild(loc_indices=frozendict(loc_indices),


maybe the cast to frozendict should be moved to inside HaloSchemeEntry.__init__

should it be for all attributes?

devito/passes/iet/mpi.py

FabioLuporini · 2024-11-12T08:33:29Z

devito/passes/iet/mpi.py

-                   for hs, v in cond_mapper.items()}
+    cond_mapper = _make_cond_mapper(iet)
+
+    mapper = {}

    iter_mapper = MapNodes(Iteration, HaloSpot, 'immediate').visit(iet)


you may introudce a private function to get this iter_mapper too, since you're doing the same exact filtering that appears in hoist_invariant

now that I think about it, you can probably also move the cond_mapper inside such a new utility function, and do the conditional filtering in there directly, rather than repeating it in both passes (it's basically the same stuff repeated at the moment)

FabioLuporini · 2024-11-12T08:34:12Z

devito/passes/iet/mpi.py

+            for hs, v in cond_mapper.items()}
+
+
+def ensure_control_flow(hs0, hs1, cond_mapper):


FabioLuporini · 2024-11-12T08:34:18Z

devito/passes/iet/mpi.py

+
+def ensure_control_flow(hs0, hs1, cond_mapper):
+    """
+    # If there are Conditionals involved, both `hs0` and `hs1` must be


FabioLuporini · 2024-11-12T08:34:29Z

devito/passes/iet/mpi.py

+def ensure_control_flow(hs0, hs1, cond_mapper):
+    """
+    # If there are Conditionals involved, both `hs0` and `hs1` must be
+    # within the same Conditional, otherwise we would break the control


control flow

FabioLuporini · 2024-11-12T08:39:09Z

devito/passes/iet/mpi.py

+    return cond0 != cond1
+
+
+def motion_rules():


And I'd call this simply "rules"

def _rule0(...): """....""" <impl> def _rule1(...): """....""" <impl> rules = (_rule0, _rule1)

tests/test_mpi.py

mloubout · 2024-11-12T13:13:05Z

tests/test_mpi.py

+    @pytest.mark.parallel(mode=1)
+    def test_halo_structure(self, mode):
+
+        mytest = TestTTI()


remove, oneline TestTTI().tti_operator

mloubout · 2024-11-12T13:14:58Z

devito/passes/iet/mpi.py

-        if i is None or len(halo_spots) <= 1:
-            continue
+    # Drop pairs that have keys that are None
+    iter_mapper = {k: v for k, v in iter_mapper.items() if k is not None}


This whole section looks like a plain copy paste of 89-119 especially the loop next, iter_maper filtering and conditions, can't this be. merged or lifted in some common piece?

simplified a bit

georgebisbas · 2024-11-13T19:02:32Z

devito/mpi/halo_scheme.py

@@ -95,8 +127,16 @@ def __init__(self, exprs, ispace):
        self._honored = frozendict(self._honored)

    def __repr__(self):


I was also thinking of this, but fmapper is the problem, we need to drop it from the entry? and thus be a class itself?
needs some thought, I may restructure more

georgebisbas · 2024-11-14T10:33:00Z

devito/mpi/halo_scheme.py

-HaloSchemeEntry = namedtuple('HaloSchemeEntry', 'loc_indices loc_dirs halos dims')
+class HaloSchemeEntry(Reconstructable):
+
+    __rkwargs__ = ('loc_indices', 'loc_dirs', 'halos', 'dims')


georgebisbas · 2024-11-14T10:41:04Z

devito/mpi/halo_scheme.py

-                hse = HaloSchemeEntry(frozendict(loc_indices),
-                                      frozendict(loc_dirs),
-                                      hse0.halos, hse0.dims)
+                hse = hse0._rebuild(loc_indices=frozendict(loc_indices),


georgebisbas · 2024-11-14T10:41:13Z

devito/mpi/halo_scheme.py

-                hse = HaloSchemeEntry(frozendict(loc_indices),
-                                      frozendict(loc_dirs),
-                                      hse0.halos, hse0.dims)
+                hse = hse0._rebuild(loc_indices=frozendict(loc_indices),


should it be for all attributes?

devito/passes/iet/mpi.py

georgebisbas · 2024-11-14T10:42:47Z

devito/passes/iet/mpi.py

+    iter_mapper = {k: v for k, v in iter_mapper.items() if len(v) > 1}
+
+    for it, halo_spots in iter_mapper.items():
+


georgebisbas · 2024-11-14T10:43:03Z

devito/passes/iet/mpi.py

+            if ensure_control_flow(hs0, hs1, cond_mapper):
+                continue
+
+            # If there are overlapping time accesses, skip


cahnged, is it better?

devito/passes/iet/mpi.py

georgebisbas · 2024-11-14T10:54:27Z

examples/seismic/viscoacoustic/operators.py

@@ -524,7 +524,7 @@ def ForwardOperator(model, geometry, space_order=4, kernel='sls', time_order=2,

    # Substitute spacing terms to reduce flops
    return Operator(eqn + src_term + rec_term, subs=model.spacing_map,
-                    name='Forward', **kwargs)
+                    name='ViscoAcForward', **kwargs)


@mloubout @FabioLuporini are you ok with this?

mloubout · 2024-11-15T15:03:58Z

devito/passes/iet/mpi.py

+                    for d in hse.loc_indices:
+                        md = hse.loc_indices[d]
+                        if md in it.uindices:
+                            if it.direction is Backward:


can this be moved into it as it.start ?

EdCaunt

Few minor comments. Looks very good to me

EdCaunt · 2024-11-18T11:12:47Z

devito/passes/iet/mpi.py

+
+# *** Utilities
+
+def filter_iter_mapper(iet):


Nitpick: if these aren't going to be exposed, do they need to be private? Either way, they should probably all be consistent with one another

EdCaunt · 2024-11-18T11:13:29Z

devito/passes/iet/mpi.py

+               for d, v in loc_indices.items())
+
+
+rules = [_rule0, _rule1]


Nitpick: Making this a tuple would be safer

EdCaunt · 2024-11-18T11:17:41Z

tests/test_mpi.py

+
+        # Test two variants of receiver interpolation
+        nrec = 1
+        rec = SparseTimeFunction(name="rec", grid=grid, npoint=nrec, nt=tn)


Nitpick: You could remove some lines by having npoint=1, nt=30 in here. You can do similar stuff throughout this test (for example Grid(shape=(2,)).

EdCaunt · 2024-11-18T11:18:51Z

tests/test_mpi.py

+        rec_term3 = rec2.interpolate(expr=v2.forward)
+
+        # Test receiver interpolation 0, here we have a halo exchange hoisted
+        op2 = Operator([u_v] + [u_v2] + [u_tau] + [u_tau2] + rec_term0 + rec_term2)


Should this and op3 be split off in their own tests, or is it necessary to have them in this test?

georgebisbas added MPI mpi-related compiler labels Nov 4, 2024

georgebisbas requested review from FabioLuporini, mloubout and EdCaunt November 4, 2024 11:16

georgebisbas self-assigned this Nov 4, 2024

georgebisbas linked an issue Nov 4, 2024 that may be closed by this pull request

compiler: Redundant haloupdate #2448

Open

georgebisbas mentioned this pull request Nov 4, 2024

[RFC] compiler: Add new HaloSpot optimization #2475

Closed

georgebisbas commented Nov 4, 2024

View reviewed changes

devito/passes/iet/mpi.py Show resolved Hide resolved

tests/test_mpi.py Show resolved Hide resolved

EdCaunt reviewed Nov 4, 2024

View reviewed changes

FabioLuporini reviewed Nov 5, 2024

View reviewed changes

georgebisbas force-pushed the halo_opt_revamp_II branch 4 times, most recently from ca5251a to 5951eaf Compare November 8, 2024 11:04

georgebisbas commented Nov 8, 2024

View reviewed changes

georgebisbas mentioned this pull request Nov 8, 2024

tests: Add a test-only init for issue 2448 #2465

Closed

EdCaunt reviewed Nov 11, 2024

View reviewed changes

FabioLuporini reviewed Nov 12, 2024

View reviewed changes

mloubout reviewed Nov 12, 2024

View reviewed changes

georgebisbas force-pushed the halo_opt_revamp_II branch from 9f23d51 to 57b0464 Compare November 14, 2024 11:04

georgebisbas commented Nov 14, 2024

View reviewed changes

mloubout reviewed Nov 15, 2024

View reviewed changes

EdCaunt reviewed Nov 18, 2024

View reviewed changes

georgebisbas added 4 commits November 18, 2024 14:48

compiler: Initialize additional halo opt pass

1280293

compiler: Improve candidate selection and mapper application

68ab524

compiler: Drop duplicated code

9cdc9c7

compiler: Concretize halo_opt_revamp, tests, drop redundant code

130d803

georgebisbas added 7 commits November 18, 2024 14:48

examples: Update MPI notebook output

53ab530

examples: Update with edge case of 03_PML notebook

62f5ba1

compiler: Inline halo_to_halo

d681b70

compiler: class HaloSchemeEntry(Reconstructable)

546ce9c

compiler: Rework terminology, minor cleanup, comments

4f2cc1d

compiler: add frozendict to __init__

c084c22

tests: Add backward test and drop redundancies

eb1435a

georgebisbas force-pushed the halo_opt_revamp_II branch from d6ed1a7 to eb1435a Compare November 18, 2024 12:53

		@@ -95,8 +127,16 @@ def __init__(self, exprs, ispace):
		self._honored = frozendict(self._honored)

		def __repr__(self):


		candidates = [i.dim._defines for i in iters[n:]]
		for hs1 in halo_spots[i+1:]:


		for it, halo_spots in iter_mapper.items():

		for hs0, hs1 in combinations(halo_spots, r=2):

		iet = Transformer(mapper, nested=True).visit(iet)

		# Clean up: de-nest HaloSpots if necessary

		for hs, v in cond_mapper.items()}


		def ensure_control_flow(hs0, hs1, cond_mapper):

		iter_mapper = {k: v for k, v in iter_mapper.items() if len(v) > 1}

		for it, halo_spots in iter_mapper.items():

compiler: Rework HaloSpot optimization #2477

Are you sure you want to change the base?

compiler: Rework HaloSpot optimization #2477

Conversation

georgebisbas commented Nov 4, 2024

codecov bot commented Nov 4, 2024 • edited Loading

Codecov Report

review-notebook-app bot commented Nov 4, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

EdCaunt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Nov 4, 2024 •

edited

Loading