Merge branch 'main' into fix_caching_issue

.github/pull_request_template.md

@@ -15,7 +15,7 @@ Delete this comment and add a proper description of the changes contained in thi
         - test: Adding missing tests or correcting existing tests
 
     <scope>: cartesian | eve | next | storage
-    # ONLY if changes are limited to a specific subsytem
+    # ONLY if changes are limited to a specific subsystem
 
 - PR Description:
 
@@ -27,7 +27,7 @@ Delete this comment and add a proper description of the changes contained in thi
 ## Requirements
 
 - [ ] All fixes and/or new features come with corresponding tests.
-- [ ] Important design decisions have been documented in the approriate ADR inside the [docs/development/ADRs/](docs/development/ADRs/Index.md) folder.
+- [ ] Important design decisions have been documented in the appropriate ADR inside the [docs/development/ADRs/](docs/development/ADRs/Index.md) folder.
 
 If this PR contains code authored by new contributors please make sure:
 

src/gt4py/cartesian/frontend/gtscript_frontend.py

@@ -1460,6 +1460,13 @@ def visit_Assign(self, node: ast.Assign) -> list:
                                 loc=nodes.Location.from_ast_node(t),
                             )
 
+                    if self.backend_name in ["gt:gpu"]:
+                        raise GTScriptSyntaxError(
+                            message=f"Assignment to non-zero offsets in K is not available in {self.backend_name} as an unsolved bug remains."
+                            "Please refer to https://github.com/GridTools/gt4py/issues/1754.",
+                            loc=nodes.Location.from_ast_node(t),
+                        )
+
             if not self._is_known(name):
                 if name in self.temp_decls:
                     field_decl = self.temp_decls[name]

src/gt4py/eve/__init__.py

@@ -24,8 +24,7 @@
 
 """
 
-from __future__ import annotations  # isort:skip
-
+from __future__ import annotations
 
 from .concepts import (
     AnnexManager,
@@ -89,15 +88,6 @@
     "SymbolRef",
     "VType",
     "register_annex_user",
-    "# datamodels" "Coerced",
-    "DataModel",
-    "FrozenModel",
-    "GenericDataModel",
-    "Unchecked",
-    "concretize",
-    "datamodel",
-    "field",
-    "frozenmodel",
     # datamodels
     "Coerced",
     "DataModel",
@@ -122,7 +112,7 @@
     "pre_walk_values",
     "walk_items",
     "walk_values",
-    "# type_definition",
+    # type_definitions
     "NOTHING",
     "ConstrainedStr",
     "Enum",

src/gt4py/eve/codegen.py

@@ -347,7 +347,7 @@ def __str__(self) -> str:
 class Template(Protocol):
     """Protocol (abstract base class) defining the Template interface.
 
-    Direct subclassess of this base class only need to implement the
+    Direct subclasses of this base class only need to implement the
     abstract methods to adapt different template engines to this
     interface.
 
