[Orchestration] Common cast operation replacements

ndsl/dsl/dace/orchestration.py

@@ -18,6 +18,7 @@
 from dace.transformation.passes.simplify import SimplifyPass
 from gt4py import storage
 
+import ndsl.dsl.dace.replacements  # noqa # We load in the DaCe replacements
 from ndsl.comm.mpi import MPI
 from ndsl.dsl.dace.build import get_sdfg_path, write_build_info
 from ndsl.dsl.dace.dace_config import (

ndsl/dsl/dace/replacements.py

@@ -0,0 +1,32 @@
+"""This module uses DaCe's op_repository feature to override symbols/AST object
+during parsing and replace them with an SDFG compatible representation. This
+allow custom NDSL system to be natively orchestratable."""
+
+from dace import SDFG, SDFGState, dtypes
+from dace.frontend.common import op_repository as oprepo
+from dace.frontend.python.newast import ProgramVisitor
+from dace.frontend.python.replacements import UfuncInput, _datatype_converter
+
+from ndsl.dsl.typing import Float, Int
+
+
+@oprepo.replaces("Float")
+def _convert_Float(_pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, arg: UfuncInput):
+    """Replace `Float(x)` with a typecast of `x` to the proper floating precision type"""
+    return _datatype_converter(
+        sdfg,
+        state,
+        arg,
+        dtype=dtypes.dtype_to_typeclass(Float),
+    )
+
+
+@oprepo.replaces("Int")
+def _convert_Int(_pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, arg: UfuncInput):
+    """Replace `Int(x)` with a typecast of `x` to the proper integer precision type"""
+    return _datatype_converter(
+        sdfg,
+        state,
+        arg,
+        dtype=dtypes.dtype_to_typeclass(Int),
+    )