Vertical loop fusion and demotion of temporaries #374

MichaelSt98 · 2024-09-10T11:12:47Z

Vertical loop fusion and demotion of temporaries in the vertical dimension if possible based on
pragmas in the input source code and heuristics.

This is intended to be a starting point for a more automated and sophisticated approach (including dataflow analysis) as well as further optimisations that require this as a first step.

To test with CLOUDSC, checkout branch nams-vertical-loop-fusion.
This results in the following performance improvements (executed with 1 262144 128):

VARIANT	ORIG/WITHOUT	WITH NEW TRAFO
SCC	ca. 200 GF/s	ca. 460 GF/s
SCC-HOIST	ca. 320 GF/s	ca. 500 GF/s
SCC-STACK	ca. 320 GF/s	ca. 520 GF/s

github-actions · 2024-09-10T11:17:45Z

Documentation for this branch can be viewed at https://sites.ecmwf.int/docs/loki/374/index.html

codecov · 2024-09-10T15:45:22Z

Codecov Report

Attention: Patch coverage is 98.48943% with 5 lines in your changes missing coverage. Please review.

Project coverage is 95.56%. Comparing base (1d95896) to head (7a65ebe).
Report is 36 commits behind head on main.

Files with missing lines	Patch %	Lines
loki/transformations/transform_loop.py	96.29%	3 Missing ⚠️
...ormations/single_column/tests/test_scc_vertical.py	97.97%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #374      +/-   ##
==========================================
+ Coverage   95.54%   95.56%   +0.02%     
==========================================
  Files         188      190       +2     
  Lines       39468    39691     +223     
==========================================
+ Hits        37709    37930     +221     
- Misses       1759     1761       +2

Flag	Coverage Δ
lint_rules	`96.39% <ø> (ø)`
loki	`95.54% <98.48%> (+0.02%)`	⬆️

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.

mlange05

Overall this looks really great, and is a first step towards some very advanced optimisations. I've left a few comments, but they are mostly cosmetic. Once addressed, I think this should be good to go.

The direct inclusion in SCC also ok for now, as it will basically become a no-op unless any pragma markers are inserted. In the future, we might want to add this stuff to a more generic non-SCC loop-change transformation that can be prepended to SCC in config, or run in it's own right (eg. for CPU opts). But that is beyond this PR.

mlange05 · 2024-09-28T03:45:30Z

loki/transformations/single_column/vertical.py

+        routine : :any:`Subroutine`
+            The subroutine in the vertical loops should be fused and
+            temporaries be demoted.
+        horizontal : :any:`Dimension`


I think the horizontal is not actually used here.

mlange05 · 2024-09-28T03:50:07Z

loki/transformations/single_column/vertical.py

+
+    .. note::
+        This transfomation currently relies on pragmas being inserted in the input
+        source files.


It might be good to briefly explain the pragma annotations needed. For example, listing and linking loop-interchange and loop-fusion notation, as well as the custom "-init" and "ignore" notations.

mlange05 · 2024-09-28T03:54:41Z

loki/transformations/single_column/vertical.py

+        loop_var_map = {}
+        with pragmas_attached(routine, ir.Loop):
+            for loop in FindNodes(ir.Loop).visit(routine.body):
+                if loop.variable.name.lower() == self.vertical.index.lower():


Is the .lower() really needed here? Are there edge-cases where expression-string comparison not trigger?

mlange05 · 2024-09-28T04:00:57Z

loki/transformations/single_column/vertical.py

+                        parameters = get_pragma_parameters(loop.pragma, starts_with='fused-loop')
+                        group = parameters.get('group', 'default')
+                        if group == 'ignore':
+                            ignore = True


Instead of tracking ignore, I think you could simply continue here, no?

mlange05 · 2024-09-28T04:02:48Z

loki/transformations/single_column/vertical.py

