DataOffload: Field offload of driver regions #457

mlange05 · 2024-12-05T19:37:25Z

This PR refactors the current Field-API offload transformation and its utilities and adds two previously missing features:

The ability to detect aliased array pointer usage, so CALL F(obj%a(:, 1), obj%a(:, 2)) only declares and offloads obj_a once.
The ability to run FieldOffloadTransformation over compute regions in drivers without kernel calls

Moreover, it turns I've started moving things around a little and ended up putting most of the heavy lifting in the FieldPointerMap utility. To avoid future cyclic import and other issues, I've then moved this out to a pure helper sub-module loki.transformations.field_api and renamed the using transformations data_offload.field_offload and parallel.field_views.

In refactoring the existing transformation and utilities to achieve the above, I moved quite a few things around - sorry for the large diff. I'd be happy to split this up, if this will aid the review process.

In a little more detail, the other changes in this PR are:

Refactoring Field API offload tests to use a common fixture for state_type, where appropriate
Filtering duplicates at alls stages via dict.fromkeys instead of set, as this retains ordering and thus ensures reproducible orders of variables/utility calls at all stages
The FieldPointerMap not converts from field view array symbols to general data pointers (with block dimension) on-the-fly. This allows the data pointers to be given as a mere property on the pointer map object.
Similarly the FieldPointerMap now provides the lists of boilerplate calls for host-to-device copies and host-sync calls as properties and derives them from the three internally stored argument/array symbol lists.
Apply symbol substitution over all array symbols that match the pointer/array symbols stored in the field pointer map over an entire region, not just calls
Derive the array symbols for which to generate Field API offload semantics via dataflow analysis, so that we can run this over any marked pragma region

Sets are unordered, so using them to filter creates an arbitrary ordering, which in turn yields an unpredictable order of declarations.

This allows us to re-use the key pieces without being hooked into the `!$loki data` regions semantics.

Instead of subbing just on calls, we apply the remapping over the whole routine body.

Without tying the transformation to calls, explicit no-target skipping becomes virtually impossible; hence removing the test for it.

github-actions · 2024-12-05T19:40:22Z

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

codecov · 2024-12-05T20:07:05Z

Codecov Report

Attention: Patch coverage is 99.06977% with 2 lines in your changes missing coverage. Please review.

Project coverage is 93.28%. Comparing base (e9760e8) to head (9d572ee).

Files with missing lines	Patch %	Lines
loki/transformations/data_offload/field_offload.py	98.68%	1 Missing ⚠️
loki/transformations/field_api.py	98.52%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #457   +/-   ##
=======================================
  Coverage   93.28%   93.28%           
=======================================
  Files         220      221    +1     
  Lines       41224    41249   +25     
=======================================
+ Hits        38457    38481   +24     
- Misses       2767     2768    +1

Flag	Coverage Δ
lint_rules	`96.39% <ø> (ø)`
loki	`93.24% <99.06%> (+<0.01%)`	⬆️

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.

wertysas

It's very nice to get these extensions of the field offload capabilities. It makes a lot of sense to separate the field utilities and offload transformations like this and I think it's very nice to have the offload map tracking the offload variables like this👍.

wertysas · 2024-12-06T10:36:10Z

loki/transformations/field_api.py

