Repository: GridTools/gt4py
Pull request #1139 (2023-01-25 20:52:38)

Feature/literal parametrized dtype
## Description

In order to run a 32-bit float version of the model, we need at minima to be able to control literal dtype. This PR implements option 4 of (the gtscript concept workshop)[https://github.com/GridTools/concepts/blob/master/collaboration/gtscript-workshop/GTScript-Syntax-Discussion.md#gtscript-syntax-discussed-issues-20200829]

## Requirements

Before submitting this PR, please make sure:

- [x] You have run the code checks, tests and documentation build successfully
- [x] All fixes and all new functionality are tested and documentation is up to date
- [x] You looked at the [review checklist](https://github.com/GridTools/gt4py/blob/main/REVIEW_CHECKLIST.txt)

Additionally, if this PR contains code authored by new contributors:

- [x] All the authors are covered by a valid contributor assignment agreement,
signed by the employer if needed, provided to ETH Zurich

 





Patch:
diff --git a/AUTHORS.rst b/AUTHORS.rst
index 98eaedb7..10bac07c 100644
--- a/AUTHORS.rst
+++ b/AUTHORS.rst
@@ -8,6 +8,7 @@ Authors
 * Cruz, Felipe. ETH Zurich - CSCS
 * Dahm, Johann. Allen Institute for AI
 * Davis, Eddie. Allen Institute for AI
+* Deconinck, Florian. SSAI/NASA-GSFC
 * Ehrengruber, Till. ETH Zurich - CSCS
 * Farabullini, Nicoletta. ETH Zurich - C2SM
 * George, Rhea. Allen Institute for AI
diff --git a/src/gt4py/cartesian/frontend/base.py b/src/gt4py/cartesian/frontend/base.py
index bff7fc3c..28af28a8 100644
--- a/src/gt4py/cartesian/frontend/base.py
+++ b/src/gt4py/cartesian/frontend/base.py
@@ -74,7 +74,11 @@ class Frontend(abc.ABC):
     @classmethod
     @abc.abstractmethod
     def generate(
-        cls, definition: AnyStencilFunc, externals: Dict[str, Any], options: BuildOptions
+        cls,
+        definition: AnyStencilFunc,
+        externals: Dict[str, Any],
+        dtypes: Dict[Type, Type],
+        options: BuildOptions,
     ) -> gtir.Stencil:
         """
         Generate a StencilDefinition from a stencil Python function.
diff --git a/src/gt4py/cartesian/frontend/defir_to_gtir.py b/src/gt4py/cartesian/frontend/defir_to_gtir.py
index 78b5b71f..7326b0ca 100644
--- a/src/gt4py/cartesian/frontend/defir_to_gtir.py
+++ b/src/gt4py/cartesian/frontend/defir_to_gtir.py
@@ -18,6 +18,8 @@ import itertools
 import numbers
 from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
+import numpy as np
+
 from gt4py.cartesian.frontend.node_util import (
     IRNodeMapper,
     IRNodeVisitor,
@@ -73,20 +75,29 @@ def _convert_dtype(data_type) -> common.DataType:
 
 def _make_literal(v: numbers.Number) -> gtir.Literal:
     value: Union[BuiltinLiteral, str]
-    if isinstance(v, bool):
+    if isinstance(v, (bool, np.bool_)):
         dtype = common.DataType.BOOL
         value = common.BuiltInLiteral.TRUE if v else common.BuiltInLiteral.FALSE
     else:
-        if isinstance(v, int):
-            # note: could be extended to support 32 bit integers by checking the values magnitude
+        if isinstance(v, (int, np.int64)):
             dtype = common.DataType.INT64
-        elif isinstance(v, float):
+        elif isinstance(v, (float, np.float64)):
             dtype = common.DataType.FLOAT64
+        elif isinstance(v, np.float32):
+            dtype = common.DataType.FLOAT32
+        elif isinstance(v, np.int64):
+            dtype = common.DataType.INT32
+        elif isinstance(v, np.int32):
+            dtype = common.DataType.INT32
+        elif isinstance(v, np.int16):
+            dtype = common.DataType.INT16
+        elif isinstance(v, np.int8):
+            dtype = common.DataType.INT8
         else:
-            print(
-                f"Warning: Only INT64 and FLOAT64 literals are supported currently. Implicitly upcasting `{v}` to FLOAT64"
+            raise TypeError(
+                "Warning: Only INTs, FLOAT64 and FLOAT32 (via np.X types) literals are "
+                f"supported currently. Type {type(v)} unsupported."
             )
-            dtype = common.DataType.FLOAT64
         value = str(v)
     return gtir.Literal(dtype=dtype, value=value)
 
diff --git a/src/gt4py/cartesian/frontend/gtscript_frontend.py b/src/gt4py/cartesian/frontend/gtscript_frontend.py
index 8e0fd656..a8280387 100644
--- a/src/gt4py/cartesian/frontend/gtscript_frontend.py
+++ b/src/gt4py/cartesian/frontend/gtscript_frontend.py
@@ -22,7 +22,7 @@ import textwrap
 import time
 import types
 import warnings
-from typing import Any, Dict, Final, List, Literal, Optional, Sequence, Set, Tuple, Union
+from typing import Any, Dict, Final, List, Literal, Optional, Sequence, Set, Tuple, Type, Union
 
 import numpy as np
 
@@ -701,6 +701,7 @@ class IRMaker(ast.NodeVisitor):
         *,
         domain: nodes.Domain,
         temp_decls: Optional[Dict[str, nodes.FieldDecl]] = None,
+        dtypes: Optional[Dict[Type, Type]] = None,
     ):
         fields = fields or {}
         parameters = parameters or {}
@@ -720,6 +721,7 @@ class IRMaker(ast.NodeVisitor):
         self.decls_stack = []
         self.parsing_horizontal_region = False
         self.written_vars: Set[str] = set()
+        self.dtypes = dtypes
         nodes.NativeFunction.PYTHON_SYMBOL_TO_IR_OP = {
             "abs": nodes.NativeFunction.ABS,
             "min": nodes.NativeFunction.MIN,
@@ -996,7 +998,12 @@ class IRMaker(ast.NodeVisitor):
                 loc=nodes.Location.from_ast_node(node),
             )
         elif isinstance(value, numbers.Number):
-            data_type = nodes.DataType.from_dtype(np.dtype(type(value)))
+            value_type = (
+                self.dtypes[type(value)]
+                if self.dtypes and type(value) in self.dtypes.keys()
+                else np.dtype(type(value))
+            )
+            data_type = nodes.DataType.from_dtype(value_type)
             return nodes.ScalarLiteral(value=value, data_type=data_type)
         else:
             raise GTScriptSyntaxError(
@@ -1612,7 +1619,7 @@ class GTScriptParser(ast.NodeVisitor):
         gtscript.Axis,
     )
 
-    def __init__(self, definition, *, options, externals=None):
+    def __init__(self, definition, *, options, externals=None, dtypes=None):
         assert isinstance(definition, types.FunctionType)
         self.definition = definition
         self.filename = inspect.getfile(definition)
@@ -1625,6 +1632,7 @@ class GTScriptParser(ast.NodeVisitor):
         self.external_context = externals or {}
         self.resolved_externals = {}
         self.block = None
+        self.dtypes = dtypes
 
     def __str__(self):
         result = "<GT4Py.GTScriptParser> {\n"
@@ -2025,6 +2033,7 @@ class GTScriptParser(ast.NodeVisitor):
             local_symbols={},  # Not used
             domain=domain,
             temp_decls=temp_decls,
+            dtypes=self.dtypes,
         )(self.ast_root)
 
         self.definition_ir = nodes.StencilDefinition(
@@ -2077,13 +2086,13 @@ class GTScriptFrontend(Frontend):
         return GTScriptParser.annotate_definition(definition, externals)
 
     @classmethod
-    def generate(cls, definition, externals, options):
+    def generate(cls, definition, externals, dtypes, options):
         if options.build_info is not None:
             start_time = time.perf_counter()
 
         if not hasattr(definition, "_gtscript_"):
             cls.prepare_stencil_definition(definition, externals)
-        translator = GTScriptParser(definition, externals=externals, options=options)
+        translator = GTScriptParser(definition, externals=externals, dtypes=dtypes, options=options)
         definition_ir = translator.run()
 
         # GTIR only supports LatLonGrids
diff --git a/src/gt4py/cartesian/gtscript.py b/src/gt4py/cartesian/gtscript.py
index dbe033a1..a56df74b 100644
--- a/src/gt4py/cartesian/gtscript.py
+++ b/src/gt4py/cartesian/gtscript.py
@@ -307,6 +307,7 @@ def stencil(
             backend=backend,
             build_options=build_options,
             externals=externals or {},
+            dtypes=dtypes or {},
         )
         definition_func.__annotations__ = original_annotations
         return out
@@ -449,16 +450,16 @@ def lazy_stencil(
             build_options.name = f"{definition_func.__name__}"
         if backend and "dace" in backend:
             stencil = DaCeLazyStencil(
-                StencilBuilder(
-                    definition_func, backend=backend, options=build_options
-                ).with_externals(externals or {})
+                StencilBuilder(definition_func, backend=backend, options=build_options)
+                .with_externals(externals or {})
+                .with_dtypes(dtypes or {})
             )
 
         else:
             stencil = LazyStencil(
-                StencilBuilder(
-                    definition_func, backend=backend, options=build_options
-                ).with_externals(externals or {})
+                StencilBuilder(definition_func, backend=backend, options=build_options)
+                .with_externals(externals or {})
+                .with_dtypes(dtypes or {})
             )
         if eager:
             stencil = stencil.implementation
diff --git a/src/gt4py/cartesian/loader.py b/src/gt4py/cartesian/loader.py
index 74a74951..83d7c6aa 100644
--- a/src/gt4py/cartesian/loader.py
+++ b/src/gt4py/cartesian/loader.py
@@ -36,6 +36,7 @@ def load_stencil(
     backend_name: str,
     definition_func: StencilFunc,
     externals: Dict[str, Any],
+    dtypes: Dict[Type, Type],
     build_options: "BuildOptions",
 ) -> Type["StencilObject"]:
     """Generate a new class object implementing the provided definition."""
@@ -48,9 +49,13 @@ def load_stencil(
     if frontend is None:
         raise ValueError(f"Invalid frontend name ({frontend_name})")
 
-    builder = StencilBuilder(
-        definition_func, options=build_options, backend=backend_name, frontend=frontend
-    ).with_externals(externals)
+    builder = (
+        StencilBuilder(
+            definition_func, options=build_options, backend=backend_name, frontend=frontend
+        )
+        .with_externals(externals)
+        .with_dtypes(dtypes)
+    )
 
     return builder.build()
 
@@ -60,12 +65,15 @@ def gtscript_loader(
     backend: str,
     build_options: "BuildOptions",
     externals: Dict[str, Any],
+    dtypes: Dict[Type, Type],
 ) -> "StencilObject":
     if not isinstance(definition_func, types.FunctionType):
         raise ValueError("Invalid stencil definition object ({obj})".format(obj=definition_func))
 
     if not build_options.name:
         build_options.name = f"{definition_func.__name__}"
-    stencil_class = load_stencil("gtscript", backend, definition_func, externals, build_options)
+    stencil_class = load_stencil(
+        "gtscript", backend, definition_func, externals, dtypes, build_options
+    )
 
     return stencil_class()
diff --git a/src/gt4py/cartesian/stencil_builder.py b/src/gt4py/cartesian/stencil_builder.py
index e43e4f5b..eb015aa2 100644
--- a/src/gt4py/cartesian/stencil_builder.py
+++ b/src/gt4py/cartesian/stencil_builder.py
@@ -79,6 +79,7 @@ class StencilBuilder:
         self.frontend: Type["FrontendType"] = frontend
         self.with_caching("jit")
         self._externals: Dict[str, Any] = {}
+        self._dtypes: Dict[Type, Type] = {}
 
     def build(self) -> Type["StencilObject"]:
         """Generate, compile and/or load everything necessary to provide a usable stencil class."""
@@ -218,6 +219,12 @@ class StencilBuilder:
             "externals", self._externals.copy()
         )
 
+    @property
+    def dtypes(self) -> Dict[Type, Type]:
+        return self._build_data.get("dtypes") or self._build_data.setdefault(
+            "dtypes", self._dtypes.copy()
+        )
+
     def with_externals(self: "StencilBuilder", externals: Dict[str, Any]) -> "StencilBuilder":
         """
         Fluidly set externals for this build.
@@ -229,6 +236,12 @@ class StencilBuilder:
         self.with_caching(self.caching.name)
         return self
 
+    def with_dtypes(self: "StencilBuilder", dtypes: Dict[Type, Type]) -> "StencilBuilder":
+        self._build_data = {}
+        self._dtypes = dtypes
+        self.with_caching(self.caching.name)
+        return self
+
     @property
     def backend_data(self) -> Dict[str, Any]:
         return self._build_data.get("backend_data", {}).copy()
@@ -268,7 +281,7 @@ class StencilBuilder:
         return self._build_data.get("gtir_pipeline") or self._build_data.setdefault(
             "gtir_pipeline",
             GtirPipeline(
-                self.frontend.generate(self.definition, self.externals, self.options),
+                self.frontend.generate(self.definition, self.externals, self.dtypes, self.options),
                 self.stencil_id,
             ),
         )


Test patch:
diff --git a/tests/cartesian_tests/unit_tests/test_gtc/test_gtir.py b/tests/cartesian_tests/unit_tests/test_gtc/test_gtir.py
index 748028af..61bb99a4 100644
--- a/tests/cartesian_tests/unit_tests/test_gtc/test_gtir.py
+++ b/tests/cartesian_tests/unit_tests/test_gtc/test_gtir.py
@@ -181,5 +181,5 @@ def test_variable_k_offset_in_access():
 
 
 def test_visit_ScalarIf_HorizontalMask_fail():
-    with pytest.raises(TypeError):
+    with pytest.raises(TypeError, match="HorizontalMask"):
         ScalarIfStmtFactory(cond=HorizontalMaskFactory())
diff --git a/tests/cartesian_tests/unit_tests/test_gtc_backend/test_defir_to_gtir.py b/tests/cartesian_tests/unit_tests/test_gtc_backend/test_defir_to_gtir.py
index 9e95f9cc..e40f62d7 100644
--- a/tests/cartesian_tests/unit_tests/test_gtc_backend/test_defir_to_gtir.py
+++ b/tests/cartesian_tests/unit_tests/test_gtc_backend/test_defir_to_gtir.py
@@ -12,9 +12,10 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+import numpy as np
 import pytest
 
-from gt4py.cartesian.frontend.defir_to_gtir import DefIRToGTIR
+from gt4py.cartesian.frontend.defir_to_gtir import DefIRToGTIR, _make_literal
 from gt4py.cartesian.frontend.nodes import (
     AxisBound,
     AxisInterval,
@@ -170,3 +171,30 @@ def test_field_decl_dims(defir_to_gtir, axes, expected_mask):
     field_decl = FieldDecl(name="a", data_type=DataType.INT64, axes=axes, is_api=True)
     gtir_decl = defir_to_gtir.visit_FieldDecl(field_decl)
     assert gtir_decl.dimensions == expected_mask
+
+
+def test_make_literal():
+    # All of those are o.k.
+    gtir_lit = _make_literal(10.10)
+    assert gtir_lit.dtype == common.DataType.FLOAT64
+    gtir_lit = _make_literal(np.float64(10.10))
+    assert gtir_lit.dtype == common.DataType.FLOAT64
+    gtir_lit = _make_literal(np.float32(10.10))
+    assert gtir_lit.dtype == common.DataType.FLOAT32
+    gtir_lit = _make_literal(10)
+    assert gtir_lit.dtype == common.DataType.INT64
+    gtir_lit = _make_literal(np.int64(10))
+    assert gtir_lit.dtype == common.DataType.INT64
+    gtir_lit = _make_literal(np.int32(10))
+    assert gtir_lit.dtype == common.DataType.INT32
+    gtir_lit = _make_literal(np.int16(10))
+    assert gtir_lit.dtype == common.DataType.INT16
+    gtir_lit = _make_literal(np.int8(10))
+    assert gtir_lit.dtype == common.DataType.INT8
+    gtir_lit = _make_literal(True)
+    assert gtir_lit.dtype == common.DataType.BOOL
+    gtir_lit = _make_literal(np.bool_(True))
+    assert gtir_lit.dtype == common.DataType.BOOL
+    # Not allowed
+    with pytest.raises(TypeError):
+        gtir_lit = _make_literal("a")
diff --git a/tests/cartesian_tests/unit_tests/test_gtscript_frontend.py b/tests/cartesian_tests/unit_tests/test_gtscript_frontend.py
index eb3dac82..c37b2d8b 100644
--- a/tests/cartesian_tests/unit_tests/test_gtscript_frontend.py
+++ b/tests/cartesian_tests/unit_tests/test_gtscript_frontend.py
@@ -71,7 +71,7 @@ def parse_definition(
         definition_func, externals=externals or {}
     )
     definition_ir = gt_frontend.GTScriptParser(
-        definition_func, externals=externals or {}, options=build_options
+        definition_func, externals=externals or {}, options=build_options, dtypes=dtypes
     ).run()
 
     setattr(definition_func, "__annotations__", original_annotations)
@@ -1318,6 +1318,32 @@ class TestDTypes:
             )
 
 
+class TestBuiltinDTypes:
+    @staticmethod
+    def literal_add_func(
+        in_field: gtscript.Field[float],
+        out_field: gtscript.Field["my_float"],
+    ):
+        with computation(PARALLEL), interval(...):
+            out_field = in_field + 42.0
+
+    @pytest.mark.parametrize("the_float", [np.float32, np.float64])
+    def test_literal_floating_parametrization(self, the_float):
+        def_ir = parse_definition(
+            definition_func=self.literal_add_func,
+            name=inspect.stack()[0][3],
+            module=self.__class__.__name__,
+            dtypes={float: the_float, "my_float": the_float},
+        )
+        # Check fields dtype are replaced only when str is used
+        def_ir.api_fields[0].data_type == nodes.DataType.FLOAT64
+        def_ir.api_fields[1].data_type == nodes.DataType.from_dtype(the_float)
+        # Check scalar literal dtype is always replaced
+        assert def_ir.computations[0].body.stmts[
+            0
+        ].value.rhs.data_type == nodes.DataType.from_dtype(the_float)
+
+
 class TestAssignmentSyntax:
     def test_ellipsis(self):
         def func(in_field: gtscript.Field[np.float_], out_field: gtscript.Field[np.float_]):