+                        if group == 'ignore':
+                            ignore = True
+                    if not ignore:
+                        loop_var_map[loop] = as_tuple(var.name.lower() for var in FindVariables().visit(loop.body)


Would building the inverse of this map not make the logic much simple? As in var_loop_map[var.name] = loop would allow you to then simply check for len(var_loop_map[var.name]) > 1 in the next part? Or am I misreading this?

Also, a CaseInsensitiveDict might be useful here for var_loop_map.

mlange05 · 2024-09-28T04:09:24Z

loki/transformations/single_column/vertical.py

+        # look for 'loki k-caching ignore(var1, var2, ...)' pragmas within routine and ignore those vars
+        for pragma in pragmas:
+            if is_loki_pragma(pragma, starts_with='k-caching'):
+                pragma_ignore = get_pragma_parameters(pragma, starts_with='k-caching').get('ignore', None)


Minor nitpick: I wonder if the is_loki_pragma and get_pragma_parameters can be combined to something like if pragma_ignore := get_pragma_parameters(pragma, starts_with='k-caching').get(...). Not strictly necessary, just wondering...

mlange05 · 2024-09-28T04:11:39Z

loki/transformations/single_column/vertical.py

+
+    def correct_init_of_multilevel_arrays(self, routine, multilevel_local_arrays):
+        """
+        Possibly handle initaliztion of those multilevel local arrays via


typo: "initialization"

loki/transformations/single_column/vertical.py

mlange05 · 2024-09-28T04:21:11Z

loki/transformations/single_column/vertical.py

+                            new_pragmas = [pragma.clone(content=pragma.content.replace('-init', '')) if '-init'
+                                    in pragma.content else pragma for pragma in pragmas]
+                            loop_map[loop] = (ir.Comment('! Loki generated loop for init ...'),
+                                    Transformer(node_map_init).visit(loop.clone(\


Two inline Transformer invocations makes this very hard to read. Can we apply them in sequence to build the tuple, before putting them into the loop_map?

Not sure what you mean.
I need to apply the Transformer twice with two different maps since the key would be overwritten otherwise ...
If you propose to introduce temporary variables: sure this is possible!

Oh, I just meant the cosmetics of it - the use of two transformers is fine, it just gets really hard to read if everything is applied in a single line. If you build the tuple incrementally, each Transformer.visit(..) invocation could be a single line - but this is really just a cosmetic suggestion and not strictly necessary.

loki/transformations/transform_loop.py

MichaelSt98 · 2024-10-02T06:36:56Z

Thanks for reviewing @mlange05. I addressed all your requested changes.
I am thinking about introducing a flag to activate or deactivate that transformation. It is not completely true that this is a no-op unless any pragma markers are inserted. Without pragma markers no loop-interchange and loop-fusion is performed, however, it would still look for temporaries that could be demoted (so in principle this transformation could optimise even without inserted pragmas if some temporaries unnecessarily have a vertical dimension). I don't know how much time it would take to execute this performance for e.g., ecwam or ecphys and/or whether the result is correct. Therefore, I should at least check whether it works with ecwam/ecphys and/or introduce a flag to activate or deactivate. What do you think?

mlange05 · 2024-10-06T11:04:26Z

Ok, nitpicks and code stye comments have all been addressed. Many thanks for this and over to @reuterbal for final approval.

As for the on/off switch, I see you argument, and I think having a safety switch built in is a good idea. Would you mind adding this please to the constructor? I'll check offline if having it enabled affects EC-physics. Once that switch is added, I think this is ready.

On a wider note, I think that once we refactor this a little, we can have the loop-pragma steps as a separate "pre-processing Trafo" that can be added to the pipeline in the config files, rather than be part of SCC. Demotion of the vertical would then simply require additional parameters to the SCCDemote step - but all of this if beyond this initial PR.

reuterbal

Thanks for a very comprehensive and nicely structured implementation and the round of review and fixes already performed! This looks already very good to me. I've left a few (minor) questions and suggestions, and a comment for one potential source of problems.

reuterbal · 2024-10-07T10:32:00Z

loki/transformations/single_column/vertical.py

+        1. loop interchange to expose vertical loops
+        2. fuse vertical loops (possibly into multiple groups)
+        3. find local arrays to be demoted and apply heuristics to check whether this is safe
+        4. demote those arrays which are safe to be demoted


Very nice docstring! 👏

loki/transformations/single_column/vertical.py

loki/transformations/transform_loop.py

reuterbal

Very nice, many thanks!

mlange05 · 2024-10-14T11:31:16Z

Ok, sorry for the delay. I've just run a locally rebased version over this with ecphs-regression, and I can confirm that it works and is performance neutral. I'll trigger one more rebase, and then GTG from me.

…some kind of ordering via potentially loking for 'insert' in the fuse pragma

…hat fuses loops and demotes temporaries in the vertical dimension

…is now necessary due to 'SCCFuseVerticalLoops')

…implemented check whether demotion is safe

…ct routines this transformation is applied to

…r functions defined by the Fortran standard (#390)

mlange05

Looking great now, and confirmed safe and fast with ecphys-regression. GTG.

MichaelSt98 marked this pull request as ready for review September 10, 2024 15:51

MichaelSt98 requested review from reuterbal and mlange05 September 10, 2024 16:26

mlange05 requested changes Sep 28, 2024

View reviewed changes

MichaelSt98 requested a review from mlange05 October 2, 2024 06:29

reuterbal requested changes Oct 7, 2024

View reviewed changes

MichaelSt98 force-pushed the nams-vertical-loop-fusion branch from 13b73b5 to 829e214 Compare October 8, 2024 07:42

reuterbal approved these changes Oct 8, 2024

View reviewed changes

reuterbal added the ready for merge This PR has been approved and is ready to be merged label Oct 8, 2024

MichaelSt98 added 15 commits October 14, 2024 13:31

Loop fusion to account for pragmas when loop was fused and allow for …

780b91b

…some kind of ordering via potentially loking for 'insert' in the fuse pragma

Introduce new transformation 'SCCFuseVerticalLoops' transformation, t…

4ac235a

…hat fuses loops and demotes temporaries in the vertical dimension

Include 'SCCFuseVerticalLoops' in SCC pipelines

94d7668

Pass the vertical dimension to SCC pipelines (since this information …

12c9f03

…is now necessary due to 'SCCFuseVerticalLoops')

[SCCFuseVerticalLoops] bail if vertical dimension is not defined, re-…

72f1164

…implemented check whether demotion is safe

loop_fusion: remove debug print statement

2cde419

Scheduler test: account for new trafo in SCCVectorPipeline

1c4b44b

SCCFuseVerticalLoops: minor refactoring and simplifications

b40ae8a

(vertical) loop fusion: more strongly guard of '-init' ending

6538ac2

loop fusion: use 'insert-loc' instead of 'insert'

6b3bac8

please linter

a9f7287

fix imports: 'FindVariables' has been moved to loki/ir

29e5186

SCCFuseVerticalLoops: introduce new arg 'apply_to' to possibly restri…

b12a33d

…ct routines this transformation is applied to

fix imports (again): 'FindVariables' has been moved to loki/ir

38a0a9b

loop fusion: mention issue regarding provision of type information fo…

7a65ebe

…r functions defined by the Fortran standard (#390)

mlange05 force-pushed the nams-vertical-loop-fusion branch from be63803 to 7a65ebe Compare October 14, 2024 11:31

mlange05 approved these changes Oct 14, 2024

View reviewed changes

reuterbal merged commit b6eb817 into main Oct 14, 2024
13 checks passed

reuterbal deleted the nams-vertical-loop-fusion branch October 14, 2024 14:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Vertical loop fusion and demotion of temporaries #374

Vertical loop fusion and demotion of temporaries #374

MichaelSt98 commented Sep 10, 2024 •

edited

Loading

github-actions bot commented Sep 10, 2024

codecov bot commented Sep 10, 2024 •

edited

Loading

mlange05 left a comment

mlange05 Sep 28, 2024

mlange05 Sep 28, 2024

mlange05 Sep 28, 2024

mlange05 Sep 28, 2024

mlange05 Sep 28, 2024

mlange05 Sep 28, 2024

mlange05 Sep 28, 2024

mlange05 Sep 28, 2024

MichaelSt98 Oct 1, 2024

mlange05 Oct 1, 2024

MichaelSt98 commented Oct 2, 2024

mlange05 commented Oct 6, 2024

reuterbal left a comment

reuterbal Oct 7, 2024

reuterbal left a comment

mlange05 commented Oct 14, 2024

mlange05 left a comment

Vertical loop fusion and demotion of temporaries #374

Vertical loop fusion and demotion of temporaries #374

Conversation

MichaelSt98 commented Sep 10, 2024 • edited Loading

github-actions bot commented Sep 10, 2024

codecov bot commented Sep 10, 2024 • edited Loading

Codecov Report

mlange05 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

MichaelSt98 commented Oct 2, 2024

mlange05 commented Oct 6, 2024

reuterbal left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

reuterbal left a comment

Choose a reason for hiding this comment

mlange05 commented Oct 14, 2024

mlange05 left a comment

Choose a reason for hiding this comment

MichaelSt98 commented Sep 10, 2024 •

edited

Loading

codecov bot commented Sep 10, 2024 •

edited

Loading