DataOffload: Add new FieldOffloadTransformation for FIELD API boilerplate injection #437

wertysas · 2024-11-11T12:47:39Z

This adds a new transformation, FieldOffloadTransformation, for offloading F-API fields to GPU. Specifically this transformation is meant to operate on CPU driver code that passes view pointers to the kernels. The transformation

Removes the view update calls
Adds new device pointers for the arrays to offload
Adds F-API data offload and sync calls
Replaces the old view-based arguments in kernel calls with device pointer slices

The corresponding dwarf-p-cloudsc implementation is on the branch dwarf-p-cloudsc/je-field-api-offload-v2. This adds a new transformation pipeline for converting view based CPU driver code to GPU offloaded F-API code, [pipelines.scc-field], and a new Loki target dwarf-cloudsc-loki-scc-field.

github-actions · 2024-11-11T12:49:41Z

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

…allStatement pre_init method

…pi.py

codecov · 2024-11-15T13:09:09Z

Codecov Report

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

Project coverage is 93.31%. Comparing base (00d4290) to head (ac0aa75).
Report is 25 commits behind head on main.

Files with missing lines	Patch %	Lines
loki/transformations/data_offload.py	98.49%	2 Missing ⚠️
loki/transformations/parallel/field_api.py	96.42%	1 Missing ⚠️
loki/transformations/tests/test_data_offload.py	99.32%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #437      +/-   ##
==========================================
+ Coverage   93.23%   93.31%   +0.08%     
==========================================
  Files         212      212              
  Lines       40259    40728     +469     
==========================================
+ Hits        37535    38007     +472     
+ Misses       2724     2721       -3

Flag	Coverage Δ
lint_rules	`96.39% <ø> (ø)`
loki	`93.27% <98.82%> (+0.08%)`	⬆️

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.

🚨 Try these New Features:

Flaky Tests Detection - Detect and resolve failed and flaky tests

mlange05

Excellent contribution and very nicely structured and test; many thanks! This is a great start towards a very comprehensive set of tools for dealing with FIELD API boilerplate.

I've left a bunch of cosmetic comments (mostly docstrings, and minor style things), and a few notes to myself. The reason for the latter is that this one does not yet cover more general cases with compute code in the driver loop, and it will also likely require some restructuring for the utilities to be used outside the DataOffloadTransformation - but all of these will happen in a follow-on consolidation effort and should not stop this.

loki/transformations/data_offload.py

loki/transformations/parallel/field_api.py

mlange05 · 2024-11-19T08:20:14Z

loki/transformations/tests/test_data_offload.py

+    return Module.from_source(fcode, frontend=frontend, xmods=[tmp_path])
+
+@pytest.fixture(name="field_module")
+def fixture_field_module(tmp_path, frontend):


Yeah! It was a good suggestion to use a fixture for this to reduce boilerplate in the tests.

loki/transformations/tests/test_data_offload.py

mlange05 · 2024-11-19T15:59:03Z

loki/transformations/data_offload.py

+
+    def __init__(self, devptr_prefix=None, field_group_types=None, offload_index=None):
+        self.deviceptr_prefix = 'loki_devptr_' if devptr_prefix is None else devptr_prefix
+        field_group_types = [''] if field_group_types is None else field_group_types


[no action] For future use, as_tuple is an internal utility that we use for this pattern everywhere.

mlange05

Very nice! And many thanks for addressing this so promptly. GTG from me.

…kwargs to named kwargs in FieldOffload init

awnawab

Many thanks for this superb piece of work 👌 The implementation is really neat and well tested 🙏 I especially love your dedication to type hints (even though I am too lazy to adhere to it myself 😇 ).

I've left a few comments, mostly reminders for us/myself for the future but there are a couple which I would like you to please address.

awnawab · 2024-11-19T16:32:38Z

loki/transformations/data_offload.py

@@ -49,6 +52,12 @@ def __init__(self, **kwargs):
        self.has_data_regions = False
        self.remove_openmp = kwargs.get('remove_openmp', False)
        self.assume_deviceptr = kwargs.get('assume_deviceptr', False)
