feat[next][dace]: GTIR-to-DaCe lowering of map-reduce (only full connectivity) #1683

edopao · 2024-10-09T07:19:21Z

This PR adds support for lowering of map_ and make_const_list builtin functions. However, the current implementation only supports neighbor tables with full connectivity (no skip values). The support for skip values will be added in next PR.

To be noted:

This PR generalizes the handling of tasklets without arguments inside a map scope. The return type for input_connections is extended to contain a TaskletConnection variant, which is lowered to an empty edge from map entry node to the tasklet node.
The result of make_const_list is a scalar value to be broadcasted on a local field. However, in order to keep the lowering simple, this value is represented as a 1D 1-element array (shape=(1,)).

…map_reduce

philip-paul-mueller

There are some points that needs further clarifications but it does not look that bad.
There are also some suggestions that are not directly related to the PR per se.

philip-paul-mueller · 2024-10-09T07:41:40Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

-            [gtx_common.Dimension("")] if isinstance(arg.data_type.dtype, gtir_ts.ListType) else []
+            [gtx_common.Dimension("")] if isinstance(arg.data_type.dtype, itir_ts.ListType) else []


Could it happen that you have two anonymous dimensions in a field?
I mean something like Field(KDim, "", EdgeDim, "") wouldn't this be a problem.
I would propose to give these dimensions a name like __anonymous_dim_{UNIQUE_STUFF}.

In theory it can happen, but the current IR does not support it. I agree that we should give it a name, not just leave an empty string. However, the name of dimension does not have to be unique. My idea is to use the offset name (e.g. V2E). So you could have something like Field(KDim, "V2E", EdgeDim, "V2E"), and this should not be a problem (I would have to ensure that map indices are unique, though). This will be addressed in next iteration (I was thinking about a separate PR).

Okay it is a bit off topic, but Rule 11 implies that the name of a dimensions and the iteration variable name is liked.

I think it is not a problem, but we should probably clarify the rule a bit.

philip-paul-mueller · 2024-10-09T07:44:03Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

        _, _, reduce_identity = gtir_to_tasklet.get_reduce_params(stencil_expr.expr)
+        reduce_identity_for_args = reduce_identity


Suggested change

_, _, reduce_identity = gtir_to_tasklet.get_reduce_params(stencil_expr.expr)

reduce_identity_for_args = reduce_identity

_, _, reduce_identity_for_args = gtir_to_tasklet.get_reduce_params(stencil_expr.expr)

I also need to set reduce_identity, because it is passed to LambdaToTasklet to lower the current node. Note that reduce_identity_for_args and reduce_identity could be different (e.g. the first None, the latter not None). That is what I tried to explain below in the branch for neighbors nodes: we do not support nested reduction, so the reduction identity value is used (consumed) by the current neighbors expression to fill the skip values.

Okay, I would make this aspect more clearer.

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_to_tasklet.py

.../next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py

edopao · 2024-10-09T16:17:36Z

Thank you @philip-paul-mueller for the review comments. I have pushed a commit that should clarify how InputConnection was used. I will address the remaining comments tomorrow.

Implements some refactoring changes discussed during review of #1683.

…-gtir-map_reduce

philip-paul-mueller

There is nothing serious, but some aspects need a bit more work.
But I expect that we can merge in the next round.

src/gt4py/next/iterator/transforms/fuse_maps.py

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

src/gt4py/next/program_processors/runners/dace_fieldview/utility.py

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

.../next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py

philip-paul-mueller

There are some things that needs some discussion but they are rather minor points.

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

philip-paul-mueller · 2024-10-15T06:20:03Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

+    # list (see the second if-branch), we stop carrying the reduce identity further
+    # (`reduce_identity_for_args = None`) because it is not needed aymore;
+    # the reason being that reduce operates along a single axis, which in this case
+    # corresponds to the local dimension of the list of neighbor values.


Excuse my ignorance and slow understanding, but why, you still have to initialize the reduction.

You are right, I do not need to initialize reduce_identity in this case.

I think the whole problem is that there are different reduce identities:

self.reduce_identity

reduce_identity_for_args

I do not get why they are split.
Also I do not fully understand the explanation.
I think it is that if you hit the neighbor built in, then you would reduce over the list of neighbor hood values of a point, and because you are currently ignoring skip values, i.e. assume that all are defined, you do not have to fuill something.
So technically speaking it is not needed but not wrong to do.
If there is no reason to threat is as special case then I would not treat it as special case.
Except you want to do some prep work for the skip value case, but then you have to clearly indicate this.

