diff --git a/docs/user/next/advanced/HackTheToolchain.md b/docs/user/next/advanced/HackTheToolchain.md index 74cf66e1da..15e2e98ff1 100644 --- a/docs/user/next/advanced/HackTheToolchain.md +++ b/docs/user/next/advanced/HackTheToolchain.md @@ -3,7 +3,7 @@ import dataclasses import typing from gt4py import next as gtx -from gt4py.next.otf import toolchain, workflow +from gt4py.next.otf import workflow from gt4py.next.ffront import field_operator_ast as foast, stages as ff_stages from gt4py import eve ``` @@ -22,7 +22,7 @@ cached_lowering_toolchain = gtx.backend.DEFAULT_TRANSFORMS.replace( ## Skip Steps / Change Order ```python -DUMMY_FOP = toolchain.ConcreteArtifact( +DUMMY_FOP = workflow.ConcreteArtifact( data=ff_stages.DSLFieldOperatorDef(definition=None), args=None ) ``` @@ -57,11 +57,11 @@ class Cpp2BindingsGen: ... class PureCpp2WorkflowFactory(gtx.program_processors.runners.gtfn.GTFNCompileWorkflowFactory): translation: workflow.Workflow[ - gtx.otf.definitions.CompilableProgramDef, gtx.otf.stages.ProgramSource + gtx.otf.stages.CompilableProgramDef, gtx.otf.artifacts.ProgramSource ] = MyCodeGen() - bindings: workflow.Workflow[gtx.otf.stages.ProgramSource, gtx.otf.stages.ExtensionSource] = ( - Cpp2BindingsGen() - ) + bindings: workflow.Workflow[ + gtx.otf.artifacts.ProgramSource, gtx.otf.artifacts.ExtensionSource + ] = Cpp2BindingsGen() PureCpp2WorkflowFactory(cmake_build_type=gtx.config.CMAKE_BUILD_TYPE.DEBUG) diff --git a/src/gt4py/next/backend.py b/src/gt4py/next/backend.py index 9063fe09cb..eae12981f0 100644 --- a/src/gt4py/next/backend.py +++ b/src/gt4py/next/backend.py @@ -24,7 +24,7 @@ ) from gt4py.next.ffront.past_passes import linters as past_linters from gt4py.next.iterator import ir as itir -from gt4py.next.otf import arguments, definitions, stages, toolchain, workflow +from gt4py.next.otf import arguments, artifacts, stages, toolchain, workflow def jit_to_aot_args( @@ -34,8 +34,8 @@ def jit_to_aot_args( def adapted_jit_to_aot_args_factory() -> workflow.Workflow[ - definitions.ConcreteProgramDef[definitions.IRDefinitionT, arguments.JITArgs], - definitions.ConcreteProgramDef[definitions.IRDefinitionT, arguments.CompileTimeArgs], + stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.JITArgs], + stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.CompileTimeArgs], ]: """Wrap `jit_to_aot` into a workflow adapter to fit into backend transform workflows.""" return toolchain.ArgsOnlyAdapter(jit_to_aot_args) @@ -44,8 +44,8 @@ def adapted_jit_to_aot_args_factory() -> workflow.Workflow[ @dataclasses.dataclass(frozen=True) class Transforms( workflow.MultiWorkflow[ - definitions.ConcreteProgramDef[definitions.IRDefinitionT, definitions.ArgsDefinitionT], - definitions.CompilableProgramDef, + stages.ConcreteProgramDef[stages.IRDefinitionT, stages.ArgsDefinitionT], + stages.CompilableProgramDef, ] ): """ @@ -63,8 +63,8 @@ class Transforms( """ aotify_args: workflow.Workflow[ - definitions.ConcreteProgramDef[definitions.IRDefinitionT, arguments.JITArgs], - definitions.ConcreteProgramDef[definitions.IRDefinitionT, arguments.CompileTimeArgs], + stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.JITArgs], + stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.CompileTimeArgs], ] = dataclasses.field(default_factory=adapted_jit_to_aot_args_factory) func_to_foast: workflow.Workflow[ @@ -92,10 +92,10 @@ class Transforms( ] = dataclasses.field(default_factory=past_process_args.transform_program_args_factory) past_to_itir: workflow.Workflow[ - ffront_stages.ConcretePASTProgramDef, definitions.CompilableProgramDef + ffront_stages.ConcretePASTProgramDef, stages.CompilableProgramDef ] = dataclasses.field(default_factory=past_to_itir.past_to_gtir_factory) - def step_order(self, inp: definitions.ConcreteProgramDef) -> list[str]: + def step_order(self, inp: stages.ConcreteProgramDef) -> list[str]: steps: list[str] = [] if isinstance(inp.args, arguments.JITArgs): steps.append("aotify_args") @@ -147,19 +147,19 @@ def step_order(self, inp: definitions.ConcreteProgramDef) -> list[str]: @dataclasses.dataclass(frozen=True) class Backend(Generic[core_defs.DeviceTypeT]): name: str - executor: workflow.Workflow[definitions.CompilableProgramDef, stages.CompilationArtifact] + executor: workflow.Workflow[stages.CompilableProgramDef, artifacts.CompilationArtifact] allocator: next_allocators.FieldBufferAllocatorProtocol[core_defs.DeviceTypeT] - transforms: workflow.Workflow[definitions.ConcreteProgramDef, definitions.CompilableProgramDef] + transforms: workflow.Workflow[stages.ConcreteProgramDef, stages.CompilableProgramDef] def compile( - self, program: definitions.IRDefinitionT, compile_time_args: arguments.CompileTimeArgs - ) -> stages.ExecutableProgram: + self, program: stages.IRDefinitionT, compile_time_args: arguments.CompileTimeArgs + ) -> artifacts.ExecutableProgram: artifact = self.executor( - self.transforms(definitions.ConcreteProgramDef(data=program, args=compile_time_args)) + self.transforms(stages.ConcreteProgramDef(data=program, args=compile_time_args)) ) return self.load_artifact(artifact) - def load_artifact(self, artifact: stages.CompilationArtifact) -> stages.ExecutableProgram: + def load_artifact(self, artifact: artifacts.CompilationArtifact) -> artifacts.ExecutableProgram: """Load an artifact into an executable program. Backends may override this method to inject backend-specific runtime data diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 0df00aeb58..0a6db1aabf 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -47,7 +47,7 @@ from gt4py.next.ffront.gtcallable import GTCallable from gt4py.next.instrumentation import hook_machinery, metrics from gt4py.next.iterator import ir as itir -from gt4py.next.otf import arguments, compiled_program, options, toolchain +from gt4py.next.otf import arguments, compiled_program, options, workflow from gt4py.next.type_system import type_info, type_specifications as ts, type_translation @@ -264,9 +264,7 @@ def __gt_type__(self) -> ts_ffront.ProgramType: # TODO(ricoh): linting should become optional, up to the backend. def __post_init__(self) -> None: - no_args_past = toolchain.ConcreteArtifact( - self.past_stage, arguments.CompileTimeArgs.empty() - ) + no_args_past = workflow.ConcreteArtifact(self.past_stage, arguments.CompileTimeArgs.empty()) _ = self._frontend_transforms.past_lint(no_args_past).data @property @@ -289,7 +287,7 @@ def definition(self) -> types.FunctionType: @functools.cached_property def past_stage(self) -> ffront_stages.PASTProgramDef: # backwards compatibility for backends that do not support the full toolchain - no_args_def = toolchain.ConcreteArtifact( + no_args_def = workflow.ConcreteArtifact( self.definition_stage, arguments.CompileTimeArgs.empty() ) return self._frontend_transforms.func_to_past(no_args_def).data @@ -309,7 +307,7 @@ def _all_closure_vars(self) -> dict[str, Any]: @functools.cached_property def gtir(self) -> itir.Program: - no_args_past = toolchain.ConcreteArtifact( + no_args_past = workflow.ConcreteArtifact( data=ffront_stages.PASTProgramDef( past_node=self.past_stage.past_node, closure_vars=self.past_stage.closure_vars, @@ -609,7 +607,7 @@ def __post_init__(self) -> None: @functools.cached_property def foast_stage(self) -> ffront_stages.FOASTOperatorDef: return self._frontend_transforms.func_to_foast( - toolchain.ConcreteArtifact( + workflow.ConcreteArtifact( data=self.definition_stage, args=arguments.CompileTimeArgs.empty() ) ).data diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index 9a560b7ff8..beabe53c71 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -20,7 +20,7 @@ from gt4py.next.ffront.past_passes import closure_var_type_deduction, type_deduction from gt4py.next.ffront.stages import ConcreteFOASTOperatorDef, ConcretePASTProgramDef from gt4py.next.iterator import ir as itir -from gt4py.next.otf import toolchain, workflow +from gt4py.next.otf import workflow from gt4py.next.type_system import type_info, type_specifications as ts @@ -62,7 +62,7 @@ class OperatorToProgram(workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePAST Example: >>> from gt4py import next as gtx - >>> from gt4py.next.otf import arguments, toolchain + >>> from gt4py.next.otf import arguments, workflow >>> IDim = gtx.Dimension("I") >>> @gtx.field_operator @@ -83,7 +83,7 @@ class OperatorToProgram(workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePAST ... ) >>> copy_program = op_to_prog( - ... toolchain.ConcreteArtifact(copy.foast_stage, compile_time_args) + ... workflow.ConcreteArtifact(copy.foast_stage, compile_time_args) ... ) >>> print(copy_program.data.past_node.id) @@ -169,7 +169,7 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: ) past_node = type_deduction.ProgramTypeDeduction.apply(untyped_past_node) - return toolchain.ConcreteArtifact( + return workflow.ConcreteArtifact( data=ffront_stages.PASTProgramDef( past_node=past_node, closure_vars=fieldop_itir_closure_vars, # type: ignore[arg-type] diff --git a/src/gt4py/next/ffront/past_process_args.py b/src/gt4py/next/ffront/past_process_args.py index 13ce743e86..859776ccc7 100644 --- a/src/gt4py/next/ffront/past_process_args.py +++ b/src/gt4py/next/ffront/past_process_args.py @@ -14,7 +14,7 @@ stages as ffront_stages, type_specifications as ts_ffront, ) -from gt4py.next.otf import arguments, toolchain, workflow +from gt4py.next.otf import arguments, workflow from gt4py.next.type_system import type_info, type_specifications as ts @@ -24,7 +24,7 @@ def transform_program_args( rewritten_args, rewritten_kwargs = _process_args( past_node=inp.data.past_node, args=inp.args.args, kwargs=inp.args.kwargs ) - return toolchain.ConcreteArtifact( + return workflow.ConcreteArtifact( data=inp.data, args=arguments.CompileTimeArgs( args=rewritten_args, diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index 3febb910ef..cd908fca14 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -29,19 +29,19 @@ from gt4py.next.iterator import ir as itir from gt4py.next.iterator.ir_utils import ir_makers as im from gt4py.next.iterator.transforms import remap_symbols, replace_get_domain_range_with_constants -from gt4py.next.otf import arguments, definitions, workflow +from gt4py.next.otf import arguments, stages, workflow from gt4py.next.type_system import type_info, type_specifications as ts # FIXME[#1582](tehrengruber): This should only depend on the program not the arguments. Remove # dependency as soon as column axis can be deduced from ITIR in consumers of the CompilableProgram. -def past_to_gtir(inp: ConcretePASTProgramDef) -> definitions.CompilableProgramDef: +def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgramDef: """ Lower a PAST program definition to Iterator IR. Example: >>> from gt4py import next as gtx - >>> from gt4py.next.otf import arguments, toolchain + >>> from gt4py.next.otf import arguments, workflow >>> IDim = gtx.Dimension("I") >>> @gtx.field_operator @@ -63,7 +63,7 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> definitions.CompilableProgramDe ... ) >>> itir_copy = past_to_gtir( - ... toolchain.ConcreteArtifact(copy_program.past_stage, compile_time_args) + ... workflow.ConcreteArtifact(copy_program.past_stage, compile_time_args) ... ) >>> print(itir_copy.data.id) @@ -144,12 +144,12 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> definitions.CompilableProgramDe if config.DEBUG or inp.data.debug: devtools.debug(itir_program) - return definitions.CompilableProgramDef(data=itir_program, args=compile_time_args) + return stages.CompilableProgramDef(data=itir_program, args=compile_time_args) def past_to_gtir_factory( cached: bool = True, -) -> workflow.Workflow[ConcretePASTProgramDef, definitions.CompilableProgramDef]: +) -> workflow.Workflow[ConcretePASTProgramDef, stages.CompilableProgramDef]: wf = workflow.make_step(past_to_gtir) if cached: wf = workflow.CachedStep.in_memory( diff --git a/src/gt4py/next/ffront/stages.py b/src/gt4py/next/ffront/stages.py index 0651a69739..46f59d3551 100644 --- a/src/gt4py/next/ffront/stages.py +++ b/src/gt4py/next/ffront/stages.py @@ -29,7 +29,7 @@ from gt4py.next import common, fingerprinting from gt4py.next.ffront import field_operator_ast as foast, program_ast as past, source_utils -from gt4py.next.otf import arguments, toolchain +from gt4py.next.otf import arguments, workflow @dataclasses.dataclass(frozen=True) @@ -79,7 +79,7 @@ class DSLFieldOperatorDef(BaseStage): debug: bool = False -ConcreteDSLFieldOperatorDef: typing.TypeAlias = toolchain.ConcreteArtifact[ +ConcreteDSLFieldOperatorDef: typing.TypeAlias = workflow.ConcreteArtifact[ DSLFieldOperatorDef, arguments.CompileTimeArgs ] @@ -93,7 +93,7 @@ class FOASTOperatorDef(BaseStage): debug: bool = False -ConcreteFOASTOperatorDef: typing.TypeAlias = toolchain.ConcreteArtifact[ +ConcreteFOASTOperatorDef: typing.TypeAlias = workflow.ConcreteArtifact[ FOASTOperatorDef, arguments.CompileTimeArgs ] @@ -105,7 +105,7 @@ class DSLProgramDef(BaseStage): debug: bool = False -ConcreteDSLProgramDef: typing.TypeAlias = toolchain.ConcreteArtifact[ +ConcreteDSLProgramDef: typing.TypeAlias = workflow.ConcreteArtifact[ DSLProgramDef, arguments.CompileTimeArgs ] @@ -118,7 +118,7 @@ class PASTProgramDef(BaseStage): debug: bool = False -ConcretePASTProgramDef: typing.TypeAlias = toolchain.ConcreteArtifact[ +ConcretePASTProgramDef: typing.TypeAlias = workflow.ConcreteArtifact[ PASTProgramDef, arguments.CompileTimeArgs ] diff --git a/src/gt4py/next/otf/artifacts.py b/src/gt4py/next/otf/artifacts.py new file mode 100644 index 0000000000..204e2e75d0 --- /dev/null +++ b/src/gt4py/next/otf/artifacts.py @@ -0,0 +1,243 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Data models of the artifacts flowing through a build pipeline. + +This module hosts the DSL-agnostic vocabulary of on-the-fly compilation: +source code specifications (`SourceCodeSpec` and friends), the source +containers produced and consumed by the pipeline steps (`ProgramSource`, +`BindingSource`, `ExtensionSource`) and the contracts of its end products +(`CompilationArtifact`, `ExecutableProgram`, `BuildSystemProject`). + +Nothing in here knows about the GT4Py IRs or the DSL frontend; what a +program *is* stays in `gt4py.next.otf.stages`. The `next` type system is +not excluded: an entry point is an `otf.binding.interface.Function`, whose +parameters carry `TypeSpec`s, so the signature of a generated program stays +describable without reaching for an IR. +""" + +from __future__ import annotations + +import dataclasses +import functools +from collections.abc import Callable, Mapping +from typing import Any, Generic, Optional, Protocol, TypeAlias, TypeVar, runtime_checkable + +from gt4py.eve import codegen +from gt4py.next.otf.binding import interface + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class SourceCodeSpec: + """ + Basic settings for any source programming language. + + Formatting will happen through ``eve.codegen.format_source``. + For available formatting options, check the options of the + specific formatter used depending on ``.formatter_key``. + """ + + source_language: str + file_extension: str + formatter_key: str | None = None + formatter_options: Mapping[str, Any] | None = None + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class HeaderAndSourceCodeSpec(SourceCodeSpec): + """Add a header file extension setting on top of the basic set.""" + + header_extension: str + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class PythonCodeSpec(SourceCodeSpec): + """Settings for Python language.""" + + source_language: str = "python" + file_extension: str = "py" + formatter_key: str = "python" + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class SDFGCodeSpec(SourceCodeSpec): + """Settings for SDFGs.""" + + source_language: str = "SDFG" + file_extension: str = "sdfg" + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class CPPLikeCodeSpec(HeaderAndSourceCodeSpec): + """Settings for C++-like language.""" + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class CPPCodeSpec(CPPLikeCodeSpec): + """Settings for C++ language.""" + + source_language: str = "CXX" + file_extension: str = "cpp" + header_extension: str = "hpp" + formatter_key: str = "cpp" + formatter_options: Mapping[str, Any] = dataclasses.field( + default_factory=functools.partial(dict, style="LLVM") + ) + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class CUDACodeSpec(CPPLikeCodeSpec): + """Settings for CUDA language.""" + + source_language: str = "CUDA" + file_extension: str = "cu" + header_extension: str = "cuh" + formatter_key: str = "cpp" + formatter_options: Mapping[str, Any] = dataclasses.field( + default_factory=functools.partial(dict, style="LLVM") + ) + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class HIPCodeSpec(CPPLikeCodeSpec): + """Settings for HIP language.""" + + source_language: str = "HIP" + file_extension: str = "hip" + header_extension: str = "h" + formatter_key: str = "cpp" + formatter_options: Mapping[str, Any] = dataclasses.field( + default_factory=functools.partial(dict, style="LLVM") + ) + + +def format_source(source_code_spec: SourceCodeSpec, source: str) -> str: + assert source_code_spec.formatter_key is not None, ( + "No formatter key specified in source code specification." + ) + return codegen.format_source( + source_code_spec.formatter_key, source, **(source_code_spec.formatter_options or {}) + ) + + +CodeSpecT = TypeVar("CodeSpecT", bound=SourceCodeSpec) +TargetCodeSpecT = TypeVar("TargetCodeSpecT", bound=SourceCodeSpec) + + +@dataclasses.dataclass(frozen=True) +class ProgramSource(Generic[CodeSpecT]): + """ + Standalone source code translated from an IR along with information relevant for OTF compilation. + + Contains additional information required for further OTF steps, such as + - implementation language and language conventions + - dependencies on implementation language libraries + - how to call the program + """ + + entry_point: interface.Function + source_code: str + library_deps: tuple[interface.LibraryDependency, ...] + code_spec: CodeSpecT + + +@dataclasses.dataclass(frozen=True) +class BindingSource(Generic[CodeSpecT, TargetCodeSpecT]): + """ + Companion source code for translated program source code. + + This is only needed for OTF compilation if the translated program source code is + not directly callable from python and therefore requires bindings. + This can also optionally be added to compile bindings for other languages than python + when using GT4Py as part of the build for a non-python driver project. + """ + + source_code: str + library_deps: tuple[interface.LibraryDependency, ...] + + +@dataclasses.dataclass(frozen=True) +class ExtensionSource(Generic[CodeSpecT, TargetCodeSpecT]): + """ + Encapsulate all the source code required for OTF compilation. + + The bindings module is optional if and only if the program_source is directly callable. + This should only be the case if the source language / framework supports this out of the box. + If bindings are required, it is recommended to create them in a separate step to ensure reusability. + """ + + program_source: ProgramSource[CodeSpecT] + binding_source: Optional[BindingSource[CodeSpecT, TargetCodeSpecT]] + + @property + def library_deps(self) -> tuple[interface.LibraryDependency, ...]: + if not self.binding_source: + return self.program_source.library_deps + return _unique_libs(*self.program_source.library_deps, *self.binding_source.library_deps) + + +CodeSpecT_co = TypeVar("CodeSpecT_co", bound=SourceCodeSpec, covariant=True) +TargetCodeSpecT_co = TypeVar("TargetCodeSpecT_co", bound=SourceCodeSpec, covariant=True) + + +class BuildSystemProject(Protocol[CodeSpecT_co, TargetCodeSpecT_co]): + """ + Use source code extracted from an ``ExtensionSource`` to configure and build a GT4Py program. + + Should only be considered an OTF stage if used as an endpoint, as this only runs commands on source files + and is not responsible for importing the results into Python. + """ + + def build(self) -> None: ... + + +ExecutableProgram: TypeAlias = Callable + + +@runtime_checkable +class CompilationArtifact(Protocol): + """The output of an ``OTFCompileWorkflow``. + + Each backend defines its own concrete artifact dataclass; all share this + Protocol. Implementations are frozen dataclasses, picklable, and carry no + live process-bound state — that is reconstructed by ``load``, which + returns a directly-callable ``ExecutableProgram`` taking gt4py-shaped + arguments. + + The one current exception is ``RoundtripArtifact`` when it is configured + with a ``dispatch_backend``: that field holds a ``Backend`` reference + whose role belongs at the runner / load-time seam, not in the artifact + itself. + """ + + def load(self) -> ExecutableProgram: ... + + +def _unique_libs(*args: interface.LibraryDependency) -> tuple[interface.LibraryDependency, ...]: + """ + Filter out multiple occurrences of the same ``interface.LibraryDependency``. + + Examples: + --------- + >>> libs_a = ( + ... interface.LibraryDependency("foo", "1.2.3"), + ... interface.LibraryDependency("common", "1.0.0"), + ... ) + >>> libs_b = ( + ... interface.LibraryDependency("common", "1.0.0"), + ... interface.LibraryDependency("bar", "1.2.3"), + ... ) + >>> _unique_libs(*libs_a, *libs_b) + (LibraryDependency(name='foo', version='1.2.3'), LibraryDependency(name='common', version='1.0.0'), LibraryDependency(name='bar', version='1.2.3')) + """ + unique: list[interface.LibraryDependency] = [] + for lib in args: + if lib not in unique: + unique.append(lib) + return tuple(unique) diff --git a/src/gt4py/next/otf/binding/interface.py b/src/gt4py/next/otf/binding/interface.py index 96cab18c8a..2a5cee42a8 100644 --- a/src/gt4py/next/otf/binding/interface.py +++ b/src/gt4py/next/otf/binding/interface.py @@ -11,17 +11,6 @@ import dataclasses import gt4py.next.type_system.type_specifications as ts -from gt4py.eve import codegen -from gt4py.next.otf import code_specs - - -def format_source(source_code_spec: code_specs.SourceCodeSpec, source: str) -> str: - assert source_code_spec.formatter_key is not None, ( - "No formatter key specified in source code specification." - ) - return codegen.format_source( - source_code_spec.formatter_key, source, **(source_code_spec.formatter_options or {}) - ) @dataclasses.dataclass(frozen=True) diff --git a/src/gt4py/next/otf/binding/nanobind.py b/src/gt4py/next/otf/binding/nanobind.py index a353dcff3e..4c9297e691 100644 --- a/src/gt4py/next/otf/binding/nanobind.py +++ b/src/gt4py/next/otf/binding/nanobind.py @@ -17,12 +17,12 @@ import gt4py.eve as eve from gt4py.eve.codegen import JinjaTemplate as as_jinja, TemplatedGenerator from gt4py.next import common, config -from gt4py.next.otf import code_specs, cpp_utils, stages +from gt4py.next.otf import artifacts, cpp_utils from gt4py.next.otf.binding import cpp_interface, interface from gt4py.next.type_system import type_specifications as ts -CodeSpecT = TypeVar("CodeSpecT", bound=code_specs.CPPLikeCodeSpec, covariant=True) +CodeSpecT = TypeVar("CodeSpecT", bound=artifacts.CPPLikeCodeSpec, covariant=True) class Expr(eve.Node): @@ -234,8 +234,9 @@ def make_argument( def create_bindings( - program_source: stages.ProgramSource[CodeSpecT], unstructured_horizontal_has_unit_stride: bool -) -> stages.BindingSource[CodeSpecT, code_specs.PythonCodeSpec]: + program_source: artifacts.ProgramSource[CodeSpecT], + unstructured_horizontal_has_unit_stride: bool, +) -> artifacts.BindingSource[CodeSpecT, artifacts.PythonCodeSpec]: """ Generate Python bindings through which a C++ function can be called. @@ -244,7 +245,7 @@ def create_bindings( program_source The program source for which the bindings are created """ - if not isinstance(program_source.code_spec, code_specs.CPPLikeCodeSpec): + if not isinstance(program_source.code_spec, artifacts.CPPLikeCodeSpec): raise ValueError( f"Can only create bindings for C++ program sources, received '{program_source.code_spec.source_language}'." ) @@ -289,7 +290,7 @@ def create_bindings( ), on_device=isinstance( program_source.code_spec, - (code_specs.CUDACodeSpec, code_specs.HIPCodeSpec), + (artifacts.CUDACodeSpec, artifacts.HIPCodeSpec), ), ), binding_module=BindingModule( @@ -306,11 +307,11 @@ def create_bindings( ), ) - src = interface.format_source( + src = artifacts.format_source( program_source.code_spec, BindingCodeGenerator.apply(file_binding) ) - return stages.BindingSource(src, (interface.LibraryDependency("nanobind", "2.0.0"),)) + return artifacts.BindingSource(src, (interface.LibraryDependency("nanobind", "2.0.0"),)) @dataclasses.dataclass(frozen=True) @@ -322,9 +323,11 @@ class ExtensionGenerator: unstructured_horizontal_has_unit_stride: bool = config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE def __call__( - self, program_source: stages.ProgramSource[CodeSpecT] - ) -> stages.ExtensionSource[CodeSpecT, code_specs.PythonCodeSpec]: + self, program_source: artifacts.ProgramSource[CodeSpecT] + ) -> artifacts.ExtensionSource[CodeSpecT, artifacts.PythonCodeSpec]: binding_source = create_bindings( program_source, self.unstructured_horizontal_has_unit_stride ) - return stages.ExtensionSource(program_source=program_source, binding_source=binding_source) + return artifacts.ExtensionSource( + program_source=program_source, binding_source=binding_source + ) diff --git a/src/gt4py/next/otf/code_specs.py b/src/gt4py/next/otf/code_specs.py deleted file mode 100644 index f5c215bc3e..0000000000 --- a/src/gt4py/next/otf/code_specs.py +++ /dev/null @@ -1,98 +0,0 @@ -# GT4Py - GridTools Framework -# -# Copyright (c) 2014-2024, ETH Zurich -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import dataclasses -import functools -from collections.abc import Mapping -from typing import Any - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class SourceCodeSpec: - """ - Basic settings for any source programming language. - - Formatting will happen through ``eve.codegen.format_source``. - For available formatting options, check the options of the - specific formatter used depending on ``.formatter_key``. - """ - - source_language: str - file_extension: str - formatter_key: str | None = None - formatter_options: Mapping[str, Any] | None = None - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class HeaderAndSourceCodeSpec(SourceCodeSpec): - """Add a header file extension setting on top of the basic set.""" - - header_extension: str - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class PythonCodeSpec(SourceCodeSpec): - """Settings for Python language.""" - - source_language: str = "python" - file_extension: str = "py" - formatter_key: str = "python" - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class SDFGCodeSpec(SourceCodeSpec): - """Settings for SDFGs.""" - - source_language: str = "SDFG" - file_extension: str = "sdfg" - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class CPPLikeCodeSpec(HeaderAndSourceCodeSpec): - """Settings for C++-like language.""" - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class CPPCodeSpec(CPPLikeCodeSpec): - """Settings for C++ language.""" - - source_language: str = "CXX" - file_extension: str = "cpp" - header_extension: str = "hpp" - formatter_key: str = "cpp" - formatter_options: Mapping[str, Any] = dataclasses.field( - default_factory=functools.partial(dict, style="LLVM") - ) - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class CUDACodeSpec(CPPLikeCodeSpec): - """Settings for CUDA language.""" - - source_language: str = "CUDA" - file_extension: str = "cu" - header_extension: str = "cuh" - formatter_key: str = "cpp" - formatter_options: Mapping[str, Any] = dataclasses.field( - default_factory=functools.partial(dict, style="LLVM") - ) - - -@dataclasses.dataclass(frozen=True, kw_only=True) -class HIPCodeSpec(CPPLikeCodeSpec): - """Settings for HIP language.""" - - source_language: str = "HIP" - file_extension: str = "hip" - header_extension: str = "h" - formatter_key: str = "cpp" - formatter_options: Mapping[str, Any] = dataclasses.field( - default_factory=functools.partial(dict, style="LLVM") - ) diff --git a/src/gt4py/next/otf/compilation/build_systems/cmake.py b/src/gt4py/next/otf/compilation/build_systems/cmake.py index 4c21273524..79341d4a72 100644 --- a/src/gt4py/next/otf/compilation/build_systems/cmake.py +++ b/src/gt4py/next/otf/compilation/build_systems/cmake.py @@ -15,7 +15,7 @@ from gt4py._core import definitions as core_defs from gt4py.next import config, errors -from gt4py.next.otf import code_specs, stages +from gt4py.next.otf import artifacts from gt4py.next.otf.compilation import build_data, cache, common, compiler from gt4py.next.otf.compilation.build_systems import cmake_lists @@ -33,12 +33,12 @@ def get_cmake_device_arch_option() -> str: return cmake_flag_template.format(device_archs=device_archs) if device_archs else "" -CPPLikeCodeSpecT = TypeVar("CPPLikeCodeSpecT", bound=code_specs.CPPLikeCodeSpec) +CPPLikeCodeSpecT = TypeVar("CPPLikeCodeSpecT", bound=artifacts.CPPLikeCodeSpec) @dataclasses.dataclass class CMakeFactory( - compiler.BuildSystemProjectGenerator[CPPLikeCodeSpecT, code_specs.PythonCodeSpec] + compiler.BuildSystemProjectGenerator[CPPLikeCodeSpecT, artifacts.PythonCodeSpec] ): """Create a CMakeProject from an ``ExtensionSource`` stage object with given CMake settings.""" @@ -48,7 +48,7 @@ class CMakeFactory( def __call__( self, - source: stages.ExtensionSource[CPPLikeCodeSpecT, code_specs.PythonCodeSpec], + source: artifacts.ExtensionSource[CPPLikeCodeSpecT, artifacts.PythonCodeSpec], cache_lifetime: config.BuildCacheLifetime, ) -> CMakeProject: if not source.binding_source: @@ -60,8 +60,8 @@ def __call__( bindings_name = f"{name}_bindings.{source.program_source.code_spec.file_extension}" cmake_languages = [cmake_lists.Language(name="CXX")] if (src_lang_name := source.program_source.code_spec.source_language) in { - code_specs.CUDACodeSpec.source_language, - code_specs.HIPCodeSpec.source_language, + artifacts.CUDACodeSpec.source_language, + artifacts.HIPCodeSpec.source_language, }: cmake_languages = [*cmake_languages, cmake_lists.Language(name=src_lang_name)] if device_arch_flag := get_cmake_device_arch_option(): @@ -88,7 +88,7 @@ def __call__( @dataclasses.dataclass -class CMakeProject(stages.BuildSystemProject[CPPLikeCodeSpecT, code_specs.PythonCodeSpec]): +class CMakeProject(artifacts.BuildSystemProject[CPPLikeCodeSpecT, artifacts.PythonCodeSpec]): """ CMake build system for gt4py programs. diff --git a/src/gt4py/next/otf/compilation/build_systems/compiledb.py b/src/gt4py/next/otf/compilation/build_systems/compiledb.py index 487d43afa5..92a9facb08 100644 --- a/src/gt4py/next/otf/compilation/build_systems/compiledb.py +++ b/src/gt4py/next/otf/compilation/build_systems/compiledb.py @@ -18,13 +18,13 @@ from gt4py._core import file_utils, locking from gt4py.next import config, errors, fingerprinting -from gt4py.next.otf import code_specs, stages +from gt4py.next.otf import artifacts from gt4py.next.otf.binding import interface from gt4py.next.otf.compilation import build_data, cache, compiler from gt4py.next.otf.compilation.build_systems import cmake -CPPLikeCodeSpecT = TypeVar("CPPLikeCodeSpecT", bound=code_specs.CPPLikeCodeSpec) +CPPLikeCodeSpecT = TypeVar("CPPLikeCodeSpecT", bound=artifacts.CPPLikeCodeSpec) #: Name prefix of the synthetic program under which the shared compiledb is cached. #: Its cache folder sits next to the folders of real programs, so tools that scan @@ -34,7 +34,7 @@ @dataclasses.dataclass class CompiledbFactory( - compiler.BuildSystemProjectGenerator[CPPLikeCodeSpecT, code_specs.PythonCodeSpec] + compiler.BuildSystemProjectGenerator[CPPLikeCodeSpecT, artifacts.PythonCodeSpec] ): """ Create a CompiledbProject from an ``ExtensionSource`` stage object with given CMake settings. @@ -50,7 +50,7 @@ class CompiledbFactory( def __call__( self, - source: stages.ExtensionSource[CPPLikeCodeSpecT, code_specs.PythonCodeSpec], + source: artifacts.ExtensionSource[CPPLikeCodeSpecT, artifacts.PythonCodeSpec], cache_lifetime: config.BuildCacheLifetime, ) -> CompiledbProject: if not source.binding_source: @@ -109,7 +109,7 @@ def _relative_path_to_parent(current: pathlib.Path, parent: pathlib.Path) -> str @dataclasses.dataclass() -class CompiledbProject(stages.BuildSystemProject[CPPLikeCodeSpecT, code_specs.PythonCodeSpec]): +class CompiledbProject(artifacts.BuildSystemProject[CPPLikeCodeSpecT, artifacts.PythonCodeSpec]): """ Compiledb build system for gt4py programs. @@ -255,10 +255,10 @@ def _cc_prototype_program_source( deps: tuple[interface.LibraryDependency, ...], build_type: config.CMakeBuildType, cmake_flags: list[str], - code_spec: code_specs.CPPLikeCodeSpec, -) -> stages.ProgramSource: + code_spec: artifacts.CPPLikeCodeSpec, +) -> artifacts.ProgramSource: name = _cc_prototype_program_name(deps, build_type.value, cmake_flags) - return stages.ProgramSource( + return artifacts.ProgramSource( entry_point=interface.Function(name=name, parameters=()), source_code="", library_deps=deps, @@ -268,15 +268,15 @@ def _cc_prototype_program_source( def _cc_get_compiledb( renew_compiledb: bool, - prototype_program_source: stages.ProgramSource, + prototype_program_source: artifacts.ProgramSource, build_type: config.CMakeBuildType, cmake_flags: list[str], cache_lifetime: config.BuildCacheLifetime, ) -> pathlib.Path: # Use the same prototype source (with empty bindings) for both locating and creating the # compiledb, so `get_cache_folder` names the same folder in either path. - prototype_source: stages.ExtensionSource = stages.ExtensionSource( - prototype_program_source, stages.BindingSource(source_code="", library_deps=()) + prototype_source: artifacts.ExtensionSource = artifacts.ExtensionSource( + prototype_program_source, artifacts.BindingSource(source_code="", library_deps=()) ) cache_path = cache.get_cache_folder(prototype_source, cache_lifetime) @@ -312,7 +312,7 @@ def _cc_find_compiledb(path: pathlib.Path) -> Optional[pathlib.Path]: def _cc_create_compiledb( - prototype_source: stages.ExtensionSource, + prototype_source: artifacts.ExtensionSource, build_type: config.CMakeBuildType, cmake_flags: list[str], cache_lifetime: config.BuildCacheLifetime, diff --git a/src/gt4py/next/otf/compilation/cache.py b/src/gt4py/next/otf/compilation/cache.py index b279b1d2f6..eac9c07dc2 100644 --- a/src/gt4py/next/otf/compilation/cache.py +++ b/src/gt4py/next/otf/compilation/cache.py @@ -13,7 +13,7 @@ from typing import Final from gt4py.next import config, fingerprinting -from gt4py.next.otf import stages +from gt4py.next.otf import artifacts #: Regex describing the folder names produced by `get_cache_folder` (use @@ -74,7 +74,7 @@ def get_cache_base_path(lifetime: config.BuildCacheLifetime) -> pathlib.Path: def get_cache_folder( - ext_source: stages.ExtensionSource, + ext_source: artifacts.ExtensionSource, lifetime: config.BuildCacheLifetime, build_context_id: str = "", ) -> pathlib.Path: diff --git a/src/gt4py/next/otf/compilation/compiler.py b/src/gt4py/next/otf/compilation/compiler.py index fa9d767cb9..e2151e9745 100644 --- a/src/gt4py/next/otf/compilation/compiler.py +++ b/src/gt4py/next/otf/compilation/compiler.py @@ -14,7 +14,7 @@ from gt4py._core import definitions as core_defs, locking from gt4py.next import config, fingerprinting -from gt4py.next.otf import code_specs, definitions, stages, workflow +from gt4py.next.otf import artifacts, workflow from gt4py.next.otf.compilation import build_data, cache, importer @@ -37,17 +37,17 @@ def is_usable( return data is not None and is_compiled(data) and module_exists(data, src_dir) -CodeSpecT = TypeVar("CodeSpecT", bound=code_specs.SourceCodeSpec) -TargetCodeSpecT = TypeVar("TargetCodeSpecT", bound=code_specs.SourceCodeSpec) -CPPLikeCodeSpecT = TypeVar("CPPLikeCodeSpecT", bound=code_specs.CPPLikeCodeSpec) +CodeSpecT = TypeVar("CodeSpecT", bound=artifacts.SourceCodeSpec) +TargetCodeSpecT = TypeVar("TargetCodeSpecT", bound=artifacts.SourceCodeSpec) +CPPLikeCodeSpecT = TypeVar("CPPLikeCodeSpecT", bound=artifacts.CPPLikeCodeSpec) class BuildSystemProjectGenerator(Protocol[CodeSpecT, TargetCodeSpecT]): def __call__( self, - source: stages.ExtensionSource[CodeSpecT, TargetCodeSpecT], + source: artifacts.ExtensionSource[CodeSpecT, TargetCodeSpecT], cache_lifetime: config.BuildCacheLifetime, - ) -> stages.BuildSystemProject[CodeSpecT, TargetCodeSpecT]: ... + ) -> artifacts.BuildSystemProject[CodeSpecT, TargetCodeSpecT]: ... @dataclasses.dataclass(frozen=True) @@ -63,7 +63,7 @@ class CPPCompilationArtifact: entry_point_name: str device_type: core_defs.DeviceType - def load(self) -> stages.ExecutableProgram: + def load(self) -> artifacts.ExecutableProgram: """Import the .so and return the raw entry point. Must run in the process that will call the returned program: @@ -78,14 +78,13 @@ def load(self) -> stages.ExecutableProgram: @dataclasses.dataclass(frozen=True) class CPPCompiler( workflow.ChainableWorkflowMixin[ - stages.ExtensionSource[CPPLikeCodeSpecT, code_specs.PythonCodeSpec], + artifacts.ExtensionSource[CPPLikeCodeSpecT, artifacts.PythonCodeSpec], CPPCompilationArtifact, ], workflow.ReplaceEnabledWorkflowMixin[ - stages.ExtensionSource[CPPLikeCodeSpecT, code_specs.PythonCodeSpec], + artifacts.ExtensionSource[CPPLikeCodeSpecT, artifacts.PythonCodeSpec], CPPCompilationArtifact, ], - definitions.CompilationStep[CPPLikeCodeSpecT, code_specs.PythonCodeSpec], ): """Drive a CPP-style build system into a ``CPPCompilationArtifact``. @@ -93,14 +92,14 @@ class CPPCompiler( """ cache_lifetime: config.BuildCacheLifetime - builder_factory: BuildSystemProjectGenerator[CPPLikeCodeSpecT, code_specs.PythonCodeSpec] + builder_factory: BuildSystemProjectGenerator[CPPLikeCodeSpecT, artifacts.PythonCodeSpec] device_type: core_defs.DeviceType fingerprint_builder_factory: bool = True force_recompile: bool = False def __call__( self, - inp: stages.ExtensionSource[CPPLikeCodeSpecT, code_specs.PythonCodeSpec], + inp: artifacts.ExtensionSource[CPPLikeCodeSpecT, artifacts.PythonCodeSpec], ) -> CPPCompilationArtifact: build_context_id = ( fingerprinting.strict_fingerprinter(self.builder_factory) diff --git a/src/gt4py/next/otf/compilation_tasks.py b/src/gt4py/next/otf/compilation_tasks.py index 25bb9b5060..e172cf370f 100644 --- a/src/gt4py/next/otf/compilation_tasks.py +++ b/src/gt4py/next/otf/compilation_tasks.py @@ -29,7 +29,7 @@ from gt4py._core import definitions as core_defs from gt4py.eve import extended_typing as xtyping from gt4py.next import backend as gtx_backend, common, constructors -from gt4py.next.otf import arguments, definitions as otf_definitions, runners, stages +from gt4py.next.otf import arguments, artifacts, runners, stages def _connectivity_from_file( @@ -127,9 +127,9 @@ def _offset_provider_with_file_refs( class _PreloadedArtifact: """Wraps the already-loaded program of a backend with a customized ``compile``.""" - program: stages.ExecutableProgram + program: artifacts.ExecutableProgram - def load(self) -> stages.ExecutableProgram: + def load(self) -> artifacts.ExecutableProgram: return self.program @@ -158,10 +158,10 @@ def make_compilation_task( # function module attribute, so the raw `types.FunctionType` must not cross # a process boundary; the lowered `CompilableProgramDef` is pickle-safe. compilable = backend.transforms( - otf_definitions.ConcreteProgramDef(data=definition_stage, args=compile_time_args) + stages.ConcreteProgramDef(data=definition_stage, args=compile_time_args) ) - def construct_compilable(with_refs: bool) -> otf_definitions.CompilableProgramDef: + def construct_compilable(with_refs: bool) -> stages.CompilableProgramDef: if not with_refs or not compilable.args.offset_provider: return compilable # The shipped copy must not carry the connectivity buffers: they may diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 50f372f5a9..0b49f9ec0c 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -32,7 +32,7 @@ type_translation, ) from gt4py.next.instrumentation import hook_machinery, metrics -from gt4py.next.otf import arguments, compilation_tasks, runners, stages +from gt4py.next.otf import arguments, artifacts, compilation_tasks, runners from gt4py.next.type_system import type_info, type_specifications as ts from gt4py.next.utils import tree_map @@ -157,7 +157,7 @@ def compiled_program_call_context( # In-flight compilations (future -> "program (backend)" label). Weak keys: a # future disappears here once its pool consumed or dropped it. _ongoing_compilations: weakref.WeakKeyDictionary[ - concurrent.futures.Future[stages.CompilationArtifact], str + concurrent.futures.Future[artifacts.CompilationArtifact], str ] = weakref.WeakKeyDictionary() @@ -360,12 +360,12 @@ class CompiledProgramsPool(Generic[ffront_stages.DSLDefinitionT]): #: Note: The list is not ordered. argument_descriptor_mapping: dict[type[arguments.ArgStaticDescriptor], Sequence[str]] | None # store for the compiled programs - compiled_programs: dict[CompiledProgramsKey, stages.ExecutableProgram] = dataclasses.field( + compiled_programs: dict[CompiledProgramsKey, artifacts.ExecutableProgram] = dataclasses.field( default_factory=dict, init=False ) _compilation_jobs: dict[ - CompiledProgramsKey, concurrent.futures.Future[stages.CompilationArtifact] + CompiledProgramsKey, concurrent.futures.Future[artifacts.CompilationArtifact] ] = dataclasses.field(default_factory=dict, init=False) @functools.cached_property @@ -487,8 +487,8 @@ def _describe_argument_descriptors(self, descriptor_values: tuple[Hashable, ...] ) def _load_artifact( - self, artifact_future: concurrent.futures.Future[stages.CompilationArtifact] - ) -> stages.ExecutableProgram: + self, artifact_future: concurrent.futures.Future[artifacts.CompilationArtifact] + ) -> artifacts.ExecutableProgram: artifact = artifact_future.result() # re-raises errors from the compilation worker try: return self.backend.load_artifact(artifact) diff --git a/src/gt4py/next/otf/definitions.py b/src/gt4py/next/otf/definitions.py deleted file mode 100644 index 5a30abafcb..0000000000 --- a/src/gt4py/next/otf/definitions.py +++ /dev/null @@ -1,60 +0,0 @@ -# GT4Py - GridTools Framework -# -# Copyright (c) 2014-2024, ETH Zurich -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -from typing import Protocol, TypeAlias, TypeVar - -from gt4py.next.ffront import stages as ffront_stages -from gt4py.next.iterator import ir as itir -from gt4py.next.otf import arguments, code_specs, stages, toolchain, workflow - - -CodeSpecT = TypeVar("CodeSpecT", bound=code_specs.SourceCodeSpec) -TargetCodeSpecT = TypeVar("TargetCodeSpecT", bound=code_specs.SourceCodeSpec) - - -IRDefinitionT = TypeVar( - "IRDefinitionT", - ffront_stages.DSLFieldOperatorDef, - ffront_stages.DSLProgramDef, - ffront_stages.FOASTOperatorDef, - ffront_stages.PASTProgramDef, - itir.Program, -) -ArgsDefinitionT = TypeVar("ArgsDefinitionT", arguments.JITArgs, arguments.CompileTimeArgs) - -ConcreteProgramDef: TypeAlias = toolchain.ConcreteArtifact[IRDefinitionT, ArgsDefinitionT] -CompilableProgramDef: TypeAlias = ConcreteProgramDef[itir.Program, arguments.CompileTimeArgs] - - -class TranslationStep( - workflow.ReplaceEnabledWorkflowMixin[CompilableProgramDef, stages.ProgramSource[CodeSpecT]], - Protocol[CodeSpecT], -): - """Translate a GT4Py program to source code (ProgramCall -> ProgramSource).""" - - ... - - -class CompilationStep( - workflow.Workflow[ - stages.ExtensionSource[CodeSpecT, TargetCodeSpecT], stages.CompilationArtifact - ], - Protocol[CodeSpecT, TargetCodeSpecT], -): - """Run the build system and produce a ``stages.CompilationArtifact``. - - Each backend defines its own concrete artifact dataclass (frozen, - picklable, with a ``load`` method); they all satisfy the - ``stages.CompilationArtifact`` Protocol structurally. - """ - - def __call__( - self, source: stages.ExtensionSource[CodeSpecT, TargetCodeSpecT] - ) -> stages.CompilationArtifact: ... diff --git a/src/gt4py/next/otf/recipes.py b/src/gt4py/next/otf/recipes.py index 1c8b7853a7..b905ee0bd3 100644 --- a/src/gt4py/next/otf/recipes.py +++ b/src/gt4py/next/otf/recipes.py @@ -10,15 +10,15 @@ import dataclasses -from gt4py.next.otf import definitions, stages, workflow +from gt4py.next.otf import artifacts, stages, workflow @dataclasses.dataclass(frozen=True) class OTFCompileWorkflow( - workflow.NamedStepSequence[definitions.CompilableProgramDef, stages.CompilationArtifact] + workflow.NamedStepSequence[stages.CompilableProgramDef, artifacts.CompilationArtifact] ): """The typical compiled backend steps composed into a workflow.""" - translation: definitions.TranslationStep - bindings: workflow.Workflow[stages.ProgramSource, stages.ExtensionSource] - compilation: workflow.Workflow[stages.ExtensionSource, stages.CompilationArtifact] + translation: stages.TranslationStep + bindings: workflow.Workflow[artifacts.ProgramSource, artifacts.ExtensionSource] + compilation: workflow.Workflow[artifacts.ExtensionSource, artifacts.CompilationArtifact] diff --git a/src/gt4py/next/otf/runners.py b/src/gt4py/next/otf/runners.py index c23c48061e..f260a33061 100644 --- a/src/gt4py/next/otf/runners.py +++ b/src/gt4py/next/otf/runners.py @@ -24,7 +24,7 @@ from gt4py._core import definitions as core_defs from gt4py.next import config -from gt4py.next.otf import stages +from gt4py.next.otf import artifacts from gt4py.next.otf.compilation import cache as _cache, common as compilation_common @@ -42,11 +42,11 @@ class CompilationTask: construct_compilable: Callable[[bool], Any] #: The artifact-producing step. A runner may execute it anywhere, including #: another process; its picklability is the runner's concern. - executor: Callable[[Any], stages.CompilationArtifact] + executor: Callable[[Any], artifacts.CompilationArtifact] #: Reason this task is known not to be shippable to another process, if any. no_offload_reason: str | None = None - def compile(self, with_refs: bool = False) -> stages.CompilationArtifact: + def compile(self, with_refs: bool = False) -> artifacts.CompilationArtifact: return self.executor(self.construct_compilable(with_refs)) @@ -54,7 +54,7 @@ def compile(self, with_refs: bool = False) -> stages.CompilationArtifact: class Runner(Protocol): def submit( self, task: CompilationTask - ) -> concurrent.futures.Future[stages.CompilationArtifact]: + ) -> concurrent.futures.Future[artifacts.CompilationArtifact]: """Schedule `task`. Returns: @@ -69,8 +69,8 @@ def shutdown(self, wait: bool = True) -> None: def _run_in_calling_thread( task: CompilationTask, -) -> concurrent.futures.Future[stages.CompilationArtifact]: - future: concurrent.futures.Future[stages.CompilationArtifact] = concurrent.futures.Future() +) -> concurrent.futures.Future[artifacts.CompilationArtifact]: + future: concurrent.futures.Future[artifacts.CompilationArtifact] = concurrent.futures.Future() try: future.set_result(task.compile()) except BaseException as exception: # re-raised via the future @@ -83,7 +83,7 @@ class SerialRunner: def submit( self, task: CompilationTask - ) -> concurrent.futures.Future[stages.CompilationArtifact]: + ) -> concurrent.futures.Future[artifacts.CompilationArtifact]: return _run_in_calling_thread(task) def shutdown(self, wait: bool = True) -> None: @@ -98,7 +98,7 @@ def __init__(self, max_workers: int) -> None: def submit( self, task: CompilationTask - ) -> concurrent.futures.Future[stages.CompilationArtifact]: + ) -> concurrent.futures.Future[artifacts.CompilationArtifact]: return self._pool.submit(task.compile) def shutdown(self, wait: bool = True) -> None: @@ -169,7 +169,7 @@ def _run_compilation_task_in_worker( compilable: Any, config_overrides: dict[str, Any], recursion_limit: int, -) -> stages.CompilationArtifact: +) -> artifacts.CompilationArtifact: """Worker entry point: deserialize the executor and run it.""" # `sys.setrecursionlimit` is per-process: a limit the parent raised for # deeply nested IR would silently reset to the interpreter default in a @@ -210,7 +210,7 @@ def __init__(self, max_workers: int, shared_session_cache_dir: str) -> None: def submit( self, task: CompilationTask - ) -> concurrent.futures.Future[stages.CompilationArtifact]: + ) -> concurrent.futures.Future[artifacts.CompilationArtifact]: reason = task.no_offload_reason executor_blob: bytes | None = None if reason is None: diff --git a/src/gt4py/next/otf/stages.py b/src/gt4py/next/otf/stages.py index 17a2961013..4781907c1b 100644 --- a/src/gt4py/next/otf/stages.py +++ b/src/gt4py/next/otf/stages.py @@ -6,134 +6,65 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from __future__ import annotations - -import dataclasses -from collections.abc import Callable -from typing import Final, Generic, Optional, Protocol, TypeAlias, TypeVar, runtime_checkable - -from gt4py.next import fingerprinting -from gt4py.next.otf import code_specs -from gt4py.next.otf.binding import interface - - -compilable_program_fingerprinter: Final[fingerprinting.Fingerprinter] = ( - fingerprinting.strict_fingerprinter -) - - -CodeSpecT = TypeVar("CodeSpecT", bound=code_specs.SourceCodeSpec) -TargetCodeSpecT = TypeVar("TargetCodeSpecT", bound=code_specs.SourceCodeSpec) - - -@dataclasses.dataclass(frozen=True) -class ProgramSource(Generic[CodeSpecT]): - """ - Standalone source code translated from an IR along with information relevant for OTF compilation. - - Contains additional information required for further OTF steps, such as - - implementation language and language conventions - - dependencies on implementation language libraries - - how to call the program - """ - - entry_point: interface.Function - source_code: str - library_deps: tuple[interface.LibraryDependency, ...] - code_spec: CodeSpecT - - -@dataclasses.dataclass(frozen=True) -class BindingSource(Generic[CodeSpecT, TargetCodeSpecT]): - """ - Companion source code for translated program source code. - - This is only needed for OTF compilation if the translated program source code is - not directly callable from python and therefore requires bindings. - This can also optionally be added to compile bindings for other languages than python - when using GT4Py as part of the build for a non-python driver project. - """ - - source_code: str - library_deps: tuple[interface.LibraryDependency, ...] +""" +Definition stages of the toolchain and the contracts of the steps between them. +This module hosts the DSL-aware half of the on-the-fly compilation +vocabulary: which forms a program definition can take on its way from DSL +source to a compilable IR program, and the typed contracts of the steps a +compile pipeline is composed of. The DSL-agnostic artifact models these +contracts produce live in `gt4py.next.otf.artifacts`. +""" -@dataclasses.dataclass(frozen=True) -class ExtensionSource(Generic[CodeSpecT, TargetCodeSpecT]): - """ - Encapsulate all the source code required for OTF compilation. - - The bindings module is optional if and only if the program_source is directly callable. - This should only be the case if the source language / framework supports this out of the box. - If bindings are required, it is recommended to create them in a separate step to ensure reusability. - """ - - program_source: ProgramSource[CodeSpecT] - binding_source: Optional[BindingSource[CodeSpecT, TargetCodeSpecT]] +from __future__ import annotations - @property - def library_deps(self) -> tuple[interface.LibraryDependency, ...]: - if not self.binding_source: - return self.program_source.library_deps - return _unique_libs(*self.program_source.library_deps, *self.binding_source.library_deps) +from typing import Protocol, TypeAlias, TypeVar +from gt4py.next.ffront import stages as ffront_stages +from gt4py.next.iterator import ir as itir +from gt4py.next.otf import arguments, artifacts, workflow -CodeSpecT_co = TypeVar("CodeSpecT_co", bound=code_specs.SourceCodeSpec, covariant=True) -TargetCodeSpecT_co = TypeVar("TargetCodeSpecT_co", bound=code_specs.SourceCodeSpec, covariant=True) +CodeSpecT = TypeVar("CodeSpecT", bound=artifacts.SourceCodeSpec) +TargetCodeSpecT = TypeVar("TargetCodeSpecT", bound=artifacts.SourceCodeSpec) -class BuildSystemProject(Protocol[CodeSpecT_co, TargetCodeSpecT_co]): - """ - Use source code extracted from an ``ExtensionSource`` to configure and build a GT4Py program. - Should only be considered an OTF stage if used as an endpoint, as this only runs commands on source files - and is not responsible for importing the results into Python. - """ +IRDefinitionT = TypeVar( + "IRDefinitionT", + ffront_stages.DSLFieldOperatorDef, + ffront_stages.DSLProgramDef, + ffront_stages.FOASTOperatorDef, + ffront_stages.PASTProgramDef, + itir.Program, +) +ArgsDefinitionT = TypeVar("ArgsDefinitionT", arguments.JITArgs, arguments.CompileTimeArgs) - def build(self) -> None: ... +ConcreteProgramDef: TypeAlias = workflow.ConcreteArtifact[IRDefinitionT, ArgsDefinitionT] +CompilableProgramDef: TypeAlias = ConcreteProgramDef[itir.Program, arguments.CompileTimeArgs] -ExecutableProgram: TypeAlias = Callable +class TranslationStep( + workflow.ReplaceEnabledWorkflowMixin[CompilableProgramDef, artifacts.ProgramSource[CodeSpecT]], + Protocol[CodeSpecT], +): + """Translate a GT4Py program to source code (ProgramCall -> ProgramSource).""" + ... -@runtime_checkable -class CompilationArtifact(Protocol): - """The output of an ``OTFCompileWorkflow``. - Each backend defines its own concrete artifact dataclass; all share this - Protocol. Implementations are frozen dataclasses, picklable, and carry no - live process-bound state — that is reconstructed by ``load``, which - returns a directly-callable ``ExecutableProgram`` taking gt4py-shaped - arguments. +class CompilationStep( + workflow.Workflow[ + artifacts.ExtensionSource[CodeSpecT, TargetCodeSpecT], artifacts.CompilationArtifact + ], + Protocol[CodeSpecT, TargetCodeSpecT], +): + """Run the build system and produce an ``artifacts.CompilationArtifact``. - The one current exception is ``RoundtripArtifact`` when it is configured - with a ``dispatch_backend``: that field holds a ``Backend`` reference - whose role belongs at the runner / load-time seam, not in the artifact - itself. + Each backend defines its own concrete artifact dataclass (frozen, + picklable, with a ``load`` method); they all satisfy the + ``artifacts.CompilationArtifact`` Protocol structurally. """ - def load(self) -> ExecutableProgram: ... - - -def _unique_libs(*args: interface.LibraryDependency) -> tuple[interface.LibraryDependency, ...]: - """ - Filter out multiple occurrences of the same ``interface.LibraryDependency``. - - Examples: - --------- - >>> libs_a = ( - ... interface.LibraryDependency("foo", "1.2.3"), - ... interface.LibraryDependency("common", "1.0.0"), - ... ) - >>> libs_b = ( - ... interface.LibraryDependency("common", "1.0.0"), - ... interface.LibraryDependency("bar", "1.2.3"), - ... ) - >>> _unique_libs(*libs_a, *libs_b) - (LibraryDependency(name='foo', version='1.2.3'), LibraryDependency(name='common', version='1.0.0'), LibraryDependency(name='bar', version='1.2.3')) - """ - unique: list[interface.LibraryDependency] = [] - for lib in args: - if lib not in unique: - unique.append(lib) - return tuple(unique) + def __call__( + self, source: artifacts.ExtensionSource[CodeSpecT, TargetCodeSpecT] + ) -> artifacts.CompilationArtifact: ... diff --git a/src/gt4py/next/otf/toolchain.py b/src/gt4py/next/otf/toolchain.py index 0c816759ff..14d1b1cdf8 100644 --- a/src/gt4py/next/otf/toolchain.py +++ b/src/gt4py/next/otf/toolchain.py @@ -21,46 +21,44 @@ ArgsT = typing.TypeVar("ArgsT") -@dataclasses.dataclass -class ConcreteArtifact(Generic[DefT, ArgsT]): - data: DefT - args: ArgsT - - @dataclasses.dataclass(frozen=True) class DataOnlyAdapter( workflow.ChainableWorkflowMixin, workflow.ReplaceEnabledWorkflowMixin, - workflow.Workflow[ConcreteArtifact[S, ArgsT], ConcreteArtifact[T, ArgsT]], + workflow.Workflow[workflow.ConcreteArtifact[S, ArgsT], workflow.ConcreteArtifact[T, ArgsT]], Generic[ArgsT, S, T], ): step: workflow.Workflow[S, T] - def __call__(self, inp: ConcreteArtifact[S, ArgsT]) -> ConcreteArtifact[T, ArgsT]: - return ConcreteArtifact(data=self.step(inp.data), args=inp.args) + def __call__( + self, inp: workflow.ConcreteArtifact[S, ArgsT] + ) -> workflow.ConcreteArtifact[T, ArgsT]: + return workflow.ConcreteArtifact(data=self.step(inp.data), args=inp.args) @dataclasses.dataclass(frozen=True) class ArgsOnlyAdapter( workflow.ChainableWorkflowMixin, workflow.ReplaceEnabledWorkflowMixin, - workflow.Workflow[ConcreteArtifact[DefT, S], ConcreteArtifact[DefT, T]], + workflow.Workflow[workflow.ConcreteArtifact[DefT, S], workflow.ConcreteArtifact[DefT, T]], Generic[DefT, S, T], ): step: workflow.Workflow[S, T] - def __call__(self, inp: ConcreteArtifact[DefT, S]) -> ConcreteArtifact[DefT, T]: - return ConcreteArtifact(data=inp.data, args=self.step(inp.args)) + def __call__( + self, inp: workflow.ConcreteArtifact[DefT, S] + ) -> workflow.ConcreteArtifact[DefT, T]: + return workflow.ConcreteArtifact(data=inp.data, args=self.step(inp.args)) @dataclasses.dataclass(frozen=True) class StripArgsAdapter( workflow.ChainableWorkflowMixin, workflow.ReplaceEnabledWorkflowMixin, - workflow.Workflow[ConcreteArtifact[S, ArgsT], T], + workflow.Workflow[workflow.ConcreteArtifact[S, ArgsT], T], Generic[ArgsT, S, T], ): step: workflow.Workflow[S, T] - def __call__(self, inp: ConcreteArtifact[S, ArgsT]) -> T: + def __call__(self, inp: workflow.ConcreteArtifact[S, ArgsT]) -> T: return self.step(inp.data) diff --git a/src/gt4py/next/otf/workflow.py b/src/gt4py/next/otf/workflow.py index 274e8da029..6e5bf42837 100644 --- a/src/gt4py/next/otf/workflow.py +++ b/src/gt4py/next/otf/workflow.py @@ -31,6 +31,26 @@ HashT = TypeVar("HashT") DataT = TypeVar("DataT") ArgT = TypeVar("ArgT") +DefT = TypeVar("DefT") +ArgsT = TypeVar("ArgsT") + + +@dataclasses.dataclass +class ConcreteArtifact(Generic[DefT, ArgsT]): + """Pair of a program definition in any stage with the arguments it is compiled for. + + This is the envelope threaded through the definition-transforming half of + the toolchain, so it is generic over DSL-frontend types. It nevertheless + lives in this DSL-neutral bottom module rather than beside the stage + definitions that use it: those parameterize it while they are being + imported, so hosting it any higher would close an import cycle. The + invariant that keeps this working -- this module reaching no IR or + frontend module at import time -- is checked by the OTF import-boundary + test in `tests/next_tests/unit_tests/otf_tests/`. + """ + + data: DefT + args: ArgsT def make_step(function: Workflow[StartT, EndT]) -> ChainableWorkflowMixin[StartT, EndT]: 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 a452bf53fe..c32c1c15d1 100644 --- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py +++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py @@ -21,7 +21,7 @@ from gt4py.next.ffront import fbuiltins from gt4py.next.iterator import ir as itir from gt4py.next.iterator.transforms import pass_manager -from gt4py.next.otf import code_specs, definitions, stages, workflow +from gt4py.next.otf import artifacts, stages, workflow from gt4py.next.otf.binding import cpp_interface, interface from gt4py.next.program_processors.codegens.gtfn.codegen import GTFNCodegen, GTFNIMCodegen from gt4py.next.program_processors.codegens.gtfn.gtfn_ir_to_gtfn_im_ir import GTFN_IM_lowering @@ -39,15 +39,15 @@ def get_param_description(name: str, type_: Any) -> interface.Parameter: @dataclasses.dataclass(frozen=True) class GTFNTranslationStep( workflow.ReplaceEnabledWorkflowMixin[ - definitions.CompilableProgramDef, - stages.ProgramSource[code_specs.HeaderAndSourceCodeSpec], + stages.CompilableProgramDef, + artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec], ], workflow.ChainableWorkflowMixin[ - definitions.CompilableProgramDef, - stages.ProgramSource[code_specs.HeaderAndSourceCodeSpec], + stages.CompilableProgramDef, + artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec], ], ): - code_spec: Optional[code_specs.HeaderAndSourceCodeSpec] = None + code_spec: Optional[artifacts.HeaderAndSourceCodeSpec] = None # TODO replace by more general mechanism, see https://github.com/GridTools/gt4py/issues/1135 enable_itir_transforms: bool = True use_imperative_backend: bool = False @@ -55,14 +55,14 @@ class GTFNTranslationStep( symbolic_domain_sizes: dict[str, itir.Expr] | None = None use_max_domain_range_on_unstructured_shift: bool | None = None - def _default_code_spec(self) -> code_specs.HeaderAndSourceCodeSpec: + def _default_code_spec(self) -> artifacts.HeaderAndSourceCodeSpec: match self.device_type: case core_defs.DeviceType.CUDA: - return code_specs.CUDACodeSpec() + return artifacts.CUDACodeSpec() case core_defs.DeviceType.ROCM: - return code_specs.HIPCodeSpec() + return artifacts.HIPCodeSpec() case core_defs.DeviceType.CPU: - return code_specs.CPPCodeSpec() + return artifacts.CPPCodeSpec() case _: raise self._not_implemented_for_device_type() @@ -196,8 +196,8 @@ def generate_stencil_source( return codegen.format_source("cpp", generated_code, style="LLVM") def __call__( - self, inp: definitions.CompilableProgramDef - ) -> stages.ProgramSource[code_specs.HeaderAndSourceCodeSpec]: + self, inp: stages.CompilableProgramDef + ) -> artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec]: """Generate GTFN C++ code from the ITIR definition.""" program: itir.Program = inp.data @@ -230,7 +230,7 @@ def __call__( inp.args.offset_provider, inp.args.column_axis, ) - source_code = interface.format_source( + source_code = artifacts.format_source( self._code_spec(), f""" #include <{self._backend_header()}> @@ -240,11 +240,13 @@ def __call__( """.strip(), ) - module: stages.ProgramSource[code_specs.HeaderAndSourceCodeSpec] = stages.ProgramSource( - entry_point=function, - library_deps=(interface.LibraryDependency(self._library_name(), "master"),), - source_code=source_code, - code_spec=self._code_spec(), + module: artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec] = ( + artifacts.ProgramSource( + entry_point=function, + library_deps=(interface.LibraryDependency(self._library_name(), "master"),), + source_code=source_code, + code_spec=self._code_spec(), + ) ) return module @@ -266,7 +268,7 @@ def _backend_type(self) -> str: case _: raise self._not_implemented_for_device_type() - def _code_spec(self) -> code_specs.HeaderAndSourceCodeSpec: + def _code_spec(self) -> artifacts.HeaderAndSourceCodeSpec: return self.code_spec if self.code_spec is not None else self._default_code_spec() def _library_name(self) -> str: @@ -289,8 +291,8 @@ class Meta: model = GTFNTranslationStep -translate_program_cpu: Final[definitions.TranslationStep] = GTFNTranslationStepFactory() # type: ignore[assignment] # factory-boy typing not precise enough +translate_program_cpu: Final[stages.TranslationStep] = GTFNTranslationStepFactory() # type: ignore[assignment] # factory-boy typing not precise enough -translate_program_gpu: Final[definitions.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 device_type=core_defs.DeviceType.CUDA ) diff --git a/src/gt4py/next/program_processors/runners/dace/program.py b/src/gt4py/next/program_processors/runners/dace/program.py index 7626d7b3b5..f8f8c7e1a4 100644 --- a/src/gt4py/next/program_processors/runners/dace/program.py +++ b/src/gt4py/next/program_processors/runners/dace/program.py @@ -18,7 +18,7 @@ from gt4py.next.ffront import decorator from gt4py.next.iterator import ir as itir, transforms as itir_transforms from gt4py.next.iterator.transforms import extractors as extractors -from gt4py.next.otf import arguments, toolchain, workflow +from gt4py.next.otf import arguments, workflow from gt4py.next.program_processors.runners.dace import sdfg_args as gtx_dace_args from gt4py.next.type_system import type_specifications as ts @@ -44,7 +44,7 @@ def __sdfg__(self, *args: Any, **kwargs: Any) -> dace.sdfg.sdfg.SDFG: # TODO(ricoh): connectivity tables required here for now. gtir_stage = typing.cast(gtx_backend.Transforms, self.backend.transforms).past_to_itir( - toolchain.ConcreteArtifact( + workflow.ConcreteArtifact( data=self.past_stage, args=arguments.CompileTimeArgs( args=tuple(p.type for p in self.past_stage.past_node.params), 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 34ac8e6f4e..c0ea33daf0 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -17,7 +17,7 @@ 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.otf import stages +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 ( common as gtx_wfdcommon, @@ -32,7 +32,7 @@ class DaCeBackend(backend.Backend[Any]): external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None - def load_artifact(self, artifact: stages.CompilationArtifact) -> stages.ExecutableProgram: + def load_artifact(self, artifact: artifacts.CompilationArtifact) -> artifacts.ExecutableProgram: program = super().load_artifact(artifact) assert isinstance(program, gtx_wfddecoration.DaCeDecoratedProgram) # Inject the backend-level workspace so it is used when arguments are constructed. 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 311ce428cf..82e6b3fc3c 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/bindings.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/bindings.py @@ -13,7 +13,7 @@ import dace from gt4py.eve import codegen -from gt4py.next.otf import code_specs, stages +from gt4py.next.otf import artifacts from gt4py.next.program_processors.runners.dace import ( sdfg_args as gtx_dace_args, sdfg_utils as gtx_dace_utils, @@ -226,9 +226,9 @@ def _parse_gt_connectivities( def _create_sdfg_bindings( - program_source: stages.ProgramSource[code_specs.SDFGCodeSpec], + program_source: artifacts.ProgramSource[artifacts.SDFGCodeSpec], bind_func_name: str, -) -> stages.BindingSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec]: +) -> artifacts.BindingSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec]: """ Creates a Python translation function to convert the GT4Py arguments list to the SDFG calling convention. @@ -285,19 +285,19 @@ def _create_sdfg_bindings( # arrays as well in SDFG fastcall. _parse_gt_connectivities(code, sdfg_arglist) - return stages.BindingSource(code.text, library_deps=tuple()) + return artifacts.BindingSource(code.text, library_deps=tuple()) def bind_sdfg( - inp: stages.ProgramSource[code_specs.SDFGCodeSpec], + inp: artifacts.ProgramSource[artifacts.SDFGCodeSpec], bind_func_name: str, -) -> stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec]: +) -> artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec]: """ Method to be used as workflow stage for generation of SDFG bindings. Refer to `_create_sdfg_bindings` documentation. """ - return stages.ExtensionSource( + return artifacts.ExtensionSource( program_source=inp, binding_source=_create_sdfg_bindings(inp, bind_func_name), ) 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 87abce5cf6..6b2683b7ef 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -23,7 +23,7 @@ from gt4py._core import definitions as core_defs, locking from gt4py.eve import extended_typing as xtyping from gt4py.next import common, config, fingerprinting -from gt4py.next.otf import code_specs, definitions, stages, workflow +from gt4py.next.otf import artifacts, workflow from gt4py.next.otf.compilation import cache as gtx_cache from gt4py.next.program_processors.runners.dace.workflow import ( common as gtx_wfdcommon, @@ -34,8 +34,8 @@ _COMPILE_COMPLETE_MARKER: Final = ".gt4py_compile_complete" -SDFGExtensionSource: TypeAlias = stages.ExtensionSource[ - code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec +SDFGExtensionSource: TypeAlias = artifacts.ExtensionSource[ + artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec ] @@ -257,7 +257,7 @@ class DaCeCompilationArtifact: bind_func_name: str device_type: core_defs.DeviceType - def load(self) -> stages.ExecutableProgram: + def load(self) -> artifacts.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``. @@ -270,14 +270,13 @@ def load(self) -> stages.ExecutableProgram: @dataclasses.dataclass(frozen=True) class DaCeCompiler( workflow.ChainableWorkflowMixin[ - stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], + artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec], DaCeCompilationArtifact, ], workflow.ReplaceEnabledWorkflowMixin[ - stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], + artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec], DaCeCompilationArtifact, ], - definitions.CompilationStep[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], ): """Run the DaCe build system and produce an on-disk ``DaCeCompilationArtifact``.""" 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 51dcc93553..2f37f90cd7 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py @@ -14,8 +14,8 @@ import factory from gt4py._core import definitions as core_defs, filecache -from gt4py.next import config -from gt4py.next.otf import recipes, stages, workflow +from gt4py.next import config, fingerprinting +from gt4py.next.otf import recipes, workflow 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 ( @@ -44,7 +44,7 @@ class Params: translation=factory.LazyAttribute( lambda o: workflow.CachedStep.persistent( o.bare_translation, - input_fingerprinter=stages.compilable_program_fingerprinter, + input_fingerprinter=fingerprinting.strict_fingerprinter, cache=filecache.FileCache( cache.get_translation_cache_folder( cache.get_cache_base_path(config.BUILD_CACHE_LIFETIME), "dace" 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 e1fac0600a..f26e982eff 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py @@ -18,7 +18,7 @@ from gt4py.next import common from gt4py.next.instrumentation import metrics from gt4py.next.iterator import ir as itir, transforms as itir_transforms -from gt4py.next.otf import code_specs, definitions, stages, workflow +from gt4py.next.otf import artifacts, stages, workflow from gt4py.next.otf.binding import interface from gt4py.next.program_processors.runners.dace import ( lowering as gtx_dace_lowering, @@ -343,10 +343,13 @@ def make_sdfg_call_sync(sdfg: dace.SDFG, gpu: bool) -> None: @dataclasses.dataclass(frozen=True) class DaCeTranslator( workflow.ChainableWorkflowMixin[ - definitions.CompilableProgramDef, - stages.ProgramSource[code_specs.SDFGCodeSpec], + stages.CompilableProgramDef, + artifacts.ProgramSource[artifacts.SDFGCodeSpec], + ], + workflow.ReplaceEnabledWorkflowMixin[ + stages.CompilableProgramDef, + artifacts.ProgramSource[artifacts.SDFGCodeSpec], ], - definitions.TranslationStep[code_specs.SDFGCodeSpec], ): device_type: core_defs.DeviceType auto_optimize: bool @@ -438,8 +441,8 @@ def _generate_sdfg_without_configuring_dace( return sdfg def __call__( - self, inp: definitions.CompilableProgramDef - ) -> stages.ProgramSource[code_specs.SDFGCodeSpec]: + self, inp: stages.CompilableProgramDef + ) -> artifacts.ProgramSource[artifacts.SDFGCodeSpec]: """Generate DaCe SDFG file from the GTIR definition.""" program: itir.Program = inp.data assert isinstance(program, itir.Program) @@ -457,11 +460,11 @@ def __call__( for param, arg_type in zip(program.params, arg_types) ) - module: stages.ProgramSource[code_specs.SDFGCodeSpec] = stages.ProgramSource( + module: artifacts.ProgramSource[artifacts.SDFGCodeSpec] = artifacts.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. library_deps=tuple(), - code_spec=code_specs.SDFGCodeSpec(), + code_spec=artifacts.SDFGCodeSpec(), ) return module diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py index d8ad151749..c48caab305 100644 --- a/src/gt4py/next/program_processors/runners/gtfn.py +++ b/src/gt4py/next/program_processors/runners/gtfn.py @@ -16,10 +16,10 @@ 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 +from gt4py.next import backend, common, config, field_utils, fingerprinting from gt4py.next.embedded import nd_array_field from gt4py.next.instrumentation import metrics -from gt4py.next.otf import recipes, stages, workflow +from gt4py.next.otf import artifacts, recipes, workflow from gt4py.next.otf.binding import nanobind from gt4py.next.otf.compilation import cache, compiler from gt4py.next.otf.compilation.build_systems import compiledb @@ -45,8 +45,8 @@ def convert_arg(arg: Any) -> Any: def convert_args( - inp: stages.ExecutableProgram, device: core_defs.DeviceType = core_defs.DeviceType.CPU -) -> stages.ExecutableProgram: + inp: artifacts.ExecutableProgram, device: core_defs.DeviceType = core_defs.DeviceType.CPU +) -> artifacts.ExecutableProgram: def decorated_program( *args: Any, offset_provider: dict[str, common.OffsetProviderElem], @@ -106,7 +106,7 @@ def extract_connectivity_args( @dataclasses.dataclass(frozen=True) class GTFNCompilationArtifact(compiler.CPPCompilationArtifact): - def load(self) -> stages.ExecutableProgram: + def load(self) -> artifacts.ExecutableProgram: return convert_args(super().load(), device=self.device_type) @@ -148,7 +148,7 @@ class Params: translation=factory.LazyAttribute( lambda o: workflow.CachedStep.persistent( o.bare_translation, - input_fingerprinter=stages.compilable_program_fingerprinter, + input_fingerprinter=fingerprinting.strict_fingerprinter, cache=filecache.FileCache( cache.get_translation_cache_folder( cache.get_cache_base_path(config.BUILD_CACHE_LIFETIME), "gtfn" @@ -164,7 +164,7 @@ class Params: ) translation = factory.LazyAttribute(lambda o: o.bare_translation) - bindings: workflow.Workflow[stages.ProgramSource, stages.ExtensionSource] = ( + 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 diff --git a/src/gt4py/next/program_processors/runners/roundtrip.py b/src/gt4py/next/program_processors/runners/roundtrip.py index 5c015cd69c..09f173d3f9 100644 --- a/src/gt4py/next/program_processors/runners/roundtrip.py +++ b/src/gt4py/next/program_processors/runners/roundtrip.py @@ -27,7 +27,7 @@ ) from gt4py.next.ffront import foast_to_gtir, foast_to_past, past_to_itir from gt4py.next.iterator import ir as itir, transforms as itir_transforms -from gt4py.next.otf import definitions, stages, workflow +from gt4py.next.otf import artifacts, stages, workflow from gt4py.next.type_system import type_info, type_specifications as ts @@ -224,7 +224,7 @@ class RoundtripArtifact: dispatch_backend: next_backend.Backend | None debug: bool - def load(self) -> stages.ExecutableProgram: + def load(self) -> artifacts.ExecutableProgram: mod = _load_module(self.source_code, self.debug) fencil = getattr(mod, self.entry_point_name) captured_column_axis = self.column_axis @@ -253,13 +253,13 @@ def decorated_fencil( @dataclasses.dataclass(frozen=True) -class Roundtrip(workflow.Workflow[definitions.CompilableProgramDef, RoundtripArtifact]): +class Roundtrip(workflow.Workflow[stages.CompilableProgramDef, RoundtripArtifact]): debug: Optional[bool] = None use_embedded: bool = True dispatch_backend: Optional[next_backend.Backend] = None transforms: itir_transforms.GTIRTransform = itir_transforms.apply_common_transforms # type: ignore[assignment] # TODO(havogt): cleanup interface of `apply_common_transforms` - def __call__(self, inp: definitions.CompilableProgramDef) -> RoundtripArtifact: + def __call__(self, inp: stages.CompilableProgramDef) -> RoundtripArtifact: debug = config.DEBUG if self.debug is None else self.debug source_code, entry_point_name = _generate_source( diff --git a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py index 9e524fcfeb..ba83a07b2d 100644 --- a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py +++ b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/conftest.py @@ -14,12 +14,12 @@ import gt4py.next as gtx import gt4py.next.type_system.type_specifications as ts from gt4py.next import config -from gt4py.next.otf import code_specs, stages +from gt4py.next.otf import artifacts from gt4py.next.otf.binding import cpp_interface, interface, nanobind from gt4py.next.otf.compilation import cache -def make_program_source(name: str) -> stages.ProgramSource: +def make_program_source(name: str) -> artifacts.ProgramSource: entry_point = interface.Function( name, parameters=( @@ -69,11 +69,11 @@ def make_program_source(name: str) -> stages.ProgramSource: """ ).render(func=func) - return stages.ProgramSource( + return artifacts.ProgramSource( entry_point=entry_point, source_code=src, library_deps=(interface.LibraryDependency("gridtools_cpu", "master"),), - code_spec=code_specs.CPPCodeSpec(), + code_spec=artifacts.CPPCodeSpec(), ) @@ -89,7 +89,7 @@ def program_source_example(): @pytest.fixture def extension_source_example(program_source_example): - return stages.ExtensionSource( + return artifacts.ExtensionSource( program_source=program_source_example, binding_source=nanobind.create_bindings( program_source_example, config.UNSTRUCTURED_HORIZONTAL_HAS_UNIT_STRIDE diff --git a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/test_cache_consistency.py b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/test_cache_consistency.py index 81d76e6cbd..f17517c42e 100644 --- a/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/test_cache_consistency.py +++ b/tests/next_tests/unit_tests/otf_tests/compilation_tests/build_systems_tests/test_cache_consistency.py @@ -23,7 +23,7 @@ from gt4py._core import definitions as core_defs from gt4py.next import config, fingerprinting -from gt4py.next.otf import stages +from gt4py.next.otf import artifacts from gt4py.next.otf.compilation import build_data, cache, compiler from gt4py.next.otf.compilation.build_systems import compiledb @@ -36,7 +36,9 @@ def _compiler() -> compiler.CPPCompiler: ) -def _src_dir(comp: compiler.CPPCompiler, extension_source: stages.ExtensionSource) -> pathlib.Path: +def _src_dir( + comp: compiler.CPPCompiler, extension_source: artifacts.ExtensionSource +) -> pathlib.Path: return cache.get_cache_folder( extension_source, config.BuildCacheLifetime.SESSION, 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 2f9120bda5..9b3f4bc2cb 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 @@ -20,7 +20,7 @@ from gt4py.next import errors, backend, broadcast, common from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple from gt4py.next.iterator.ir_utils import ir_makers as im -from gt4py.next.otf import toolchain, arguments, compiled_program +from gt4py.next.otf import arguments, compiled_program, workflow from gt4py.next.type_system import type_specifications as ts from gt4py.next.iterator import ir as itir from gt4py.next.program_processors.runners import gtfn @@ -90,7 +90,7 @@ def _verify_program_has_expected_true_value(program: itir.Program): def test_inlining_of_scalars_works(testee_prog): - input_pair = toolchain.ConcreteArtifact( + input_pair = workflow.ConcreteArtifact( data=testee_prog.definition_stage, args=arguments.CompileTimeArgs( args=list(testee_prog.past_stage.past_node.type.definition.pos_or_kw_args.values()), @@ -123,7 +123,7 @@ class _NoOpArtifact: def load(self): return lambda *args, **kwargs: None - def pirate(program: toolchain.ConcreteArtifact): + def pirate(program: workflow.ConcreteArtifact): # Replaces the gtfn otf_workflow: steals the compilable program, then # returns a dummy artifact whose materialization is a no-op callable. nonlocal hijacked_program @@ -298,7 +298,7 @@ def _verify_program_has_expected_domain( def test_inlining_of_static_domain_works(testee_prog, uids: utils.IDGeneratorPool): domain = gtx.Domain(dims=(TDim,), ranges=(gtx.UnitRange(0, 1),)) - input_pair = toolchain.ConcreteArtifact( + input_pair = workflow.ConcreteArtifact( data=testee_prog.definition_stage, args=arguments.CompileTimeArgs( args=list(testee_prog.past_stage.past_node.type.definition.pos_or_kw_args.values()), diff --git a/tests/next_tests/unit_tests/otf_tests/test_import_boundaries.py b/tests/next_tests/unit_tests/otf_tests/test_import_boundaries.py new file mode 100644 index 0000000000..72a40a6b52 --- /dev/null +++ b/tests/next_tests/unit_tests/otf_tests/test_import_boundaries.py @@ -0,0 +1,253 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +""" +Executable check of the layering inside `gt4py.next.otf`. + +The DSL-agnostic half of the toolchain (artifact models, build systems, +compilers, binding interface) must not pull in the GT4Py IRs. That boundary +cannot be observed by importing the modules: every one of them does +``from gt4py.next import config``, which executes the `gt4py.next` package +`__init__` and therefore loads the whole frontend regardless of what the +module itself declares. So the graph is reconstructed from the AST instead, +counting only the imports that a real interpreter would execute when the +module is loaded. Two kinds of import are therefore not edges: + +- the `if TYPE_CHECKING:` branch, which never executes (its `else:` branch, + the one a real interpreter takes, does count). +- imports inside a function body, which execute only if that function is + called. Deferring an import into a function is the sanctioned way to use a + higher layer from a lower one without taking an import-time dependency on + it, so counting those would forbid the very escape hatch this boundary + relies on (`instrumentation.stage_dump` formats FOAST nodes that way). + +Class bodies and module-level conditionals *are* entered: both execute while +the module is being loaded. +""" + +from __future__ import annotations + +import ast +import collections +import pathlib +from collections.abc import Iterable, Iterator, Mapping, Set + +import pytest + +import gt4py + + +#: Root of the `gt4py` package, and the path its dotted module names are +#: relative to. Only the package itself is scanned: with a non-editable +#: install `SRC_ROOT` is `site-packages`, and walking that would both parse +#: every unrelated installed package and let their modules into the graph. +PACKAGE_ROOT = pathlib.Path(gt4py.__file__).resolve().parent +SRC_ROOT = PACKAGE_ROOT.parent + +#: Packages that hold the GT4Py IRs and the DSL frontend. +IR_PACKAGES = ("gt4py.next.iterator", "gt4py.next.ffront") + +#: Modules that must stay clear of the IRs. `otf.compilation.*` is expanded +#: to the package and all its submodules, so that a newly added file is +#: covered automatically. +IR_FREE_MODULES = ( + "gt4py.next.otf.workflow", + "gt4py.next.otf.artifacts", + "gt4py.next.otf.binding.interface", + "gt4py.next.otf.runners", + "gt4py.next.otf.compilation.*", +) + +#: `binding.cpp_interface` uses `type_system.type_info.is_compatible_type`, +#: which still special-cases iterator types. See the TODO(tehrengruber) on +#: that function in `src/gt4py/next/type_system/type_info.py` (line 429): +#: until it is split, `otf.binding` cannot be IR-free. Pinned exactly so that +#: the debt neither grows nor silently outlives its fix. +PINNED_IR_DEPENDENCIES = { + "gt4py.next.otf.binding.cpp_interface": {"gt4py.next.iterator.type_system.type_specifications"}, + "gt4py.next.otf.binding.nanobind": {"gt4py.next.iterator.type_system.type_specifications"}, +} + + +def _is_type_checking_guard(node: ast.stmt) -> bool: + if not isinstance(node, ast.If): + return False + test = node.test + return (isinstance(test, ast.Name) and test.id == "TYPE_CHECKING") or ( + isinstance(test, ast.Attribute) and test.attr == "TYPE_CHECKING" + ) + + +def _import_time_nodes(nodes: Iterable[ast.AST]) -> Iterator[ast.AST]: + """ + Yield the nodes that are evaluated while the module is being loaded. + + Descends through class bodies and conditionals, which execute at load + time, but not into function bodies or the `if TYPE_CHECKING:` branch, + which do not. The `else:` branch of a `TYPE_CHECKING` guard *is* the + branch a real interpreter takes, so it is descended into. + """ + for node in nodes: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + continue + if _is_type_checking_guard(node): + assert isinstance(node, ast.If) + yield from _import_time_nodes(node.orelse) + continue + yield node + yield from _import_time_nodes(ast.iter_child_nodes(node)) + + +def _imported_names(path: pathlib.Path, module: str) -> Iterator[str]: + """Yield the dotted names imported when `module` is executed at import time.""" + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + # The package a relative import is resolved against: for a package + # `__init__` that is the module's own dotted name, for a plain module its + # parent. Getting this wrong silently drops the edge, because the + # bogus dotted name resolves back to `module` itself. + self_parts = module.split(".") + if path.name != "__init__.py": + self_parts = self_parts[:-1] + for node in _import_time_nodes(tree.body): + if isinstance(node, ast.Import): + for alias in node.names: + yield alias.name + elif isinstance(node, ast.ImportFrom): + if node.level: # relative import + package = ".".join(self_parts[: len(self_parts) - (node.level - 1)]) + base = f"{package}.{node.module}" if node.module else package + else: + base = node.module or "" + for alias in node.names: + yield f"{base}.{alias.name}" + + +def _collect_modules() -> Mapping[str, pathlib.Path]: + modules = {} + for path in sorted(PACKAGE_ROOT.rglob("*.py")): + parts = list(path.relative_to(SRC_ROOT).with_suffix("").parts) + if parts[-1] == "__init__": + parts = parts[:-1] + modules[".".join(parts)] = path + return modules + + +def _build_import_graph() -> Mapping[str, Set[str]]: + """Map each module to the modules it imports at module-execution time.""" + modules = _collect_modules() + graph: dict[str, set[str]] = {module: set() for module in modules} + for module, path in modules.items(): + for name in _imported_names(path, module): + # Resolve to the most specific existing module, so that + # `from gt4py.next import config` is an edge to `gt4py.next.config` + # and not to the `gt4py.next` package (whose `__init__` re-exports + # the whole DSL and would make every edge reach everything). + parts = name.split(".") + while parts: + candidate = ".".join(parts) + if candidate in modules: + if candidate != module: + graph[module].add(candidate) + break + parts.pop() + return graph + + +def _is_ir_module(module: str) -> bool: + return any(module == pkg or module.startswith(f"{pkg}.") for pkg in IR_PACKAGES) + + +def _shortest_path_to_ir(graph: Mapping[str, Set[str]], start: str) -> list[str] | None: + """Return the shortest import chain from `start` into an IR package, if any.""" + queue = collections.deque([(start, [start])]) + seen = {start} + while queue: + module, chain = queue.popleft() + for imported in sorted(graph.get(module, ())): + if imported in seen: + continue + seen.add(imported) + extended = [*chain, imported] + if _is_ir_module(imported): + return extended + queue.append((imported, extended)) + return None + + +def _reachable_ir_modules(graph: Mapping[str, Set[str]], start: str) -> set[str]: + reachable: set[str] = set() + stack = [start] + while stack: + module = stack.pop() + for imported in graph.get(module, ()): + if imported not in reachable: + reachable.add(imported) + stack.append(imported) + return {module for module in reachable if _is_ir_module(module)} + + +@pytest.fixture(scope="module") +def import_graph() -> Mapping[str, Set[str]]: + return _build_import_graph() + + +def _expand(patterns: tuple[str, ...], graph: Mapping[str, Set[str]]) -> list[str]: + expanded = [] + for pattern in patterns: + if pattern.endswith(".*"): + # The package `__init__` is covered too: it executes on every + # import of one of its submodules. + package = pattern[:-2] + prefix = f"{package}." + matched = sorted( + module for module in graph if module == package or module.startswith(prefix) + ) + assert matched, f"no modules found under '{pattern}'" + expanded.extend(matched) + else: + assert pattern in graph, f"unknown module '{pattern}'" + expanded.append(pattern) + return expanded + + +def test_source_root_is_the_working_tree(import_graph): + """Guard against the check silently passing on an unexpected install layout.""" + assert "gt4py.next.otf.artifacts" in import_graph, ( + f"'{SRC_ROOT}' does not look like the gt4py source tree" + ) + + +def test_dsl_agnostic_otf_modules_do_not_reach_the_irs(import_graph): + offenders = {} + for module in _expand(IR_FREE_MODULES, import_graph): + chain = _shortest_path_to_ir(import_graph, module) + if chain is not None: + offenders[module] = chain + assert not offenders, "DSL-agnostic OTF modules must not import the GT4Py IRs:\n" + "\n".join( + f" {module}: " + " -> ".join(chain) for module, chain in sorted(offenders.items()) + ) + + +@pytest.mark.parametrize("module", sorted(PINNED_IR_DEPENDENCIES)) +def test_binding_ir_dependency_stays_pinned(import_graph, module): + expected = PINNED_IR_DEPENDENCIES[module] + actual = _reachable_ir_modules(import_graph, module) + if actual == expected: + return + if not actual: + pytest.fail( + f"'{module}' no longer reaches the IRs. Move it to `IR_FREE_MODULES` " + f"and drop it from `PINNED_IR_DEPENDENCIES`." + ) + chain = _shortest_path_to_ir(import_graph, module) + pytest.fail( + f"'{module}' IR dependency changed.\n" + f" expected: {sorted(expected)}\n" + f" actual: {sorted(actual)}\n" + f" shortest chain: {' -> '.join(chain or [])}" + ) diff --git a/tests/next_tests/unit_tests/otf_tests/test_languages.py b/tests/next_tests/unit_tests/otf_tests/test_languages.py index 781dc73dad..08a738d017 100644 --- a/tests/next_tests/unit_tests/otf_tests/test_languages.py +++ b/tests/next_tests/unit_tests/otf_tests/test_languages.py @@ -8,14 +8,14 @@ import pytest -from gt4py.next.otf import code_specs, stages +from gt4py.next.otf import artifacts from gt4py.next.otf.binding import interface def test_header_files_settings_with_cpp_accepted(): - stages.ProgramSource( + artifacts.ProgramSource( entry_point=interface.Function(name="basic_settings_with_cpp", parameters=[]), source_code="", library_deps=(), - code_spec=code_specs.CPPCodeSpec(), + code_spec=artifacts.CPPCodeSpec(), ) 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 0a2f1a9c00..fc077b3a90 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 @@ -16,7 +16,7 @@ from gt4py.next.iterator import builtins, ir as itir from gt4py.next.iterator.ir_utils import ir_makers as im from gt4py.next import fingerprinting -from gt4py.next.otf import arguments, code_specs, stages, definitions +from gt4py.next.otf import arguments, artifacts, stages from gt4py.next.program_processors.codegens.gtfn import gtfn_module from gt4py.next.program_processors.runners import gtfn from gt4py.next.type_system import type_translation @@ -76,19 +76,19 @@ def program_example(): def test_codegen(program_example): fencil, parameters = program_example module = gtfn_module.translate_program_cpu( - definitions.CompilableProgramDef( + stages.CompilableProgramDef( data=fencil, args=arguments.CompileTimeArgs.from_concrete(*parameters, **{"offset_provider": {}}), ) ) assert module.entry_point.name == fencil.id assert any(d.name == "gridtools_cpu" for d in module.library_deps) - assert isinstance(module.code_spec, code_specs.CPPCodeSpec) + assert isinstance(module.code_spec, artifacts.CPPCodeSpec) def test_hash_and_diskcache(program_example, tmp_path): fencil, parameters = program_example - compilable_program = definitions.CompilableProgramDef( + compilable_program = stages.CompilableProgramDef( data=fencil, args=arguments.CompileTimeArgs.from_concrete(*parameters, **{"offset_provider": {}}), ) @@ -130,7 +130,7 @@ def test_hash_and_diskcache(program_example, tmp_path): def test_gtfn_file_cache(program_example): fencil, parameters = program_example - compilable_program = definitions.CompilableProgramDef( + compilable_program = stages.CompilableProgramDef( data=fencil, args=arguments.CompileTimeArgs.from_concrete(*parameters, **{"offset_provider": {}}), ) 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 4a6c7c1d7f..488364b360 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 @@ -20,7 +20,7 @@ from gt4py import next as gtx from gt4py._core import definitions as core_defs from gt4py.next import config -from gt4py.next.otf import definitions, runners +from gt4py.next.otf import runners, stages from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations from gt4py.next.program_processors.runners.dace.transformations import ( auto_optimize as gtx_auto_optimize, @@ -320,7 +320,7 @@ def testee(a: cases.IField, b: cases.IField, out: cases.IField): captured_sdfg: dace.SDFG | None = None translation_step = custom_backend.executor.translation.step - def mocked_translator(inp: definitions.CompilableProgramDef) -> dace.SDFG: + def mocked_translator(inp: stages.CompilableProgramDef) -> dace.SDFG: nonlocal captured_sdfg result = translation_step(inp) captured_sdfg = dace.SDFG.from_json(result.source_code) 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 4dbe395d83..e4ec853c23 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 @@ -16,7 +16,7 @@ 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.next.otf import artifacts 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 @@ -227,7 +227,7 @@ def {_bind_func_name}(device, sdfg_argtypes, args, sdfg_call_args, offset_provid def mocked_compile_call( self, - inp: stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], + inp: artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec], binding_source_ref: str, ): assert len(inp.library_deps) == 0 @@ -244,7 +244,7 @@ def mocked_compile_call( def mocked_compile_call_cartesian( self, - inp: stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], + inp: artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec], use_metrics: bool, use_zero_origin: bool, ): @@ -256,7 +256,7 @@ def mocked_compile_call_cartesian( def mocked_compile_call_unstructured( self, - inp: stages.ExtensionSource[code_specs.SDFGCodeSpec, code_specs.PythonCodeSpec], + inp: artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec], use_metrics: bool, use_zero_origin: bool, ): 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 ca7211ee5e..c63463d05c 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 @@ -28,7 +28,7 @@ from gt4py._core import definitions as core_defs from gt4py.next import config, fingerprinting -from gt4py.next.otf import code_specs, stages +from gt4py.next.otf import artifacts from gt4py.next.otf.binding import interface from gt4py.next.otf.compilation import cache as gtx_cache from gt4py.next.program_processors.runners.dace.workflow import compilation as dace_wf_compilation @@ -53,17 +53,17 @@ def _make_compilable_sdfg(name: str) -> dace.SDFG: return sdfg -def _make_input(name: str) -> stages.ExtensionSource: +def _make_input(name: str) -> artifacts.ExtensionSource: sdfg = _make_compilable_sdfg(name) - program_source = stages.ProgramSource( + program_source = artifacts.ProgramSource( entry_point=interface.Function(name=sdfg.name, parameters=()), source_code=sdfg.to_json(), library_deps=(), - code_spec=code_specs.SDFGCodeSpec(), + code_spec=artifacts.SDFGCodeSpec(), ) - return stages.ExtensionSource( + return artifacts.ExtensionSource( program_source=program_source, - binding_source=stages.BindingSource( + binding_source=artifacts.BindingSource( source_code="def bind(*args, **kwargs):\n return None\n", library_deps=(), ), @@ -80,7 +80,7 @@ def _compiler() -> dace_wf_compilation.DaCeCompiler: def _build_folder( - comp: dace_wf_compilation.DaCeCompiler, inp: stages.ExtensionSource + comp: dace_wf_compilation.DaCeCompiler, inp: artifacts.ExtensionSource ) -> pathlib.Path: return gtx_cache.get_cache_folder( inp, @@ -92,7 +92,7 @@ def _build_folder( @pytest.fixture def clean_build_folder(request): def factory( - comp: dace_wf_compilation.DaCeCompiler, inp: stages.ExtensionSource + comp: dace_wf_compilation.DaCeCompiler, inp: artifacts.ExtensionSource ) -> pathlib.Path: folder = _build_folder(comp, inp) shutil.rmtree(folder, ignore_errors=True) 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 2204aaa1d9..35478779e4 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 @@ -25,7 +25,7 @@ from gt4py._core import definitions as core_defs from gt4py.next import config -from gt4py.next.otf import code_specs, stages +from gt4py.next.otf import artifacts from gt4py.next.otf.binding import interface from gt4py.next.program_processors.runners.dace.workflow import ( compilation as dace_wf_compilation, @@ -88,18 +88,18 @@ def program_source() -> dace_wf_compilation.SDFGExtensionSource: Using a real source (rather than a `MagicMock`) lets the unmocked `get_cache_folder` fingerprint the program source for the build-folder name. """ - program_source = stages.ProgramSource( + program_source = artifacts.ProgramSource( entry_point=interface.Function("gpu_program", parameters=()), source_code=_make_sdfg_with_gpu_map().to_json(), library_deps=(), - code_spec=code_specs.SDFGCodeSpec(), + code_spec=artifacts.SDFGCodeSpec(), ) - binding_source = stages.BindingSource(source_code="", library_deps=()) - return stages.ExtensionSource(program_source=program_source, binding_source=binding_source) + binding_source = artifacts.BindingSource(source_code="", library_deps=()) + return artifacts.ExtensionSource(program_source=program_source, binding_source=binding_source) def _run_compiler( - inp: stages.ExtensionSource, + inp: artifacts.ExtensionSource, *, add_gpu_trace_markers: bool = False, cmake_build_type: config.CMakeBuildType = config.CMakeBuildType.RELEASE, diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py index 2c6bf03239..87e212afc9 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py @@ -20,7 +20,7 @@ from gt4py.next import common as gtx_common, fingerprinting from gt4py.next.iterator import ir as itir from gt4py.next.iterator.ir_utils import ir_makers as im -from gt4py.next.otf import arguments as otf_arguments, toolchain as otf_toolchain +from gt4py.next.otf import arguments as otf_arguments, workflow as otf_workflow from gt4py.next.program_processors.runners.dace import lowering as gtx_dace_lowering from gt4py.next.program_processors.runners.dace.workflow import ( translation as dace_wf_translation, @@ -425,7 +425,7 @@ def test_generate_sdfg_async_call_multi_state( assert not _are_streams_synchronized(sdfg) -def _make_simple_field_operator_compilable_program() -> otf_toolchain.ConcreteArtifact: +def _make_simple_field_operator_compilable_program() -> otf_workflow.ConcreteArtifact: """Return a compilable program wrapping a minimal GTIR field operator.""" ir = itir.Program( id="simple_field_operator", @@ -443,7 +443,7 @@ def _make_simple_field_operator_compilable_program() -> otf_toolchain.ConcreteAr ), ], ) - return otf_toolchain.ConcreteArtifact( + return otf_workflow.ConcreteArtifact( data=ir, args=otf_arguments.CompileTimeArgs( args=tuple(param.type for param in ir.params),