compiler: Revamp code generation for asynchronous operations #2376

FabioLuporini · 2024-05-27T09:44:47Z

This PR is the result of the last 40 days of work, more or less. It would deserve to be split over multiple PRs, as what I've done here is really really bad, but point is, 1) I didn't have the time to go back and forth and 2) it's all intertwined, somehow...

Strictly speaking, there are no new features here. "Just" a major revamp to generalise code generation for asynchronous operations, which will mostly impact PRO.

A lot of compiler passes have been rewritten from scratch, and in my opinion, the code quality has increased significantly. Some legacy code has finally gone for good...

Despite being a giant PR, it's worth saying that all edits go in the same direction and contribute to the same goal; in other words, this PR isn't the sum of N orthogonal changes, even though, I agree, the appearance may be deceptive

FabioLuporini · 2024-05-27T09:46:55Z

devito/passes/clusters/asynchrony.py

-                        writes = c1.scope.writes[target]
-                    except KeyError:
-                        # No read-write dependency, ignore
+    def _schedule_waitlocks(self, c0, d, clusters, locks, syncs):


Note for reviewers: this is mostly the same as before (with same exceptions), but split over different methods for better code readability

codecov · 2024-05-27T09:47:38Z

Codecov Report

Attention: Patch coverage is 61.75772% with 483 lines in your changes missing coverage. Please review.

Project coverage is 86.77%. Comparing base (210f899) to head (0decbb2).

Files	Patch %	Lines
devito/passes/clusters/asynchrony.py	9.77%	157 Missing ⚠️
devito/passes/iet/asynchrony.py	9.47%	86 Missing ⚠️
devito/passes/iet/engine.py	45.00%	29 Missing and 4 partials ⚠️
devito/passes/iet/orchestration.py	25.00%	33 Missing ⚠️
devito/passes/clusters/buffering.py	88.50%	15 Missing and 15 partials ⚠️
devito/passes/clusters/implicit.py	72.81%	25 Missing and 3 partials ⚠️
tests/test_gpu_common.py	0.00%	23 Missing ⚠️
devito/core/gpu.py	30.00%	14 Missing ⚠️
devito/operator/operator.py	54.16%	9 Missing and 2 partials ⚠️
devito/ir/support/syncs.py	67.74%	10 Missing ⚠️
... and 18 more

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2376      +/-   ##
==========================================
- Coverage   86.81%   86.77%   -0.05%     
==========================================
  Files         234      234              
  Lines       44004    44245     +241     
  Branches     8127     8181      +54     
==========================================
+ Hits        38204    38393     +189     
- Misses       5093     5137      +44     
- Partials      707      715       +8

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

FabioLuporini · 2024-05-27T09:49:22Z

devito/passes/clusters/buffering.py


 __all__ = ['buffering']


 @timed_pass()
-def buffering(clusters, callback, sregistry, options, **kwargs):
+def buffering(clusters, key, sregistry, options, **kwargs):


Note for reviewers: unlike passes/clusters/asynchrony.py, which for a good 85-90% is the same code but refactored, this pass has been totally revamped and the result is, imho, vastly better (more robust, easier to read, dramatically easier to understand and modify) code

devito/passes/clusters/implicit.py

devito/operator/operator.py

devito/operator/profiling.py

georgebisbas · 2024-05-27T13:35:16Z

devito/operator/profiling.py


        k = PerfKey(name, rank)

-        if not ops:
+        if not ops or any(np.isnan(i) for i in [ops, points, traffic]):


np.isfinite to catch the Inf cases too

so something like

any(not np.isfinite(i) ...)

instead of

any(np.isnan(i) ...)

?

georgebisbas · 2024-05-27T15:13:06Z

devito/passes/iet/instrument.py

devito/passes/iet/misc.py

georgebisbas

Big PR, I need some more time, to have another look and most importantly understand a few things!

review-notebook-app · 2024-05-27T15:47:03Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

devito/core/gpu.py

EdCaunt · 2024-05-30T09:31:35Z

devito/ir/iet/efunc.py

@@ -215,6 +215,13 @@ def functions(self):
            launch_args += (self.stream.function,)
        return super().functions + launch_args

+    @cached_property
+    def expr_symbols(self):
+        launch_symbols = (self.grid, self.block,)


Don't need second comma in this tuple

devito/ir/stree/algorithms.py

EdCaunt · 2024-05-30T09:57:30Z

devito/ir/support/space.py

+        if callable(maybe_callable):
+            dims = {i.dim for i in self if maybe_callable(i.dim)}
+        else:
+            dims = set(as_tuple(maybe_callable))


Suggestion: add an as_set utility to do this. They could both be thin wrappers around some as_iterable