+        self.assume_acc_mapped = kwargs.get('assume_acc_mapped', False)


The name here (and corresponding control flow) is a little confusing. If this option is not enabled, and the offload instructions are instrumented via DataOffloadTransformation as per usual, they would still be acc mapped. I suggest we use two options, present_on_device and assume_deviceptr. The control flow would look like this:

if self.present_on_device (or self.assume_deviceptr): # the "or" here is only needed if this suggested change messily breaks backwards compatibility, # resolving which shouldn't hold back this PR if self.assume_deviceptr: # add deviceptr clause else: # add present clause else: # add copy/copyin/copyout clause

The option should also be documented, it's missing from the "parameters".

I agree, updated it according to your suggestion now.

awnawab · 2024-11-19T16:35:21Z

loki/transformations/data_offload.py

                if self.assume_deviceptr:
                    offload_args = inargs + outargs + inoutargs
                    if offload_args:
                        deviceptr = f' deviceptr({", ".join(offload_args)})'
                    else:
                        deviceptr = ''
                    pragma = Pragma(keyword='acc', content=f'data{deviceptr}')
+                elif self.assume_acc_mapped:


In the short term this should be refactored using the suggestion above. In the long term, I think the arrays that are already present on device via a previous transformation should be in a trafo_data entry, and DataOffloadTransformation would then just instrument the offload for the remaining arrays. This would keep the current transformation more general.

awnawab · 2024-11-19T16:55:16Z

loki/transformations/parallel/field_api.py

+
+
+
+class FieldAPITransferType(Enum):


Very nice 👌

awnawab · 2024-11-19T20:21:15Z

loki/transformations/data_offload.py

+                if not isinstance(param, Array):
+                    continue
+                try:
+                    parent = arg.parent


Couldn't we do this with a parent = arg.getattr('parent', None) and print the warning if not parent?

Yes, but I left it as is now, since getattr does something similar under the hood.

awnawab · 2024-11-19T20:23:06Z

loki/transformations/data_offload.py

+        driver.variables += device_ptrs
+        return device_ptrs
+
+    def _devptr_from_array(self, driver, a: sym.Array):


Really love the type hints ❤️

awnawab · 2024-11-19T20:31:02Z

loki/transformations/data_offload.py

+
+    def _get_field_ptr_from_view(self, field_view):
+        type_chain = field_view.name.split('%')
+        field_type_name = 'F_' + type_chain[-1]


[noaction] this is a sensible default, which we might want to make overridable per derived-type.

awnawab · 2024-11-19T20:35:20Z

loki/transformations/data_offload.py

+        change_map = {}
+        offload_idx_expr = driver.variable_map[self.offload_index]
+        for arg, devptr in chain(offload_map.in_pairs, offload_map.inout_pairs, offload_map.out_pairs):
+            dims = (sym.RangeIndex((None, None)),) * (len(devptr.shape)-1) + (offload_idx_expr,)


Should we not clone arg.dims here and add the offload_idx_expr to it? We don't always want to pass the full array, we might for example pass FIELD_PTR(:,1,IBL) if the field is 3D but the dummy argument is 2D. These are probably mistakes that we should fix in source, but the transformation should nevertheless support this edge-case.

Nice catch! I had missed that. Fixed it now.

awnawab · 2024-11-19T20:37:17Z

loki/transformations/parallel/field_api.py

