DO NOT MERGE feat[next-dace]: Use nanobind Interface of DaCe

pyproject.toml

@@ -471,6 +471,9 @@ url = 'https://gridtools.github.io/pypi/'
 # Add the uv source below to pull dace from the gridtools index instead of PyPI:
 [tool.uv.sources]
 atlas4py = {index = "test.pypi"}
+dace = [
+  {git = "https://github.com/philip-paul-mueller/dace", branch = "nanobind-compiled-sdfg"}
+]
 
 # -- versioningit --
 [tool.versioningit]

src/gt4py/next/embedded/nd_array_field.py

@@ -189,6 +189,11 @@ def __gt_origin__(self) -> tuple[int, ...]:
         assert common.Domain.is_finite(self.domain)
         return tuple(-r.start for r in self.domain.ranges)
 
+    @functools.cached_property
+    def __dace_origin__(self) -> tuple[int, ...]:
+        assert common.Domain.is_finite(self.domain)
+        return tuple(r.start for r in self.domain.ranges)
+
     @functools.cached_property
     def __gt_buffer_info__(self) -> common.BufferInfo:
         """

src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py

@@ -1018,7 +1018,13 @@ def _add_sdfg_params(
                 transient=True,
             )
 
-        # the list of all sdfg arguments (aka non-transient arrays) which include tuple-element fields
+        # The list of the SDFG arguments, i.e. the global arrays, scalars and free symbols.
+        #  Note that tuple arguments are flattened and their name is mangled and no longer
+        #  matches the name in the signature of the field operator / program.
+        #  Also note that some scalar arguments (which are lowered to symbols) listed in
+        #  this signature are might not be part of the generated C-API, as unused symbols
+        #  are exluded from it and might have been pruned from the SDFG during optimization.
+        # NOTE: The dispatch code does not use it, instead the `user_args` are used.
         return [arg_name for arg_name, _ in sdfg_args]
 
     def visit_Program(self, node: gtir.Program) -> dace.SDFG:
@@ -1064,10 +1070,10 @@ def visit_Program(self, node: gtir.Program) -> dace.SDFG:
                 assert isinstance(nsdfg.arrays[data], dace.data.Array)
                 nsdfg.arrays.pop(data)
 
-        # Create the call signature for the SDFG.
-        #  Only the arguments required by the GT4Py program, i.e. `node.params`, are added
-        #  as positional arguments. The implicit arguments, such as the offset providers or
-        #  the arguments created by the translation process, must be passed as keyword arguments.
+        # NOTE: A program uses the `user_args` mechanism to perform the call. So,
+        #   technically these arguments is not needed. However, the orchestrator needs
+        #   it to work. Note that in the following list, tuple arguments to the
+        #   program/fieldop are expanded.
         sdfg.arg_names = sdfg_arg_names
 
         return sdfg

src/gt4py/next/program_processors/runners/dace/workflow/common.py

@@ -84,6 +84,10 @@ def set_dace_config(
     #  `gt4py.next.program_processors.runners.dace.transfromations.gpu_utils.gt_gpu_transform_non_standard_memlet()`.
     dace.Config.set("compiler.cuda.allow_implicit_memlet_to_map", value=False)
 
+    # FORCE NANOBIND
+    dace.Config.set("compiler.interface", value="nanobind")
+    dace.Config.set("compiler.nanobind_name_collision", value="error")
+
     if cmake_build_type is not None:
         dace.Config.set("compiler.build_type", value=cmake_build_type.value)
 

src/gt4py/next/program_processors/runners/dace/workflow/compilation.py

@@ -9,12 +9,10 @@
 from __future__ import annotations
 
 import dataclasses
-import json
 import os
 import pathlib
-import warnings
-from collections.abc import Callable, MutableSequence, Sequence
-from typing import Any, Final, TypeAlias
+from collections.abc import Callable
+from typing import Any, Final, Sequence, TypeAlias
 
 import dace
 import dace.codegen.compiler as dace_compiler
@@ -73,30 +71,11 @@ def _add_tx_markers(program_source: SDFGExtensionSource) -> tuple[SDFGExtensionS
 class CompiledDaceProgram:
     sdfg_program: dace.CompiledSDFG
 
-    # Sorted list of SDFG arguments as they appear in program ABI and corresponding data type;
-    # scalar arguments that are not used in the SDFG will not be present.
-    sdfg_argtypes: list[dace.dtypes.Data]
-
-    # The compiled program contains a callable object to update the SDFG arguments list.
-    update_sdfg_ctype_arglist: Callable[
-        [
-            core_defs.DeviceType,
-            Sequence[dace.dtypes.Data],
-            Sequence[Any],
-            MutableSequence[Any],
-            common.OffsetProvider,
-        ],
-        None,
+    # Callable to process the GT4Py arguments and offset providers to bring them in a form suitable for calling.
+    argument_preprocessing_function: Callable[
+        [Sequence[Any], common.OffsetProvider, int, Any], tuple[Any, ...]
     ]
 
-    # Processed argument vectors that are passed to `CompiledSDFG.fast_call()`. `None`
-    #  means that it has not been initialized, i.e. no call was ever performed.
-    #  - csdfg_argv: Arguments used for calling the actual compiled SDFG, will be updated.
-    #  - csdfg_init_argv: Arguments used for initialization; used only the first time and
-    #       never updated.
-    csdfg_argv: MutableSequence[Any] | None
-    csdfg_init_argv: Sequence[Any] | None
-
     def __init__(
         self,
         program: dace.CompiledSDFG,
@@ -105,91 +84,41 @@ def __init__(
     ):
         self.sdfg_program = program
 
-        # `dace.CompiledSDFG.arglist()` returns an ordered dictionary that maps the argument
-        # name to its data type, in the same order as arguments appear in the program ABI.
-        # This is also the same order of arguments in `dace.CompiledSDFG._lastargs[0]`.
-        self.sdfg_argtypes = list(program.sdfg.arglist().values())
-
         # The binding source code is Python tailored to this specific SDFG.
         # We dynamically compile that function and add it to the compiled program.
         global_namespace: dict[str, Any] = {}
         exec(binding_source_code, global_namespace)
-        self.update_sdfg_ctype_arglist = global_namespace[bind_func_name]
-        # For debug purpose, we set a unique module name on the compiled function.
-        self.update_sdfg_ctype_arglist.__module__ = os.path.basename(program.sdfg.build_folder)
 
-        # Since the SDFG hasn't been called yet.
-        self.csdfg_argv = None
-        self.csdfg_init_argv = None
-
-    def construct_arguments(self, **kwargs: Any) -> None:
-        """
-        This function will process the arguments and store the processed argument
-        vectors in `self.csdfg_args`, to call them use `self.fast_call()`.
-        """
-        with dace.config.set_temporary("compiler", "allow_view_arguments", value=True):
-            csdfg_argv, csdfg_init_argv = self.sdfg_program.construct_arguments(**kwargs)
-        # Note we only care about `csdfg_argv` (normal call), since we have to update it,
-        #  we ensure that it is a `list`.
-        self.csdfg_argv = [*csdfg_argv]
-        self.csdfg_init_argv = csdfg_init_argv
-
-    def fast_call(self) -> None:
-        """
-        Perform a call to the compiled SDFG using the previously generated argument
-        vectors, see `self.construct_arguments()`.
-        """
-        assert self.csdfg_argv is not None and self.csdfg_init_argv is not None, (
-            "Argument vector was not set properly."
-        )
-        self.sdfg_program.fast_call(
-            self.csdfg_argv, self.csdfg_init_argv, do_gpu_check=config.DEBUG
+        self.argument_preprocessing_function = global_namespace[bind_func_name]
+        # For debug purpose, we set a unique module name on the compiled function.
+        self.argument_preprocessing_function.__module__ = os.path.basename(
+            program.sdfg.build_folder
         )
 
     def __call__(self, **kwargs: Any) -> None:
         """Call the compiled SDFG with the given arguments.
 
-        Note that this function will not update the argument vectors stored inside
-        `self`. Furthermore, it is not recommended to use this function as it is
-        very slow.
+        A `CompiledDaceProgram` should not be called directly. Instead
+        `gt4py.next.program_processors.runners.dace.workflow.decoration.convert_args()`
+        should be used to obtain a callable.
         """
-        warnings.warn(
-            "Called an SDFG through the standard DaCe interface is not recommended, use `fast_call()` instead.",
-            stacklevel=1,
+        raise NotImplementedError(
+            "A `CompiledDaceProgram` can not be called directly. Instead use "
+            "`gt4py.next.program_processors.runners.dace.workflow.decoration.convert_args()`."
         )
-        result = self.sdfg_program(**kwargs)
-        assert result is None
 
 
 @dataclasses.dataclass(frozen=True)
 class DaCeCompilationArtifact:
-    """Result of a DaCe compilation: library path + SDFG bindings + the SDFG itself.
-
-    The SDFG is carried inline as JSON because dace's load path
-    (``get_program_handle``) needs an SDFG instance to wrap into the
-    returned ``CompiledSDFG``, and the build folder may not contain a
-    ``program.sdfg(z)`` dump under the upcoming minimal-build-dir mode.
-
-    The SDFG we store here is the one on which we called `SDFG.compile(return_program_handle=False)`.
-    Note that the `compile()` call has side effects, because it applies transformations
-    to the SDFG, in order to enable code generation for the target platform.
-    Since we pass `return_program_handle=False`, the `compile()` method does not
-    return a `CompiledSDFG` instance, therefore we cannot access `CompiledSDFG.sdfg`,
-    which would be the modified SDFG from which DaCe generates the C++/CUDA/HIP code.
-    """
+    """Result of a DaCe compilation: library path + SDFG bindings + the SDFG itself."""
 
-    library_path: pathlib.Path
-    sdfg_json: str
+    sdfg_build_folder: pathlib.Path
     binding_source_code: str
     bind_func_name: str
     device_type: core_defs.DeviceType
 
     def load(self) -> stages.ExecutableProgram:
-        # TODO(phimuell): Drop ``sdfg_json`` from the artifact once dace
-        #   exposes a load path that doesn't require an SDFG instance to wrap
-        #   into the returned ``CompiledSDFG``.
-        sdfg = dace.SDFG.from_json(json.loads(self.sdfg_json))
-        sdfg_program = dace_compiler.get_program_handle(self.library_path, sdfg)
+        sdfg_program = dace_compiler.load_precompiled_sdfg(self.sdfg_build_folder, sdfg=None)
         program = CompiledDaceProgram(sdfg_program, self.bind_func_name, self.binding_source_code)
         return gtx_wfddecoration.convert_args(program, device=self.device_type)
 
@@ -217,7 +146,8 @@ class DaCeCompiler(
     cmake_build_type: config.CMakeBuildType = dataclasses.field(
         default_factory=lambda: config.CMAKE_BUILD_TYPE
     )
-    # we store the non-default values of `dace.Config` in order to include it in the stage fingerprint
+    # We store the non-default values of `dace.Config` in order to include it in the stage fingerprint
+    # NOTE: They do not include the non default keys set through DaCe environment variables.
     dace_config_nondefaults: dict[str, Any] = dataclasses.field(init=False)
 
     def __post_init__(self) -> None:
@@ -250,8 +180,8 @@ def __call__(self, inp: SDFGExtensionSource) -> DaCeCompilationArtifact:
             # Configure the SDFG build folder
             sdfg.build_folder = sdfg_build_folder
 
-            # ``build_folder_mode`` is set by ``dace_context``; resolve the library
-            # path here so ``get_binary_name`` sees the same mode dace built under.
+            # `compiler.build_folder_mode` is set by `dace_context()`; resolve the library
+            #  path here so `get_binary_name()` sees the same mode DaCe built under.
             library_path = dace_compiler.get_binary_name(
                 object_folder=sdfg_build_folder, sdfg_name=sdfg.name
             )
@@ -275,8 +205,7 @@ def __call__(self, inp: SDFGExtensionSource) -> DaCeCompilationArtifact:
 
         assert inp.binding_source is not None
         return DaCeCompilationArtifact(
-            library_path=library_path,
-            sdfg_json=json.dumps(inp.program_source.source_code),
+            sdfg_build_folder=sdfg_build_folder,
             binding_source_code=inp.binding_source.source_code,
             bind_func_name=self.bind_func_name,
             device_type=self.device_type,

src/gt4py/next/program_processors/runners/dace/workflow/decoration.py

@@ -8,16 +8,14 @@
 
 from __future__ import annotations
 
-import functools
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 from gt4py._core import definitions as core_defs
-from gt4py.next import common as gtx_common, utils as gtx_utils
+from gt4py.next import common as gtx_common
 from gt4py.next.instrumentation import metrics
 from gt4py.next.otf import stages
-from gt4py.next.program_processors.runners.dace import sdfg_callable
 from gt4py.next.program_processors.runners.dace.workflow import common as gtx_wfdcommon
 
 
@@ -35,10 +33,7 @@ def convert_args(
     collect_time_arg = np.array(
         [1], dtype=gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME_DTYPE.as_numpy_dtype()
     )
-    # We use the callback function provided by the compiled program to update the SDFG arglist.
-    update_sdfg_call_args = functools.partial(
-        fun.update_sdfg_ctype_arglist, device, fun.sdfg_argtypes
-    )
+    argument_preprocessing_function = fun.argument_preprocessing_function
 
     def decorated_program(
         *args: Any,
@@ -48,33 +43,13 @@ def decorated_program(
         if out is not None:
             args = (*args, out)
 
-        try:
-            # Not the first call.
-            #  We will only update the argument vector  for the normal call.
-            # NOTE: If this is the first time then we will generate an exception because
-            #   `fun.csdfg_args` is `None`
-            # TODO(phimuell, edopao): Think about refactor the code such that the update
-            #   of the argument vector is a Method of the `CompiledDaceProgram`.
-            update_sdfg_call_args(args, fun.csdfg_argv, offset_provider)  # type: ignore[arg-type]  # Will error out in first call.
-
-        except TypeError:
-            # First call. Construct the initial argument vector of the `CompiledDaceProgram`.
-            assert fun.csdfg_argv is None and fun.csdfg_init_argv is None
-            flat_args: Sequence[Any] = gtx_utils.flatten_nested_tuple(args)
-            this_call_args = sdfg_callable.get_sdfg_args(
-                fun.sdfg_program.sdfg,
-                offset_provider,
-                *flat_args,
-                filter_args=False,
-            )
-            this_call_args |= {
-                gtx_wfdcommon.SDFG_ARG_METRIC_LEVEL: metrics.get_current_level(),
-                gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME: collect_time_arg,
-            }
-            fun.construct_arguments(**this_call_args)
-
-        # Perform the call to the SDFG.
-        fun.fast_call()
+        processed_args, _ = argument_preprocessing_function(
+            args,
+            offset_provider,
+            metrics.get_current_level(),
+            collect_time_arg,
+        )
+        fun.sdfg_program.user_bind_call(*processed_args)
 
         if collect_time:
             metrics.add_sample_to_current_source(metrics.COMPUTE_METRIC, collect_time_arg[0].item())