diff --git a/docs/user/next/advanced/HackTheToolchain.md b/docs/user/next/advanced/HackTheToolchain.md index 15e2e98ff1..2ba43d5a2f 100644 --- a/docs/user/next/advanced/HackTheToolchain.md +++ b/docs/user/next/advanced/HackTheToolchain.md @@ -46,25 +46,40 @@ skip_linting_transforms = SkipLinting(**same_steps) skip_linting_transforms.step_order(DUMMY_FOP) ``` -## Alternative Factory +## Alternative Workflow + +The builders take the settings shared by several steps (the device, the build +type, translation caching) as keyword arguments, and the settings of a single +step as a dict. The keys of that dict are the step's own fields, so a typo is a +type error. ```python -class MyCodeGen: ... +gtfn = gtx.program_processors.runners.gtfn +debug_gpu_no_transforms = gtfn.make_gtfn_backend( + gpu=True, + cmake_build_type=gtx.config.CMakeBuildType.DEBUG, + name_postfix="_debug_no_transforms", + translation={"enable_itir_transforms": False}, +) +``` -class Cpp2BindingsGen: ... +Compile workflows are plain frozen dataclasses, so a whole step is replaced +with `dataclasses.replace` on the workflow a builder returned. The replacement +is used as given. +```python +class MyCodeGen: ... -class PureCpp2WorkflowFactory(gtx.program_processors.runners.gtfn.GTFNCompileWorkflowFactory): - translation: workflow.Workflow[ - gtx.otf.stages.CompilableProgramDef, gtx.otf.artifacts.ProgramSource - ] = MyCodeGen() - bindings: workflow.Workflow[ - gtx.otf.artifacts.ProgramSource, gtx.otf.artifacts.ExtensionSource - ] = Cpp2BindingsGen() +class Cpp2BindingsGen: ... -PureCpp2WorkflowFactory(cmake_build_type=gtx.config.CMAKE_BUILD_TYPE.DEBUG) + +pure_cpp2_workflow = dataclasses.replace( + gtfn.make_gtfn_compile_workflow(cmake_build_type=gtx.config.CMakeBuildType.DEBUG), + translation=MyCodeGen(), + bindings=Cpp2BindingsGen(), +) ``` ## Invent new Workflow Types diff --git a/docs/user/next/advanced/WorkflowPatterns.md b/docs/user/next/advanced/WorkflowPatterns.md index 0e0abc4aea..31b67a4761 100644 --- a/docs/user/next/advanced/WorkflowPatterns.md +++ b/docs/user/next/advanced/WorkflowPatterns.md @@ -17,8 +17,6 @@ jupyter: import dataclasses import re -import factory - import gt4py.next as gtx import devtools @@ -199,7 +197,7 @@ Let's say we want to make our calculation workflow compatible with string input. ```python editable=true slideshow={"slide_type": ""} # A plain conversion step turning a string into an int, chained into the -# workflow below and reused by `StrToIntFactory(cached=True)`. +# workflow below and reused by `make_str_to_int(cached=True)`. def to_int(inp: str) -> int: assert isinstance(inp, str), "Can not work with 'int'!" # yes, this is horribly contrived return int(inp) @@ -214,9 +212,9 @@ str_calc("1") -### Step with factory (builder) +### Step with a builder -If a step can be useful with different combinations of parameters and wrappers, it should have a factory. In this case we will add a neutral wrapper around it, so we can put any combination of wrappers into that: +If a step is useful with different combinations of parameters and wrappers, give it a **builder function**: a plain function taking the cross-cutting options and returning the assembled step. Steps are frozen dataclasses, so the builder is ordinary code — no factory framework involved, and the result is fully type-checked. @@ -229,32 +227,23 @@ class AnyStrToInt(gtx.otf.workflow.ChainableWorkflowMixin[str | int, int]): return self.inner_step(inp) -class StrToIntFactory(factory.Factory): - class Meta: - model = AnyStrToInt - - class Params: - default_step = to_int - cached = factory.Trait( - inner_step=factory.LazyAttribute( - lambda o: gtx.otf.workflow.CachedStep.in_memory( - step=o.default_step, input_fingerprinter=str - ) - ) - ) - - inner_step = factory.LazyAttribute(lambda o: o.default_step) +def make_str_to_int( + *, cached: bool = False, step: gtx.otf.workflow.Workflow[str, int] = to_int +) -> AnyStrToInt: + if cached: + step = gtx.otf.workflow.CachedStep.in_memory(step=step, input_fingerprinter=str) + return AnyStrToInt(inner_step=step) -cached = StrToIntFactory(cached=True) -uncached = StrToIntFactory() +cached = make_str_to_int(cached=True) +uncached = make_str_to_int() uncached.inner_step ``` ### Example in the Wild ```python -gtx.ffront.past_passes.linters.LinterFactory?? +gtx.ffront.past_passes.linters.linter_factory?? ``` @@ -413,5 +402,5 @@ gtx.program_processors.runners.gtfn.run_gtfn_gpu.executor.otf_workflow?? ``` ```python -gtx.program_processors.runners.gtfn.GTFNBackendFactory?? +gtx.program_processors.runners.gtfn.make_gtfn_backend?? ``` diff --git a/pyproject.toml b/pyproject.toml index e5717fe037..e131ebda0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ profiling = [ ] scripts = ["pyyaml>=6.0.1", "typer>=0.16.0", "packaging"] test = [ + 'factory-boy>=3.3.3', 'hypothesis>=6.0.0', 'nbmake>=1.4.6', 'nox>=2025.02.09', @@ -99,7 +100,6 @@ dependencies = [ 'dace==2.0.0a9', 'deepdiff>=8.1.0', 'devtools>=0.6', - 'factory-boy>=3.3.3', "filelock>=3.18.0", 'frozendict>=2.3', 'gridtools-cpp>=2.3.9,==2.*', @@ -231,12 +231,6 @@ module = 'gt4py.next.iterator.*' ignore_errors = true module = 'gt4py.next.iterator.runtime' -[[tool.mypy.overrides]] -ignore_missing_imports = true -implicit_reexport = true -# factory-boy is broken, see https://github.com/FactoryBoy/factory_boy/pull/1114 -module = "factory.*" - # -- pytest -- [tool.pytest] diff --git a/src/gt4py/next/backend.py b/src/gt4py/next/backend.py index eae12981f0..2a2ef16b13 100644 --- a/src/gt4py/next/backend.py +++ b/src/gt4py/next/backend.py @@ -140,6 +140,22 @@ def step_order(self, inp: stages.ConcreteProgramDef) -> list[str]: DEFAULT_TRANSFORMS: Transforms = Transforms() +def select_device( + gpu: bool, +) -> tuple[core_defs.DeviceType, next_allocators.FieldBufferAllocatorProtocol]: + """ + Return the device type and default field allocator of a CPU or GPU backend. + + The GPU is the one CuPy was built for, or CUDA if CuPy is not available. + """ + if gpu: + return ( + core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, + next_allocators.StandardGPUFieldBufferAllocator(), + ) + return core_defs.DeviceType.CPU, next_allocators.StandardCPUFieldBufferAllocator() + + # TODO(tehrengruber): Rename class and `executor` & `transforms` attribute. Maybe: # `Backend` -> `Toolchain` # `transforms` -> `frontend_transforms` diff --git a/src/gt4py/next/otf/compilation/cache.py b/src/gt4py/next/otf/compilation/cache.py index eac9c07dc2..a8affb726a 100644 --- a/src/gt4py/next/otf/compilation/cache.py +++ b/src/gt4py/next/otf/compilation/cache.py @@ -10,10 +10,11 @@ import pathlib import tempfile -from typing import Final +from typing import Final, TypeVar +from gt4py._core import filecache from gt4py.next import config, fingerprinting -from gt4py.next.otf import artifacts +from gt4py.next.otf import artifacts, workflow #: Regex describing the folder names produced by `get_cache_folder` (use @@ -47,6 +48,9 @@ #: workflow factory enables the `cached_translation` trait. TRANSLATION_CACHE_BACKENDS: Final[tuple[str, ...]] = ("dace", "gtfn") +StartT = TypeVar("StartT") +EndT = TypeVar("EndT") + _session_cache_dir = tempfile.TemporaryDirectory(prefix="gt4py_session_") _session_cache_dir_path = pathlib.Path(_session_cache_dir.name) @@ -62,6 +66,29 @@ def get_translation_cache_folder(cache_base: pathlib.Path, backend: str) -> path return cache_base / TRANSLATION_CACHE_DIR_NAME / backend +def persistent_translation_cache( + step: workflow.Workflow[StartT, EndT], backend: str +) -> workflow.CachedStep[StartT, EndT, str]: + """ + Wrap a translation step in the persistent translation cache of `backend`. + + Args: + step: The translation step to cache. + backend: Name of the backend family, which selects the cache folder. + + Returns: + The step, cached in the translation cache folder of `backend` under the + cache base path of the configured build-cache lifetime. + """ + return workflow.CachedStep[StartT, EndT, str].persistent( + step, + input_fingerprinter=fingerprinting.strict_fingerprinter, + cache=filecache.FileCache( + get_translation_cache_folder(get_cache_base_path(config.BUILD_CACHE_LIFETIME), backend) + ), + ) + + def get_cache_base_path(lifetime: config.BuildCacheLifetime) -> pathlib.Path: """Return the base directory for cached artifacts with the given lifetime.""" match lifetime: diff --git a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py index dcf12cf509..72d1649b52 100644 --- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py +++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py @@ -11,7 +11,6 @@ import dataclasses from typing import Any, Final, Optional -import factory import numpy as np from gt4py._core import definitions as core_defs @@ -265,13 +264,8 @@ def _not_implemented_for_device_type(self) -> NotImplementedError: ) -class GTFNTranslationStepFactory(factory.Factory[GTFNTranslationStep]): - class Meta: - model = GTFNTranslationStep +translate_program_cpu: Final[stages.TranslationStep] = GTFNTranslationStep() - -translate_program_cpu: Final[stages.TranslationStep] = GTFNTranslationStepFactory() # type: ignore[assignment] # factory-boy typing not precise enough - -translate_program_gpu: Final[stages.TranslationStep] = GTFNTranslationStepFactory( # type: ignore[assignment] # factory-boy typing not precise enough +translate_program_gpu: Final[stages.TranslationStep] = GTFNTranslationStep( device_type=core_defs.DeviceType.CUDA ) diff --git a/src/gt4py/next/program_processors/formatters/gtfn.py b/src/gt4py/next/program_processors/formatters/gtfn.py index 75494a1759..cea215e55b 100644 --- a/src/gt4py/next/program_processors/formatters/gtfn.py +++ b/src/gt4py/next/program_processors/formatters/gtfn.py @@ -16,7 +16,7 @@ @program_formatter.program_formatter def format_cpp(program: itir.Program, *args: Any, **kwargs: Any) -> str: - gtfn_translation = gtfn.GTFNCompileWorkflowFactory(cached_translation=False).translation + gtfn_translation = gtfn.make_gtfn_compile_workflow().translation assert isinstance(gtfn_translation, GTFNTranslationStep) return gtfn_translation.generate_stencil_source( program, diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py b/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py index 4d825c0c9b..89746f325d 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py @@ -15,6 +15,6 @@ - `compilation` for compiling the SDFG into a program - `decoration` to parse the program arguments and pass them to the program call -The GTIR-DaCe backend factory extends `CachedBackendFactory`, thus it provides -caching of the GTIR program. +The backend builder wraps the translation step in a persistent `CachedStep`, +thus it provides caching of the GTIR program. """ diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py index c0ea33daf0..1df5db22c1 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -10,13 +10,9 @@ import dataclasses import warnings -from typing import Any, Final +from typing import Any -import factory - -import gt4py.next.custom_layout_allocators as next_allocators -from gt4py._core import definitions as core_defs -from gt4py.next import backend, common, config +from gt4py.next import backend, config from gt4py.next.otf import artifacts from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations from gt4py.next.program_processors.runners.dace.workflow import ( @@ -40,151 +36,98 @@ def load_artifact(self, artifact: artifacts.CompilationArtifact) -> artifacts.Ex return program -class DaCeBackendFactory(factory.Factory): - """ - Workflow factory for the GTIR-DaCe backend. - - Several parameters are inherithed from `backend.Backend`, see below the specific ones. - - Args: - auto_optimize: Enables the SDFG transformation pipeline. - """ - - class Meta: - model = DaCeBackend - - class Params: - name_device = "cpu" - name_postfix = "" - gpu = factory.Trait( - allocator=next_allocators.StandardGPUFieldBufferAllocator(), - device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, - name_device="gpu", - ) - device_type = core_defs.DeviceType.CPU - otf_workflow = factory.SubFactory( - gtx_wfdfactory.DaCeWorkflowFactory, - cached_translation=True, - device_type=factory.SelfAttribute("..device_type"), - auto_optimize=factory.SelfAttribute("..auto_optimize"), - ) - auto_optimize = factory.Trait(name_postfix="_opt") - - name = factory.LazyAttribute(lambda o: f"run_dace_{o.name_device}{o.name_postfix}") - executor = factory.LazyAttribute(lambda o: o.otf_workflow) - allocator = next_allocators.StandardCPUFieldBufferAllocator() - transforms = backend.DEFAULT_TRANSFORMS - external_workspace = None - - def make_dace_backend( gpu: bool, auto_optimize: bool = True, - async_sdfg_call: bool = True, - optimization_args: dict[str, Any] | None = None, + *, external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None, - unstructured_horizontal_has_unit_stride: bool = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE, - use_metrics: bool = True, - use_zero_origin: bool = False, - use_max_domain_range_on_unstructured_shift: bool | None = None, -) -> backend.Backend: + unstructured_horizontal_has_unit_stride: bool | None = None, + cached_translation: bool = True, + cmake_build_type: config.CMakeBuildType | None = None, + translation: gtx_wfdfactory.DaCeTranslationOptions | None = None, + compilation: gtx_wfdfactory.DaCeCompilationOptions | None = None, +) -> DaCeBackend: """Customize the dace backend with the given configuration parameters. + Settings shared by several steps are keyword arguments; the step-local + settings of the translation and compilation steps are passed as dicts, see + `make_dace_compile_workflow`. + Args: gpu: Enable GPU transformations and code generation. auto_optimize: Enable the SDFG auto-optimize pipeline. - async_sdfg_call: Make an asynchronous SDFG call on GPU to allow overlapping - of GPU kernel execution with the Python driver code. - optimization_args: A `dict` containing configuration parameters for - the SDFG auto-optimize pipeline, see `gt_auto_optimize()`. external_workspace: Workspace memory externally allocated, which is used for SDFG's transient arrays when `transient_memory_mode` is `EXTERNAL`. unstructured_horizontal_has_unit_stride: When the memory layout has unit stride in the horizontal dimension, replace the field stride symbol with '1'. - use_metrics: Add SDFG instrumentation to collect the metric for stencil - compute time. - use_zero_origin: Can be set to `True` when all fields passed as program - arguments have zero-based origin. This setting will skip generation - of range start-symbols `_range_0` since they can be assumed to be zero. - - Note that `gt_auto_optimize()` parameters that are derived from GT4Py configuration - cannot be overriden, and therefore cannot appear here. Thus, this function will - throw an exception if called with any argument included in `gt_optimization_args`. + Defaults to the value in `config`. + cached_translation: Wrap the translation step in a persistent cache. + cmake_build_type: Build type of the generated project. Defaults to the + value in `config`. + translation: Step-local settings of the translation step, see + `DaCeTranslator`. When an `external_workspace` is given, + `auto_optimize` is enabled and `auto_optimize_args` sets no + `transient_memory_mode`, it defaults to `EXTERNAL`; without + auto-optimize the transient memory mode is not used. + compilation: Step-local settings of the compilation step. Returns: A dace backend with custom configuration for the target device. - """ - - # The `gt_optimization_args` set contains the parameters of `gt_auto_optimize()` - # that are derived from the gt4py configuration, and therefore cannot be customized. - gt_optimization_args: Final[set[str]] = {"gpu", "constant_symbols", "unit_strides_kind"} - - if optimization_args is None: - optimization_args = {} - elif optimization_args and not auto_optimize: - warnings.warn("Optimizations args given, but auto-optimize is disabled.", stacklevel=2) - elif intersect_args := gt_optimization_args.intersection(optimization_args.keys()): - raise ValueError( - f"The following optimization arguments cannot be overriden: {intersect_args}." - ) - - # Set `unit_strides_kind` based on the gt4py env configuration. - optimization_args = optimization_args | { - "unit_strides_kind": common.DimensionKind.HORIZONTAL - if unstructured_horizontal_has_unit_stride - else None - } + Raises: + ValueError: If `auto_optimize_args` requests the `EXTERNAL` transient + memory mode without an `external_workspace`, or sets a parameter the + translation step derives itself. + """ + if unstructured_horizontal_has_unit_stride is None: + unstructured_horizontal_has_unit_stride = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE + if translation is None: + translation = gtx_wfdfactory.DaCeTranslationOptions() + + # The external workspace belongs to the backend, which injects it at load + # time, so the backend owns its coupling to the transient memory mode. + optimization_args = dict(translation.get("auto_optimize_args") or {}) + transient_memory_mode = optimization_args.get("transient_memory_mode") if external_workspace is None: - if ( - optimization_args.get("transient_memory_mode") - is gtx_transformations.TransientMemoryMode.EXTERNAL - ): + if transient_memory_mode is gtx_transformations.TransientMemoryMode.EXTERNAL: raise ValueError( "External memory workspace must be provided when 'transient_memory_mode' is 'EXTERNAL'." ) - elif transient_memory_mode := optimization_args.get("transient_memory_mode"): - if transient_memory_mode is not gtx_transformations.TransientMemoryMode.EXTERNAL: - warnings.warn( - f"External memory workspace provided but 'transient_memory_mode' is '{transient_memory_mode}', it requires '{gtx_transformations.TransientMemoryMode.EXTERNAL}'.", - stacklevel=2, + elif transient_memory_mode is None: + if auto_optimize: + optimization_args["transient_memory_mode"] = ( + gtx_transformations.TransientMemoryMode.EXTERNAL + ) + translation = translation | gtx_wfdfactory.DaCeTranslationOptions( + auto_optimize_args=optimization_args ) - else: - optimization_args["transient_memory_mode"] = ( - gtx_transformations.TransientMemoryMode.EXTERNAL + elif transient_memory_mode is not gtx_transformations.TransientMemoryMode.EXTERNAL: + warnings.warn( + f"External memory workspace provided but 'transient_memory_mode' is '{transient_memory_mode}', it requires '{gtx_transformations.TransientMemoryMode.EXTERNAL}'.", + stacklevel=2, ) - return DaCeBackendFactory( # type: ignore[return-value] # factory-boy typing not precise enough - gpu=gpu, - auto_optimize=auto_optimize, + device_type, allocator = backend.select_device(gpu) + + return DaCeBackend( + name=f"run_dace_{'gpu' if gpu else 'cpu'}{'_opt' if auto_optimize else ''}", + executor=gtx_wfdfactory.make_dace_compile_workflow( + device_type=device_type, + auto_optimize=auto_optimize, + cached_translation=cached_translation, + cmake_build_type=cmake_build_type, + unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride, + translation=translation, + compilation=compilation, + ), + allocator=allocator, + transforms=backend.DEFAULT_TRANSFORMS, external_workspace=external_workspace, - otf_workflow__bare_translation__async_sdfg_call=(async_sdfg_call if gpu else False), - otf_workflow__bare_translation__auto_optimize_args=optimization_args, - otf_workflow__bare_translation__unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride, - otf_workflow__bare_translation__use_metrics=use_metrics, - otf_workflow__bare_translation__disable_field_origin_on_program_arguments=use_zero_origin, - otf_workflow__bare_translation__use_max_domain_range_on_unstructured_shift=use_max_domain_range_on_unstructured_shift, ) -run_dace_cpu = make_dace_backend( - gpu=False, - auto_optimize=True, - async_sdfg_call=False, -) -run_dace_cpu_noopt = make_dace_backend( - gpu=False, - auto_optimize=False, - async_sdfg_call=False, -) +run_dace_cpu = make_dace_backend(gpu=False, auto_optimize=True) +run_dace_cpu_noopt = make_dace_backend(gpu=False, auto_optimize=False) -run_dace_gpu = make_dace_backend( - gpu=True, - auto_optimize=True, - async_sdfg_call=True, -) -run_dace_gpu_noopt = make_dace_backend( - gpu=True, - auto_optimize=False, - async_sdfg_call=True, -) +run_dace_gpu = make_dace_backend(gpu=True, auto_optimize=True) +run_dace_gpu_noopt = make_dace_backend(gpu=True, auto_optimize=False) 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 b202dea1cf..3a8bc9eac1 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -18,7 +18,6 @@ import dace import dace.codegen.compiler as dace_compiler -import factory from gt4py._core import definitions as core_defs, locking from gt4py.eve import xtyping @@ -367,8 +366,3 @@ def __call__(self, inp: SDFGExtensionSource) -> DaCeCompilationArtifact: bind_func_name=self.bind_func_name, device_type=self.device_type, ) - - -class DaCeCompilationStepFactory(factory.Factory): - class Meta: - model = DaCeCompiler diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py index 2f37f90cd7..f5eec4bef8 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py @@ -9,68 +9,105 @@ from __future__ import annotations import functools -from typing import Final +from typing import Any, Final, TypedDict -import factory - -from gt4py._core import definitions as core_defs, filecache -from gt4py.next import config, fingerprinting -from gt4py.next.otf import recipes, workflow +from gt4py._core import definitions as core_defs +from gt4py.next import config +from gt4py.next.otf import recipes, stages from gt4py.next.otf.compilation import cache from gt4py.next.program_processors.runners.dace.workflow import bindings as bindings_step -from gt4py.next.program_processors.runners.dace.workflow.compilation import ( - DaCeCompilationStepFactory, -) -from gt4py.next.program_processors.runners.dace.workflow.translation import ( - DaCeTranslationStepFactory, -) +from gt4py.next.program_processors.runners.dace.workflow.compilation import DaCeCompiler +from gt4py.next.program_processors.runners.dace.workflow.translation import DaCeTranslator _GT_DACE_BINDING_FUNCTION_NAME: Final[str] = "update_sdfg_args" -class DaCeWorkflowFactory(factory.Factory): - class Meta: - model = recipes.OTFCompileWorkflow - - class Params: - auto_optimize: bool = False - device_type: core_defs.DeviceType = core_defs.DeviceType.CPU - cmake_build_type: config.CMakeBuildType = factory.LazyFunction( # type: ignore[assignment] # factory-boy typing not precise enough - lambda: config.CMAKE_BUILD_TYPE - ) - - cached_translation = factory.Trait( - translation=factory.LazyAttribute( - lambda o: workflow.CachedStep.persistent( - o.bare_translation, - input_fingerprinter=fingerprinting.strict_fingerprinter, - cache=filecache.FileCache( - cache.get_translation_cache_folder( - cache.get_cache_base_path(config.BUILD_CACHE_LIFETIME), "dace" - ) - ), - ) - ), - ) - - bare_translation = factory.SubFactory( - DaCeTranslationStepFactory, - device_type=factory.SelfAttribute("..device_type"), - auto_optimize=factory.SelfAttribute("..auto_optimize"), - ) - - translation = factory.LazyAttribute(lambda o: o.bare_translation) - bindings = factory.LazyAttribute( - lambda o: functools.partial( - bindings_step.bind_sdfg, - bind_func_name=_GT_DACE_BINDING_FUNCTION_NAME, - ) +class DaCeTranslationOptions(TypedDict, total=False): + """ + Step-local settings of `DaCeTranslator`. + + The device, auto-optimize and unit-stride settings come from the builder. + """ + + auto_optimize_args: dict[str, Any] | None + async_sdfg_call: bool + use_metrics: bool + disable_itir_transforms: bool + disable_field_origin_on_program_arguments: bool + use_max_domain_range_on_unstructured_shift: bool | None + + +class DaCeCompilationOptions(TypedDict, total=False): + """Step-local settings of `DaCeCompiler`; device, build type and cache lifetime come from the builder.""" + + add_gpu_trace_markers: bool + + +#: Defaults of the `DaCeTranslator` fields that have no dataclass default. +_DEFAULT_TRANSLATION_OPTIONS: Final[DaCeTranslationOptions] = DaCeTranslationOptions( + auto_optimize_args=None, async_sdfg_call=True, use_metrics=True +) + + +def make_dace_compile_workflow( + *, + device_type: core_defs.DeviceType = core_defs.DeviceType.CPU, + auto_optimize: bool = False, + cached_translation: bool = False, + cmake_build_type: config.CMakeBuildType | None = None, + unstructured_horizontal_has_unit_stride: bool | None = None, + translation: DaCeTranslationOptions | None = None, + compilation: DaCeCompilationOptions | None = None, +) -> recipes.OTFCompileWorkflow: + """ + Build the DaCe translation -> bindings -> compilation workflow. + + Settings shared by several steps are keyword arguments, forwarded to every + step that needs them. The step-local settings of each step are passed as a + dict, unpacked into the step's constructor; the shared settings are not + part of these dicts, so they cannot be set for one step alone. To replace a + whole step, use `dataclasses.replace` on the returned workflow. + + Args: + device_type: The device the compiled program targets. + auto_optimize: Enable the SDFG auto-optimize pipeline. + cached_translation: Wrap the translation step in a persistent cache. + cmake_build_type: Build type of the generated project. Defaults to the + value in `config`. + unstructured_horizontal_has_unit_stride: Replace the field stride + symbol with '1' in the horizontal dimension. Defaults to the value + in `config`. + translation: Step-local settings of the translation step. + compilation: Step-local settings of the compilation step. + + Returns: + The composed compile workflow. + """ + if cmake_build_type is None: + cmake_build_type = config.CMAKE_BUILD_TYPE + if unstructured_horizontal_has_unit_stride is None: + unstructured_horizontal_has_unit_stride = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE + + translation_step: stages.TranslationStep = DaCeTranslator( + device_type=device_type, + auto_optimize=auto_optimize, + unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride, + **(_DEFAULT_TRANSLATION_OPTIONS | (translation or DaCeTranslationOptions())), ) - compilation = factory.SubFactory( - DaCeCompilationStepFactory, - bind_func_name=_GT_DACE_BINDING_FUNCTION_NAME, - cache_lifetime=factory.LazyFunction(lambda: config.BUILD_CACHE_LIFETIME), - device_type=factory.SelfAttribute("..device_type"), - cmake_build_type=factory.SelfAttribute("..cmake_build_type"), + if cached_translation: + translation_step = cache.persistent_translation_cache(translation_step, "dace") + + return recipes.OTFCompileWorkflow( + translation=translation_step, + bindings=functools.partial( + bindings_step.bind_sdfg, bind_func_name=_GT_DACE_BINDING_FUNCTION_NAME + ), + compilation=DaCeCompiler( + bind_func_name=_GT_DACE_BINDING_FUNCTION_NAME, + cache_lifetime=config.BUILD_CACHE_LIFETIME, + device_type=device_type, + cmake_build_type=cmake_build_type, + **(compilation or DaCeCompilationOptions()), + ), ) 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 f26e982eff..477c7703e7 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py @@ -9,11 +9,13 @@ from __future__ import annotations import dataclasses -from typing import Any, Optional +import pathlib +import warnings +from typing import Any, Final, Optional import dace -import factory +import gt4py from gt4py._core import definitions as core_defs from gt4py.next import common from gt4py.next.instrumentation import metrics @@ -340,6 +342,16 @@ def make_sdfg_call_sync(sdfg: dace.SDFG, gpu: bool) -> None: ) +#: The parameters of `gt_auto_optimize()` that the translation step derives from +#: its own configuration, and which therefore cannot be customized. +_DERIVED_OPTIMIZATION_ARGS: Final[frozenset[str]] = frozenset( + {"gpu", "constant_symbols", "unit_strides_kind"} +) + +#: Warnings about the step settings point at the first caller outside GT4Py. +_GT4PY_SOURCE_PREFIX: Final[str] = str(pathlib.Path(gt4py.__file__).parent) + + @dataclasses.dataclass(frozen=True) class DaCeTranslator( workflow.ChainableWorkflowMixin[ @@ -362,6 +374,19 @@ class DaCeTranslator( disable_field_origin_on_program_arguments: bool = False use_max_domain_range_on_unstructured_shift: bool | None = None + def __post_init__(self) -> None: + if not self.auto_optimize_args: + return + if derived_args := self.auto_optimize_args.keys() & _DERIVED_OPTIMIZATION_ARGS: + raise ValueError( + f"The following optimization arguments cannot be overriden: {derived_args}." + ) + if not self.auto_optimize: + warnings.warn( + "Optimizations args given, but auto-optimize is disabled.", + skip_file_prefixes=(_GT4PY_SOURCE_PREFIX,), + ) + def generate_sdfg( self, *args: Any, @@ -402,6 +427,11 @@ def _generate_sdfg_without_configuring_dace( sdfg, gpu=on_gpu, constant_symbols=constant_symbols, + unit_strides_kind=( + common.DimensionKind.HORIZONTAL + if self.unstructured_horizontal_has_unit_stride + else None + ), **auto_optimize_args, ) elif on_gpu: @@ -467,8 +497,3 @@ def __call__( code_spec=artifacts.SDFGCodeSpec(), ) return module - - -class DaCeTranslationStepFactory(factory.Factory): - class Meta: - model = DaCeTranslator diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py index 624ca269a3..8d4d5618af 100644 --- a/src/gt4py/next/program_processors/runners/gtfn.py +++ b/src/gt4py/next/program_processors/runners/gtfn.py @@ -8,18 +8,16 @@ import dataclasses import pathlib -from typing import Any +from typing import Any, TypedDict -import factory import numpy as np import gt4py._core.definitions as core_defs -import gt4py.next.custom_layout_allocators as next_allocators -from gt4py._core import filecache -from gt4py.next import backend, common, config, field_utils, fingerprinting +from gt4py.next import backend, common, config, field_utils from gt4py.next.embedded import nd_array_field from gt4py.next.instrumentation import metrics -from gt4py.next.otf import artifacts, recipes, workflow +from gt4py.next.iterator import ir as itir +from gt4py.next.otf import artifacts, recipes, stages from gt4py.next.otf.binding import nanobind from gt4py.next.otf.compilation import cache, compiler from gt4py.next.otf.compilation.build_systems import compiledb @@ -123,91 +121,141 @@ def _make_artifact( ) -class GTFNCompilerFactory(factory.Factory): - class Meta: - model = GTFNCompiler - - -class GTFNCompileWorkflowFactory(factory.Factory): - class Meta: - model = recipes.OTFCompileWorkflow - - class Params: - device_type: core_defs.DeviceType = core_defs.DeviceType.CPU - cmake_build_type: config.CMakeBuildType = factory.LazyFunction( # type: ignore[assignment] # factory-boy typing not precise enough - lambda: config.CMAKE_BUILD_TYPE - ) - unstructured_horizontal_has_unit_stride: bool = factory.LazyFunction( # type: ignore[assignment] # factory-boy typing not precise enough - lambda: config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE - ) - builder_factory: compiler.BuildSystemProjectGenerator = factory.LazyAttribute( # type: ignore[assignment] # factory-boy typing not precise enough - lambda o: compiledb.CompiledbFactory(cmake_build_type=o.cmake_build_type) - ) - - cached_translation = factory.Trait( - translation=factory.LazyAttribute( - lambda o: workflow.CachedStep.persistent( - o.bare_translation, - input_fingerprinter=fingerprinting.strict_fingerprinter, - cache=filecache.FileCache( - cache.get_translation_cache_folder( - cache.get_cache_base_path(config.BUILD_CACHE_LIFETIME), "gtfn" - ) - ), - ) - ), - ) - - bare_translation = factory.SubFactory( - gtfn_module.GTFNTranslationStepFactory, - device_type=factory.SelfAttribute("..device_type"), - ) - - translation = factory.LazyAttribute(lambda o: o.bare_translation) - bindings: workflow.Workflow[artifacts.ProgramSource, artifacts.ExtensionSource] = ( - factory.LazyAttribute( # type: ignore[assignment] # factory-boy typing not precise enough - lambda o: nanobind.ExtensionGenerator( - unstructured_horizontal_has_unit_stride=o.unstructured_horizontal_has_unit_stride - ) - ) +class GTFNTranslationOptions(TypedDict, total=False): + """Step-local settings of `GTFNTranslationStep`; the device comes from the builder.""" + + code_spec: artifacts.HeaderAndSourceCodeSpec | None + enable_itir_transforms: bool + symbolic_domain_sizes: dict[str, itir.Expr] | None + use_max_domain_range_on_unstructured_shift: bool | None + + +class GTFNBuildSystemOptions(TypedDict, total=False): + """Step-local settings of `CompiledbFactory`; the build type comes from the builder.""" + + cmake_extra_flags: list[str] + renew_compiledb: bool + + +class GTFNCompilationOptions(TypedDict, total=False): + """Step-local settings of `GTFNCompiler`; device and cache lifetime come from the builder.""" + + fingerprint_builder_factory: bool + force_recompile: bool + + +def make_gtfn_compile_workflow( + *, + device_type: core_defs.DeviceType = core_defs.DeviceType.CPU, + cached_translation: bool = False, + cmake_build_type: config.CMakeBuildType | None = None, + unstructured_horizontal_has_unit_stride: bool | None = None, + translation: GTFNTranslationOptions | None = None, + build_system: GTFNBuildSystemOptions | None = None, + compilation: GTFNCompilationOptions | None = None, +) -> recipes.OTFCompileWorkflow: + """ + Build the GTFN translation -> bindings -> compilation workflow. + + Settings shared by several steps are keyword arguments, forwarded to every + step that needs them. The step-local settings of each step are passed as a + dict, unpacked into the step's constructor; the shared settings are not + part of these dicts, so they cannot be set for one step alone. To replace a + whole step, use `dataclasses.replace` on the returned workflow. + + Args: + device_type: The device the compiled program targets. + cached_translation: Wrap the translation step in a persistent cache. + cmake_build_type: Build type of the generated CMake project. Defaults + to the value in `config`. + unstructured_horizontal_has_unit_stride: Layout assumption of the + bindings. Defaults to the value in `config`. + translation: Step-local settings of the translation step. + build_system: Step-local settings of the build system. + compilation: Step-local settings of the compilation step. + + Returns: + The composed compile workflow. + """ + if cmake_build_type is None: + cmake_build_type = config.CMAKE_BUILD_TYPE + if unstructured_horizontal_has_unit_stride is None: + unstructured_horizontal_has_unit_stride = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE + + translation_step: stages.TranslationStep = gtfn_module.GTFNTranslationStep( + device_type=device_type, **(translation or GTFNTranslationOptions()) ) - compilation = factory.SubFactory( - GTFNCompilerFactory, - cache_lifetime=factory.LazyFunction(lambda: config.BUILD_CACHE_LIFETIME), - builder_factory=factory.SelfAttribute("..builder_factory"), - device_type=factory.SelfAttribute("..device_type"), + if cached_translation: + translation_step = cache.persistent_translation_cache(translation_step, "gtfn") + + return recipes.OTFCompileWorkflow( + translation=translation_step, + # `OTFCompileWorkflow` is not parameterized over the code spec, so its + # `bindings` field is typed for `ProgramSource[Any]` while + # `ExtensionGenerator` accepts only C++-like specs. + bindings=nanobind.ExtensionGenerator( # type: ignore[arg-type] # see comment above + unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride + ), + compilation=GTFNCompiler( + cache_lifetime=config.BUILD_CACHE_LIFETIME, + builder_factory=compiledb.CompiledbFactory( + cmake_build_type=cmake_build_type, **(build_system or GTFNBuildSystemOptions()) + ), + device_type=device_type, + **(compilation or GTFNCompilationOptions()), + ), ) -class GTFNBackendFactory(factory.Factory): - class Meta: - model = backend.Backend - - class Params: - name_device = "cpu" - name_postfix = "" - gpu = factory.Trait( - allocator=next_allocators.StandardGPUFieldBufferAllocator(), - device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, - name_device="gpu", - ) - device_type = core_defs.DeviceType.CPU - otf_workflow = factory.SubFactory( - GTFNCompileWorkflowFactory, - cached_translation=True, - device_type=factory.SelfAttribute("..device_type"), - ) - - name = factory.LazyAttribute(lambda o: f"run_gtfn_{o.name_device}{o.name_postfix}") - executor = factory.LazyAttribute(lambda o: o.otf_workflow) - allocator = next_allocators.StandardCPUFieldBufferAllocator() - transforms = backend.DEFAULT_TRANSFORMS +def make_gtfn_backend( + *, + gpu: bool = False, + name_postfix: str = "", + cached_translation: bool = True, + cmake_build_type: config.CMakeBuildType | None = None, + unstructured_horizontal_has_unit_stride: bool | None = None, + translation: GTFNTranslationOptions | None = None, + build_system: GTFNBuildSystemOptions | None = None, + compilation: GTFNCompilationOptions | None = None, +) -> backend.Backend: + """ + Build a GTFN backend. + + Args: + gpu: Target the GPU instead of the CPU. + name_postfix: Appended to the backend name, which must stay unique. + cached_translation: Wrap the translation step in a persistent cache. + cmake_build_type: See `make_gtfn_compile_workflow`. + unstructured_horizontal_has_unit_stride: See `make_gtfn_compile_workflow`. + translation: Step-local settings of the translation step. + build_system: Step-local settings of the build system. + compilation: Step-local settings of the compilation step. + + Returns: + The configured backend. + """ + device_type, allocator = backend.select_device(gpu) + + return backend.Backend( + name=f"run_gtfn_{'gpu' if gpu else 'cpu'}{name_postfix}", + executor=make_gtfn_compile_workflow( + device_type=device_type, + cached_translation=cached_translation, + cmake_build_type=cmake_build_type, + unstructured_horizontal_has_unit_stride=unstructured_horizontal_has_unit_stride, + translation=translation, + build_system=build_system, + compilation=compilation, + ), + allocator=allocator, + transforms=backend.DEFAULT_TRANSFORMS, + ) -run_gtfn = GTFNBackendFactory() +run_gtfn = make_gtfn_backend() -run_gtfn_gpu = GTFNBackendFactory(gpu=True) +run_gtfn_gpu = make_gtfn_backend(gpu=True) -run_gtfn_no_transforms = GTFNBackendFactory( - otf_workflow__bare_translation__enable_itir_transforms=False +run_gtfn_no_transforms = make_gtfn_backend( + name_postfix="_no_transforms", translation={"enable_itir_transforms": False} ) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py index 80963d83ae..e1e3adb1ea 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py @@ -34,14 +34,14 @@ def exec_alloc_descriptor(): return backend.Backend( name="run_gtfn_with_temporaries_and_sizes", transforms=backend.DEFAULT_TRANSFORMS, - executor=gtfn.GTFNCompileWorkflowFactory( - translation=gtfn.gtfn_module.GTFNTranslationStepFactory( - symbolic_domain_sizes={ + executor=gtfn.make_gtfn_compile_workflow( + translation={ + "symbolic_domain_sizes": { "Cell": "num_cells", "Edge": "num_edges", "Vertex": "num_vertices", } - ) + } ), allocator=next_allocators.StandardCPUFieldBufferAllocator(), ) diff --git a/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py b/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py index 9b3f4bc2cb..a7df05ead8 100644 --- a/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py +++ b/tests/next_tests/unit_tests/otf_tests/test_compiled_program.py @@ -130,7 +130,9 @@ def pirate(program: workflow.ConcreteArtifact): hijacked_program = program return _NoOpArtifact() - hacked_gtfn_backend = gtfn.GTFNBackendFactory(name_postfix="_custom", executor=pirate) + hacked_gtfn_backend = dataclasses.replace( + gtfn.make_gtfn_backend(name_postfix="_custom"), executor=pirate + ) testee = testee_prog.with_backend(hacked_gtfn_backend).compile(cond=[True], offset_provider={}) testee( diff --git a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py index fc077b3a90..2f05f8debd 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py +++ b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py @@ -134,11 +134,11 @@ def test_gtfn_file_cache(program_example): data=fencil, args=arguments.CompileTimeArgs.from_concrete(*parameters, **{"offset_provider": {}}), ) - cached_gtfn_translation_step = gtfn.GTFNCompileWorkflowFactory( + cached_gtfn_translation_step = gtfn.make_gtfn_compile_workflow( cached_translation=True ).translation - bare_gtfn_translation_step = gtfn.GTFNCompileWorkflowFactory( + bare_gtfn_translation_step = gtfn.make_gtfn_compile_workflow( cached_translation=False ).translation diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py index 488364b360..1785e9cba1 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py @@ -29,6 +29,7 @@ backend as dace_wf_backend, common as dace_wf_common, decoration as dace_wf_decoration, + factory as dace_wf_factory, ) from next_tests.integration_tests import cases, cases_utils @@ -104,10 +105,12 @@ def mocked_gpu_transformation(*args, **kwargs) -> dace.SDFG: custom_backend = dace_wf_backend.make_dace_backend( gpu=on_gpu, auto_optimize=auto_optimize, - async_sdfg_call=True, - optimization_args=optimization_args, unstructured_horizontal_has_unit_stride=on_gpu, - use_metrics=True, + translation={ + "async_sdfg_call": True, + "auto_optimize_args": optimization_args, + "use_metrics": True, + }, ) # The monkeypatched transformation functions exist only in this process, so # compilation must not be offloaded to a worker. @@ -187,11 +190,13 @@ def test_make_backend_accepts_external_workspace_with_external_mode(): backend = dace_wf_backend.make_dace_backend( gpu=False, auto_optimize=True, - async_sdfg_call=False, - optimization_args={ - "transient_memory_mode": gtx_transformations.TransientMemoryMode.EXTERNAL, - }, external_workspace={core_defs.DeviceType.CPU: workspace}, + translation={ + "async_sdfg_call": False, + "auto_optimize_args": { + "transient_memory_mode": gtx_transformations.TransientMemoryMode.EXTERNAL, + }, + }, ) assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace @@ -203,8 +208,8 @@ def test_make_backend_infers_external_mode_when_workspace_is_provided(): backend = dace_wf_backend.make_dace_backend( gpu=False, auto_optimize=True, - async_sdfg_call=False, external_workspace={core_defs.DeviceType.CPU: workspace}, + translation={"async_sdfg_call": False}, ) assert ( @@ -221,11 +226,13 @@ def test_make_backend_warns_external_workspace_without_external_mode(): backend = dace_wf_backend.make_dace_backend( gpu=False, auto_optimize=True, - async_sdfg_call=False, - optimization_args={ - "transient_memory_mode": gtx_transformations.TransientMemoryMode.POOL, - }, external_workspace={core_defs.DeviceType.CPU: workspace}, + translation={ + "async_sdfg_call": False, + "auto_optimize_args": { + "transient_memory_mode": gtx_transformations.TransientMemoryMode.POOL, + }, + }, ) # Explicit mode stays as requested by the caller; backend only warns. @@ -236,6 +243,72 @@ def test_make_backend_warns_external_workspace_without_external_mode(): assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace +def test_make_backend_forwards_step_options(): + backend = dace_wf_backend.make_dace_backend( + gpu=False, + translation={"use_metrics": False}, + compilation={"add_gpu_trace_markers": True}, + ) + + assert backend.executor.translation.step.use_metrics is False + assert backend.executor.translation.step.device_type is core_defs.DeviceType.CPU + assert backend.executor.compilation.add_gpu_trace_markers is True + + +def test_make_backend_rejects_derived_optimization_args(): + with pytest.raises(ValueError, match="cannot be overriden"): + dace_wf_backend.make_dace_backend( + gpu=False, translation={"auto_optimize_args": {"unit_strides_kind": None}} + ) + + +def test_compile_workflow_rejects_derived_optimization_args(): + # The translation step validates its own settings, so the check also covers + # workflows built without `make_dace_backend`. + with pytest.raises(ValueError, match="cannot be overriden"): + dace_wf_factory.make_dace_compile_workflow( + auto_optimize=True, translation={"auto_optimize_args": {"gpu": True}} + ) + + +def test_compile_workflow_rejects_derived_optimization_args_without_auto_optimize(): + with pytest.raises(ValueError, match="cannot be overriden"): + dace_wf_factory.make_dace_compile_workflow( + auto_optimize=False, translation={"auto_optimize_args": {"gpu": True}} + ) + + +def test_compile_workflow_defaults_follow_config(monkeypatch): + monkeypatch.setattr(config, "UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE", True) + monkeypatch.setattr(config, "CMAKE_BUILD_TYPE", config.CMakeBuildType.DEBUG) + + workflow = dace_wf_factory.make_dace_compile_workflow() + + assert workflow.translation.unstructured_horizontal_has_unit_stride is True + assert workflow.compilation.cmake_build_type is config.CMakeBuildType.DEBUG + + +def test_compile_workflow_warns_on_unused_optimization_args(): + with pytest.warns(UserWarning, match="auto-optimize is disabled"): + dace_wf_factory.make_dace_compile_workflow( + auto_optimize=False, translation={"auto_optimize_args": {"blocking_size": 10}} + ) + + +def test_compile_workflow_matches_backend_translation(): + """Building the workflow directly gives the translation step of the backend.""" + workflow = dace_wf_factory.make_dace_compile_workflow( + device_type=core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA, + auto_optimize=True, + unstructured_horizontal_has_unit_stride=True, + ) + backend = dace_wf_backend.make_dace_backend( + gpu=True, unstructured_horizontal_has_unit_stride=True, cached_translation=False + ) + + assert workflow.translation == backend.executor.translation + + def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str: # Helper function to ignore the GPU device initialization code in the generated # cuda code, which is not relevant to the test. @@ -292,11 +365,13 @@ def test_transient_memory_mode(device_type, transient_memory_mode, monkeypatch): custom_backend = dace_wf_backend.make_dace_backend( gpu=on_gpu, auto_optimize=True, - async_sdfg_call=False, - optimization_args={ - "transient_memory_mode": transient_memory_mode, - }, external_workspace=external_workspace, + translation={ + "async_sdfg_call": False, + "auto_optimize_args": { + "transient_memory_mode": transient_memory_mode, + }, + }, ) @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 e4ec853c23..fd46d2672b 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 @@ -296,8 +296,10 @@ def testee( backend = dace_runner.make_dace_backend( gpu=False, - use_metrics=use_metrics, - use_zero_origin=use_zero_origin, + translation={ + "use_metrics": use_metrics, + "disable_field_origin_on_program_arguments": use_zero_origin, + }, ) monkeypatch.setattr( dace_workflow.compilation.DaCeCompiler, @@ -350,8 +352,10 @@ def testee(a: cases.VField, b: cases.VField): backend = dace_runner.make_dace_backend( gpu=False, - use_metrics=use_metrics, - use_zero_origin=use_zero_origin, + translation={ + "use_metrics": use_metrics, + "disable_field_origin_on_program_arguments": use_zero_origin, + }, ) monkeypatch.setattr( dace_workflow.compilation.DaCeCompiler, diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py index cb72c80f42..73c8e35ee6 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py @@ -22,6 +22,8 @@ import pathlib import unittest.mock +import pytest + import gt4py._core.definitions as core_defs from gt4py.next import config, custom_layout_allocators from gt4py.next.otf import workflow @@ -29,34 +31,35 @@ from gt4py.next.program_processors.runners import gtfn -def test_backend_factory_trait_device(): - cpu_version = gtfn.GTFNBackendFactory(gpu=False) - gpu_version = gtfn.GTFNBackendFactory(gpu=True) +_GPU_DEVICE_TYPE = core_defs.CUPY_DEVICE_TYPE or core_defs.DeviceType.CUDA + + +def test_make_gtfn_backend_trait_device(): + cpu_version = gtfn.make_gtfn_backend(gpu=False) + gpu_version = gtfn.make_gtfn_backend(gpu=True) assert cpu_version.name == "run_gtfn_cpu" assert isinstance(cpu_version.executor.translation, workflow.CachedStep) assert cpu_version.executor.translation.step.device_type is core_defs.DeviceType.CPU assert gpu_version.name == "run_gtfn_gpu" assert isinstance(gpu_version.executor.translation, workflow.CachedStep) - assert gpu_version.executor.translation.step.device_type is core_defs.DeviceType.CUDA + assert gpu_version.executor.translation.step.device_type is _GPU_DEVICE_TYPE # The compilation step now also carries device_type so it can stamp the artifact. assert cpu_version.executor.compilation.device_type is core_defs.DeviceType.CPU - assert gpu_version.executor.compilation.device_type is core_defs.DeviceType.CUDA + assert gpu_version.executor.compilation.device_type is _GPU_DEVICE_TYPE assert custom_layout_allocators.is_field_allocator_for( cpu_version.allocator, core_defs.DeviceType.CPU ) - assert custom_layout_allocators.is_field_allocator_for( - gpu_version.allocator, core_defs.DeviceType.CUDA - ) + assert custom_layout_allocators.is_field_allocator_for(gpu_version.allocator, _GPU_DEVICE_TYPE) -def test_backend_factory_build_cache_config(monkeypatch): +def test_make_gtfn_backend_build_cache_config(monkeypatch): monkeypatch.setattr(config, "BUILD_CACHE_LIFETIME", config.BuildCacheLifetime.SESSION) - session_version = gtfn.GTFNBackendFactory() + session_version = gtfn.make_gtfn_backend() monkeypatch.setattr(config, "BUILD_CACHE_LIFETIME", config.BuildCacheLifetime.PERSISTENT) - persistent_version = gtfn.GTFNBackendFactory() + persistent_version = gtfn.make_gtfn_backend() assert session_version.executor.compilation.cache_lifetime is config.BuildCacheLifetime.SESSION assert ( @@ -65,11 +68,11 @@ def test_backend_factory_build_cache_config(monkeypatch): ) -def test_backend_factory_build_type_config(monkeypatch): +def test_make_gtfn_backend_build_type_config(monkeypatch): monkeypatch.setattr(config, "CMAKE_BUILD_TYPE", config.CMakeBuildType.RELEASE) - release_version = gtfn.GTFNBackendFactory() + release_version = gtfn.make_gtfn_backend() monkeypatch.setattr(config, "CMAKE_BUILD_TYPE", config.CMakeBuildType.MIN_SIZE_REL) - min_size_version = gtfn.GTFNBackendFactory() + min_size_version = gtfn.make_gtfn_backend() assert ( release_version.executor.compilation.builder_factory.cmake_build_type @@ -90,9 +93,9 @@ def test_cmake_build_type_changes_build_folder(monkeypatch, tmp_path): land in different cache folders. """ monkeypatch.setattr(config, "CMAKE_BUILD_TYPE", config.CMakeBuildType.RELEASE) - release_version = gtfn.GTFNBackendFactory() + release_version = gtfn.make_gtfn_backend() monkeypatch.setattr(config, "CMAKE_BUILD_TYPE", config.CMakeBuildType.DEBUG) - debug_version = gtfn.GTFNBackendFactory() + debug_version = gtfn.make_gtfn_backend() release_compiler = release_version.executor.compilation debug_compiler = debug_version.executor.compilation @@ -126,3 +129,43 @@ def fake_get_cache_folder( assert len(build_context_ids) == 2 assert build_context_ids[0] != build_context_ids[1] + + +def test_step_options_reach_their_step_with_shared_settings(): + toolchain = gtfn.make_gtfn_backend( + gpu=True, + cmake_build_type=config.CMakeBuildType.DEBUG, + translation={"enable_itir_transforms": False}, + build_system={"cmake_extra_flags": ["-DEXTRA=ON"]}, + compilation={"force_recompile": True}, + ) + + translation = toolchain.executor.translation + assert isinstance(translation, workflow.CachedStep) + assert translation.step.enable_itir_transforms is False + assert translation.step.device_type is _GPU_DEVICE_TYPE + compilation = toolchain.executor.compilation + assert compilation.force_recompile is True + assert compilation.device_type is _GPU_DEVICE_TYPE + assert compilation.builder_factory.cmake_extra_flags == ["-DEXTRA=ON"] + assert compilation.builder_factory.cmake_build_type is config.CMakeBuildType.DEBUG + + +def test_shared_setting_cannot_be_set_in_step_options(): + with pytest.raises(TypeError, match="device_type"): + gtfn.make_gtfn_backend(gpu=True, translation={"device_type": core_defs.DeviceType.CPU}) + + +def test_prebuilt_backend_names_are_unique(): + names = [gtfn.run_gtfn.name, gtfn.run_gtfn_gpu.name, gtfn.run_gtfn_no_transforms.name] + + assert gtfn.run_gtfn_no_transforms.name == "run_gtfn_cpu_no_transforms" + assert len(set(names)) == len(names) + + +def test_compile_workflow_unit_stride_default_follows_config(monkeypatch): + monkeypatch.setattr(config, "UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE", True) + + workflow = gtfn.make_gtfn_compile_workflow() + + assert workflow.bindings.unstructured_horizontal_has_unit_stride is True diff --git a/uv.lock b/uv.lock index d024752ae3..ef2d5ca914 100644 --- a/uv.lock +++ b/uv.lock @@ -1278,7 +1278,6 @@ dependencies = [ { name = "dace" }, { name = "deepdiff" }, { name = "devtools" }, - { name = "factory-boy" }, { name = "filelock" }, { name = "frozendict" }, { name = "gridtools-cpp" }, @@ -1357,6 +1356,7 @@ dev = [ { name = "atlas4py" }, { name = "cython" }, { name = "esbonio" }, + { name = "factory-boy" }, { name = "hypothesis" }, { name = "jupytext" }, { name = "matplotlib" }, @@ -1421,6 +1421,7 @@ scripts = [ { name = "typer" }, ] test = [ + { name = "factory-boy" }, { name = "hypothesis" }, { name = "nbmake" }, { name = "nox" }, @@ -1469,7 +1470,6 @@ requires-dist = [ { name = "dace", specifier = "==2.0.0a9" }, { name = "deepdiff", specifier = ">=8.1.0" }, { name = "devtools", specifier = ">=0.6" }, - { name = "factory-boy", specifier = ">=3.3.3" }, { name = "filelock", specifier = ">=3.18.0" }, { name = "frozendict", specifier = ">=2.3" }, { name = "gridtools-cpp", specifier = "==2.*,>=2.3.9" }, @@ -1514,6 +1514,7 @@ dev = [ { name = "atlas4py", specifier = ">=0.41", index = "https://test.pypi.org/simple" }, { name = "cython", specifier = ">=3.0.0" }, { name = "esbonio", specifier = ">=0.16.0" }, + { name = "factory-boy", specifier = ">=3.3.3" }, { name = "hypothesis", specifier = ">=6.0.0" }, { name = "jupytext", specifier = ">=1.14" }, { name = "matplotlib", specifier = ">=3.9.0" }, @@ -1576,6 +1577,7 @@ scripts = [ { name = "typer", specifier = ">=0.16.0" }, ] test = [ + { name = "factory-boy", specifier = ">=3.3.3" }, { name = "hypothesis", specifier = ">=6.0.0" }, { name = "nbmake", specifier = ">=1.4.6" }, { name = "nox", specifier = ">=2025.2.9" },