feat[next]: Index builtin #1699

SF-N · 2024-10-22T14:05:11Z

Adds index builtin for embedded and gtfn backends.

SF-N · 2024-10-22T14:06:23Z

src/gt4py/next/iterator/pretty_parser.py

@@ -31,9 +31,9 @@
    INT_LITERAL: SIGNED_INT
    FLOAT_LITERAL: SIGNED_FLOAT
    OFFSET_LITERAL: ( INT_LITERAL | CNAME ) "ₒ"
-    _literal: INT_LITERAL | FLOAT_LITERAL | OFFSET_LITERAL
+    AXIS_LITERAL: CNAME ("ᵥ" | "ₕ")


Renamed for consistency.

SF-N · 2024-10-22T14:07:34Z

src/gt4py/next/iterator/pretty_parser.py

@@ -31,9 +31,9 @@
    INT_LITERAL: SIGNED_INT
    FLOAT_LITERAL: SIGNED_FLOAT
    OFFSET_LITERAL: ( INT_LITERAL | CNAME ) "ₒ"
-    _literal: INT_LITERAL | FLOAT_LITERAL | OFFSET_LITERAL
+    AXIS_LITERAL: CNAME ("ᵥ" | "ₕ")
+    _literal: INT_LITERAL | FLOAT_LITERAL | OFFSET_LITERAL | AXIS_LITERAL


Otherwise index(OFFSET_LITERAL) isn't parsed properly.

Do you mean index(AXIS_LITERAL)?

SF-N · 2024-10-22T14:11:22Z

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

@@ -611,7 +611,6 @@ def convert_el_to_sid(el_expr: Expr, el_type: ts.ScalarType | ts.FieldType) -> E
                tuple_constructor=lambda *elements: SidComposite(values=list(elements)),
            )

-            assert isinstance(lowered_input_as_sid, (SidComposite, SidFromScalar, SymRef))


This check was superfluous from the beginning since the format is checked anyway when the SetAt is constructed. We just removed it instead of allowing an index FunCall.

havogt

a few comments

havogt · 2024-10-23T11:59:58Z

src/gt4py/next/iterator/embedded.py

@@ -1156,6 +1156,7 @@ def ndarray(self) -> core_defs.NDArrayObject:
    def asnumpy(self) -> np.ndarray:
        raise NotImplementedError()

+    # TODO(tehrengruber): Use a regular zero dimensional field instead.


I don't understand the comment.

We moved the comment to a better place. With respect to the comment itself: The IndexField class is rather strange. It has two modes of operation: Either it is a field with (conceptually) field(domain) == domain or it is a zero-dimensional field. Both modes don't share any implementation similarities, but are mushed into the same class. The way the class behaves is then controlled using _cur_index. It would be much simpler to just make field[index] return a zero-dimensional field which is exactly what we want instead of re-implementing it here.

I don't understand what you are saying...

havogt · 2024-10-23T12:02:59Z

src/gt4py/next/iterator/pretty_parser.py

@@ -31,9 +31,9 @@
    INT_LITERAL: SIGNED_INT
    FLOAT_LITERAL: SIGNED_FLOAT
    OFFSET_LITERAL: ( INT_LITERAL | CNAME ) "ₒ"
-    _literal: INT_LITERAL | FLOAT_LITERAL | OFFSET_LITERAL
+    AXIS_LITERAL: CNAME ("ᵥ" | "ₕ")
+    _literal: INT_LITERAL | FLOAT_LITERAL | OFFSET_LITERAL | AXIS_LITERAL


Do you mean index(AXIS_LITERAL)?

havogt · 2024-10-23T12:07:37Z

src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py

+        assert all(
+            isinstance(inp, SymRef) for inp in node.inputs
+        )  # backend only supports SymRef inputs, not `index` calls
+        input_names = [str(inp.id) for inp in node.inputs]  # type: ignore[union-attr]  # backend only supports SymRef inputs, not `index` calls


Suggested change

input_names = [str(inp.id) for inp in node.inputs] # type: ignore[union-attr] # backend only supports SymRef inputs, not `index` calls

input_names = [str(inp.id) for inp in node.inputs] # type: ignore[union-attr] # `inp`s are only `SymRef`s

or similar