+# nor does it submit to any jurisdiction.
+
+"""
+A set utility classes for dealing with FIELD API boilerplate in


Suggested change

A set utility classes for dealing with FIELD API boilerplate in

A set of utility classes for dealing with FIELD API boilerplate in

wertysas · 2024-12-06T14:50:10Z

loki/transformations/field_api.py

This is much cleaner now with the separation of the general f-api utilities from the offloading transformations.

wertysas · 2024-12-06T16:04:23Z

loki/transformations/field_api.py

+        """
+        Returns a tuple of :any:`CallStatement` for host-to-device transfers on fields.
+        """
+        READ_ONLY, READ_WRITE = FieldAPITransferType.READ_ONLY, FieldAPITransferType.READ_WRITE


This is neat

wertysas · 2024-12-06T16:28:25Z

loki/transformations/data_offload/field_offload.py

+                    replace_kernel_args(driver, offload_map, self.offload_index)
+
+
+def find_offload_variables(driver, region, field_group_types):


Perhaps this can be split into two functions in the future. One more general utility in transformations/data_offload/offload.py that finds the inargs/inoutargs/outargs using dataflow analysis. And the logic of the second method could even be a helper function here or a method of the FieldOffloadTransformation.

In the meantime I think it would be good to have a docstring that specifies that this has to be run with dataflow analysis attached.

wertysas · 2024-12-06T16:32:41Z

loki/transformations/data_offload/field_offload.py

+    return inargs, inoutargs, outargs
+
+
+def declare_device_ptrs(driver, deviceptrs):


I didn't realise this when I wrote this helper, but there is nothing field or device ptr specific about this function. I suggest we name it something else, perhaps add_variables, and add it as member function of the ProgramUnit base class.

wertysas · 2024-12-09T11:06:12Z

loki/transformations/data_offload/field_offload.py

+    change_map = {}
+    offload_idx_expr = driver.variable_map[offload_index]
+
+    args = tuple(chain(offload_map.inargs, offload_map.inoutargs, offload_map.outargs))


I think it would be useful to expose this chained tuple as a property, args, of the offload map.

wertysas · 2024-12-09T11:08:32Z

loki/transformations/field_api.py

+        return field_view.parent.get_derived_type_member(field_type_name)
+
+    @property
+    def dataptrs(self):


The same comment with the name applies here.

wertysas · 2024-12-09T12:38:59Z

loki/transformations/field_api.py

+    Helper class to map FIELD API pointers to intents and access descriptors.
+
+    This utility is used to store arrays passed to target kernel calls
+    and the corresponding device pointers added by the transformation.


Suggested change

and the corresponding device pointers added by the transformation.

and easily access corresponding device pointers added by the transformation.

wertysas · 2024-12-09T12:55:25Z

loki/transformations/field_api.py

+        return tuple(dict.fromkeys(
+            self.dataptr_from_array(a)
+            for a in chain(*(self.inargs, self.inoutargs, self.outargs))
+        ))


Suggested change

return tuple(dict.fromkeys(

self.dataptr_from_array(a)

for a in chain(*(self.inargs, self.inoutargs, self.outargs))

))

return tuple(

self.dataptr_from_array(a)

for a in chain(self.inargs, self.inoutargs, self.outargs)

)

If there are no duplicates in self.inargs/inoutargs/outargs, then there shouldn't be any in the ones created either.

wertysas · 2024-12-09T13:15:03Z

loki/transformations/field_api.py

+
+    The pointer/array variable pairs are exposed through the class
+    properties, based on the intent of the kernel argument.


Suggested change

The pointer/array variable pairs are exposed through the class

properties, based on the intent of the kernel argument.

mlange05 added 13 commits December 3, 2024 13:02

DataOffload: Move FieldPointerMap outside of transformer class

595e048

DataOffload: Filter duplicates with dict to retain ordering

d7d4308

Sets are unordered, so using them to filter creates an arbitrary ordering, which in turn yields an unpredictable order of declarations.

DataOffload: Expose the core functionalities of FieldOffloadTrafo

7713209

This allows us to re-use the key pieces without being hooked into the `!$loki data` regions semantics.

DataOffload: Use single state_module fixture in field api tests

f4fddd3

DataOffload: Add test for aliasing argument pointers in kernel calls

3c1bb91

FieldAPI: Move utility classes and methods to dedicated sub-package

214e2ff

FieldAPI: Rename Field API modules in parallel and data_offload

d5a0667

FieldAPI: Internalise data pointer generation to FieldPointerMapp

752fbcb

FieldAPI: Internalise data transfer call generation and drop pairs

a4a446e

FieldAPI: Filter duplicates in FieldPointerMap when generating calls

33429be

DataOffload: Apply symbol substitution over driver in field offload

7bfbf38

Instead of subbing just on calls, we apply the remapping over the whole routine body.

DataOffload: Derive offload variables via dataflow analysis for regions

cce8b93

DataOffload: Tidy up imports and remove obsolete utility and test

9d572ee

Without tying the transformation to calls, explicit no-target skipping becomes virtually impossible; hence removing the test for it.

mlange05 requested review from wertysas and awnawab December 5, 2024 19:37

mlange05 marked this pull request as ready for review December 6, 2024 03:25

wertysas requested changes Dec 9, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

DataOffload: Field offload of driver regions #457

DataOffload: Field offload of driver regions #457

mlange05 commented Dec 5, 2024

github-actions bot commented Dec 5, 2024

codecov bot commented Dec 5, 2024 •

edited

Loading

wertysas left a comment

wertysas Dec 6, 2024

wertysas Dec 6, 2024

wertysas Dec 6, 2024

wertysas Dec 6, 2024

wertysas Dec 6, 2024

wertysas Dec 9, 2024

wertysas Dec 9, 2024

wertysas Dec 9, 2024

wertysas Dec 9, 2024

wertysas Dec 9, 2024

	A set utility classes for dealing with FIELD API boilerplate in
	A set of utility classes for dealing with FIELD API boilerplate in

		replace_kernel_args(driver, offload_map, self.offload_index)


		def find_offload_variables(driver, region, field_group_types):

		return inargs, inoutargs, outargs


		def declare_device_ptrs(driver, deviceptrs):

	and the corresponding device pointers added by the transformation.
	and easily access corresponding device pointers added by the transformation.


		The pointer/array variable pairs are exposed through the class
		properties, based on the intent of the kernel argument.

DataOffload: Field offload of driver regions #457

Are you sure you want to change the base?

DataOffload: Field offload of driver regions #457

Conversation

mlange05 commented Dec 5, 2024

github-actions bot commented Dec 5, 2024

codecov bot commented Dec 5, 2024 • edited Loading

Codecov Report

wertysas 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

codecov bot commented Dec 5, 2024 •

edited

Loading