feature[next]: Nested scalars args & cleanup #1540

tehrengruber · 2024-05-02T19:46:18Z

Cleanup handling of tuple arguments in GTFN backend. This enables using scalar arguments in both the domain expression of a program and as an argument to a field operator simultaneously. Scalar arguments to an applied as_fieldop are detected using the new type inference.
Adds support for nested scalar args in the GTFN backend
Adds support for nested scalar & field args in GTFN and Embedded

havogt · 2024-06-28T07:05:25Z

src/gt4py/next/ffront/past_process_args.py

-                    " tuple) need to have the same shape and dimensions."
-                )
-            size_args.extend(shape if shape else [None] * len(dims))
+            if shapes_and_dims:  # scalar or zero-dim field otherwise


Does it mean we allow writing to a scalar?

Looking at the implementation _field_constituents_shape_and_dims, it seems the comment is wrong as 0d fields probably returns a tuple of empty tuples.

Does it mean we allow writing to a scalar?

Not sure what motivated this question. This code block only extracts size arguments, the change here merely avoids errors for cases like (scalar, (field1, field2)). I've added a comment that explains this a little more.

Looking at the implementation _field_constituents_shape_and_dims, it seems the comment is wrong as 0d fields probably returns a tuple of empty tuples.

The comment was alright though not helpful^^, I've fixed _field_constituents_shape_and_dims which was wrong and added a test test_zero_dim_tuple_arg.

havogt · 2024-06-28T07:09:28Z

src/gt4py/next/ffront/past_to_itir.py

+                .filter(lambda dims: len(dims) > 0)
+                .to_list()
+            )
+            if len(fields_dims) > 0:  # param has no field constituent otherwise


what is a "field constituent"?

A constituent of a composite with type field, e.g. here (scalar, (field, scalar)) the composite is a tuple and field is a field constituent. I have introduced the nomenclature composite and constituent as a generalization of tuples and structs where the constituents are the elements and members respectively. Given the amount of question marks this caused we should probably revisit this naming.

havogt · 2024-06-28T07:16:32Z

src/gt4py/next/iterator/embedded.py