@@ -654,8 +654,8 @@ def apply(  # redefinition of symbol
 
         Args:
             root: An IR node.
-            node_templates (optiona): see :class:`NodeDumper`.
-            dump_function (optiona): see :class:`NodeDumper`.
+            node_templates (optional): see :class:`NodeDumper`.
+            dump_function (optional): see :class:`NodeDumper`.
             ``**kwargs`` (optional): custom extra parameters forwarded to `visit_NODE_TYPE_NAME()`.
 
         Returns:

src/gt4py/eve/datamodels/__init__.py

@@ -11,7 +11,7 @@
 Data Models can be considered as enhanced  `attrs <https://www.attrs.org>`_
 / `dataclasses <https://docs.python.org/3/library/dataclasses.html>`_ providing
 additional features like automatic run-time type validation. Values assigned to fields
-at initialization can be validated with automatic type checkings using the
+at initialization can be validated with automatic type checking using the
 field type definition. Custom field validation methods can also be added with
 the :func:`validator` decorator, and global instance validation methods with
 :func:`root_validator`.
@@ -33,7 +33,7 @@
     1. ``__init__()``.
 
         a.  If a custom ``__init__`` already exists in the class, it will not be overwritten.
-            It is your responsability to call ``__auto_init__`` from there to obtain
+            It is your responsibility to call ``__auto_init__`` from there to obtain
             the described behavior.
         b.  If there is not custom ``__init__``, the one generated by datamodels
             will be called first.

src/gt4py/eve/datamodels/core.py

@@ -24,7 +24,7 @@
 
 
 try:
-    # For perfomance reasons, try to use cytoolz when possible (using cython)
+    # For performance reasons, try to use cytoolz when possible (using cython)
     import cytoolz as toolz
 except ModuleNotFoundError:
     # Fall back to pure Python toolz
@@ -270,7 +270,7 @@ def datamodel(
 
 
 @overload
-def datamodel(  # redefinion of unused symbol
+def datamodel(  # redefinition of unused symbol
     cls: Type[_T],
     /,
     *,
@@ -289,7 +289,7 @@ def datamodel(  # redefinion of unused symbol
 
 
 # TODO(egparedes): Use @dataclass_transform(eq_default=True, field_specifiers=("field",))
-def datamodel(  # redefinion of unused symbol
+def datamodel(  # redefinition of unused symbol
     cls: Optional[Type[_T]] = None,
     /,
     *,
@@ -867,7 +867,7 @@ def _substitute_typevars(
 
 def _make_counting_attr_from_attribute(
     field_attrib: Attribute, *, include_type: bool = False, **kwargs: Any
-) -> Any:  # attr.s lies a bit in some typing definitons
+) -> Any:  # attr.s lies a bit in some typing definitions
     args = [
         "default",
         "validator",
@@ -965,7 +965,7 @@ def _type_converter(value: Any) -> _T:
                 return value if isinstance(value, type_annotation) else type_annotation(value)
             except Exception as error:
                 raise TypeError(
-                    f"Error during coertion of given value '{value}' for field '{name}'."
+                    f"Error during coercion of given value '{value}' for field '{name}'."
                 ) from error
 
         return _type_converter
@@ -996,7 +996,7 @@ def _type_converter(value: Any) -> _T:
         return _make_type_converter(origin_type, name)
 
     raise exceptions.EveTypeError(
-        f"Automatic type coertion for {type_annotation} types is not supported."
+        f"Automatic type coercion for {type_annotation} types is not supported."
     )
 
 
@@ -1085,7 +1085,7 @@ def _make_datamodel(
                     )
 
         else:
-            # Create field converter if automatic coertion is enabled
+            # Create field converter if automatic coercion is enabled
             converter: TypeConverter = cast(
                 TypeConverter,
                 _make_type_converter(type_hint, qualified_field_name) if coerce_field else None,
@@ -1099,7 +1099,7 @@ def _make_datamodel(
                 if isinstance(attr_value_in_cls, _KNOWN_MUTABLE_TYPES):
                     warnings.warn(
                         f"'{attr_value_in_cls.__class__.__name__}' value used as default in '{cls.__name__}.{key}'.\n"
-                        "Mutable types should not defbe normally used as field defaults (use 'default_factory' instead).",
+                        "Mutable types should not be used as field defaults (use 'default_factory' instead).",
                         stacklevel=_stacklevel_offset + 2,
                     )
                 setattr(

src/gt4py/eve/extended_typing.py

@@ -14,12 +14,8 @@
 
 from __future__ import annotations
 
-import abc as _abc
 import array as _array
-import collections.abc as _collections_abc
-import ctypes as _ctypes
 import dataclasses as _dataclasses
-import enum as _enum
 import functools as _functools
 import inspect as _inspect
 import mmap as _mmap

src/gt4py/eve/trees.py

@@ -31,14 +31,6 @@
 from .type_definitions import Enum
 
 
-try:
-    # For performance reasons, try to use cytoolz when possible (using cython)
-    import cytoolz as toolz
-except ModuleNotFoundError:
-    # Fall back to pure Python toolz
-    import toolz  # noqa: F401 [unused-import]
-
-
 TreeKey = Union[int, str]
 
 

src/gt4py/eve/type_validation.py

@@ -311,7 +311,7 @@ def __call__(
                 #          ...
                 #
                 # Since this can be an arbitrary type (not something regular like a collection) there is
-                # no way to check if the type parameter is verifed in the actual instance.
+                # no way to check if the type parameter is verified in the actual instance.
                 # The only check can be done at run-time is to verify that the value is an instance of
                 # the original type, completely ignoring the annotation. Ideally, the static type checker
                 # can do a better job to try figure out if the type parameter is ok ...

src/gt4py/eve/utils.py

@@ -69,7 +69,7 @@
 
 
 try:
-    # For perfomance reasons, try to use cytoolz when possible (using cython)
+    # For performance reasons, try to use cytoolz when possible (using cython)
     import cytoolz as toolz
 except ModuleNotFoundError:
     # Fall back to pure Python toolz

src/gt4py/next/config.py

@@ -11,7 +11,6 @@
 import enum
 import os
 import pathlib
-import tempfile
 from typing import Final
 
 
@@ -51,37 +50,34 @@ def env_flag_to_bool(name: str, default: bool) -> bool:
             )
 
 
-_PREFIX: Final[str] = "GT4PY"
-
 #: 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("GT4PY_DEBUG", default=False)
 
 
 #: Verbose flag for DSL compilation errors
 VERBOSE_EXCEPTIONS: bool = env_flag_to_bool(
-    f"{_PREFIX}_VERBOSE_EXCEPTIONS", default=True if DEBUG else False
+    "GT4PY_VERBOSE_EXCEPTIONS", default=True if DEBUG else False
 )
 
 
 #: Where generated code projects should be persisted.
 #: Only active if BUILD_CACHE_LIFETIME is set to PERSISTENT
 BUILD_CACHE_DIR: pathlib.Path = (
-    pathlib.Path(os.environ.get(f"{_PREFIX}_BUILD_CACHE_DIR", tempfile.gettempdir()))
-    / "gt4py_cache"
+    pathlib.Path(os.environ.get("GT4PY_BUILD_CACHE_DIR", pathlib.Path.cwd())) / ".gt4py_cache"
 )
 
 
 #: Whether generated code projects should be kept around between runs.
 #: - SESSION: generated code projects get destroyed when the interpreter shuts down
 #: - PERSISTENT: generated code projects are written to BUILD_CACHE_DIR and persist between runs
 BUILD_CACHE_LIFETIME: BuildCacheLifetime = BuildCacheLifetime[
-    os.environ.get(f"{_PREFIX}_BUILD_CACHE_LIFETIME", "persistent" if DEBUG else "session").upper()
+    os.environ.get("GT4PY_BUILD_CACHE_LIFETIME", "persistent" if DEBUG else "session").upper()
 ]
 
 #: Build type to be used when CMake is used to compile generated code.
 #: Might have no effect when CMake is not used as part of the toolchain.
 CMAKE_BUILD_TYPE: CMakeBuildType = CMakeBuildType[
-    os.environ.get(f"{_PREFIX}_CMAKE_BUILD_TYPE", "debug" if DEBUG else "release").upper()
+    os.environ.get("GT4PY_CMAKE_BUILD_TYPE", "debug" if DEBUG else "release").upper()
 ]

src/gt4py/next/ffront/decorator.py

@@ -230,7 +230,7 @@ def __call__(self, *args: Any, offset_provider: common.OffsetProvider, **kwargs:
         if self.backend is None:
             warnings.warn(
                 UserWarning(
-                    f"Field View Program '{self.definition_stage.definition.__name__}': Using Python execution, consider selecting a perfomance backend."
+                    f"Field View Program '{self.definition_stage.definition.__name__}': Using Python execution, consider selecting a performance backend."
                 ),
                 stacklevel=2,
             )

src/gt4py/next/iterator/ir_utils/domain_utils.py

@@ -16,6 +16,7 @@
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.iterator.transforms import trace_shifts
+from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
 
 
 def _max_domain_sizes_by_location_type(offset_provider: Mapping[str, Any]) -> dict[str, int]:
@@ -168,6 +169,8 @@ def domain_union(*domains: SymbolicDomain) -> SymbolicDomain:
             lambda current_expr, el_expr: im.call("maximum")(current_expr, el_expr),
             [domain.ranges[dim].stop for domain in domains],
         )
+        # constant fold expression to keep the tree small
+        start, stop = ConstantFolding.apply(start), ConstantFolding.apply(stop)  # type: ignore[assignment]  # always an itir.Expr
         new_domain_ranges[dim] = SymbolicRange(start, stop)
 
     return SymbolicDomain(domains[0].grid_type, new_domain_ranges)

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

@@ -97,7 +97,6 @@ def new_name(name):
 
     if all(eligible_params):
         new_expr.location = node.location
-        return new_expr
     else:
         new_expr = ir.FunCall(
             fun=ir.Lambda(
@@ -111,11 +110,11 @@ def new_name(name):
             args=[arg for arg, eligible in zip(node.args, eligible_params) if not eligible],
             location=node.location,
         )
-        for attr in ("type", "recorded_shifts", "domain"):
-            if hasattr(node.annex, attr):
-                setattr(new_expr.annex, attr, getattr(node.annex, attr))
-        itir_inference.copy_type(from_=node, to=new_expr, allow_untyped=True)
-        return new_expr
+    for attr in ("type", "recorded_shifts", "domain"):
+        if hasattr(node.annex, attr):
+            setattr(new_expr.annex, attr, getattr(node.annex, attr))
+    itir_inference.copy_type(from_=node, to=new_expr, allow_untyped=True)
+    return new_expr
 
 
 @dataclasses.dataclass

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

@@ -10,6 +10,7 @@
 
 from gt4py.eve import NodeTranslator, PreserveLocationVisitor, SymbolTableTrait
 from gt4py.next.iterator import ir
+from gt4py.next.iterator.type_system import inference as type_inference
 
 
 class RemapSymbolRefs(PreserveLocationVisitor, NodeTranslator):
@@ -46,7 +47,9 @@ def visit_SymRef(
         self, node: ir.SymRef, *, name_map: Dict[str, str], active: Optional[Set[str]] = None
     ):
         if active and node.id in active:
-            return ir.SymRef(id=name_map.get(node.id, node.id))
+            new_ref = ir.SymRef(id=name_map.get(node.id, node.id))
+            type_inference.copy_type(from_=node, to=new_ref, allow_untyped=True)
+            return new_ref
         return node
 
     def generic_visit(  # type: ignore[override]

src/gt4py/next/iterator/type_system/inference.py

@@ -95,14 +95,17 @@ def _set_node_type(node: itir.Node, type_: ts.TypeSpec) -> None:
     node.type = type_
 
 
-def copy_type(from_: itir.Node, to: itir.Node, allow_untyped=False) -> None:
+def copy_type(from_: itir.Node, to: itir.Node, allow_untyped: bool = False) -> None:
     """
     Copy type from one node to another.
 
     This function mainly exists for readability reasons.
     """
     assert allow_untyped is not None or isinstance(from_.type, ts.TypeSpec)
-    _set_node_type(to, from_.type)  # type: ignore[arg-type]
+    if from_.type is None:
+        assert allow_untyped
+        return
+    _set_node_type(to, from_.type)
 
 
 def on_inferred(callback: Callable, *args: Union[ts.TypeSpec, ObservableTypeSynthesizer]) -> None: