diff --git a/ndsl/dsl/dace/orchestration.py b/ndsl/dsl/dace/orchestration.py index 38e7be09..00a2bd9a 100644 --- a/ndsl/dsl/dace/orchestration.py +++ b/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 ( diff --git a/ndsl/dsl/dace/replacements.py b/ndsl/dsl/dace/replacements.py new file mode 100644 index 00000000..4dec53ab --- /dev/null +++ b/ndsl/dsl/dace/replacements.py @@ -0,0 +1,32 @@ +"""This module uses DaCe's op_repository feature to override symbols/AST objects +during parsing and replace them with an SDFG compatible representation. This +allows custom NDSL syntax, objects and symbols 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), + )