self.reduce_identity is defined in LambdaToDataflow, not here. Here reduce_identity is an input to the function. By visiting the node arguments with reduce_identity_for_args = None, their reduce_identity will be None.

Setting reduce_identity_for_args = None will also work in case of skip values. The reduce_identity value is filled in in place of skip values in the context of neighbors itself, not in the arguments context.

You are right that it would not be wrong to set reduce_identity_for_args also in case of neighbors arguments. I do it because it enables a sanity check that the sequence reduce(V2E) -> neighbors(V2E) -> reduce(C2E) -> neighbors(C2E) is accepted, while the sequence reduce(V2E) -> reduce(C2E) -> neighbors(V2E) -> neighbors(C2E) is not. The latter sequence would raise the NotImplementedError("nested reductions not supported.") exception.

Side note. There is a simple alternative design where I can avoid passing reduce_identity around. It would be to write a dummy value for skip values instead of doing array access (for neighbors) or executing the map operation (for map - I cannot safely apply the map operation in case for example of division by 0). Then, in the lowering of reduce I would put -- before the reduce library node -- another local map that fills the skip values with the reduce identity. I really like this approach (separate PR) but it will result in one extra write - replacing the dummy value with the reduce identity.

I would add part of the explanation you write here as comments.
Because, you get from my question that this is really hard to understand.

philip-paul-mueller · 2024-10-15T07:12:29Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

-            [gtx_common.Dimension("")] if isinstance(arg.data_type.dtype, gtir_ts.ListType) else []
+            [gtx_common.Dimension("")] if isinstance(arg.data_type.dtype, itir_ts.ListType) else []


Okay it is a bit off topic, but Rule 11 implies that the name of a dimensions and the iteration variable name is liked.

I think it is not a problem, but we should probably clarify the rule a bit.

philip-paul-mueller · 2024-10-15T07:13:47Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

        _, _, reduce_identity = gtir_to_tasklet.get_reduce_params(stencil_expr.expr)
+        reduce_identity_for_args = reduce_identity


Okay, I would make this aspect more clearer.

philip-paul-mueller · 2024-10-15T07:14:57Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_python_codegen.py

+def make_const_list(arg: str) -> str:
+    return arg


So you can say it is like a "lazy fill"?

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

philip-paul-mueller

There are some smallish things.
But It generally looks good.

philip-paul-mueller · 2024-10-16T06:20:28Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

+    # list (see the second if-branch), we stop carrying the reduce identity further
+    # (`reduce_identity_for_args = None`) because it is not needed aymore;
+    # the reason being that reduce operates along a single axis, which in this case
+    # corresponds to the local dimension of the list of neighbor values.


I think the whole problem is that there are different reduce identities:

self.reduce_identity

reduce_identity_for_args

I do not get why they are split.
Also I do not fully understand the explanation.
I think it is that if you hit the neighbor built in, then you would reduce over the list of neighbor hood values of a point, and because you are currently ignoring skip values, i.e. assume that all are defined, you do not have to fuill something.
So technically speaking it is not needed but not wrong to do.
If there is no reason to threat is as special case then I would not treat it as special case.
Except you want to do some prep work for the skip value case, but then you have to clearly indicate this.

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

philip-paul-mueller · 2024-10-16T06:33:33Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py

