diff --git a/pyproject.toml b/pyproject.toml index 6b5a08f78a..42892bc2ad 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index 01f933fde3..2c86b44a2a 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/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: """ diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py index 8a8d42f9ba..8e9e13fac9 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py +++ b/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 diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/bindings.py b/src/gt4py/next/program_processors/runners/dace/workflow/bindings.py index 172668af65..59a7e2e363 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/bindings.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/bindings.py @@ -8,281 +8,140 @@ from __future__ import annotations -from typing import Final +from typing import Any, Literal, Optional, Sequence import dace -from gt4py.eve import codegen +from gt4py.next import common as gtx_common from gt4py.next.otf import code_specs, stages -from gt4py.next.program_processors.runners.dace import sdfg_args as gtx_dace_args +from gt4py.next.otf.binding import interface +from gt4py.next.program_processors.runners.dace import sdfg_args as gtx_sdfg_args +from gt4py.next.program_processors.runners.dace.workflow import common as gtx_wfdcommon from gt4py.next.type_system import type_specifications as ts -_cb_args: Final[str] = "args" -_cb_device: Final[str] = "device" -_cb_sdfg_argtypes: Final[str] = "sdfg_argtypes" -_cb_sdfg_call_args: Final[str] = "sdfg_call_args" -_cb_neighbor_table: Final[str] = "table" -_cb_offset_provider: Final[str] = "offset_provider" - - -def _update_sdfg_array_ptr(code: codegen.TextBlock, arg: str, sdfg_arg_index: int) -> None: - code.append(f"assert field_utils.verify_device_field_type({arg}, {_cb_device})") - code.append(f"assert isinstance({_cb_sdfg_call_args}[{sdfg_arg_index}], ctypes.c_void_p)") - code.append(f"{_cb_sdfg_call_args}[{sdfg_arg_index}].value = {arg}.__gt_buffer_info__.data_ptr") - - -def _update_sdfg_array_strides( - code: codegen.TextBlock, - sdfg_arglist: dict[str, dace.data.Data], - arg: str, - sdfg_arg_desc: dace.data.Array, - sdfg_arg_index: int, -) -> None: - for i, array_stride in enumerate(sdfg_arg_desc.strides): - arg_stride = f"{arg}.__gt_buffer_info__.elem_strides[{i}]" - if isinstance(array_stride, int) or str(array_stride).isdigit(): - # The array stride is set to constant value in this dimension. - code.append( - f"assert {_cb_sdfg_argtypes}[{sdfg_arg_index}].strides[{i}] == {arg_stride}" - ) - else: - # The strides of a global array are defined by a sequence of SDFG symbols. - _parse_gt_param( - param_name=array_stride.name, - param_type=gtx_dace_args.as_itir_type(array_stride.dtype), - arg=arg_stride, - code=code, - sdfg_arglist=sdfg_arglist, - ) +def argument_processing_function( + args: tuple[Any, ...], + offset_provider: gtx_common.OffsetProvider, + metrics_level: int, + runtime_return_value: Any, +) -> tuple[tuple[Any, ...], dict[str, Any]]: + """Signature generated by `_create_sdfg_bindings()` and used to process the arguments for an SDFG call. + See `_create_sdfg_bindings()` and `bind_sdfg()` for more. -def _update_sdfg_scalar_arg( - code: codegen.TextBlock, - sdfg_arg_desc: dace.data.Data, - sdfg_arg_index: int, - call_arg: str, -) -> None: - """ - Emit Python code to update a scalar argument in the SDFG arglist - with the argument value passed to the gt4py program call. - """ - assert isinstance(sdfg_arg_desc, dace.data.Scalar) - actype = sdfg_arg_desc.dtype.as_ctypes() - actype_call = f"{actype.__module__}.{actype.__name__}" - code.append(f"assert isinstance({_cb_sdfg_call_args}[{sdfg_arg_index}], ctypes._SimpleCData)") - code.append(f"{_cb_sdfg_call_args}[{sdfg_arg_index}] = {actype_call}({call_arg})") - - -def _unpack_args(code: codegen.TextBlock, num_args: int, arg_name: str) -> list[str]: - """Unpack a sequence of arguments (either a list or a tuple) into variables. - - Each element of the given sequence gets a name 'arg_name' with an index-based suffix. - - >>> code = codegen.TextBlock() - >>> _unpack_args(code, 3, "var") - ['var_0', 'var_1', 'var_2'] - >>> code.lines - ['var_0, var_1, var_2, = var'] - >>> _unpack_args(code, 1, "var_2") - ['var_2_0'] - >>> code.lines - ['var_0, var_1, var_2, = var', 'var_2_0, = var_2'] + Args: + args: The arguments that are passed to the SDFG. A tuple of fields and nested tuples. + offset_provider: The offset providers (order is used verbatim). + metrics_level: The level of the metric that is active. + runtime_return_value: NumPy array used to write back the runtime spend in the call. """ - tuple_args = [f"{arg_name}_{i}" for i in range(num_args)] - if num_args == 0: - raise ValueError("Cannot unpack argument with length zero.") - else: - # The trailing comma is needed to unpack single-element tuples - code.append(f"{', '.join(tuple_args)}, = {arg_name}") - return tuple_args + raise NotImplementedError("`argument_processing_function()` is only for documenting purposes.`") -def _parse_gt_param( - param_name: str, - param_type: ts.DataType, - arg: str, - code: codegen.TextBlock, - sdfg_arglist: dict[str, dace.data.Data], -) -> None: - """Emit Python code to parse a program argument and set the required fields in the SDFG arglist. +def _create_sdfg_bindings( + program_parameters: Sequence[interface.Parameter], + bind_func_name: str, + use_metrics: bool, + sdfg: Optional[dace.SDFG], + backend: Literal["gtfn", "dace"], +) -> str: + """ + Creates a Python function that translates GT4Py arguments into arguments suitable for a SDFG call. - For scalar arguments, a single field is set in the SDFG arglist. + The returned string can be passed to exec to then dynamically create a function to + perform the translation. The function will have the same signature as + `argument_processing_function()` and have the name `bind_func_name`. + The function that is generated targets the interface defined through the + `user_args` defined interface of the SDFG. - For array arguments, in addition to the data pointer, the fields for array shape - and strides are also set in SDFG arglist. This results in nested calls to - `_parse_gt_param()` with the scalar values of array shape and strides. + The function has experimental support for GTFN. - For tuple arguments, this function is recursively called on all elements of the tuple. + Args: + prog: The program definition. + bind_func_name: The name of the function to perform the translation. + use_metrics: If metric support was added to the underlying compiled code. + In that case the last two arguments of the generated function signature + are ignored and also not included in the output. + sdfg: If provided and `backend` is `dace` then only the offset providers that + are used in the SDFG are processed, otherwise `None` is passed. + This is compatible with the `user_args` signature. + backend: For which backend the generated function should be used. """ - if isinstance(param_type, ts.TupleType): - # Each element of a tuple gets a name with an index-based suffix and it is recursively visited. - tuple_args = _unpack_args(code=code, num_args=len(param_type.types), arg_name=arg) - for i, (tuple_arg, tuple_arg_type) in enumerate(zip(tuple_args, param_type.types)): - assert isinstance(tuple_arg_type, ts.DataType) - _parse_gt_param( - param_name=f"{param_name}_{i}", - param_type=tuple_arg_type, - arg=tuple_arg, - code=code, - sdfg_arglist=sdfg_arglist, - ) - - elif param_name not in sdfg_arglist: - # There are two reasons for this case: - # 1) The argument is a symbol/scalar that is not used in the generated code. - # 2) The argument was demoted, see `demote_fields` argument of `gt_auto_optimize()` - # and was not put back. - pass - - else: - sdfg_arg_desc = sdfg_arglist[param_name] - sdfg_arg_index = list(sdfg_arglist.keys()).index(param_name) - - if isinstance(param_type, ts.FieldType): - if len(param_type.dims) == 0: - # Pass zero-dimensional fields as scalars. - assert isinstance(sdfg_arg_desc, dace.data.Scalar) - _update_sdfg_scalar_arg( - code=code, - sdfg_arg_desc=sdfg_arg_desc, - sdfg_arg_index=sdfg_arg_index, - call_arg=f"{arg}.as_scalar()", - ) - else: - assert isinstance(sdfg_arg_desc, dace.data.Array) - _update_sdfg_array_ptr(code, arg, sdfg_arg_index) - for i, (dim, array_size) in enumerate( - zip(param_type.dims, sdfg_arg_desc.shape, strict=True) - ): - if isinstance(array_size, int) or str(array_size).isdigit(): - # The array shape in this dimension is set at compile-time. - code.append( - f"assert {_cb_sdfg_argtypes}[{sdfg_arg_index}].shape[{i}] == {arg}.__gt_buffer_info__.shape[{i}]" - ) - else: - # The array shape is defined as a sequence of expressions - # like 'range_stop - range_start', where 'range_start' and - # 'range_stop' are the SDFG symbols for the domain range. - arg_range = f"{arg}.domain.ranges[{i}]" - rstart = gtx_dace_args.range_start_symbol(param_name, dim) - rstop = gtx_dace_args.range_stop_symbol(param_name, dim) - for suffix, sdfg_range_symbol in [("start", rstart), ("stop", rstop)]: - _parse_gt_param( - param_name=sdfg_range_symbol.name, - param_type=gtx_dace_args.as_itir_type(sdfg_range_symbol.dtype), - arg=f"{arg_range}.{suffix}", - code=code, - sdfg_arglist=sdfg_arglist, - ) - _update_sdfg_array_strides(code, sdfg_arglist, arg, sdfg_arg_desc, sdfg_arg_index) - - elif isinstance(param_type, ts.ScalarType): - assert isinstance(sdfg_arg_desc, dace.data.Scalar) - _update_sdfg_scalar_arg( - code=code, - sdfg_arg_desc=sdfg_arg_desc, - sdfg_arg_index=sdfg_arg_index, - call_arg=arg, + assert backend in ["gtfn", "dace"] + assert backend == "dace" or sdfg is None + + if backend == "dace": + is_offset_arg_name = lambda arg_name: arg_name.startswith("@") # noqa: E731 [lambda-assignment] + get_offset_name_from_prog_arg_name = lambda arg_name: arg_name[1:] # noqa: E731 [lambda-assignment] + if sdfg is not None: + sdfg_arglist = sdfg.arglist() + is_offset_used = lambda offset_name: ( # noqa: E731 [lambda-assignment] + gtx_sdfg_args.connectivity_identifier(offset_name) in sdfg_arglist ) - else: - raise ValueError(f"Unexpected paramter type {param_type}") - - -def _parse_gt_connectivities( - code: codegen.TextBlock, sdfg_arglist: dict[str, dace.data.Data] -) -> None: - for sdfg_arg_index, (arg_name, arg_desc) in enumerate(sdfg_arglist.items()): - if gtx_dace_args.is_connectivity_identifier(arg_name): - assert isinstance(arg_desc, dace.data.Array) - assert len(arg_desc.shape) == 2 - assert isinstance(arg_desc.shape[1], int) or str(arg_desc.shape[1]).isdigit() - origin_size_arg = arg_desc.shape[0] - assert len(origin_size_arg.free_symbols) == 1 - origin_size_param = next(iter(origin_size_arg.free_symbols)) - m = gtx_dace_args.CONNECTIVITY_INDENTIFIER_RE.match(arg_name) - assert m is not None - conn_arg = f"{_cb_neighbor_table}_{m[1]}" - code.append(f'{conn_arg} = {_cb_offset_provider}["{m[1]}"]') - _update_sdfg_array_ptr(code, conn_arg, sdfg_arg_index) - _parse_gt_param( # set the size in the horizontal dimension - param_name=origin_size_param, - param_type=gtx_dace_args.as_itir_type(gtx_dace_args.FIELD_SYMBOL_DTYPE), - arg=f"{conn_arg}.__gt_buffer_info__.shape[0]", - code=code, - sdfg_arglist=sdfg_arglist, - ) - _update_sdfg_array_strides( - code, - sdfg_arglist, - conn_arg, - arg_desc, - sdfg_arg_index, - ) + is_offset_used = lambda _: True # noqa: E731 [lambda-assignment] + else: + raise NotImplementedError() + code_lines: list[str] = [] + code_lines.append("_PAIR_OF_ZEROS = (0, 0)") + code_lines.append( + f"def {bind_func_name}(args, offset_provider, metrics_level, runtime_return_value):" + ) -def _create_sdfg_bindings( - program_source: stages.ProgramSource[code_specs.SDFGCodeSpec], - bind_func_name: str, -) -> stages.BindingSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec]: - """ - Creates a Python translation function to convert the GT4Py arguments list - to the SDFG calling convention. + unpacked_program_arguments: list[Optional[str]] = [ + None if is_offset_arg_name(param.name) else f"__gtx_expanded_names_{param.name}" + for param in program_parameters + ] - Args: - program_source: The json representation of the SDFG. - bind_func_name: Name to use for the translation function. + unpacked_variables = [ + ", ".join(arg for arg in unpacked_program_arguments if arg is not None) + ", = args" + ] - Returns: - The Python code to convert call arguments from gt4py canonical form to the - SDFG canonical form. - """ - sdfg = dace.SDFG.from_json(program_source.source_code) - - # `dace.SDFG.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]`. - sdfg_arglist = sdfg.arglist() - - code = codegen.TextBlock() - - code.append("import ctypes") - code.empty_line() - code.append("from gt4py.next import common as gtx_common, field_utils") - code.empty_line() - code.append( - "def {funname}({arg0}, {arg1}, {arg2}, {arg3}, {arg4}):".format( - funname=bind_func_name, - arg0=_cb_device, - arg1=_cb_sdfg_argtypes, - arg2=_cb_args, - arg3=_cb_sdfg_call_args, - arg4=_cb_offset_provider, - ) - ) + positional_arguments = "" + kwargs_arguments: dict[str, str] = {} - # The SDFG binding function is used with fast-call, to update the SDFG arguments - # list, therefore it is only used from the second time the SDFG is called. - # On the first time, we use the regular SDFG call, which constructs the SDFG - # arguments list and validates that all data containers and free symbols are set. - with code.indented(): - arg_vars = _unpack_args( - code=code, num_args=len(program_source.entry_point.parameters), arg_name=_cb_args - ) - for param, arg in zip(program_source.entry_point.parameters, arg_vars): - assert isinstance(param.type_, ts.DataType) - _parse_gt_param(param.name, param.type_, arg, code, sdfg_arglist) + for arg_name, param in zip(unpacked_program_arguments, program_parameters, strict=True): + assert param.type_ is not None + real_arg_name = param.name - # In the regular case, the connectivity fields are allocated at the beginning - # of the application and then used during its entire lifetime and never reallocated. - # However, this might not be the case all the time, for example in unit tests - # where, due to limited lifetime of the fixtures, the connectivity fields - # might get reallocated. In order to avoid problems, we update the connectivity - # arrays as well in SDFG fastcall. - _parse_gt_connectivities(code, sdfg_arglist) + if is_offset_arg_name(real_arg_name): + assert arg_name is None + offset_name = get_offset_name_from_prog_arg_name(real_arg_name) + if is_offset_used(offset_name): + positional_arguments += ( + f"(offset_provider['{offset_name}'].ndarray, _PAIR_OF_ZEROS), " + ) + else: + positional_arguments += "None, " + else: + assert isinstance(arg_name, str) + positional_arguments += _process_argument( + argument=arg_name, + param_type=param.type_, + unpacked_variables=unpacked_variables, + backend=backend, + ) + assert positional_arguments.strip().endswith(",") + assert positional_arguments.strip().endswith(",") + + positional_arguments, kwargs_arguments = _process_metric_arguments( + metric_level_arg_name="metrics_level", + compute_time_arg_name="runtime_return_value", + positional_arguments=positional_arguments, + kwargs_arguments=kwargs_arguments, + backend=backend, + use_metrics=use_metrics, + ) + code_lines.extend(("\t" + s for s in unpacked_variables)) + code_lines.append( + f"\treturn ({positional_arguments}), {{" + + ", ".join(f"'{k}': {v}" for k, v in kwargs_arguments.items()) + + "}" + ) - return stages.BindingSource(code.text, library_deps=tuple()) + return "\n".join(code_lines) def bind_sdfg( @@ -292,9 +151,98 @@ def bind_sdfg( """ Method to be used as workflow stage for generation of SDFG bindings. - Refer to `_create_sdfg_bindings` documentation. + Refer to `_create_sdfg_bindings()` documentation. """ + assert isinstance(inp.source_code, dict) + + sdfg = dace.SDFG.from_json(inp.source_code) + use_metrics = gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME in sdfg.arrays + + bind_source = _create_sdfg_bindings( + program_parameters=inp.entry_point.parameters, + bind_func_name=bind_func_name, + use_metrics=use_metrics, + sdfg=sdfg, + backend="dace", + ) + binding_source = stages.BindingSource(bind_source, library_deps=tuple()) + return stages.ExtensionSource( program_source=inp, - binding_source=_create_sdfg_bindings(inp, bind_func_name), + binding_source=binding_source, ) + + +def _perfom_tuple_unpacking( + tuple_argument: str, + tuple_len: int, + unpacked_variables: list[str], +) -> list[str]: + assert tuple_len > 0 + unpacked_arguments = [f"{tuple_argument}_{i}" for i in range(tuple_len)] + # Trailing comma is needed to handle 1 element tuples correctly. + unpacked_variables.append(f"{', '.join(unpacked_arguments)}, = {tuple_argument}") + return unpacked_arguments + + +def _process_argument( + argument: str, + param_type: ts.TypeSpec, + unpacked_variables: list[str], + backend: Literal["gtfn", "dace"], +) -> str: + processed_argument = "" + if isinstance(param_type, ts.TupleType): + tuple_member_types = param_type.types + tuple_members = _perfom_tuple_unpacking( + argument, len(tuple_member_types), unpacked_variables + ) + processed_argument = "(" + for tuple_member, tuple_member_type in zip(tuple_members, param_type.types): + processed_argument += _process_argument( + argument=tuple_member, + param_type=tuple_member_type, + unpacked_variables=unpacked_variables, + backend=backend, + ) + processed_argument += ")" + + elif isinstance(param_type, ts.FieldType): + if len(param_type.dims) == 0: + # GTFN wants a field, dace a scalar. + processed_argument = argument if backend == "gtfn" else f"{argument}.as_scalar()" + else: + # Full array, the differences is where the origin comes from. + origin_method = "__gt_origin__" if backend == "gtfn" else "__dace_origin__" + processed_argument = f"({argument}.ndarray, ({argument}.{origin_method}))" + + elif isinstance(param_type, ts.ScalarType): + processed_argument = ( + f"bool({argument})" + if param_type.kind == ts.ScalarKind.BOOL and backend == "gtfn" + else argument + ) + else: + raise ValueError(f"Parameter `{argument}` had unexpected type `{param_type}`") + + assert processed_argument + return processed_argument + ", " + + +def _process_metric_arguments( + metric_level_arg_name: str, + compute_time_arg_name: str, + positional_arguments: str, + kwargs_arguments: dict[str, str], + backend: Literal["gtfn", "dace"], + use_metrics: bool, +) -> tuple[str, dict[str, str]]: + + if use_metrics: + if backend == "gtfn": + raise NotImplementedError() + else: + assert positional_arguments.strip().endswith(",") + positional_arguments += f"{metric_level_arg_name}, {compute_time_arg_name}, " + + return positional_arguments, kwargs_arguments diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/common.py b/src/gt4py/next/program_processors/runners/dace/workflow/common.py index 07f053bbc9..ea40522446 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/common.py +++ b/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) diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py index 7b78942090..0733617d2e 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/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, diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py b/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py index 8c559a6d02..c18118798b 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py +++ b/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()) diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py index 42cc1f15a4..5dcc7dffa0 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py @@ -9,10 +9,11 @@ from __future__ import annotations import dataclasses -from typing import Any, Optional +from typing import Any, Mapping, Optional, Sequence, Union import dace import factory +import numpy as np from gt4py._core import definitions as core_defs from gt4py.next import common @@ -26,7 +27,7 @@ transformations as gtx_transformations, ) from gt4py.next.program_processors.runners.dace.workflow import common as gtx_wfdcommon -from gt4py.next.type_system import type_specifications as ts +from gt4py.next.type_system import type_specifications as ts, type_translation def find_constant_symbols( @@ -170,6 +171,9 @@ def add_instrumentation(sdfg: dace.SDFG, gpu: bool) -> None: The execution time is measured in seconds and represented as a 'float64' value. It is written to the global array 'SDFG_ARG_METRIC_COMPUTE_TIME'. + + Note: + This function will add the arguments related to the metric to `user_args`. """ output, _ = sdfg.add_array(gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME, [1], dace.float64) start_time, _ = sdfg.add_scalar("gt_start_time", dace.int64, transient=True) @@ -340,6 +344,148 @@ def make_sdfg_call_sync(sdfg: dace.SDFG, gpu: bool) -> None: ) +def make_user_args( + program_parameters: Sequence[interface.Parameter], + sdfg: dace.SDFG, +) -> list[Union[str, tuple[Union[str, tuple], ...]]]: + """Generates the `user_args` field for `sdfg`. + + The generated signature is compatible with the one generated by the bindings + function. See `_create_sdfg_bindings()` for more information. + + The metric arguments are detected by inspecting the SDFG's argument list. + For that reason this function must be called _after_ they have been added. + + Args: + program_parameters: The parameter list of the entry point function. + sdfg: The generated SDFG. + """ + + def _make_user_args( + sdfg: dace.SDFG, + sdfg_arglist: Mapping[str, dace.data.Data], + param_name: str, + param_type: ts.TypeSpec, + ) -> Union[str, tuple[Union[str, tuple], ...]]: + if param_type is None: + ValueError(f"Expected that parameter `{param_name}` carries a type.") + + if param_name.startswith("@"): + # We found an connectivity argument so we add it. Note that we add it + # regardless if it is needed or not, we do this for compatibility with GTFN. + # NOTE: If changed the bindings function must be updated. + assert isinstance(param_type, ts.FieldType) + sdfg_connectivity_name = gtx_dace_args.connectivity_identifier(param_name[1:]) + return ( + (sdfg_connectivity_name, ("", "")) if sdfg_connectivity_name in sdfg_arglist else "" + ) + + elif isinstance(param_type, ts.TupleType): + return tuple( + _make_user_args( + sdfg=sdfg, + sdfg_arglist=sdfg_arglist, + param_name=f"{param_name}_{i}", + param_type=tuple_arg_type, + ) + for i, tuple_arg_type in enumerate(param_type.types) + ) + + elif param_name not in sdfg_arglist: + # There are two reasons for this case: + # 1) The argument is a symbol/scalar that is not used in the generated code. + # 2) The argument was demoted, see `demote_fields` argument of `gt_auto_optimize()` + # and was not put back. + + # This will not only work for a scalar argument, but also for fields/tuples, as it + # will ignore everything beneath it. It will still be passed, but the bindings + # will not consider it. + # TODO(phimuell): Consider updating the bindings generator such that it also + # does not forward them. + return "" + + elif isinstance(param_type, ts.FieldType): + if param_name not in sdfg.arrays: + ValueError( + f"Not not find array parameter `{param_name}` in the SDFG array registry." + ) + if param_name not in sdfg_arglist: + ValueError( + f"Did not find array parameter `{param_name}` in the SDFG argument list." + ) + + sdfg_arg_desc: dace.data.Data = sdfg.arrays[param_name] + if sdfg_arg_desc.transient: + ValueError(f"GT4Py parameter `{param_name}` is a transient.") + + if len(param_type.dims) == 0: + # Zero Dimensional Array: + # Will be passed as scalar, i.e. `zero_dim_array.as_scalar()`, this "unpacking" is + # performed by the argument pre-processing function, see `bindings.py`. This is + # different from GTFN where it is passed as a normal array. + assert isinstance(sdfg_arg_desc, dace.data.Scalar) + return param_name + + else: + # Full array: + # It will be passed as `(array_name_in_sdfg, (domain_start...))`. Note that the + # bindings will extract the `domain_stop...` symbols from the shape of the array, + # which is `domain_stop - domain_start` and the explicitly passed `domain_start` + # symbols. + # This is signature compatible with GTFN, but GTFN passes `-domain_sytart`. + assert isinstance(sdfg_arg_desc, dace.data.Array) + origins: list[str] = [] + found_needed_symbol = False + for dim in param_type.dims: + rstart = str(gtx_dace_args.range_start_symbol(param_name, dim)) + if rstart in sdfg_arglist: + assert rstart in sdfg.symbols + assert rstart not in sdfg.arrays + origins.append(rstart) + found_needed_symbol = True + else: + # For certain reason the dimension parameter is not needed and thus not + # included. For compatibility with GTFN we have to provide it, but ignore it. + # Note that it could still be inside the symbols table. + origins.append("") + + return (param_name, tuple(origins) if found_needed_symbol else "") + + elif isinstance(param_type, ts.ScalarType): + # A scalar name, so simply return the name of the parameter. + if not (param_name in sdfg.arrays or param_name in sdfg.symbols): + ValueError( + f"Not not find scalar parameter `{param_name}` in the SDFG array/symbol registry." + ) + if param_name not in sdfg_arglist: + ValueError( + f"Did not find scalar parameter `{param_name}` in the SDFG argument list." + ) + return param_name + + else: + raise ValueError(f"Parameter `{param_name}` had unexpected type `{param_type}`") + + # Now process the input arguments (fields and scalars). + sdfg_arglist = sdfg.arglist().copy() + user_args = [ + _make_user_args( + sdfg=sdfg, + sdfg_arglist=sdfg_arglist, + param_name=param.name, + param_type=param.type_, + ) + for param in program_parameters + ] + + if gtx_wfdcommon.SDFG_ARG_METRIC_LEVEL in sdfg_arglist: + assert gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME in sdfg_arglist + user_args.append(gtx_wfdcommon.SDFG_ARG_METRIC_LEVEL) + user_args.append(gtx_wfdcommon.SDFG_ARG_METRIC_COMPUTE_TIME) + + return user_args + + @dataclasses.dataclass(frozen=True) class DaCeTranslator( workflow.ChainableWorkflowMixin[ @@ -450,21 +596,74 @@ def __call__( inp.args.column_axis, ) - arg_types = inp.args.args + program_parameters = self._make_entry_point_parameters( + program, inp.args.offset_provider_type + ) - program_parameters = tuple( - interface.Parameter(param.id, arg_type) - for param, arg_type in zip(program.params, arg_types) + # This function must be called _after_ the optimization of the SDFG has been finished. + sdfg.user_args = make_user_args( + program_parameters=program_parameters, + sdfg=sdfg, ) + # NOTE: The source code is typed as `str` but we will put the json representation + # of the SDFG there. module: stages.ProgramSource[code_specs.SDFGCodeSpec] = stages.ProgramSource( entry_point=interface.Function(program.id, program_parameters), - source_code=gtx_wfdcommon.serialize_sdfg_as_json(sdfg), # type: ignore[arg-type] # The source code is typed as a `str`, but we assign a JSON dictionary. + source_code=gtx_wfdcommon.serialize_sdfg_as_json(sdfg), # type: ignore[arg-type] library_deps=tuple(), code_spec=code_specs.SDFGCodeSpec(), ) return module + def _make_entry_point_parameters( + self, + program: itir.Program, + offset_provider_type: common.OffsetProviderType, + ) -> tuple[interface.Parameter, ...]: + """Generates the parameter list for the entry point. + + The generated parameters are calling compatible with GTFN, this means: + - Normal program arguments, i.e. `program.params`, come first. + - Their name is unmodified, thus the name of argument `i` is `program.params[i].id`. + - Tuple arguments are not flattened/expanded. + - Followed by all `NeighborConnectivityType` offset providers, i.e. unused are + not filtered out. + - Offset providers are listed in the same order they appear in `offset_provider`. + - The name of offset providers name is the string used as key in `offset_provider` + prefixed by an `@`. + """ + # Regular arguments + assert all(param.type is not None for param in program.params) + program_parameters: list[interface.Parameter] = [ + interface.Parameter(param.id, param.type) # type: ignore[arg-type] + for param in program.params + ] + + # Offset provider arguments + # NOTE: The prefixing is important for the binding stages and used as the sole + # criteria to determine if a parameter is a connection or not. Furthermore, + # the name, i.e. the thing after `@`, is used to find out where to find the + # table in the offset provider. This means we are no longer dependent on the + # ordering of `offset_providers`. + for name, connectivity_type in offset_provider_type.items(): + if isinstance(connectivity_type, common.NeighborConnectivityType): + if connectivity_type.dtype.scalar_type not in [np.int32, np.int64]: + raise ValueError( + "Neighbor table indices must be of type 'np.int32' or 'np.int64'." + ) + program_parameters.append( + interface.Parameter( + name="@" + name, + type_=ts.FieldType( + dims=list(connectivity_type.domain), + dtype=type_translation.from_dtype(connectivity_type.dtype), + ), + ) + ) + + return tuple(program_parameters) + class DaCeTranslationStepFactory(factory.Factory): class Meta: diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace.py index 560b8320b2..89de6a605c 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace.py @@ -139,6 +139,8 @@ def run_and_verify_fastcall( def test_dace_fastcall(cartesian_case, monkeypatch): """Test reuse of SDFG arguments between program calls by means of SDFG fastcall API.""" + pytest.skip("`fast_call` is no longer supported.") + @gtx.field_operator def testee( a: cases.IField, @@ -209,6 +211,8 @@ def verify_testee(): def test_dace_fastcall_with_connectivity(unstructured_case, monkeypatch): """Test reuse of SDFG arguments between program calls by means of SDFG fastcall API.""" + pytest.skip("`fast_call` is no longer supported.") + connectivity_E2V = unstructured_case.offset_provider["E2V"].asnumpy() @gtx.field_operator diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_bindings.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_bindings.py index 2c4811dc7d..6dc536519a 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_bindings.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_bindings.py @@ -6,274 +6,322 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -"""Test the bindings stage of the dace backend workflow.""" - -import functools import numpy as np import pytest -from gt4py.eve import codegen + dace = pytest.importorskip("dace") from gt4py import next as gtx -from gt4py.next import common as gtx_common, int32 -from gt4py.next.otf import code_specs, stages +from gt4py.eve import codegen +from gt4py.next import common as gtx_common, config as gtx_config, int32, neighbor_sum +from gt4py.next.otf import code_specs, runners as otf_runners, stages +from gt4py.next.otf.binding import interface from gt4py.next.program_processors.runners import dace as dace_runner -from gt4py.next.program_processors.runners.dace import workflow as dace_workflow -from gt4py.next import neighbor_sum -from next_tests.integration_tests.cases import E2V, E2VDim, V2E, V2EDim +from gt4py.next.program_processors.runners.dace import ( + sdfg_args as gtx_sdfg_args, + workflow as dace_workflow, +) +from gt4py.next.program_processors.runners.dace.workflow import ( + bindings as dace_wf_bindings, + common as dace_wf_common, +) +from gt4py.next.type_system import type_specifications as ts -from next_tests.integration_tests import cases -from next_tests.integration_tests import cases_utils +from next_tests.integration_tests import cases, cases_utils +from next_tests.integration_tests.cases import E2V, V2E, E2VDim, V2EDim from next_tests.unit_tests.test_common import IDim, JDim, KDim _bind_func_name = "update_sdfg_args" +_float64_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64) +_int32_type = ts.ScalarType(kind=ts.ScalarKind.INT32) +_bool_type = ts.ScalarType(kind=ts.ScalarKind.BOOL) +_ij_field_type = ts.FieldType(dims=[IDim, JDim], dtype=_float64_type) +_zero_dim_field_type = ts.FieldType(dims=[], dtype=_float64_type) +_v2e_field_type = ts.FieldType(dims=[cases.Vertex, V2EDim], dtype=_int32_type) +_e2v_field_type = ts.FieldType(dims=[cases.Edge, E2VDim], dtype=_int32_type) + + +def _make_testee_parameters() -> tuple[interface.Parameter, ...]: + """An entry-point parameter list covering fields, scalars, bools, zero-dimensional + fields and nested tuples, followed by the offset providers as `@` prefixed + parameters, as generated by `DaCeTranslator._make_entry_point_parameters()`. + """ + return ( + interface.Parameter("a", _ij_field_type), + interface.Parameter("s", _int32_type), + interface.Parameter("flag", _bool_type), + interface.Parameter("zd", _zero_dim_field_type), + interface.Parameter( + "t", + ts.TupleType(types=[_int32_type, ts.TupleType(types=[_ij_field_type]), _ij_field_type]), + ), + interface.Parameter("@V2E", _v2e_field_type), # used inside the SDFG + interface.Parameter("@E2V", _e2v_field_type), # not used inside the SDFG + ) + -_bind_header = """\ -import ctypes +def _make_testee_sdfg() -> dace.SDFG: + """A minimal SDFG matching `_make_testee_parameters()`. -from gt4py.next import common as gtx_common, field_utils + The bindings generator only inspects the argument list of the SDFG to decide + which offset providers are used; only the 'V2E' connectivity table is present. + """ + sdfg = dace.SDFG("testee") + sdfg.add_array(gtx_sdfg_args.connectivity_identifier("V2E"), shape=(2, 4), dtype=dace.int32) + return sdfg -""" +def _make_testee_arguments() -> tuple[tuple, gtx_common.OffsetProvider]: + """Runtime arguments matching `_make_testee_parameters()`. + Note that the offset providers are not part of the argument tuple; the generated + function looks them up in the offset provider mapping by name. + """ -def _binding_source_cartesian(use_metrics: bool) -> str: - # In this SDFG 'sdfg_call_args[2]' is used to collect the stencil compute time. - # Note that the position of 'gt_compute_time' in the SDFG arguments list is - # defined by dace, based on alphabetical order from index 0 ('a', 'b', 'gt_compute_time'). - metrics_arg_index = 2 - idx = [21, 22, 0, 3, 5, 7, 4, 6, 8, 1, 9, 11, 13, 10, 12, 14, 23, 2, 15, 17, 19, 16, 18, 20] - if use_metrics: - idx = [idx + 1 if idx >= metrics_arg_index else idx for idx in idx] - return ( - _bind_header - + f"""\ -def {_bind_func_name}(device, sdfg_argtypes, args, sdfg_call_args, offset_provider): - ( - args_0, - args_1, - args_2, - args_3, - args_4, - args_5, - ) = args - ( - args_0_0, - args_0_1, - ) = args_0 - sdfg_call_args[{idx[0]}] = ctypes.c_int(args_0_0) - ( - args_0_1_0, - args_0_1_1, - args_0_1_2, - ) = args_0_1 - sdfg_call_args[{idx[1]}] = ctypes.c_int(args_0_1_0) - sdfg_call_args[{idx[2]}].value = args_0_1_1.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[3]}] = ctypes.c_int(args_0_1_1.domain.ranges[0].start) - sdfg_call_args[{idx[4]}] = ctypes.c_int(args_0_1_1.domain.ranges[1].start) - sdfg_call_args[{idx[5]}] = ctypes.c_int(args_0_1_1.domain.ranges[2].start) - sdfg_call_args[{idx[6]}] = ctypes.c_int(args_0_1_1.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[7]}] = ctypes.c_int(args_0_1_1.__gt_buffer_info__.elem_strides[1]) - sdfg_call_args[{idx[8]}] = ctypes.c_int(args_0_1_1.__gt_buffer_info__.elem_strides[2]) - ( - args_1_0, - args_1_1, - ) = args_1 - (args_1_0_0,) = args_1_0 - sdfg_call_args[{idx[9]}].value = args_1_0_0.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[10]}] = ctypes.c_int(args_1_0_0.domain.ranges[0].start) - sdfg_call_args[{idx[11]}] = ctypes.c_int(args_1_0_0.domain.ranges[1].start) - sdfg_call_args[{idx[12]}] = ctypes.c_int(args_1_0_0.domain.ranges[2].start) - sdfg_call_args[{idx[13]}] = ctypes.c_int(args_1_0_0.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[14]}] = ctypes.c_int(args_1_0_0.__gt_buffer_info__.elem_strides[1]) - sdfg_call_args[{idx[15]}] = ctypes.c_int(args_1_0_0.__gt_buffer_info__.elem_strides[2]) - sdfg_call_args[{idx[16]}] = ctypes.c_int(args_1_1) - sdfg_call_args[{idx[17]}].value = args_5.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[18]}] = ctypes.c_int(args_5.domain.ranges[0].start) - sdfg_call_args[{idx[19]}] = ctypes.c_int(args_5.domain.ranges[1].start) - sdfg_call_args[{idx[20]}] = ctypes.c_int(args_5.domain.ranges[2].start) - sdfg_call_args[{idx[21]}] = ctypes.c_int(args_5.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[22]}] = ctypes.c_int(args_5.__gt_buffer_info__.elem_strides[1]) - sdfg_call_args[{idx[23]}] = ctypes.c_int(args_5.__gt_buffer_info__.elem_strides[2]) -""" + def make_ij_field(seed: int) -> gtx.Field: + # Use a non-zero origin such that we can check that the correct one is passed on. + domain = gtx_common.domain({IDim: (1, 5), JDim: (2, 8)}) + return gtx.as_field(domain, np.arange(seed, seed + 24, dtype=np.float64).reshape(4, 6)) + + args = ( + make_ij_field(0), + np.int32(42), + np.bool_(True), + gtx.as_field([], np.asarray(41.5)), + (np.int32(7), (make_ij_field(100),), make_ij_field(200)), ) + offset_provider = { + "V2E": gtx.as_connectivity( + domain={cases.Vertex: 2, V2EDim: 4}, + codomain=cases.Edge, + data=np.array([[0, 1, 2, 3], [3, 2, 1, 0]], dtype=gtx.IndexType), + skip_value=None, + ), + "E2V": gtx.as_connectivity( + domain={cases.Edge: 4, E2VDim: 2}, + codomain=cases.Vertex, + data=np.array([[0, 1], [1, 0], [0, 1], [1, 0]], dtype=gtx.IndexType), + skip_value=None, + ), + } + return args, offset_provider -def _binding_source_cartesian_with_zero_origin(use_metrics: bool) -> str: - # In this SDFG 'sdfg_call_args[2]' is used to collect the stencil compute time. - # Note that the position of 'gt_compute_time' in the SDFG arguments list is - # defined by dace, based on alphabetical order from index 0 ('a', 'b', 'gt_compute_time'). - metrics_arg_index = 2 - idx = [12, 13, 0, 3, 4, 5, 1, 6, 7, 8, 14, 2, 9, 10, 11] - if use_metrics: - idx = [idx + 1 if idx >= metrics_arg_index else idx for idx in idx] - return ( - _bind_header - + f"""\ -def {_bind_func_name}(device, sdfg_argtypes, args, sdfg_call_args, offset_provider): - ( - args_0, - args_1, - args_2, - args_3, - args_4, - args_5, - ) = args - ( - args_0_0, - args_0_1, - ) = args_0 - sdfg_call_args[{idx[0]}] = ctypes.c_int(args_0_0) - ( - args_0_1_0, - args_0_1_1, - args_0_1_2, - ) = args_0_1 - sdfg_call_args[{idx[1]}] = ctypes.c_int(args_0_1_0) - sdfg_call_args[{idx[2]}].value = args_0_1_1.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[3]}] = ctypes.c_int(args_0_1_1.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[4]}] = ctypes.c_int(args_0_1_1.__gt_buffer_info__.elem_strides[1]) - sdfg_call_args[{idx[5]}] = ctypes.c_int(args_0_1_1.__gt_buffer_info__.elem_strides[2]) - ( - args_1_0, - args_1_1, - ) = args_1 - (args_1_0_0,) = args_1_0 - sdfg_call_args[{idx[6]}].value = args_1_0_0.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[7]}] = ctypes.c_int(args_1_0_0.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[8]}] = ctypes.c_int(args_1_0_0.__gt_buffer_info__.elem_strides[1]) - sdfg_call_args[{idx[9]}] = ctypes.c_int(args_1_0_0.__gt_buffer_info__.elem_strides[2]) - sdfg_call_args[{idx[10]}] = ctypes.c_int(args_1_1) - sdfg_call_args[{idx[11]}].value = args_5.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[12]}] = ctypes.c_int(args_5.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[13]}] = ctypes.c_int(args_5.__gt_buffer_info__.elem_strides[1]) - sdfg_call_args[{idx[14]}] = ctypes.c_int(args_5.__gt_buffer_info__.elem_strides[2]) -""" +def _create_testee_bindings( + use_metrics: bool, backend: str = "dace", with_sdfg: bool = True +) -> str: + return dace_wf_bindings._create_sdfg_bindings( + program_parameters=_make_testee_parameters(), + bind_func_name=_bind_func_name, + use_metrics=use_metrics, + sdfg=_make_testee_sdfg() if with_sdfg else None, + backend=backend, ) -def _binding_source_unstructured(use_metrics: bool) -> str: - metrics_arg_index = 2 - idx = [0, 4, 5, 1, 6, 7, 8, 2, 10, 9, 3, 12, 11] - if use_metrics: - idx = [idx + 1 if idx >= metrics_arg_index else idx for idx in idx] +def _compile_bindings(binding_source: str): + """Turn the generated source into a callable, as `CompiledDaceProgram` does.""" + namespace: dict = {} + exec(binding_source, namespace) + return namespace[_bind_func_name] + + +def _expected_testee_binding_source(use_metrics: bool, e2v_used: bool = False) -> str: + metric_args = "metrics_level, runtime_return_value, " if use_metrics else "" + e2v_arg = "(offset_provider['E2V'].ndarray, _PAIR_OF_ZEROS)" if e2v_used else "None" + return f"""\ +_PAIR_OF_ZEROS = (0, 0) +def {_bind_func_name}(args, offset_provider, metrics_level, runtime_return_value): + __gtx_expanded_names_a, __gtx_expanded_names_s, __gtx_expanded_names_flag, __gtx_expanded_names_zd, __gtx_expanded_names_t, = args + __gtx_expanded_names_t_0, __gtx_expanded_names_t_1, __gtx_expanded_names_t_2, = __gtx_expanded_names_t + __gtx_expanded_names_t_1_0, = __gtx_expanded_names_t_1 return ( - _bind_header - + f"""\ -def {_bind_func_name}(device, sdfg_argtypes, args, sdfg_call_args, offset_provider): - ( - args_0, - args_1, - ) = args - sdfg_call_args[{idx[0]}].value = args_0.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[1]}] = ctypes.c_int(args_0.domain.ranges[0].start) - sdfg_call_args[{idx[2]}] = ctypes.c_int(args_0.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[3]}].value = args_1.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[4]}] = ctypes.c_int(args_1.domain.ranges[0].start) - sdfg_call_args[{idx[5]}] = ctypes.c_int(args_1.domain.ranges[0].stop) - sdfg_call_args[{idx[6]}] = ctypes.c_int(args_1.__gt_buffer_info__.elem_strides[0]) - table_E2V = offset_provider["E2V"] - sdfg_call_args[{idx[7]}].value = table_E2V.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[8]}] = ctypes.c_int(table_E2V.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[9]}] = ctypes.c_int(table_E2V.__gt_buffer_info__.elem_strides[1]) - table_V2E = offset_provider["V2E"] - sdfg_call_args[{idx[10]}].value = table_V2E.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[11]}] = ctypes.c_int(table_V2E.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[12]}] = ctypes.c_int(table_V2E.__gt_buffer_info__.elem_strides[1]) + (__gtx_expanded_names_a.ndarray, (__gtx_expanded_names_a.__dace_origin__)), + __gtx_expanded_names_s, + __gtx_expanded_names_flag, + __gtx_expanded_names_zd.as_scalar(), + ( + __gtx_expanded_names_t_0, + ((__gtx_expanded_names_t_1_0.ndarray, (__gtx_expanded_names_t_1_0.__dace_origin__)),), + (__gtx_expanded_names_t_2.ndarray, (__gtx_expanded_names_t_2.__dace_origin__)), + ), + (offset_provider['V2E'].ndarray, _PAIR_OF_ZEROS), + {e2v_arg}, + {metric_args} + ), {{}} """ + + +@pytest.mark.parametrize("use_metrics", [False, True], ids=["no_metrics", "use_metrics"]) +def test_create_sdfg_bindings_source(use_metrics): + """The generated source is fully unrolled from the entry-point parameters. + + Offset providers appear at their parameter position: as a `(buffer, origin)` + lookup when their table is in the SDFG argument list, and as an ignored `None` + placeholder otherwise. + """ + binding_source = _create_testee_bindings(use_metrics) + + assert codegen.format_python_source(binding_source) == codegen.format_python_source( + _expected_testee_binding_source(use_metrics) ) -def _binding_source_unstructured_with_zero_origin(use_metrics: bool) -> str: - metrics_arg_index = 2 - idx = [0, 4, 1, 5, 6, 2, 8, 7, 3, 10, 9] - if use_metrics: - idx = [idx + 1 if idx >= metrics_arg_index else idx for idx in idx] - return ( - _bind_header - + f"""\ -def {_bind_func_name}(device, sdfg_argtypes, args, sdfg_call_args, offset_provider): - ( - args_0, - args_1, - ) = args - sdfg_call_args[{idx[0]}].value = args_0.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[1]}] = ctypes.c_int(args_0.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[2]}].value = args_1.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[3]}] = ctypes.c_int(args_1.domain.ranges[0].stop) - sdfg_call_args[{idx[4]}] = ctypes.c_int(args_1.__gt_buffer_info__.elem_strides[0]) - table_E2V = offset_provider["E2V"] - sdfg_call_args[{idx[5]}].value = table_E2V.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[6]}] = ctypes.c_int(table_E2V.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[7]}] = ctypes.c_int(table_E2V.__gt_buffer_info__.elem_strides[1]) - table_V2E = offset_provider["V2E"] - sdfg_call_args[{idx[8]}].value = table_V2E.__gt_buffer_info__.data_ptr - sdfg_call_args[{idx[9]}] = ctypes.c_int(table_V2E.__gt_buffer_info__.elem_strides[0]) - sdfg_call_args[{idx[10]}] = ctypes.c_int(table_V2E.__gt_buffer_info__.elem_strides[1]) -""" +@pytest.mark.parametrize("use_metrics", [False, True], ids=["no_metrics", "use_metrics"]) +def test_binding_function_processes_arguments(use_metrics): + """The compiled binding function translates the arguments for `user_bind_call()`.""" + args, offset_provider = _make_testee_arguments() + a, s, flag, zd, t = args + binding_function = _compile_bindings(_create_testee_bindings(use_metrics)) + + metric_level = 2 + compute_time = np.zeros(1, dtype=np.float64) + processed, processed_kwargs = binding_function( + args, offset_provider, metric_level, compute_time ) + assert processed_kwargs == {} + assert isinstance(processed, tuple) + assert len(processed) == (9 if use_metrics else 7) -# The difference between the two bindings versions is that one uses field domain -# with zero origin, therefore the range-start symbols are not present in the SDFG. -assert _binding_source_cartesian_with_zero_origin != _binding_source_cartesian -assert _binding_source_unstructured_with_zero_origin != _binding_source_unstructured + # A field is passed as a `(buffer, origin)` pair; the buffer is forwarded, not copied. + assert processed[0][0] is a.ndarray + assert processed[0][1] == a.__dace_origin__ == (1, 2) + # Scalars are passed through unchanged. + assert processed[1] == s + assert processed[2] == flag -_dace_compile_call = dace_workflow.compilation.DaCeCompiler.__call__ + # A zero-dimensional field is passed as a scalar. + assert processed[3] == zd.as_scalar() + # A tuple argument keeps its structure, with every member processed. + processed_t = processed[4] + assert processed_t[0] == t[0] + assert processed_t[1][0][0] is t[1][0].ndarray and processed_t[1][0][1] == (1, 2) + assert processed_t[2][0] is t[2].ndarray and processed_t[2][1] == (1, 2) -def mocked_compile_call( - self, - inp: stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], - binding_source_ref: str, -): - assert len(inp.library_deps) == 0 + # A used offset provider is passed as a `(buffer, origin)` pair with zero origin. + assert processed[5][0] is offset_provider["V2E"].ndarray + assert processed[5][1] == (0, 0) - # ignore assert statements - binding_source_pruned = "\n".join( - line - for line in inp.binding_source.source_code.splitlines() - if not line.lstrip().startswith("assert") - ) - assert codegen.format_python_source(binding_source_pruned) == binding_source_ref - return _dace_compile_call(self, inp) - - -def mocked_compile_call_cartesian( - self, - inp: stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], - use_metrics: bool, - use_zero_origin: bool, -): - binding_ref_fun = ( - _binding_source_cartesian_with_zero_origin if use_zero_origin else _binding_source_cartesian + # An unused offset provider is passed as an ignored placeholder. + assert processed[6] is None + + if use_metrics: + assert processed[7] == metric_level + assert processed[8] is compute_time + + +def test_create_sdfg_bindings_without_sdfg(): + """Without an SDFG every offset provider is treated as used.""" + binding_source = _create_testee_bindings(use_metrics=False, with_sdfg=False) + + assert codegen.format_python_source(binding_source) == codegen.format_python_source( + _expected_testee_binding_source(use_metrics=False, e2v_used=True) ) - return mocked_compile_call(self, inp, binding_ref_fun(use_metrics)) - - -def mocked_compile_call_unstructured( - self, - inp: stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], - use_metrics: bool, - use_zero_origin: bool, -): - binding_ref_fun = ( - _binding_source_unstructured_with_zero_origin - if use_zero_origin - else _binding_source_unstructured + + +def test_create_sdfg_bindings_gtfn_not_supported(): + """The GTFN flavour of the bindings generator is not implemented (yet).""" + with pytest.raises(NotImplementedError): + _create_testee_bindings(use_metrics=False, backend="gtfn", with_sdfg=False) + + +@pytest.mark.parametrize("use_metrics", [False, True], ids=["no_metrics", "use_metrics"]) +def test_bind_sdfg_stage(use_metrics): + """`bind_sdfg()` generates the bindings from the SDFG JSON and the entry point. + + Whether the metric arguments are needed is detected from the presence of the + compute-time array in the SDFG. + """ + sdfg = _make_testee_sdfg() + if use_metrics: + sdfg.add_array(dace_wf_common.SDFG_ARG_METRIC_COMPUTE_TIME, shape=(1,), dtype=dace.float64) + inp = stages.ProgramSource( + entry_point=interface.Function("testee", _make_testee_parameters()), + source_code=sdfg.to_json(), # type: ignore[arg-type] # The translation stage also puts the JSON dict on the `str` typed field. + library_deps=(), + code_spec=code_specs.SDFGCodeSpec(), ) - return mocked_compile_call(self, inp, binding_ref_fun(use_metrics)) + + ext = dace_wf_bindings.bind_sdfg(inp, _bind_func_name) + + assert ext.program_source is inp + assert ext.binding_source.library_deps == tuple() + assert codegen.format_python_source( + ext.binding_source.source_code + ) == codegen.format_python_source(_expected_testee_binding_source(use_metrics)) + + +_dace_compile_call = dace_workflow.compilation.DaCeCompiler.__call__ + + +@pytest.fixture +def captured_binding_source(monkeypatch) -> dict: + """Monkeypatch `DaCeCompiler.__call__` to capture the binding source it receives. + + Also force in-process compilation, otherwise the compilation would run in a + worker process where the monkeypatch is not effective. Since the default + runner is a process-wide singleton created on first use, it must be discarded + as well: an earlier test may already have created a process pool, which would + ignore the `BUILD_JOBS` override. + """ + monkeypatch.setattr(gtx_config, "BUILD_JOBS", 0) + otf_runners.reset_default_runner() + captured: dict = {} + + def mocked_compile_call(self, inp): + captured["binding_source"] = inp.binding_source.source_code + return _dace_compile_call(self, inp) + + monkeypatch.setattr(dace_workflow.compilation.DaCeCompiler, "__call__", mocked_compile_call) + yield captured + # Discard the serial runner created under the `BUILD_JOBS` override, so that + # later tests resolve a runner from the restored configuration again. + otf_runners.reset_default_runner() + + +def _expected_cartesian_binding_source(use_metrics: bool) -> str: + metric_args = "metrics_level, runtime_return_value, " if use_metrics else "" + return f"""\ +_PAIR_OF_ZEROS = (0, 0) +def {_bind_func_name}(args, offset_provider, metrics_level, runtime_return_value): + __gtx_expanded_names_a, __gtx_expanded_names_b, __gtx_expanded_names_M, __gtx_expanded_names_N, __gtx_expanded_names_K, __gtx_expanded_names_out, = args + __gtx_expanded_names_a_0, __gtx_expanded_names_a_1, = __gtx_expanded_names_a + __gtx_expanded_names_a_1_0, __gtx_expanded_names_a_1_1, __gtx_expanded_names_a_1_2, = __gtx_expanded_names_a_1 + __gtx_expanded_names_b_0, __gtx_expanded_names_b_1, = __gtx_expanded_names_b + __gtx_expanded_names_b_0_0, = __gtx_expanded_names_b_0 + return ( + ( + __gtx_expanded_names_a_0, + ( + __gtx_expanded_names_a_1_0, + (__gtx_expanded_names_a_1_1.ndarray, (__gtx_expanded_names_a_1_1.__dace_origin__)), + __gtx_expanded_names_a_1_2, + ), + ), + ( + ((__gtx_expanded_names_b_0_0.ndarray, (__gtx_expanded_names_b_0_0.__dace_origin__)),), + __gtx_expanded_names_b_1, + ), + __gtx_expanded_names_M, + __gtx_expanded_names_N, + __gtx_expanded_names_K, + (__gtx_expanded_names_out.ndarray, (__gtx_expanded_names_out.__dace_origin__)), + {metric_args} + ), {{}} +""" @pytest.mark.parametrize("use_metrics", [False, True], ids=["no_metrics", "use_metrics"]) @pytest.mark.parametrize( "use_zero_origin", [False, True], ids=["no_zero_origin", "use_zero_origin"] ) -def test_cartesian_bind_sdfg(use_metrics, use_zero_origin, monkeypatch): +def test_cartesian_bind_sdfg(use_metrics, use_zero_origin, captured_binding_source): M, N, K = (30, 20, 10) @gtx.field_operator @@ -300,13 +348,6 @@ def testee( use_metrics=use_metrics, use_zero_origin=use_zero_origin, ) - monkeypatch.setattr( - dace_workflow.compilation.DaCeCompiler, - "__call__", - functools.partialmethod( - mocked_compile_call_cartesian, use_metrics=use_metrics, use_zero_origin=use_zero_origin - ), - ) test_case = cases.Case.from_cartesian_grid_descriptor( cases_utils.simple_cartesian_grid(), @@ -333,12 +374,44 @@ def testee( program(a, b, out=c, M=M, N=N, K=K) assert np.all(c.asnumpy() == ref) + # The binding source only depends on the entry-point parameters; in particular + # it is independent of `use_zero_origin`, since the origin is always passed + # and, if not needed, ignored on the SDFG side. + assert codegen.format_python_source( + captured_binding_source["binding_source"] + ) == codegen.format_python_source(_expected_cartesian_binding_source(use_metrics)) + + +def _expected_unstructured_binding_source( + use_metrics: bool, offset_provider: gtx_common.OffsetProvider +) -> str: + metric_args = "metrics_level, runtime_return_value, " if use_metrics else "" + # The connectivities appear in offset-provider order; only 'E2V' and 'V2E' are + # used by the program, all others are passed as ignored placeholders. + offset_provider_args = "".join( + f"(offset_provider['{name}'].ndarray, _PAIR_OF_ZEROS), " + if name in ("E2V", "V2E") + else "None, " + for name in offset_provider + ) + return f"""\ +_PAIR_OF_ZEROS = (0, 0) +def {_bind_func_name}(args, offset_provider, metrics_level, runtime_return_value): + __gtx_expanded_names_a, __gtx_expanded_names_b, = args + return ( + (__gtx_expanded_names_a.ndarray, (__gtx_expanded_names_a.__dace_origin__)), + (__gtx_expanded_names_b.ndarray, (__gtx_expanded_names_b.__dace_origin__)), + {offset_provider_args} + {metric_args} + ), {{}} +""" + @pytest.mark.parametrize("use_metrics", [False, True], ids=["no_metrics", "use_metrics"]) @pytest.mark.parametrize( "use_zero_origin", [False, True], ids=["no_zero_origin", "use_zero_origin"] ) -def test_unstructured_bind_sdfg(use_metrics, use_zero_origin, monkeypatch): +def test_unstructured_bind_sdfg(use_metrics, use_zero_origin, captured_binding_source): @gtx.field_operator def testee_op(a: cases.VField) -> cases.VField: tmp = neighbor_sum(a(E2V), axis=E2VDim) @@ -354,15 +427,6 @@ def testee(a: cases.VField, b: cases.VField): use_metrics=use_metrics, use_zero_origin=use_zero_origin, ) - monkeypatch.setattr( - dace_workflow.compilation.DaCeCompiler, - "__call__", - functools.partialmethod( - mocked_compile_call_unstructured, - use_metrics=use_metrics, - use_zero_origin=use_zero_origin, - ), - ) SIMPLE_MESH = cases_utils.simple_mesh(None) offset_provider = SIMPLE_MESH.offset_provider @@ -379,11 +443,16 @@ def testee(a: cases.VField, b: cases.VField): axis=1, ) - static_args = {} program = ( testee.with_grid_type(gtx_common.GridType.UNSTRUCTURED) .with_backend(backend) - .compile(offset_provider=offset_provider, **static_args) + .compile(offset_provider=offset_provider) ) program(a, b, offset_provider=offset_provider) assert np.all(b.asnumpy() == ref) + + assert codegen.format_python_source( + captured_binding_source["binding_source"] + ) == codegen.format_python_source( + _expected_unstructured_binding_source(use_metrics, offset_provider) + ) diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_cache_consistency.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_cache_consistency.py index 4cb2b00ba3..087b97a584 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_cache_consistency.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_cache_consistency.py @@ -25,8 +25,8 @@ import pytest - dace = pytest.importorskip("dace") +from dace.codegen import compiler as dace_compiler from gt4py._core import definitions as core_defs from gt4py.next import config, fingerprinting @@ -114,12 +114,28 @@ def test_dace_recovers_from_truncated_library(clean_build_folder): build_folder = clean_build_folder(comp, inp) artifact = comp(inp) - assert artifact.library_path.is_file() + assert artifact.sdfg_build_folder.is_dir() + + sdfg = dace.SDFG.from_file(artifact.sdfg_build_folder / "program.sdfgz") + # Resolve the folder mode from the build folder itself: the compile step runs + # inside `dace_context()` where `compiler.build_folder_mode` may differ from + # the ambient configuration of this process. + library_path = dace_compiler.get_binary_name( + object_folder=artifact.sdfg_build_folder, + sdfg_name=sdfg.name, + folder_mode=dace_compiler.get_folder_mode(artifact.sdfg_build_folder), + ) # Simulate a build interrupted mid-link: truncated library, no completion marker. - artifact.library_path.write_bytes(b"\x00" * 64) + library_path.write_bytes(b"\x00" * 64) (build_folder / dace_wf_compilation._COMPILE_COMPLETE_MARKER).unlink() recovered = comp(inp) - - ctypes.CDLL(str(recovered.library_path)) # raises OSError if still truncated + recovered_sdfg = dace.SDFG.from_file(recovered.sdfg_build_folder / "program.sdfgz") + assert artifact.sdfg_build_folder == recovered.sdfg_build_folder + recovered_library_path = dace_compiler.get_binary_name( + object_folder=recovered.sdfg_build_folder, + sdfg_name=recovered_sdfg.name, + folder_mode=dace_compiler.get_folder_mode(recovered.sdfg_build_folder), + ) + ctypes.CDLL(str(recovered_library_path)) # raises OSError if still truncated diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py index 610567f6b2..1c5dcbf741 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py @@ -166,8 +166,7 @@ def test_compiler_skips_tx_markers_for_non_gpu_device(program_source): def test_dace_compilation_artifact_pickle_round_trip(tmp_path: pathlib.Path): artifact = dace_wf_compilation.DaCeCompilationArtifact( - library_path=tmp_path / "build" / "libprogram.so", - sdfg_json="{}", + sdfg_build_folder=tmp_path, binding_source_code="def update_sdfg_args(*a, **k): ...", bind_func_name="update_sdfg_args", device_type=core_defs.DeviceType.CPU, @@ -197,7 +196,8 @@ def test_same_artifact(add_gpu_trace_markers, program_source): device_type=core_defs.DeviceType.CUDA, ) - assert artifact_1.library_path == artifact_2.library_path + assert artifact_1.sdfg_build_folder == artifact_2.sdfg_build_folder + assert artifact_1 == artifact_2 assert ( sdfg_1.hash_sdfg() == sdfg_2.hash_sdfg() ) # might contain different GUIDs, `hash_sdfg()` ignores them @@ -212,7 +212,7 @@ def test_apply_tx_markers_changes_artifact(program_source): program_source, device_type=core_defs.DeviceType.CUDA, add_gpu_trace_markers=True ) - assert artifact_base.library_path != artifact_with_markers.library_path + assert artifact_base.sdfg_build_folder != artifact_with_markers.sdfg_build_folder # `CXXFLAGS`, `CUDAFLAGS` and `HIPFLAGS` feed `compiler.cpu.args`, `compiler.cuda.args` @@ -244,7 +244,7 @@ def test_compiler_flags_change_artifact( # The differing `dace_config_nondefaults` make the two compilers fingerprint differently, # so `get_cache_folder` names two distinct artifacts. - assert artifact_default.library_path != artifact_custom.library_path + assert artifact_default.sdfg_build_folder != artifact_custom.sdfg_build_folder def test_cmake_build_type_changes_artifact(program_source): @@ -261,4 +261,4 @@ def test_cmake_build_type_changes_artifact(program_source): ) artifact_debug, _ = _run_compiler(program_source, cmake_build_type=config.CMakeBuildType.DEBUG) - assert artifact_release.library_path != artifact_debug.library_path + assert artifact_release.sdfg_build_folder != artifact_debug.sdfg_build_folder diff --git a/uv.lock b/uv.lock index 2211528abd..47770762a9 100644 --- a/uv.lock +++ b/uv.lock @@ -908,11 +908,13 @@ wheels = [ [[package]] name = "dace" version = "2.0.0a5" -source = { registry = "https://pypi.org/simple" } +source = { git = "https://github.com/philip-paul-mueller/dace?branch=nanobind-compiled-sdfg#5c9d291c4b90699482d691be2987d806dfdd4934" } dependencies = [ { name = "astunparse" }, { name = "dill" }, { name = "fparser" }, + { name = "ml-dtypes" }, + { name = "nanobind" }, { name = "networkx" }, { name = "numpy" }, { name = "packaging" }, @@ -922,7 +924,6 @@ dependencies = [ { name = "sympy" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/55/b1/6ecb926fad1efcb99afc9f148493f4bd33a788c3c7079b90ad764e7f56f3/dace-2.0.0a5.tar.gz", hash = "sha256:09d9185c9d0dd663ce00c3cfd11d0c9048f80415db7080d09d109fce0e044a02", size = 6185419, upload-time = "2026-07-23T17:57:33.557Z" } [[package]] name = "debugpy" @@ -1442,7 +1443,7 @@ requires-dist = [ { name = "cupy-cuda13x", marker = "extra == 'cuda13'", specifier = ">=14.0" }, { name = "cupy-rocm-7-0", marker = "extra == 'rocm7'", specifier = ">=14.0" }, { name = "cytoolz", specifier = ">=1.0.1" }, - { name = "dace", specifier = ">=2.0.0a5" }, + { name = "dace", git = "https://github.com/philip-paul-mueller/dace?branch=nanobind-compiled-sdfg" }, { name = "deepdiff", specifier = ">=8.1.0" }, { name = "devtools", specifier = ">=0.6" }, { name = "factory-boy", specifier = ">=3.3.3" },