devito/operator/operator.py

devito/passes/clusters/implicit.py

EdCaunt · 2024-05-30T11:15:53Z

devito/passes/clusters/buffering.py

-        intervals = IntervalGroup(intervals, relations=relations)
-
-        return IterationSpace(intervals, sub_iterators, directions)
+        return self.ispace


Could this be dropped then?

maybe, given the major change I prefer to keep it like that for now, it also helps me remember what .ispace actually represents (and why) when obtained via .written

devito/tools/utils.py

devito/types/misc.py

mloubout

Looks mostly good

mloubout · 2024-05-30T12:48:24Z

devito/core/cpu.py

+        options0 = kwargs0.pop('options')
+
+        def wrapper(expressions, **kwargs1):
+            options = {**options0, **kwargs1.pop('options', {})}


just wrapper(expressions, options=None, **kwargs1): and skip the pop

mloubout · 2024-05-30T12:48:49Z

devito/core/cpu.py

+            kwargs = {'options': options, **kwargs0, **kwargs1}
+
+            # User-provided openmp flag has precedence over defaults
+            if not options['openmp']:


options.get ?

no square brackets is fine because by construction openmp must be there

devito/core/cpu.py

devito/core/gpu.py

mloubout · 2024-05-30T12:56:41Z

devito/ir/stree/algorithms.py

-            processed.append(c.rebuild(exprs=None, guards=c.guards, syncs=c.syncs))
+        elif c.is_critical_region:
+            if c.is_wait:
+                processed.append(c.rebuild(exprs=[], syncs=c.syncs))


can probably drop c.sync it's unchanged

mloubout · 2024-05-30T13:45:58Z

devito/symbolics/search.py

+    """Shorthand to retrieve the DiscreteFunctions in `exprs`."""
+    indexeds = search(exprs, q_indexed, mode, 'dfs')
+
+    functions = search(exprs, q_function, mode, 'dfs')


should this be in a deep mode like for retrieve_dimensions

devito/types/basic.py

devito/types/dimension.py

mloubout · 2024-05-30T13:53:16Z

tests/test_subdomains.py

@@ -672,29 +672,6 @@ class Dummy(SubDomainSet):
        assert z.is_Parallel


-class TestMultiSubDimension:


Removed test?

There were issues rebuilding MultiSubDimension because of how the thickness worked, so it has a bunch of extra gubbins to handle this. This PR removes the thickness as a kwarg and all the extra machinery, so this test isn't really needed anymore

tests/test_buffering.py

devito/ir/iet/visitors.py

devito/ir/stree/algorithms.py

devito/ir/support/guards.py

georgebisbas · 2024-06-06T14:43:13Z

devito/passes/clusters/buffering.py


    @cached_property
-    def readfrom(self):
+    def last_idx(self):


That is actually useful for the xdsl lowering, I had no idea it existed

georgebisbas · 2024-06-06T15:15:53Z

devito/passes/iet/orchestration.py

+    def _make_snapout(self, iet, *args):
+        return self._make_callable('write_snapshot', iet, *args)
+
+    def _make_snapin(self, iet, *args):
+        return self._make_callable('read_snapshot', iet, *args)


I would keep the func name as the callable arg for being homogeneous to 'init_array'

examples/performance/01_gpu.ipynb

tests/test_linearize.py

mloubout · 2024-06-17T12:39:56Z

devito/ir/iet/visitors.py

+        # 2-tuples representing ranges
+        imask = []
+        for v in o.imask:
+            try:


The single dispatch of uxrteplace doesn't handle these two cases?

mloubout · 2024-06-17T12:43:13Z

devito/operator/profiling.py


        k = PerfKey(name, rank)

-        if not ops:
+        if not ops or any(np.isnan(i) for i in [ops, points, traffic]):


georgebisbas · 2024-06-17T14:43:25Z

devito/operator/profiling.py


        k = PerfKey(name, rank)

-        if not ops:
+        if not ops or any(np.isnan(i) for i in [ops, points, traffic]):


georgebisbas · 2024-06-17T14:45:08Z

devito/passes/__init__.py

@@ -36,6 +36,16 @@ def is_on_device(obj, gpu_fit):
    return all(f in gpu_fit for f in fsave)


+def stream_dimensions(f):


Not really related to this PR but is it good that these functions live in init?
Maybe some better place? other/new file?

georgebisbas · 2024-06-17T15:27:57Z

devito/passes/clusters/implicit.py

+            return clusters
+
+        try:
+            pd = prefix[-2].dim


This is not equivalent to dim.parent?

georgebisbas · 2024-06-17T15:48:07Z

devito/types/grid.py