havogt · 2024-10-24T18:52:34Z

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

@@ -260,6 +260,17 @@ def visit_Program(self, node: gtfn_ir.Program, **kwargs: Any) -> Union[str, Coll
    #include <gridtools/fn/${grid_type_str}.hpp>
    #include <gridtools/fn/sid_neighbor_table.hpp>
    #include <gridtools/stencil/global_parameter.hpp>
+    #include <gridtools/stencil/positional.hpp>
+
+    // TODO(havogt): move to gtfn?


I'd propose to wait until GridTools/gridtools#1806 is merged next week.

Once #1720 is in, we can remove this.

Thanks, I removed that code.

updates minimum gridtools_cpp version for #1699

havogt · 2024-11-04T10:14:48Z

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

+    #include <gridtools/stencil/positional.hpp>
+
+    // TODO(havogt): move to gtfn?
+    namespace gridtools{
+    namespace fn{
+    template <class T>
+    auto index(T){
+    return stencil::positional<std::decay_t<T>>();}
+    }
+    }
+


Suggested change

#include <gridtools/stencil/positional.hpp>

// TODO(havogt): move to gtfn?

namespace gridtools{

namespace fn{

template <class T>

auto index(T){

return stencil::positional<std::decay_t<T>>();}

}

}

havogt

lgtm

SF-N added 2 commits October 22, 2024 16:03

Add index builtin

e1e30b2

Cleanup

181b912

SF-N commented Oct 22, 2024

View reviewed changes

Skip 'STARTS_FROM_GTIR_PROGRAM' only in dace backends

60d2100

SF-N changed the title ~~Index builtin~~ feat[next] Index builtin Oct 22, 2024

SF-N changed the title ~~feat[next] Index builtin~~ feat[next]: Index builtin Oct 22, 2024

SF-N requested a review from havogt October 22, 2024 14:58

havogt reviewed Oct 23, 2024

View reviewed changes

Address review comments

19998af

SF-N requested a review from havogt October 23, 2024 15:51

SF-N added 2 commits October 23, 2024 17:51

Merge branch 'main' into index_builtin

a8c99b8

Merge branch 'main' into index_builtin

8031bf9

havogt reviewed Oct 24, 2024

View reviewed changes

SF-N added 2 commits October 28, 2024 11:13

Merge branch 'main' into index_builtin

655e260

Merge branch 'main' into index_builtin

b53897a

havogt mentioned this pull request Nov 1, 2024

build: update dependencies #1720

Merged

havogt added a commit that referenced this pull request Nov 4, 2024

build: update dependencies (#1720)

162a512

updates minimum gridtools_cpp version for #1699

Merge branch 'main' into index_builtin

1540e8d

havogt reviewed Nov 4, 2024

View reviewed changes

SF-N added 2 commits November 4, 2024 11:47

Merge branch 'main' into index_builtin

a61ef9a

Remove GTFN C++ code as PR#1720 is merged

f558e35

SF-N requested a review from havogt November 4, 2024 10:51

Merge branch 'main' into index_builtin

0c50850

havogt approved these changes Nov 4, 2024

View reviewed changes

SF-N merged commit 9f3b0a7 into GridTools:main Nov 4, 2024
31 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat[next]: Index builtin #1699

feat[next]: Index builtin #1699

SF-N commented Oct 22, 2024

SF-N Oct 22, 2024

SF-N Oct 22, 2024

havogt Oct 23, 2024

SF-N Oct 23, 2024

SF-N Oct 22, 2024

havogt left a comment

havogt Oct 23, 2024

tehrengruber Oct 23, 2024

havogt Oct 23, 2024

havogt Oct 23, 2024

havogt Oct 23, 2024

havogt Oct 24, 2024

havogt Nov 1, 2024

SF-N Nov 4, 2024

havogt Nov 4, 2024

havogt left a comment

	input_names = [str(inp.id) for inp in node.inputs] # type: ignore[union-attr] # backend only supports SymRef inputs, not `index` calls
	input_names = [str(inp.id) for inp in node.inputs] # type: ignore[union-attr] # `inp`s are only `SymRef`s

feat[next]: Index builtin #1699

feat[next]: Index builtin #1699

Conversation

SF-N commented Oct 22, 2024

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

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