@@ -676,11 +676,23 @@ def _is_concrete_position(pos: Position) -> TypeGuard[ConcretePosition]:

 def _get_axes(
    field_or_tuple: LocatedField | tuple,
+    *,
+    ignore_zero_dims=False,


From the PR it's not obvious to me how this is used. Please add an itir test.

I've added doctests here. The motivation is similar as in the frontend for (scalar, (field1, field2)) we sometimes only want to check that field1 and field2 have the same dimensions, but don't care about scalar.

Are the changes in this file related to the newly added tests?

Yes, they are required for "mixed" tuple args, e.g. an argument with type (scalar, (field, scalar)). This is the code path in make_in_iterator where this function is called.

src/gt4py/next/otf/binding/nanobind.py

src/gt4py/next/program_processors/codegens/gtfn/codegen.py

src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py

havogt · 2024-06-28T07:30:08Z

src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py

+        for input_ in inputs:
+            lowered_input = self.visit(input_, **kwargs)
+
+            def _convert_input(


unclear why this should be a closure

_convert_input -> better name

Refactored this section and the convert_output case below to use a common utility function _process_elements.

havogt · 2024-06-28T07:31:39Z

src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py

-        raise ValueError("Expected 'SymRef' or 'make_tuple' in output argument.")
+        lowered_output = self.visit(node)
+
+        def _convert_output(el_type: ts.ScalarType | ts.FieldType, path: tuple[int, ...]) -> Expr:


see comment at _convert_input

Let's discuss this.

See comment above.

src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir_to_gtfn_im_ir.py

havogt · 2024-08-02T09:52:59Z

src/gt4py/next/config.py

@@ -61,7 +61,7 @@ def env_flag_to_bool(name: str, default: bool) -> bool:
 #: Master debug flag
 #: Changes defaults for all the other options to be as helpful for debugging as possible.
 #: Does not override values set in environment variables.
-DEBUG: Final[bool] = env_flag_to_bool(f"{_PREFIX}_DEBUG", default=False)
+DEBUG: Final[bool] = env_flag_to_bool(f"{_PREFIX}_DEBUG", default=True)


What motivated you to change this here instead of setting the env var? Anyway, this is a reminder to revert.

I haven't figured out how to have default env vars for test execution in a project in PyCharm. Reverted.

havogt · 2024-08-02T10:06:53Z

src/gt4py/next/ffront/past_process_args.py

+            if dims:
+                assert hasattr(arg, "shape") and len(arg.shape) == len(dims)
                yield (arg.shape, dims)
            else:
-                yield (tuple(), dims)
+                pass
        case ts.ScalarType():
-            yield (tuple(), [])
+            pass


This constructs looks weird to me, I guess it's equivalent to replacing pass by return which might be slightly more expressive, but still ugly. Alternatives could be yield from []

I used yield from [] now.

src/gt4py/next/ffront/past_to_itir.py

havogt · 2024-08-02T11:13:04Z

src/gt4py/next/iterator/embedded.py

@@ -676,11 +676,23 @@ def _is_concrete_position(pos: Position) -> TypeGuard[ConcretePosition]:

 def _get_axes(
    field_or_tuple: LocatedField | tuple,
+    *,
+    ignore_zero_dims=False,


Are the changes in this file related to the newly added tests?

src/gt4py/next/otf/binding/nanobind.py

src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir_to_gtfn_im_ir.py

tehrengruber · 2024-08-06T11:30:33Z

src/gt4py/next/iterator/embedded.py

@@ -895,7 +935,7 @@ def deref(self) -> Any:

        assert self.pos is not None
        shifted_pos = self.pos.copy()
-        axes = _get_axes(self.field)
+        axes = _get_axes(self.field, ignore_zero_dims=True)


TODO:

write iterator tests

check if (vertex_field, vertex_k_field) should also be valid here.

I've added two tests.

I've added two test named test_tuple_arg_with_unpromotable_dims and test_scalar_arg_with_field. The implementation is beyond the scope of this PR.

…ple args with different dims

havogt

2 remarks

havogt · 2024-09-16T14:23:03Z

src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir.py

+    if (
+        isinstance(expr, FunCall)
+        and isinstance(expr.fun, SymRef)
+        and expr.fun.id == "tuple_get"
+        and len(expr.args) == 2
+        and _is_ref_or_tuple_expr_of_ref(expr.args[1])
+    ):
+        return True


Is this case tested?

This is indeed tested, the _process_elements creates such gtfn_ir expressions for example in the test_multicopy case.

tests/next_tests/integration_tests/cases.py

tehrengruber added 30 commits April 15, 2024 07:45

Use type specification for itir.Literal

aa937da

Small cleanup

095ae13

Small fix

cace41a

Fix failing doctests

b22f929

Address review comments

288ea3c

New type inference first draft

c695935

Merge origin/main

72b50ff

Fix some tests

beabac2

Cleanup

ed4bf27

Use is_call_to instead of equality comparison with itir.Ref.

1c94bf0

Cleanup

3799b7a

Cleanup

f75f0e4

Address reviewer comments

779ba3d

Cleanup

c7e1e93

Merge origin_tehrengruber/fix_im_ref_comp

172982f

Cleanup

79bee51

Cleanup

dd5e601

Merge origin_tehrengruber/fix_im_ref_comp

1b803b7

Merge origin_tehrengruber/fix_im_ref_comp

f41e032

Merge origin/main

520c805

Add test for neighbor / sparse input field

129f6b3

Multiple fixes

a99ca9a

Multiple fixes

896582e

Cleanup

c8fb2d9

Cleanup

cb0ebbb

Fix dace

873c60e

Formatting

d42fa12

Formatting

8c632de

Don't use temporaries in dace

1d043d1

Try fixing test failures by removing types

b792a53

tehrengruber added 7 commits May 2, 2024 23:22

Fix gtfn_module test

18ab7a4

Address review comments

5947934

Cleanup

fb6fac5

Small fix

7971fff

Fix doctest

876d3df

Address reviewer comments

88c4f0d

Merge remote-tracking branch 'origin/main' into itir_type_inference

6bcc38a

havogt requested changes Jun 28, 2024

View reviewed changes

tehrengruber added 6 commits July 24, 2024 09:11

Address review comments

810d3d9

Merge branch 'itir_type_inference' into nested_scalars_args

803b5b1

Merge origin/main

45bd6bc

Fix headers

f1efa25

Format

770fb4a

Merge remote-tracking branch 'origin/main' into nested_scalars_args

83f62c9

tehrengruber changed the base branch from itir_type_inference to main July 31, 2024 07:00

havogt requested changes Aug 2, 2024

View reviewed changes

tehrengruber commented Aug 6, 2024

View reviewed changes

tehrengruber added 5 commits September 11, 2024 12:29

Address review comments

799ada5

Merge remote-tracking branch 'origin/main' into nested_scalars_args

b22a39d

Cleanup GTFN lowering of composite args and returns, add tests for tu…

fe1fd16

…ple args with different dims

Add iterator tests

404d2a9

Merge remote-tracking branch 'origin/main' into nested_scalars_args

0b62b2c

tehrengruber requested a review from havogt September 16, 2024 13:42

havogt approved these changes Sep 16, 2024

View reviewed changes

tehrengruber added 5 commits September 16, 2024 18:28

Fix failing tests

357b1aa

Address review comments

6b32cb9

Fix failing tests

940bb2c

Fix failing tests

f8d38b1

Fix format

bbca010

tehrengruber merged commit 15baa37 into GridTools:main Sep 16, 2024
31 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feature[next]: Nested scalars args & cleanup #1540

feature[next]: Nested scalars args & cleanup #1540

tehrengruber commented May 2, 2024 •

edited

Loading

havogt Jun 28, 2024

havogt Jun 28, 2024

tehrengruber Jul 18, 2024

havogt Jun 28, 2024

tehrengruber Jul 18, 2024

havogt Jun 28, 2024

tehrengruber Jul 18, 2024

havogt Aug 2, 2024

tehrengruber Sep 16, 2024

havogt Jun 28, 2024

tehrengruber Sep 16, 2024

havogt Jun 28, 2024

tehrengruber Jul 31, 2024

tehrengruber Sep 16, 2024

havogt Aug 2, 2024

tehrengruber Sep 11, 2024

havogt Aug 2, 2024

tehrengruber Sep 11, 2024

havogt Aug 2, 2024

tehrengruber Aug 6, 2024

tehrengruber Sep 16, 2024

havogt left a comment

havogt Sep 16, 2024

tehrengruber Sep 16, 2024

feature[next]: Nested scalars args & cleanup #1540

feature[next]: Nested scalars args & cleanup #1540

Conversation

tehrengruber commented May 2, 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

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

havogt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tehrengruber commented May 2, 2024 •

edited

Loading