-        right = parent.symbolic_max - rst
+    def __init_finalize__(self, name, parent, msd):
+        # NOTE: a MultiSubDimension stashes a reference to the originating
+        # MultiSubDomain.  This creates a circular pattern as the `msd` itself


two spaces?

FabioLuporini added the compiler label May 27, 2024

FabioLuporini requested review from mloubout, georgebisbas and EdCaunt May 27, 2024 09:44

FabioLuporini commented May 27, 2024

View reviewed changes

devito/passes/clusters/implicit.py Show resolved Hide resolved

georgebisbas reviewed May 27, 2024

View reviewed changes

EdCaunt reviewed May 30, 2024

View reviewed changes

mloubout reviewed May 30, 2024

View reviewed changes

georgebisbas reviewed May 31, 2024

View reviewed changes

devito/ir/iet/visitors.py Show resolved Hide resolved

devito/ir/stree/algorithms.py Show resolved Hide resolved

devito/ir/support/guards.py Show resolved Hide resolved

FabioLuporini force-pushed the revamp-async-codegen-final branch 2 times, most recently from 57c7acd to 2e5d1eb Compare June 6, 2024 14:00

georgebisbas reviewed Jun 6, 2024

View reviewed changes

mloubout approved these changes Jun 17, 2024

View reviewed changes

georgebisbas reviewed Jun 17, 2024

View reviewed changes

FabioLuporini added 10 commits June 18, 2024 09:49

compile: Reuse BlockDimensions when it makes sense

716a377

api: Accept frozendict for opt-options

953590a

compiler: Tweak profiling of subsections

ce8673c

compiler: Trim down Operator apply profiling

e257a09

compiler: Add Operator._sanitize_exprs

05a409f

compiler: Add VirtualDimension

1688c5a

compiler: Add IterationSpace.{split,insert}

5ab35e3

compiler: Simplify gpu callbacks

1744f04

compiler: Simplify tasking

2995c69

compiler: Simplify streaming

756f204

FabioLuporini added 26 commits June 18, 2024 09:49

misc: pep8 happiness

ef839a4

compiler: Refactor wraps_memcpy

e5534d4

compiler: Tweak double buffering

0c018b5

compiler: Revamp generation of sub-iterators

9541d84

compiler: Introduce InitArray and SyncArray ops

f91cbf7

compiler: Fix SyncArray codegen

bb6f54b

compiler: Switch tasking back to being a Queue

85714be

compiler: Improve Section placement

e0d1343

compiler: Misc cleanup

5359c8d

compiler: Patch stree construction

7659338

compiler: Fix deterministic code gen

83deb16

compiler: Revamp async codegen

951dda4

misc: pep8 happiness

3773164

compiler: Patch Uxreplace for ParallelTree

bffe9d3

compiler: Revamp linearization

1b465d6

tools: Reinstate erroneously dropped filter_ordered implementation

c08f3d0

compiler: Fixup orchestration

809bcb8

compiler: Fix Uxreplace.PragmaTransfer

e01d67d

compiler: Patch KernelLaunch.expr_symbols

7a7bf29

example: Update expected output

77cd2ed

tests: Drop --no-summary from custom markers

84faf27

compiler: Fixup rcompile post-rebasing

ad4b477

compiler: Polishing

17e8512

compiler: Add stream_dimensions()

60a2ee7

compiler: Refactor _generate_macros

1905fb6

misc: Improve performance summary

0decbb2

FabioLuporini force-pushed the revamp-async-codegen-final branch from 2e5d1eb to 0decbb2 Compare June 18, 2024 09:51

georgebisbas approved these changes Jun 18, 2024

View reviewed changes

FabioLuporini merged commit be7c403 into master Jun 18, 2024
31 checks passed

FabioLuporini deleted the revamp-async-codegen-final branch June 18, 2024 12:57

		@@ -672,29 +672,6 @@ class Dummy(SubDomainSet):
		assert z.is_Parallel


		class TestMultiSubDimension:

		@@ -36,6 +36,16 @@ def is_on_device(obj, gpu_fit):
		return all(f in gpu_fit for f in fsave)


		def stream_dimensions(f):

compiler: Revamp code generation for asynchronous operations #2376

compiler: Revamp code generation for asynchronous operations #2376

Conversation

FabioLuporini commented May 27, 2024

Choose a reason for hiding this comment

codecov bot commented May 27, 2024 • edited Loading

Codecov Report

FabioLuporini May 27, 2024 • edited Loading

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

georgebisbas left a comment

Choose a reason for hiding this comment

review-notebook-app bot commented May 27, 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

mloubout 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

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

codecov bot commented May 27, 2024 •

edited

Loading

FabioLuporini May 27, 2024 •

edited

Loading