Repository: GridTools/gt4py
Pull request #1182 (2023-02-22 21:10:59)

Fix typing of `cast_` builtin in ITIR type inference
The typing inference rule for `cast_` was not correctly implemented when the builtin was introduced. This PR fixes that by making `cast_` part of the grammar (similar to how this is handled in the frontend). This restriction could be lifted in the future (e.g. by introducting a type of a type).

tehrengruber: @havogt Please also take a quick look at the implication of making `cast_` part of the grammar (see the test for an example).

Patch:
diff --git a/src/gt4py/next/iterator/type_inference.py b/src/gt4py/next/iterator/type_inference.py
index 6ddbd807..237fe334 100644
--- a/src/gt4py/next/iterator/type_inference.py
+++ b/src/gt4py/next/iterator/type_inference.py
@@ -293,7 +293,6 @@ BUILTIN_TYPES: typing.Final[dict[str, Type]] = {
         ret=Val_BOOL_T1,
     ),
     "if_": FunctionType(args=Tuple.from_elems(Val_BOOL_T1, Val_T0_T1, Val_T0_T1), ret=Val_T0_T1),
-    "cast_": FunctionType(args=Tuple.from_elems(Val_T0_T1, Val_T0_T1), ret=Val_T0_T1),
     "lift": FunctionType(
         args=Tuple.from_elems(
             FunctionType(
@@ -402,6 +401,7 @@ class _TypeInferrer(eve.NodeTranslator):
             "shift",
             "cartesian_domain",
             "unstructured_domain",
+            "cast_",
         ):
             raise TypeError(
                 f"Builtin '{node.id}' is only supported as applied/called function by the type checker"
@@ -477,6 +477,32 @@ class _TypeInferrer(eve.NodeTranslator):
                 )
                 constraints.add((tup, val))
                 return Val(kind=kind, dtype=elem, size=size)
+            if node.fun.id == "cast_":
+                if len(node.args) != 2:
+                    raise TypeError("cast_ requires exactly two arguments.")
+                val_arg_type = self.visit(node.args[0], constraints=constraints, symtypes=symtypes)
+                type_arg = node.args[1]
+                if not isinstance(type_arg, ir.SymRef) or type_arg.id not in ir.TYPEBUILTINS:
+                    raise TypeError("The second argument to `cast_` must be a type literal.")
+
+                size = TypeVar.fresh()
+
+                constraints.add(
+                    (
+                        val_arg_type,
+                        Val(
+                            kind=Value(),
+                            dtype=TypeVar.fresh(),
+                            size=size,
+                        ),
+                    )
+                )
+
+                return Val(
+                    kind=Value(),
+                    dtype=Primitive(name=type_arg.id),
+                    size=size,
+                )
             if node.fun.id == "shift":
                 # Calls to shift are handled as being part of the grammar, not
                 # as function calls, as the type depends on the offset provider


Test patch:
diff --git a/tests/next_tests/iterator_tests/test_type_inference.py b/tests/next_tests/iterator_tests/test_type_inference.py
index 0db355c8..ee6bf711 100644
--- a/tests/next_tests/iterator_tests/test_type_inference.py
+++ b/tests/next_tests/iterator_tests/test_type_inference.py
@@ -146,6 +146,16 @@ def test_and():
     assert ti.pformat(inferred) == "(bool⁰, bool⁰) → bool⁰"
 
 
+def test_cast():
+    testee = ir.FunCall(
+        fun=ir.SymRef(id="cast_"), args=[ir.Literal(value="1.", type="float"), ir.SymRef(id="int")]
+    )
+    expected = ti.Val(kind=ti.Value(), dtype=ti.Primitive(name="int"), size=ti.TypeVar(idx=0))
+    inferred = ti.infer(testee)
+    assert inferred == expected
+    assert ti.pformat(inferred) == "int⁰"
+
+
 def test_lift():
     testee = ir.SymRef(id="lift")
     expected = ti.FunctionType(