@@ -145,7 +145,7 @@ class DataflowOutputEdge:
    """

    state: dace.SDFGState
-    expr: DataExpr


result is probably a better name as expr, but I suggest you to make it more consistent with the other two classes especially MemletInputEdge and call it destination.

It is not really a destination, it is rather the sink node of the dataflow. What about sink?

sink has a meaning in DaCe as "node with zero output degree", so it is a bit misleading, how about target?

In this case, expr is a sink node according to the dace definition. A data node the dataflow writes to.

but will it never be read from?

The sink node (typically a scalar or a local field) of the dataflow will actually be removed from the SDFG. The edge writing to the sink node will be replaced with an edge passing through the map exit node and writing to a new node, the transient field node (a multi-dimensional array with shape equal to the field operator domain).

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py

philip-paul-mueller · 2024-10-16T06:40:20Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py

+        The `plus` operation is lowered to a tasklet inside a map that computes
+        the domain of the local dimension (in this example, the number of neighbors).
+
+        The result is a 1D local field, with same size as the input local dimension.


I would also mention what the value is.
Now that I think about it, what is the size of the local field anyway.
I would say it has size 1, because $V2E_{0}$ is only one element?

The subscript is the letter 'O' and it stands for offset. I will remove it from the example in order to avoid confusion. In this example, the local size will be V2E.max_neighbors.

Again I would add what value the result has.

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py

philip-paul-mueller

I have some strong suggestions to improve the documentation of the implementation and behaviour of internals.
But the PR kind of looks good, but some documenting aspects should be improved.

philip-paul-mueller · 2024-10-16T10:56:27Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

+    elif cpm.is_call_to(stencil_expr.expr, "neighbors"):
+        reduce_identity_for_args = None


I would add the explanation you just gave me here.

philip-paul-mueller · 2024-10-16T10:57:16Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

+    # list (see the second if-branch), we stop carrying the reduce identity further
+    # (`reduce_identity_for_args = None`) because it is not needed aymore;
+    # the reason being that reduce operates along a single axis, which in this case
+    # corresponds to the local dimension of the list of neighbor values.


I would add part of the explanation you write here as comments.
Because, you get from my question that this is really hard to understand.

philip-paul-mueller · 2024-10-16T10:58:37Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py

@@ -566,7 +640,7 @@ def translate_symbol_ref(
 if TYPE_CHECKING:
    # Use type-checking to assert that all translator functions implement the `PrimitiveTranslator` protocol
    __primitive_translators: list[PrimitiveTranslator] = [


This conditional declaration does not make any sense to me.
What does it give to you what a good unit test does not give you already?

philip-paul-mueller · 2024-10-16T11:00:01Z

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py

+        The `plus` operation is lowered to a tasklet inside a map that computes
+        the domain of the local dimension (in this example, the number of neighbors).
+
+        The result is a 1D local field, with same size as the input local dimension.


Again I would add what value the result has.

#1694) This PR extends the solution for map-reduce provided in #1683 with the support for connectivity tables with skip values. The field definition is extended with a `local_offset` attribute that stores the offset provider used to build the values in the local dimension. In case the local dimension is built by the `neighbors` expression, the `local_offset` corresponds to the offset provider used to access the neighbor dimension. Since this information is carried along the data itself, whenever the data is accessed it is also possible to access the corresponding offset provider and check whether the neighbor index is valid or if there is a skip value. For local dimensions already present in the program argument, this information is retrieved from the field domain (enabled in new test case). The data is accessed in the `map_` and `reduce` expressions. Here it is now possible to check for skip values. Therefore, the main objective of this PR is the lowering of map-reduce with skip values. A secondary objective is to pave the road to simplify the lowering logic, by getting rid of the `reduce_identity` value. The current approach is propagate the `reduce_identity` value while visiting the arguments to `reduce` expressions. By introducing `local_offset`, the argument visitor will return the information needed to implement `reduce` of local values in presence of skip values.

edopao added 16 commits October 3, 2024 12:22

map-reduce for full connectivity

0272158

Merge remote-tracking branch 'origin/main' into dace-gtir-map_reduce

e62bae9

Use correct map index variable

43df0e8

enable test case

4b35572

add support for make_const_list

e8e37a4

Misc fixes for GTIR lowering

a885de2

add test case

095e28e

minor edit

8fc95a8

undo extra change

cca4fcd

review comments

1bb20cd

Merge remote-tracking branch 'origin/dace-gtir-fixes' into dace-gtir-…

be06e9f

…map_reduce

Merge remote-tracking branch 'origin/main' into dace-gtir-map_reduce

049b044

minor edit

4241b23

minor edit (1)

49f7797

remove extra change

892bca4

minor edit (2)

fa876c9

edopao marked this pull request as ready for review October 9, 2024 07:27

edopao requested a review from philip-paul-mueller October 9, 2024 07:35

philip-paul-mueller reviewed Oct 9, 2024

View reviewed changes

edopao added 2 commits October 9, 2024 17:40

fix return type of dace backend

ed7fd6e

Improved handling of map input edges for fieldop

f7d84ba

edopao added 4 commits October 9, 2024 22:26

Improved handling of map output for fieldop

c90b1cc

review comments

7f247fe

refactoring

0f897de

refactoring

397cede

edopao mentioned this pull request Oct 10, 2024

feat[next][dace]: refactoring #1685

Merged

edopao added 3 commits October 10, 2024 09:07

refactoring (1)

06fd06f

review comments

feb0951

code-style fix

9c9e7f8

undo refactoring changes

92bc1a8

edopao added a commit that referenced this pull request Oct 10, 2024

feat[next][dace]: refactoring (#1685)

92e83f3

Implements some refactoring changes discussed during review of #1683.

edopao added 5 commits October 10, 2024 14:45

Merge remote-tracking branch 'origin/dace-gtir-refactoring' into dace…

7b8155d

…-gtir-map_reduce

Merge remote-tracking branch 'origin/main' into dace-gtir-map_reduce

684e607

improve code comments

1265f93

better lowering of fieldop domain

ee972a6

rename field_op to fieldop

171288b

edopao requested a review from philip-paul-mueller October 11, 2024 12:23

philip-paul-mueller reviewed Oct 11, 2024

View reviewed changes

review comments

4334621

edopao requested a review from philip-paul-mueller October 14, 2024 10:56

philip-paul-mueller reviewed Oct 15, 2024

View reviewed changes

edopao added 2 commits October 15, 2024 09:40

minor code refactoring

b468e8a

review comments

54f5f65

edopao requested a review from philip-paul-mueller October 15, 2024 09:33

use output_connector on map tasklet

7061697

philip-paul-mueller reviewed Oct 16, 2024

View reviewed changes

review comments

e71c283

philip-paul-mueller approved these changes Oct 16, 2024

View reviewed changes

edopao added 3 commits October 16, 2024 13:36

edit code comments

083f6fe

Merge branch 'main' into dace-gtir-map_reduce

67ec6f9

extract helper method _construct_local_view

2afd8ab

edopao merged commit 0a27c7a into GridTools:main Oct 17, 2024
31 checks passed

edopao deleted the dace-gtir-map_reduce branch October 17, 2024 14:08

edopao mentioned this pull request Oct 17, 2024

feat[next][dace]: GTIR-to-DaCe lowering of map-reduce with skip values #1694

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat[next][dace]: GTIR-to-DaCe lowering of map-reduce (only full connectivity) #1683

feat[next][dace]: GTIR-to-DaCe lowering of map-reduce (only full connectivity) #1683

edopao commented Oct 9, 2024 •

edited

Loading

philip-paul-mueller left a comment

philip-paul-mueller Oct 9, 2024

edopao Oct 9, 2024 •

edited

Loading

philip-paul-mueller Oct 15, 2024

philip-paul-mueller Oct 9, 2024

edopao Oct 9, 2024

philip-paul-mueller Oct 15, 2024

edopao commented Oct 9, 2024

philip-paul-mueller left a comment

philip-paul-mueller left a comment

philip-paul-mueller Oct 15, 2024

edopao Oct 15, 2024

philip-paul-mueller Oct 16, 2024

edopao Oct 16, 2024

philip-paul-mueller Oct 16, 2024

philip-paul-mueller Oct 15, 2024

philip-paul-mueller Oct 15, 2024

philip-paul-mueller Oct 15, 2024

philip-paul-mueller left a comment

philip-paul-mueller Oct 16, 2024

philip-paul-mueller Oct 16, 2024

edopao Oct 16, 2024

philip-paul-mueller Oct 16, 2024

edopao Oct 16, 2024

philip-paul-mueller Oct 16, 2024

edopao Oct 16, 2024

philip-paul-mueller Oct 16, 2024

edopao Oct 16, 2024

philip-paul-mueller Oct 16, 2024

philip-paul-mueller left a comment

philip-paul-mueller Oct 16, 2024

philip-paul-mueller Oct 16, 2024

philip-paul-mueller Oct 16, 2024

philip-paul-mueller Oct 16, 2024

		[gtx_common.Dimension("")] if isinstance(arg.data_type.dtype, gtir_ts.ListType) else []
		[gtx_common.Dimension("")] if isinstance(arg.data_type.dtype, itir_ts.ListType) else []

		_, _, reduce_identity = gtir_to_tasklet.get_reduce_params(stencil_expr.expr)
		reduce_identity_for_args = reduce_identity

	_, _, reduce_identity = gtir_to_tasklet.get_reduce_params(stencil_expr.expr)
	reduce_identity_for_args = reduce_identity
	_, _, reduce_identity_for_args = gtir_to_tasklet.get_reduce_params(stencil_expr.expr)

		elif cpm.is_call_to(stencil_expr.expr, "neighbors"):
		reduce_identity_for_args = None

feat[next][dace]: GTIR-to-DaCe lowering of map-reduce (only full connectivity) #1683

feat[next][dace]: GTIR-to-DaCe lowering of map-reduce (only full connectivity) #1683

Conversation

edopao commented Oct 9, 2024 • edited Loading

philip-paul-mueller left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

edopao Oct 9, 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

edopao commented Oct 9, 2024

philip-paul-mueller left a comment

Choose a reason for hiding this comment

philip-paul-mueller 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

philip-paul-mueller 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

philip-paul-mueller 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

edopao commented Oct 9, 2024 •

edited

Loading

edopao Oct 9, 2024 •

edited

Loading