+        Scope of the created :any:`CallStatement`
+    """
+
+    procedure_name = 'SYNC_HOST_RDWR'


[noaction] we might want to mirror the suffix logic of field_get_device_data here.

…d requirement of present_on_device to assume_deviceptr

awnawab

Thanks a lot for addressing the changes so promptly 🙏 looks great 👌

wertysas requested review from mlange05 and awnawab November 11, 2024 13:26

wertysas added 7 commits November 15, 2024 12:52

Added new transformation for GPU data offloading with FIELD API

8530953

started adding tests for field offload functionality

60b0c9c

Added named arguments when creating Callstmts to work around bug in C…

da415b5

…allStatement pre_init method

Moved F-API helper functions from data_offload.py to parallel/field_a…

76765d1

…pi.py

Added tests for FieldOffloadTransformation

87850a1

Removed automatic inferral of offload indexes from view updates

742795d

Removed unused find_array_arg utility and FieldOffload clean

ff5ddee

wertysas force-pushed the je-field-api-offload-index-from-config branch from 4acd385 to ff5ddee Compare November 15, 2024 12:39

wertysas added 2 commits November 15, 2024 16:18

clean and linter error fixes

b7f0512

Removed unused functions and added more complete tests of data offloads

7e53319

wertysas force-pushed the je-field-api-offload-index-from-config branch from 6b9b2e4 to 7e53319 Compare November 18, 2024 16:31

wertysas marked this pull request as ready for review November 18, 2024 19:14

mlange05 requested changes Nov 19, 2024

View reviewed changes

mlange05 changed the title ~~Je field api offload index from config~~ DataOffload: Add new FieldOffloadTransformation for FIELD API boilerplate injection Nov 19, 2024

mlange05 reviewed Nov 19, 2024

View reviewed changes

mlange05 approved these changes Nov 19, 2024

View reviewed changes

Added missing docstrings, updated caplog usage in tests, and changed …

5ee7c8c

…kwargs to named kwargs in FieldOffload init

wertysas force-pushed the je-field-api-offload-index-from-config branch from e4010ba to 5ee7c8c Compare November 19, 2024 16:13

awnawab requested changes Nov 19, 2024

View reviewed changes

wertysas added 2 commits November 20, 2024 10:32

Data Offload: renamed assume_acc_mapped to present_on_device and adde…

67d1a30

…d requirement of present_on_device to assume_deviceptr

Field Offload, added missing support for sliced arguments to calls

ac0aa75

wertysas force-pushed the je-field-api-offload-index-from-config branch from 45c9e20 to ac0aa75 Compare November 20, 2024 12:56

wertysas requested a review from awnawab November 20, 2024 13:45

awnawab approved these changes Nov 20, 2024

View reviewed changes

mlange05 added the ready for merge This PR has been approved and is ready to be merged label Nov 20, 2024

mlange05 merged commit ebc7b69 into main Nov 20, 2024
13 checks passed

mlange05 deleted the je-field-api-offload-index-from-config branch November 20, 2024 15:46

wertysas mentioned this pull request Nov 21, 2024

New Field API Loki variant with state types in driver ecmwf-ifs/dwarf-p-cloudsc#100

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

DataOffload: Add new FieldOffloadTransformation for FIELD API boilerplate injection #437

DataOffload: Add new FieldOffloadTransformation for FIELD API boilerplate injection #437

wertysas commented Nov 11, 2024 •

edited

Loading

github-actions bot commented Nov 11, 2024

codecov bot commented Nov 15, 2024 •

edited

Loading

mlange05 left a comment

mlange05 Nov 19, 2024

wertysas Nov 20, 2024

mlange05 Nov 19, 2024

mlange05 left a comment

awnawab left a comment

awnawab Nov 19, 2024

awnawab Nov 19, 2024

wertysas Nov 20, 2024

awnawab Nov 19, 2024

awnawab Nov 19, 2024

awnawab Nov 19, 2024

wertysas Nov 20, 2024

awnawab Nov 19, 2024

awnawab Nov 19, 2024

awnawab Nov 19, 2024

wertysas Nov 20, 2024

awnawab Nov 19, 2024

awnawab left a comment




		class FieldAPITransferType(Enum):

DataOffload: Add new FieldOffloadTransformation for FIELD API boilerplate injection #437

DataOffload: Add new FieldOffloadTransformation for FIELD API boilerplate injection #437

Conversation

wertysas commented Nov 11, 2024 • edited Loading

github-actions bot commented Nov 11, 2024

codecov bot commented Nov 15, 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

mlange05 left a comment

Choose a reason for hiding this comment

awnawab 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

awnawab left a comment

Choose a reason for hiding this comment

wertysas commented Nov 11, 2024 •

edited

Loading

codecov bot commented Nov 15, 2024 •

edited

Loading