diff --git a/docs/development/ADRs/next/0011-On_The_Fly_Compilation.md b/docs/development/ADRs/next/0011-On_The_Fly_Compilation.md index 697fdf7253..ed43f422d0 100644 --- a/docs/development/ADRs/next/0011-On_The_Fly_Compilation.md +++ b/docs/development/ADRs/next/0011-On_The_Fly_Compilation.md @@ -11,6 +11,8 @@ tags: [backend, bindings, build, compile, otf] This supersedes [0009 - Compiled Backend Integration](0009-Compiled_Backend_Integration.md) and concentrates on the API design for on-the-fly compilation of GT4Py programs and all the steps in between IR and compiled Python extension. +Partially superseded by [0029 - Toolchain Naming and Pipeline Simplification](0029-Toolchain-Naming-and-Pipeline-Simplification.md) (workflow combinators and step-type naming). + ## Context The on-the-fly compilation (OTFC) in gt4py encompasses everything necessary to go from an IR representation of a GT4Py program to an executable Python function. Depending on the chosen route, this may include: diff --git a/docs/development/ADRs/next/0017-Toolchain-Configuration.md b/docs/development/ADRs/next/0017-Toolchain-Configuration.md index 4fc17fcb4e..367f69ad88 100644 --- a/docs/development/ADRs/next/0017-Toolchain-Configuration.md +++ b/docs/development/ADRs/next/0017-Toolchain-Configuration.md @@ -11,6 +11,8 @@ tags: [backend, otf, workflows, toolchain] In order to provide a streamlined user experience, we attempt to standardize how users of GT4Py stencils can configure how those stencils are optimized without editing GT4Py code. This describes the design of the first minimal implementation. +Refined by [0029 - Toolchain Naming and Pipeline Simplification](0029-Toolchain-Naming-and-Pipeline-Simplification.md) (the term *toolchain* becomes the name of the root object). + ## Context In this document the word toolchain is used to mean all the code components that work together to go from DSL code to an optimized, runnable python callable. diff --git a/docs/development/ADRs/next/0027-External_Workspace_Memory.md b/docs/development/ADRs/next/0027-External_Workspace_Memory.md index c2d90458ed..7807acb7bf 100644 --- a/docs/development/ADRs/next/0027-External_Workspace_Memory.md +++ b/docs/development/ADRs/next/0027-External_Workspace_Memory.md @@ -9,6 +9,8 @@ tags: [] - **Created**: 2026-07-27 - **Updated**: 2026-08-03 +Partially superseded by [0029 - Toolchain Naming and Pipeline Simplification](0029-Toolchain-Naming-and-Pipeline-Simplification.md) (the workspace is carried by an explicit `Toolchain.loading` step instead of a `DaCeBackend` subclass overriding `Backend.load_artifact`; the decision to inject the workspace at load time is unchanged). + In the context of `gt4py.next` DaCe backends that run the same compiled SDFG many times (e.g. inside a time loop), facing per-call GPU transient allocation overhead and the desire to bound workspace memory to one SDFG's transient diff --git a/docs/development/ADRs/next/0029-Toolchain-Naming-and-Pipeline-Simplification.md b/docs/development/ADRs/next/0029-Toolchain-Naming-and-Pipeline-Simplification.md new file mode 100644 index 0000000000..937d94e295 --- /dev/null +++ b/docs/development/ADRs/next/0029-Toolchain-Naming-and-Pipeline-Simplification.md @@ -0,0 +1,176 @@ +--- +tags: [backend, otf, toolchain, workflows, naming, observability] +--- + +# Toolchain Naming and Pipeline Simplification + +- **Status**: valid +- **Authors**: Enrique González Paredes (@egparedes) +- **Created**: 2026-07-30 +- **Updated**: 2026-07-30 + +In the context of the on-the-fly compilation toolchain, facing a root object +whose names (`Backend`, `transforms`, `executor`) no longer match the +established *toolchain* vocabulary and a workflow-combinator framework whose +reading cost exceeds its value, we decided to rename the root object to +`Toolchain` (with `frontend` / `backend` halves), rename the envelope and +stage types accordingly, replace the combinator tower with explicit, fully +typed pipelines, and add a sanctioned stage-observability seam. We considered +keeping the combinators and trimming only the dead ones, and accept that the +renames break the old names outright, rotate the persistent translation-cache +keys (the second such rotation in this stack), and impose a migration burden +on downstream code that subclassed the combinators. + +This partially supersedes [0011 - On The Fly Compilation](0011-On_The_Fly_Compilation.md) +(the workflow-combinator framework and the `otf.step_types` naming — its stage +vocabulary and build-system decisions remain valid) and refines +[0017 - Toolchain Configuration](0017-Toolchain-Configuration.md) (whose term +*toolchain* becomes the name of the root object; its configuration decisions +remain valid). + +This builds on [0028 - Plain Builders Instead of factory-boy Factories](0028-Plain-Builders-Instead-of-Factories.md), +which replaced the factory classes that constructed these objects, so the +renames below land in plain, type-checked builder code. + +## Context + +ADR 0017 already defines **toolchain** as "all the code components that work +together to go from DSL code to an optimized, runnable python callable" — +JIT/OTF pipelines, transformation passes, lowerings, parsers. The code grew +into exactly that shape while keeping older names: the root object is called +`Backend`, its frontend half `transforms`, its backend half `executor`, and +the pipeline pieces are spread over five vaguely-named modules (`workflow`, +`toolchain`, `definitions`, `stages`, `recipes`). + +Separately, the workflow-combinator framework introduced by ADR 0011 had grown +to thirteen abstractions (`Workflow`, two mixins, `NamedStepSequence`, +`MultiWorkflow`, `StepSequence`, `CachedStep`, `SkippableStep`, `make_step`, +the `ConcreteArtifact` envelope and three adapters) to express function +composition (some of them already deleted as dead code in the preparatory +cleanup PRs). Measured against actual use, only `CachedStep` is a deep module; +the others are shallow wrappers around `Callable[[S], T]`, and the +`NamedStepSequence.__call__` reflection loop is `Any`-typed, defeating the +static typing ADR 0011 prized. + +Finally, the seam between the toolchain and its consumers has no sanctioned +interface for observing or running intermediate stages: the one consumer that +needs an intermediate stage (the dace `__sdfg__` path) duck-types into the +default pipeline and mutates its stages. + +## Decision + +### Naming + +| Old | New | +| ---------------------------- | --------------------------------- | +| `Backend` | `Toolchain` | +| `Backend.transforms` | `Toolchain.frontend` | +| `Backend.executor` | `Toolchain.backend` | +| `Backend.load_artifact()` | `Toolchain.loading` | +| `OTFCompileWorkflow` | `CompilePipeline` | +| `ConcreteArtifact` / `.data` | `ProgramWithArgs` / `.definition` | +| `CompilableProgramDef` | `CompilableProgram` | + +These are hard renames: no deprecation aliases are kept, so the old names stop +resolving in the same release. The public `gtx.*` names (`gtx.gtfn_cpu`, +`gtx.wait_for_compilation`, ...) are unchanged, with one exception — +`gtx.typing.Backend` becomes `gtx.typing.Toolchain`. + +We considered shipping one-release read aliases for the renamed attributes and +re-export shims for the moved modules. We rejected that: an alias that silently +keeps working is exactly what lets a downstream stay on the old vocabulary past +the release where it was supposed to migrate, and the aliases would have been +partial anyway (constructing with the old keyword names could not be aliased +without hand-writing `__init__`, so `Backend(executor=...)` would have raised +`TypeError` regardless). A single loud break at a known release is easier to +act on than a half-working compatibility layer. `CompilableProgram` stays a type +alias (of the parameterized envelope) for now. + +The last row is not a pure rename. `Backend.load_artifact()` was an overridable +method, and a toolchain that needed backend-specific runtime data in the loaded +program subclassed `Toolchain` to override it (see +[0027 - External Workspace Memory for DaCe Transients](0027-External_Workspace_Memory.md)). +`Toolchain.loading` makes that seam a field like the other three — a plain step +defaulting to `artifacts.load_artifact` — so customization stays +composition-time, consistent with the rest of this ADR. It also keeps +subclassing off the load path: per +[0024 - Compilation Runners](0024-Compilation-Runners.md), a toolchain that +customizes `compile` is opaque to the process-pool runner, and a +customization mechanism built on subclassing invites exactly that. The step +lives on `Toolchain` rather than inside `CompilePipeline` because loading also +happens off the pipeline — `CompiledProgramsPool` loads artifacts straight from +a compilation future — and because monolithic backends have no `CompilePipeline` +to hang it on. + +### Pipeline, not combinators + +Named pipelines become frozen dataclasses with an explicit, fully typed +`__call__`; steps are plain callables (`Step[S, T] = Callable[[S], T]`); +customization stays composition-time via `dataclasses.replace`. The +combinator framework (both mixins, `NamedStepSequence`, `MultiWorkflow`, +`StepSequence`, `make_step`, `.chain`, the three adapters) is deleted; +`CachedStep` is kept unchanged. + +What ADR 0011's decisions become: + +| ADR 0011 requirement | Kept by | +| ------------------------------------------- | ----------------------------------------------------------------------- | +| Named steps, order visible | dataclass fields + explicit `__call__` (order literally visible) | +| Statically typed composition | explicit `__call__` — checked end-to-end, unlike the reflection loop | +| Customization at composition, not via flags | `dataclasses.replace` on frozen pipelines | +| Steps compose across backends | `Step[S, T]` is a `Callable` — every existing step already satisfies it | +| Linear workflows | unchanged (`Transforms` keeps its input-dependent step *selection*) | + +### Stage observability + +A `stage_hook(name, artifact)` event hook (on the existing +`instrumentation.hook_machinery`) is emitted by `CompilePipeline` and +`Transforms` after each step; artifacts are treated opaquely. +`GT4PY_DUMP_STAGES=` registers a subscriber that writes each stage +artifact per program. `Toolchain.translate(definition, compile_time_args)` is +the sanctioned partial run (frontend + translation only); it narrows to the +standard `CompilePipeline` shape and raises a clear error on monolithic +backends. Per-call step options are deliberately not offered — a caller +needing a variant step builds a variant pipeline with `dataclasses.replace`. + +### Phasing + +The decisions in this ADR land as a stack of PRs. The naming decisions above +land first, except the `OTFCompileWorkflow` → `CompilePipeline` rename, which +lands with the pipeline PR; the "Pipeline, not combinators" and "Stage +observability" decisions are implemented in follow-up PRs of the same stack. + +## Consequences + +- The dace `__sdfg__` reach-in (duck-typing through `executor.translation` + and mutating frozen stages) will be replaced by a dace-owned translate-only + step built with `dataclasses.replace`. +- Downstream code subclassing the deleted combinators or overriding + `Transforms.step_order` must migrate to explicit `__call__` composition. +- Downstream code must migrate to the new names in the same release: the old + ones are removed, not aliased. That includes the modules dissolved along the + way — `otf.definitions`, `otf.code_specs`, `otf.recipes` and `otf.toolchain` + cease to exist, and `gtx.typing.Backend` becomes `gtx.typing.Toolchain`. +- ADR 0011's `otf.step_types` naming section no longer describes the code. +- Cache keys rotate: fingerprints embed the qualified class and field names, + so the `ProgramWithArgs` / `definition` rename changes the keys of the + persistent translation caches (gtfn and dace). This is the *second* of two + rotations in this stack — moving the same class from `otf.toolchain` to + `otf.workflow` already rotated them once, since the qualified name alone + is enough. What makes each a cold rebuild rather than a stale hit is that + the key itself changes, so entries written under the old name are never + looked up; `BUILD_CACHE_VERSION_ID` (the gt4py version) is a separate salt + on top, which means the rotation costs nothing extra to anyone who crosses + a release — or, for a source install, any commit — boundary anyway. The two + collapse into a single cold rebuild if the whole stack lands in one + release, which is the concrete argument for not spreading it across + releases. + +## Alternatives considered + +- **Keep the combinators but trim dead ones**: rejected — the tower is a + shallow-module archetype and every recorded ADR 0011 requirement survives + its removal; trimming keeps the reading cost and the `Any`-typed + composition. +- **A per-call options API on `translate()`**: rejected — reconfiguration + stays composition-time, ADR 0011's own rule. diff --git a/docs/development/ADRs/next/README.md b/docs/development/ADRs/next/README.md index f19167ef9b..29fe79fac5 100644 --- a/docs/development/ADRs/next/README.md +++ b/docs/development/ADRs/next/README.md @@ -53,6 +53,7 @@ Writing a new ADR is simple: - [0017 - Toolchain Configuration](0017-Toolchain-Configuration.md) - [0027 - External Workspace Memory for DaCe Transients](0027-External_Workspace_Memory.md) - [0028 - Plain Builders Instead of factory-boy Factories](0028-Plain-Builders-Instead-of-Factories.md) +- [0029 - Toolchain Naming and Pipeline Simplification](0029-Toolchain-Naming-and-Pipeline-Simplification.md) ### Python Integration diff --git a/docs/user/next/advanced/HackTheToolchain.md b/docs/user/next/advanced/HackTheToolchain.md index 0d5e9ab74a..1113d8e998 100644 --- a/docs/user/next/advanced/HackTheToolchain.md +++ b/docs/user/next/advanced/HackTheToolchain.md @@ -22,8 +22,8 @@ cached_lowering_toolchain = gtx.backend.DEFAULT_TRANSFORMS.replace( ## Skip Steps / Change Order ```python -DUMMY_FOP = workflow.ConcreteArtifact( - data=ff_stages.DSLFieldOperatorDef(definition=None), args=None +DUMMY_FOP = workflow.ProgramWithArgs( + definition=ff_stages.DSLFieldOperatorDef(definition=None), args=None ) ``` diff --git a/docs/user/next/advanced/WorkflowPatterns.md b/docs/user/next/advanced/WorkflowPatterns.md index ba6fd3a104..bd63877c01 100644 --- a/docs/user/next/advanced/WorkflowPatterns.md +++ b/docs/user/next/advanced/WorkflowPatterns.md @@ -399,7 +399,7 @@ gtx.backend.DEFAULT_PROG_TRANSFORMS?? ``` ```python -gtx.program_processors.runners.gtfn.run_gtfn_gpu.executor.otf_workflow?? +gtx.program_processors.runners.gtfn.run_gtfn_gpu.backend?? ``` ```python diff --git a/src/gt4py/next/AGENTS.md b/src/gt4py/next/AGENTS.md index f0a2bfc342..ea94beaf17 100644 --- a/src/gt4py/next/AGENTS.md +++ b/src/gt4py/next/AGENTS.md @@ -31,6 +31,9 @@ field operators / programs (ffront) Distinct from `iterator/embedded.py`, which runs one level lower. - `otf/` — on-the-fly compilation toolchain (workflow steps, caching, argument descriptors). +- `backend.py` — the toolchain root object `Toolchain` (formerly `Backend`): + `frontend` (definition transforms) + `backend` (compile pipeline) + + allocator. - `program_processors/runners/` — the backends: `gtfn` (GridTools C++), `dace` (DaCe SDFG), `roundtrip` / `double_roundtrip` (pure Python). - `type_system/` — `next` type specifications and the type inference the diff --git a/src/gt4py/next/backend.py b/src/gt4py/next/backend.py index eae12981f0..1b67cb02b3 100644 --- a/src/gt4py/next/backend.py +++ b/src/gt4py/next/backend.py @@ -9,7 +9,7 @@ from __future__ import annotations import dataclasses -from typing import Generic +from typing import Callable, Generic from gt4py._core import definitions as core_defs from gt4py.next import custom_layout_allocators as next_allocators @@ -45,7 +45,7 @@ def adapted_jit_to_aot_args_factory() -> workflow.Workflow[ class Transforms( workflow.MultiWorkflow[ stages.ConcreteProgramDef[stages.IRDefinitionT, stages.ArgsDefinitionT], - stages.CompilableProgramDef, + stages.CompilableProgram, ] ): """ @@ -92,14 +92,14 @@ class Transforms( ] = dataclasses.field(default_factory=past_process_args.transform_program_args_factory) past_to_itir: workflow.Workflow[ - ffront_stages.ConcretePASTProgramDef, stages.CompilableProgramDef + ffront_stages.ConcretePASTProgramDef, stages.CompilableProgram ] = dataclasses.field(default_factory=past_to_itir.past_to_gtir_factory) def step_order(self, inp: stages.ConcreteProgramDef) -> list[str]: steps: list[str] = [] if isinstance(inp.args, arguments.JITArgs): steps.append("aotify_args") - match inp.data: + match inp.definition: case ffront_stages.DSLFieldOperatorDef(): steps.extend( [ @@ -140,32 +140,36 @@ def step_order(self, inp: stages.ConcreteProgramDef) -> list[str]: DEFAULT_TRANSFORMS: Transforms = Transforms() -# TODO(tehrengruber): Rename class and `executor` & `transforms` attribute. Maybe: -# `Backend` -> `Toolchain` -# `transforms` -> `frontend_transforms` -# `executor` -> `backend_transforms` @dataclasses.dataclass(frozen=True) -class Backend(Generic[core_defs.DeviceTypeT]): +class Toolchain(Generic[core_defs.DeviceTypeT]): + """ + Complete pipeline from a program definition to an executable program. + + The `frontend` workflow transforms any supported program definition into a + `CompilableProgram`, which the `backend` workflow then compiles into a + loadable compilation artifact. The `loading` step turns that artifact into a + directly-callable program; toolchains needing to inject backend-specific + runtime data supply their own step here. The `allocator` describes the + device the compiled program expects its buffers on. + """ + name: str - executor: workflow.Workflow[stages.CompilableProgramDef, artifacts.CompilationArtifact] + backend: workflow.Workflow[stages.CompilableProgram, artifacts.CompilationArtifact] allocator: next_allocators.FieldBufferAllocatorProtocol[core_defs.DeviceTypeT] - transforms: workflow.Workflow[stages.ConcreteProgramDef, stages.CompilableProgramDef] + frontend: workflow.Workflow[stages.ConcreteProgramDef, stages.CompilableProgram] + # A plain `Callable`, not `workflow.Workflow`: the protocol names its + # parameter `inp`, which would force every loading step to use that name. + loading: Callable[[artifacts.CompilationArtifact], artifacts.ExecutableProgram] = ( + dataclasses.field(default=artifacts.load_artifact) + ) def compile( self, program: stages.IRDefinitionT, compile_time_args: arguments.CompileTimeArgs ) -> artifacts.ExecutableProgram: - artifact = self.executor( - self.transforms(stages.ConcreteProgramDef(data=program, args=compile_time_args)) + artifact = self.backend( + self.frontend(stages.ConcreteProgramDef(definition=program, args=compile_time_args)) ) - return self.load_artifact(artifact) - - 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 - into the loaded program. - """ - return artifact.load() + return self.loading(artifact) @property def __gt_allocator__( diff --git a/src/gt4py/next/config.py b/src/gt4py/next/config.py index 56eec687a0..e62e2a1892 100644 --- a/src/gt4py/next/config.py +++ b/src/gt4py/next/config.py @@ -48,10 +48,10 @@ class BuildJobsMode(enum.Enum): #: Run compilation in a ``ThreadPoolExecutor``. THREAD = "thread" #: Run compilation in a ``ProcessPoolExecutor`` with the ``spawn`` start - #: method. Requires the backend's ``executor`` to be stdlib-picklable + #: method. Requires the toolchain's ``backend`` pipeline to be stdlib-picklable #: (standard runners and factory-constructed variants are) and to return a #: picklable ``CompilationArtifact``; backends that don't qualify (or that - #: customize ``Backend.compile``) are compiled in the calling thread + #: customize ``Toolchain.compile``) are compiled in the calling thread #: instead, with a warning. PROCESS = "process" diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 0a6db1aabf..1e5f875b4c 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -51,7 +51,7 @@ from gt4py.next.type_system import type_info, type_specifications as ts, type_translation -DEFAULT_BACKEND: next_backend.Backend | None = None +DEFAULT_BACKEND: next_backend.Toolchain | None = None ProgramCallMetricsCollector = metrics.make_collector( @@ -97,7 +97,7 @@ class _CompilableGTEntryPointMixin(Generic[ffront_stages.DSLDefinitionT]): # backend-specific compilation is keyed separately in the backend's own # caches, and fingerprinting the whole backend object graph is both wasteful # and fragile (it may hold non-importable callables, see also test doubles). - backend: Optional[next_backend.Backend] = dataclasses.field( + backend: Optional[next_backend.Toolchain] = dataclasses.field( metadata=utils.gt4py_metadata(fingerprint=False) ) compilation_options: options.CompilationOptions @@ -105,7 +105,7 @@ class _CompilableGTEntryPointMixin(Generic[ffront_stages.DSLDefinitionT]): @abc.abstractmethod def __gt_type__(self) -> ts.CallableType: ... - def with_backend(self, backend: next_backend.Backend | None) -> Self: + def with_backend(self, backend: next_backend.Toolchain | None) -> Self: return dataclasses.replace(self, backend=backend) def with_compilation_options( @@ -247,7 +247,7 @@ class Program(_CompilableGTEntryPointMixin[ffront_stages.DSLProgramDef]): def from_function( cls, definition: types.FunctionType, - backend: next_backend.Backend | None, + backend: next_backend.Toolchain | None, grid_type: common.GridType | None = None, **compilation_options: Unpack[options.CompilationOptionsArgs], ) -> Program: @@ -264,8 +264,8 @@ 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 = workflow.ConcreteArtifact(self.past_stage, arguments.CompileTimeArgs.empty()) - _ = self._frontend_transforms.past_lint(no_args_past).data + no_args_past = workflow.ProgramWithArgs(self.past_stage, arguments.CompileTimeArgs.empty()) + _ = self._frontend_transforms.past_lint(no_args_past).definition @property def __name__(self) -> str: @@ -287,19 +287,19 @@ 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 = workflow.ConcreteArtifact( + no_args_def = workflow.ProgramWithArgs( self.definition_stage, arguments.CompileTimeArgs.empty() ) - return self._frontend_transforms.func_to_past(no_args_def).data + return self._frontend_transforms.func_to_past(no_args_def).definition @property def _frontend_transforms(self) -> next_backend.Transforms: if self.backend is None: return next_backend.DEFAULT_TRANSFORMS - # TODO(tehrengruber): This class relies heavily on `self.backend.transforms` being + # TODO(tehrengruber): This class relies heavily on `self.backend.frontend` being # a `next_backend.Transforms`, but the backend type annotation does not reflect that. - assert isinstance(self.backend.transforms, next_backend.Transforms) - return self.backend.transforms + assert isinstance(self.backend.frontend, next_backend.Transforms) + return self.backend.frontend @functools.cached_property def _all_closure_vars(self) -> dict[str, Any]: @@ -307,15 +307,15 @@ def _all_closure_vars(self) -> dict[str, Any]: @functools.cached_property def gtir(self) -> itir.Program: - no_args_past = workflow.ConcreteArtifact( - data=ffront_stages.PASTProgramDef( + no_args_past = workflow.ProgramWithArgs( + definition=ffront_stages.PASTProgramDef( past_node=self.past_stage.past_node, closure_vars=self.past_stage.closure_vars, grid_type=self.definition_stage.grid_type, ), args=arguments.CompileTimeArgs.empty(), ) - return self._frontend_transforms.past_to_itir(no_args_past).data + return self._frontend_transforms.past_to_itir(no_args_past).definition def with_grid_type(self, grid_type: common.GridType) -> Program: return dataclasses.replace( @@ -504,7 +504,7 @@ def program(definition: Callable) -> Program: ... @typing.overload def program( *, - backend: next_backend.Backend | eve.NothingType | None = eve.NOTHING, + backend: next_backend.Toolchain | eve.NothingType | None = eve.NOTHING, grid_type: common.GridType | None = None, **compilation_options: Unpack[options.CompilationOptionsArgs], ) -> Callable[[Callable], Program]: ... @@ -514,7 +514,7 @@ def program( definition: Callable | None = None, *, # `NOTHING` -> default backend, `None` -> no backend (embedded execution) - backend: next_backend.Backend | eve.NothingType | None = eve.NOTHING, + backend: next_backend.Toolchain | eve.NothingType | None = eve.NOTHING, grid_type: common.GridType | None = None, **compilation_options: Unpack[options.CompilationOptionsArgs], ) -> Program | Callable[[Callable], Program]: @@ -545,7 +545,8 @@ def program_inner(definition: Callable) -> Program: program = Program.from_function( definition, backend=typing.cast( - next_backend.Backend | None, DEFAULT_BACKEND if backend is eve.NOTHING else backend + next_backend.Toolchain | None, + DEFAULT_BACKEND if backend is eve.NOTHING else backend, ), grid_type=grid_type, **compilation_options, @@ -581,7 +582,7 @@ class FieldOperator(_CompilableGTEntryPointMixin[ffront_stages.DSLFieldOperatorD def from_function( cls, definition: types.FunctionType, - backend: Optional[next_backend.Backend], + backend: Optional[next_backend.Toolchain], grid_type: Optional[common.GridType] = None, *, operator_node_cls: type[foast.OperatorNode] = foast.FieldOperator, @@ -607,10 +608,10 @@ def __post_init__(self) -> None: @functools.cached_property def foast_stage(self) -> ffront_stages.FOASTOperatorDef: return self._frontend_transforms.func_to_foast( - workflow.ConcreteArtifact( - data=self.definition_stage, args=arguments.CompileTimeArgs.empty() + workflow.ProgramWithArgs( + definition=self.definition_stage, args=arguments.CompileTimeArgs.empty() ) - ).data + ).definition @property def __name__(self) -> str: @@ -624,10 +625,10 @@ def definition(self) -> types.FunctionType: def _frontend_transforms(self) -> next_backend.Transforms: if self.backend is None: return next_backend.DEFAULT_TRANSFORMS - # TODO(tehrengruber): This class relies heavily on `self.backend.transforms` being + # TODO(tehrengruber): This class relies heavily on `self.backend.frontend` being # a `next_backend.Transforms`, but the backend type annotation does not reflect that. - assert isinstance(self.backend.transforms, next_backend.Transforms) - return self.backend.transforms + assert isinstance(self.backend.frontend, next_backend.Transforms) + return self.backend.frontend def __gt_type__(self) -> ts.CallableType: type_ = self.foast_stage.foast_node.type @@ -731,7 +732,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: def field_operator( definition: Callable, *, - backend: next_backend.Backend | eve.NothingType | None = eve.NOTHING, + backend: next_backend.Toolchain | eve.NothingType | None = eve.NOTHING, grid_type: common.GridType | None = None, ) -> FieldOperator: ... @@ -739,7 +740,7 @@ def field_operator( @typing.overload def field_operator( *, - backend: next_backend.Backend | eve.NothingType | None = eve.NOTHING, + backend: next_backend.Toolchain | eve.NothingType | None = eve.NOTHING, grid_type: common.GridType | None = None, ) -> Callable[[Callable], FieldOperator]: ... @@ -747,7 +748,7 @@ def field_operator( def field_operator( definition: Callable | None = None, *, - backend: next_backend.Backend | eve.NothingType | None = eve.NOTHING, + backend: next_backend.Toolchain | eve.NothingType | None = eve.NOTHING, grid_type: common.GridType | None = None, **compilation_options: Unpack[options.CompilationOptionsArgs], ) -> FieldOperator | Callable[[types.FunctionType], FieldOperator]: @@ -772,7 +773,8 @@ def field_operator_inner(definition: Callable) -> FieldOperator: return FieldOperator.from_function( definition, typing.cast( - next_backend.Backend | None, DEFAULT_BACKEND if backend is eve.NOTHING else backend + next_backend.Toolchain | None, + DEFAULT_BACKEND if backend is eve.NOTHING else backend, ), grid_type, **compilation_options, @@ -788,7 +790,7 @@ def scan_operator( axis: common.Dimension, forward: bool = True, init: core_defs.Scalar = 0.0, - backend: next_backend.Backend | eve.NothingType | None, + backend: next_backend.Toolchain | eve.NothingType | None, grid_type: common.GridType | None, ) -> FieldOperator: ... @@ -815,7 +817,7 @@ def scan_operator( axis: common.Dimension, forward: bool = True, init: core_defs.Scalar = 0.0, - backend: next_backend.Backend | eve.NothingType | None, + backend: next_backend.Toolchain | eve.NothingType | None, grid_type: common.GridType | None, ) -> Callable[[Callable], FieldOperator]: ... @@ -826,7 +828,7 @@ def scan_operator( axis: common.Dimension, forward: bool = True, init: core_defs.Scalar = 0.0, - backend: next_backend.Backend | eve.NothingType | None = eve.NOTHING, + backend: next_backend.Toolchain | eve.NothingType | None = eve.NOTHING, grid_type: common.GridType | None = None, ) -> FieldOperator | Callable[[Callable], FieldOperator]: """ @@ -863,7 +865,8 @@ def scan_operator_inner(definition: Callable) -> FieldOperator: return FieldOperator.from_function( definition, typing.cast( - next_backend.Backend | None, DEFAULT_BACKEND if backend is eve.NOTHING else backend + next_backend.Toolchain | None, + DEFAULT_BACKEND if backend is eve.NOTHING else backend, ), grid_type, operator_node_cls=foast.ScanOperator, diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index beabe53c71..488aa2a253 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -33,23 +33,23 @@ class ItirShim: lowering has access to the relevant information. """ - definition: ConcreteFOASTOperatorDef + operator_def: ConcreteFOASTOperatorDef foast_to_itir: workflow.Workflow[ConcreteFOASTOperatorDef, itir.FunctionDefinition] def __gt_closure_vars__(self) -> Optional[dict[str, Any]]: - return self.definition.data.closure_vars + return self.operator_def.definition.closure_vars def __gt_type__(self) -> ts.CallableType: - assert isinstance(self.definition.data.foast_node.type, ts.CallableType) - return self.definition.data.foast_node.type + assert isinstance(self.operator_def.definition.foast_node.type, ts.CallableType) + return self.operator_def.definition.foast_node.type def __gt_itir__(self) -> itir.FunctionDefinition: - return self.foast_to_itir(self.definition) + return self.foast_to_itir(self.operator_def) # FIXME[#1582](tehrengruber): remove after refactoring to GTIR def __gt_gtir__(self) -> itir.FunctionDefinition: # backend should have self.foast_to_itir set to foast_to_gtir - return self.foast_to_itir(self.definition) + return self.foast_to_itir(self.operator_def) @dataclasses.dataclass(frozen=True) @@ -82,14 +82,15 @@ class OperatorToProgram(workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePAST ... argument_descriptor_contexts={}, ... ) - >>> copy_program = op_to_prog( - ... workflow.ConcreteArtifact(copy.foast_stage, compile_time_args) - ... ) + >>> copy_program = op_to_prog(workflow.ProgramWithArgs(copy.foast_stage, compile_time_args)) - >>> print(copy_program.data.past_node.id) + >>> print(copy_program.definition.past_node.id) __field_operator_copy - >>> assert copy_program.data.closure_vars["copy"].definition.data is copy.foast_stage + >>> assert ( + ... copy_program.definition.closure_vars["copy"].operator_def.definition + ... is copy.foast_stage + ... ) """ foast_to_itir: workflow.Workflow[ConcreteFOASTOperatorDef, itir.FunctionDefinition] @@ -102,10 +103,12 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: arg_types, kwarg_types = inp.args.args, inp.args.kwargs assert not kwarg_types - type_ = inp.data.foast_node.type - loc = inp.data.foast_node.location - assert isinstance(inp.data.foast_node.type, ts.CallableType) - partial_program_type = ffront_type_info.type_in_program_context(inp.data.foast_node.type) + type_ = inp.definition.foast_node.type + loc = inp.definition.foast_node.location + assert isinstance(inp.definition.foast_node.type, ts.CallableType) + partial_program_type = ffront_type_info.type_in_program_context( + inp.definition.foast_node.type + ) assert isinstance(partial_program_type, ts_ffront.ProgramType) args_names = [ *partial_program_type.definition.pos_only_args, @@ -134,27 +137,29 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: params_ref = [past.Name(id=pdecl.id, location=loc) for pdecl in params_decl[:-1]] out_ref = past.Name(id="out", location=loc) - if inp.data.foast_node.id in inp.data.closure_vars: + if inp.definition.foast_node.id in inp.definition.closure_vars: raise RuntimeError("A closure variable has the same name as the field operator itself.") closure_symbols: list[past.Symbol] = [ past.Symbol( - id=inp.data.foast_node.id, + id=inp.definition.foast_node.id, type=ts.DeferredType(constraint=None), namespace=dialect_ast_enums.Namespace.CLOSURE, location=loc, ), ] - fieldop_itir_closure_vars = {inp.data.foast_node.id: ItirShim(inp, self.foast_to_itir)} + fieldop_itir_closure_vars = { + inp.definition.foast_node.id: ItirShim(inp, self.foast_to_itir) + } untyped_past_node = past.Program( - id=f"__field_operator_{inp.data.foast_node.id}", + id=f"__field_operator_{inp.definition.foast_node.id}", type=ts.DeferredType(constraint=ts_ffront.ProgramType), params=params_decl, body=[ past.Call( - func=past.Name(id=inp.data.foast_node.id, location=loc), + func=past.Name(id=inp.definition.foast_node.id, location=loc), args=params_ref, kwargs={"out": out_ref}, location=loc, @@ -169,11 +174,11 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: ) past_node = type_deduction.ProgramTypeDeduction.apply(untyped_past_node) - return workflow.ConcreteArtifact( - data=ffront_stages.PASTProgramDef( + return workflow.ProgramWithArgs( + definition=ffront_stages.PASTProgramDef( past_node=past_node, closure_vars=fieldop_itir_closure_vars, # type: ignore[arg-type] - grid_type=inp.data.grid_type, + grid_type=inp.definition.grid_type, ), args=inp.args, ) diff --git a/src/gt4py/next/ffront/past_process_args.py b/src/gt4py/next/ffront/past_process_args.py index 859776ccc7..ced8cf7dd8 100644 --- a/src/gt4py/next/ffront/past_process_args.py +++ b/src/gt4py/next/ffront/past_process_args.py @@ -22,10 +22,10 @@ def transform_program_args( inp: ffront_stages.ConcretePASTProgramDef, ) -> ffront_stages.ConcretePASTProgramDef: rewritten_args, rewritten_kwargs = _process_args( - past_node=inp.data.past_node, args=inp.args.args, kwargs=inp.args.kwargs + past_node=inp.definition.past_node, args=inp.args.args, kwargs=inp.args.kwargs ) - return workflow.ConcreteArtifact( - data=inp.data, + return workflow.ProgramWithArgs( + definition=inp.definition, args=arguments.CompileTimeArgs( args=rewritten_args, kwargs=rewritten_kwargs, diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index cd908fca14..5182918488 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -35,7 +35,7 @@ # 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) -> stages.CompilableProgramDef: +def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgram: """ Lower a PAST program definition to Iterator IR. @@ -63,21 +63,21 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgramDef: ... ) >>> itir_copy = past_to_gtir( - ... workflow.ConcreteArtifact(copy_program.past_stage, compile_time_args) + ... workflow.ProgramWithArgs(copy_program.past_stage, compile_time_args) ... ) - >>> print(itir_copy.data.id) + >>> print(itir_copy.definition.id) copy_program - >>> print(type(itir_copy.data)) + >>> print(type(itir_copy.definition)) """ - all_closure_vars = transform_utils._get_closure_vars_recursively(inp.data.closure_vars) + all_closure_vars = transform_utils._get_closure_vars_recursively(inp.definition.closure_vars) offsets_and_dimensions = transform_utils._filter_closure_vars_by_type( all_closure_vars, fbuiltins.FieldOffset, common.Dimension ) grid_type = transform_utils._deduce_grid_type( - inp.data.grid_type, offsets_and_dimensions.values() + inp.definition.grid_type, offsets_and_dimensions.values() ) gt_callables = transform_utils._filter_closure_vars_by_type( @@ -94,7 +94,7 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgramDef: lowered_funcs.append(gt_callable.__gt_gtir__()) itir_program = ProgramLowering.apply( - inp.data.past_node, function_definitions=lowered_funcs, grid_type=grid_type + inp.definition.past_node, function_definitions=lowered_funcs, grid_type=grid_type ) # TODO(tehrengruber): Put this in a dedicated transformation step. @@ -141,15 +141,15 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgramDef: inp.args, args=args, kwargs=kwargs, column_axis=_column_axis(all_closure_vars) ) - if config.DEBUG or inp.data.debug: + if config.DEBUG or inp.definition.debug: devtools.debug(itir_program) - return stages.CompilableProgramDef(data=itir_program, args=compile_time_args) + return stages.CompilableProgram(definition=itir_program, args=compile_time_args) def past_to_gtir_factory( cached: bool = True, -) -> workflow.Workflow[ConcretePASTProgramDef, stages.CompilableProgramDef]: +) -> workflow.Workflow[ConcretePASTProgramDef, stages.CompilableProgram]: 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 46f59d3551..0e909062a5 100644 --- a/src/gt4py/next/ffront/stages.py +++ b/src/gt4py/next/ffront/stages.py @@ -79,7 +79,7 @@ class DSLFieldOperatorDef(BaseStage): debug: bool = False -ConcreteDSLFieldOperatorDef: typing.TypeAlias = workflow.ConcreteArtifact[ +ConcreteDSLFieldOperatorDef: typing.TypeAlias = workflow.ProgramWithArgs[ DSLFieldOperatorDef, arguments.CompileTimeArgs ] @@ -93,7 +93,7 @@ class FOASTOperatorDef(BaseStage): debug: bool = False -ConcreteFOASTOperatorDef: typing.TypeAlias = workflow.ConcreteArtifact[ +ConcreteFOASTOperatorDef: typing.TypeAlias = workflow.ProgramWithArgs[ FOASTOperatorDef, arguments.CompileTimeArgs ] @@ -105,7 +105,7 @@ class DSLProgramDef(BaseStage): debug: bool = False -ConcreteDSLProgramDef: typing.TypeAlias = workflow.ConcreteArtifact[ +ConcreteDSLProgramDef: typing.TypeAlias = workflow.ProgramWithArgs[ DSLProgramDef, arguments.CompileTimeArgs ] @@ -118,7 +118,7 @@ class PASTProgramDef(BaseStage): debug: bool = False -ConcretePASTProgramDef: typing.TypeAlias = workflow.ConcreteArtifact[ +ConcretePASTProgramDef: typing.TypeAlias = workflow.ProgramWithArgs[ PASTProgramDef, arguments.CompileTimeArgs ] diff --git a/src/gt4py/next/iterator/runtime.py b/src/gt4py/next/iterator/runtime.py index 88a466229f..4907f1d7f4 100644 --- a/src/gt4py/next/iterator/runtime.py +++ b/src/gt4py/next/iterator/runtime.py @@ -74,7 +74,7 @@ def itir(self, *args): def __call__( self, *args, - backend: Optional[next_backend.Backend | program_formatter.ProgramFormatter] = None, + backend: Optional[next_backend.Toolchain | program_formatter.ProgramFormatter] = None, offset_provider=None, column_axis=None, ): @@ -87,8 +87,8 @@ def __call__( # TODO(tehrengruber): remove cirular dependency and place import at the top of the file from gt4py.next import backend as next_backend - if isinstance(backend, next_backend.Backend): - assert isinstance(backend, next_backend.Backend) + if isinstance(backend, next_backend.Toolchain): + assert isinstance(backend, next_backend.Toolchain) compiled_program = backend.compile( itir_node, arguments.CompileTimeArgs.from_concrete( @@ -102,7 +102,7 @@ def __call__( ) else: raise ValueError( - "Backend must be a 'gt4py.next.backend.Backend' or " + "Backend must be a 'gt4py.next.backend.Toolchain' or " "'gt4py.next.program_formatter.ProgramFormatter'." ) else: diff --git a/src/gt4py/next/otf/artifacts.py b/src/gt4py/next/otf/artifacts.py index 204e2e75d0..182b82b575 100644 --- a/src/gt4py/next/otf/artifacts.py +++ b/src/gt4py/next/otf/artifacts.py @@ -211,7 +211,7 @@ class CompilationArtifact(Protocol): arguments. The one current exception is ``RoundtripArtifact`` when it is configured - with a ``dispatch_backend``: that field holds a ``Backend`` reference + with a ``dispatch_backend``: that field holds a ``Toolchain`` reference whose role belongs at the runner / load-time seam, not in the artifact itself. """ @@ -219,6 +219,23 @@ class CompilationArtifact(Protocol): def load(self) -> ExecutableProgram: ... +def load_artifact(artifact: CompilationArtifact) -> ExecutableProgram: + """ + Load a compilation artifact into a directly-callable program. + + This is the default loading step of a `Toolchain`. Toolchains that need to + inject backend-specific runtime data into the loaded program supply their + own step instead of overriding a method. + + Args: + artifact: The artifact produced by the toolchain's compile pipeline. + + Returns: + The loaded, directly-callable program. + """ + return artifact.load() + + def _unique_libs(*args: interface.LibraryDependency) -> tuple[interface.LibraryDependency, ...]: """ Filter out multiple occurrences of the same ``interface.LibraryDependency``. diff --git a/src/gt4py/next/otf/compilation_tasks.py b/src/gt4py/next/otf/compilation_tasks.py index e172cf370f..b17b3f7cb6 100644 --- a/src/gt4py/next/otf/compilation_tasks.py +++ b/src/gt4py/next/otf/compilation_tasks.py @@ -134,15 +134,15 @@ def load(self) -> artifacts.ExecutableProgram: def make_compilation_task( - backend: gtx_backend.Backend, + backend: gtx_backend.Toolchain, definition_stage: Any, compile_time_args: arguments.CompileTimeArgs, ) -> runners.CompilationTask: """Prepare the compilation of `definition_stage` with `backend` as a task for a runner.""" name = getattr(backend, "name", type(backend).__name__) - if getattr(type(backend), "compile", None) is not gtx_backend.Backend.compile: + if getattr(type(backend), "compile", None) is not gtx_backend.Toolchain.compile: # A customized `compile` is opaque: it cannot be decomposed into the - # standard transforms/executor workflow (and yields an already-loaded + # standard frontend/backend workflow (and yields an already-loaded # program instead of an artifact), so the executor closes over # everything and ignores the compilable. return runners.CompilationTask( @@ -156,12 +156,12 @@ def make_compilation_task( ) # Frontend lowering happens here, main-side: decorators rebind the user's # function module attribute, so the raw `types.FunctionType` must not cross - # a process boundary; the lowered `CompilableProgramDef` is pickle-safe. - compilable = backend.transforms( - stages.ConcreteProgramDef(data=definition_stage, args=compile_time_args) + # a process boundary; the lowered `CompilableProgram` is pickle-safe. + compilable = backend.frontend( + stages.ConcreteProgramDef(definition=definition_stage, args=compile_time_args) ) - def construct_compilable(with_refs: bool) -> stages.CompilableProgramDef: + def construct_compilable(with_refs: bool) -> stages.CompilableProgram: if not with_refs or not compilable.args.offset_provider: return compilable # The shipped copy must not carry the connectivity buffers: they may @@ -176,5 +176,5 @@ def construct_compilable(with_refs: bool) -> stages.CompilableProgramDef: ) return runners.CompilationTask( - name=name, construct_compilable=construct_compilable, executor=backend.executor + name=name, construct_compilable=construct_compilable, executor=backend.backend ) diff --git a/src/gt4py/next/otf/compiled_program.py b/src/gt4py/next/otf/compiled_program.py index 0b49f9ec0c..93573ac076 100644 --- a/src/gt4py/next/otf/compiled_program.py +++ b/src/gt4py/next/otf/compiled_program.py @@ -89,7 +89,7 @@ def metrics_source_key(pool: CompiledProgramsPool, key: CompiledProgramsKey) -> def compile_variant_hook( program_pool: CompiledProgramsPool, key: CompiledProgramsKey, - backend: gtx_backend.Backend, + backend: gtx_backend.Toolchain, argument_descriptors: ArgStaticDescriptorsByType, offset_provider: common.OffsetProviderType | common.OffsetProvider, ) -> None: @@ -350,7 +350,7 @@ class CompiledProgramsPool(Generic[ffront_stages.DSLDefinitionT]): to `compile` before it can be used. """ - backend: gtx_backend.Backend + backend: gtx_backend.Toolchain definition_stage: ffront_stages.DSLDefinitionT # Note: This type can be incomplete, i.e. contain DeferredType, whenever the operator is a # scan operator. In the future it could also be the type of a generic program. @@ -491,7 +491,7 @@ def _load_artifact( ) -> artifacts.ExecutableProgram: artifact = artifact_future.result() # re-raises errors from the compilation worker try: - return self.backend.load_artifact(artifact) + return self.backend.loading(artifact) except Exception as e: raise RuntimeError( f"Failed to load the compiled program '{self.definition.__name__}'." diff --git a/src/gt4py/next/otf/recipes.py b/src/gt4py/next/otf/recipes.py index b905ee0bd3..a515fd5a98 100644 --- a/src/gt4py/next/otf/recipes.py +++ b/src/gt4py/next/otf/recipes.py @@ -15,7 +15,7 @@ @dataclasses.dataclass(frozen=True) class OTFCompileWorkflow( - workflow.NamedStepSequence[stages.CompilableProgramDef, artifacts.CompilationArtifact] + workflow.NamedStepSequence[stages.CompilableProgram, artifacts.CompilationArtifact] ): """The typical compiled backend steps composed into a workflow.""" diff --git a/src/gt4py/next/otf/stages.py b/src/gt4py/next/otf/stages.py index 4781907c1b..2d1962b4b1 100644 --- a/src/gt4py/next/otf/stages.py +++ b/src/gt4py/next/otf/stages.py @@ -39,12 +39,12 @@ ) ArgsDefinitionT = TypeVar("ArgsDefinitionT", arguments.JITArgs, arguments.CompileTimeArgs) -ConcreteProgramDef: TypeAlias = workflow.ConcreteArtifact[IRDefinitionT, ArgsDefinitionT] -CompilableProgramDef: TypeAlias = ConcreteProgramDef[itir.Program, arguments.CompileTimeArgs] +ConcreteProgramDef: TypeAlias = workflow.ProgramWithArgs[IRDefinitionT, ArgsDefinitionT] +CompilableProgram: TypeAlias = ConcreteProgramDef[itir.Program, arguments.CompileTimeArgs] class TranslationStep( - workflow.ReplaceEnabledWorkflowMixin[CompilableProgramDef, artifacts.ProgramSource[CodeSpecT]], + workflow.ReplaceEnabledWorkflowMixin[CompilableProgram, artifacts.ProgramSource[CodeSpecT]], Protocol[CodeSpecT], ): """Translate a GT4Py program to source code (ProgramCall -> ProgramSource).""" diff --git a/src/gt4py/next/otf/toolchain.py b/src/gt4py/next/otf/toolchain.py index 14d1b1cdf8..752c6e4fd9 100644 --- a/src/gt4py/next/otf/toolchain.py +++ b/src/gt4py/next/otf/toolchain.py @@ -25,40 +25,38 @@ class DataOnlyAdapter( workflow.ChainableWorkflowMixin, workflow.ReplaceEnabledWorkflowMixin, - workflow.Workflow[workflow.ConcreteArtifact[S, ArgsT], workflow.ConcreteArtifact[T, ArgsT]], + workflow.Workflow[workflow.ProgramWithArgs[S, ArgsT], workflow.ProgramWithArgs[T, ArgsT]], Generic[ArgsT, S, T], ): step: workflow.Workflow[S, T] def __call__( - self, inp: workflow.ConcreteArtifact[S, ArgsT] - ) -> workflow.ConcreteArtifact[T, ArgsT]: - return workflow.ConcreteArtifact(data=self.step(inp.data), args=inp.args) + self, inp: workflow.ProgramWithArgs[S, ArgsT] + ) -> workflow.ProgramWithArgs[T, ArgsT]: + return workflow.ProgramWithArgs(definition=self.step(inp.definition), args=inp.args) @dataclasses.dataclass(frozen=True) class ArgsOnlyAdapter( workflow.ChainableWorkflowMixin, workflow.ReplaceEnabledWorkflowMixin, - workflow.Workflow[workflow.ConcreteArtifact[DefT, S], workflow.ConcreteArtifact[DefT, T]], + workflow.Workflow[workflow.ProgramWithArgs[DefT, S], workflow.ProgramWithArgs[DefT, T]], Generic[DefT, S, T], ): step: workflow.Workflow[S, T] - def __call__( - self, inp: workflow.ConcreteArtifact[DefT, S] - ) -> workflow.ConcreteArtifact[DefT, T]: - return workflow.ConcreteArtifact(data=inp.data, args=self.step(inp.args)) + def __call__(self, inp: workflow.ProgramWithArgs[DefT, S]) -> workflow.ProgramWithArgs[DefT, T]: + return workflow.ProgramWithArgs(definition=inp.definition, args=self.step(inp.args)) @dataclasses.dataclass(frozen=True) class StripArgsAdapter( workflow.ChainableWorkflowMixin, workflow.ReplaceEnabledWorkflowMixin, - workflow.Workflow[workflow.ConcreteArtifact[S, ArgsT], T], + workflow.Workflow[workflow.ProgramWithArgs[S, ArgsT], T], Generic[ArgsT, S, T], ): step: workflow.Workflow[S, T] - def __call__(self, inp: workflow.ConcreteArtifact[S, ArgsT]) -> T: - return self.step(inp.data) + def __call__(self, inp: workflow.ProgramWithArgs[S, ArgsT]) -> T: + return self.step(inp.definition) diff --git a/src/gt4py/next/otf/workflow.py b/src/gt4py/next/otf/workflow.py index edd603f1b4..1a370037c3 100644 --- a/src/gt4py/next/otf/workflow.py +++ b/src/gt4py/next/otf/workflow.py @@ -36,7 +36,7 @@ @dataclasses.dataclass -class ConcreteArtifact(Generic[DefT, ArgsT]): +class ProgramWithArgs(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 @@ -49,7 +49,7 @@ class ConcreteArtifact(Generic[DefT, ArgsT]): test in `tests/next_tests/unit_tests/otf_tests/`. """ - data: DefT + definition: DefT args: ArgsT 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 b996be2626..d865836b75 100644 --- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py +++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py @@ -38,11 +38,11 @@ def get_param_description(name: str, type_: Any) -> interface.Parameter: @dataclasses.dataclass(frozen=True) class GTFNTranslationStep( workflow.ReplaceEnabledWorkflowMixin[ - stages.CompilableProgramDef, + stages.CompilableProgram, artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec], ], workflow.ChainableWorkflowMixin[ - stages.CompilableProgramDef, + stages.CompilableProgram, artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec], ], ): @@ -195,10 +195,10 @@ def generate_stencil_source( return codegen.format_source("cpp", generated_code, style="LLVM") def __call__( - self, inp: stages.CompilableProgramDef + self, inp: stages.CompilableProgram ) -> artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec]: """Generate GTFN C++ code from the ITIR definition.""" - program: itir.Program = inp.data + program: itir.Program = inp.definition # handle regular parameters and arguments of the program (i.e. what the user defined in # the program) diff --git a/src/gt4py/next/program_processors/runners/dace/program.py b/src/gt4py/next/program_processors/runners/dace/program.py index f8f8c7e1a4..abd88a7dad 100644 --- a/src/gt4py/next/program_processors/runners/dace/program.py +++ b/src/gt4py/next/program_processors/runners/dace/program.py @@ -43,9 +43,9 @@ def __sdfg__(self, *args: Any, **kwargs: Any) -> dace.sdfg.sdfg.SDFG: column_axis = kwargs.get("column_axis", None) # TODO(ricoh): connectivity tables required here for now. - gtir_stage = typing.cast(gtx_backend.Transforms, self.backend.transforms).past_to_itir( - workflow.ConcreteArtifact( - data=self.past_stage, + gtir_stage = typing.cast(gtx_backend.Transforms, self.backend.frontend).past_to_itir( + workflow.ProgramWithArgs( + definition=self.past_stage, args=arguments.CompileTimeArgs( args=tuple(p.type for p in self.past_stage.past_node.params), kwargs={}, @@ -55,13 +55,13 @@ def __sdfg__(self, *args: Any, **kwargs: Any) -> dace.sdfg.sdfg.SDFG: ), ) ) - program = gtir_stage.data + program = gtir_stage.definition program = itir_transforms.apply_fieldview_transforms( # run the transforms separately because they require the runtime info program, offset_provider=offset_provider ) object.__setattr__( gtir_stage, - "data", + "definition", program, ) object.__setattr__( @@ -76,7 +76,7 @@ def __sdfg__(self, *args: Any, **kwargs: Any) -> dace.sdfg.sdfg.SDFG: gt4py_program_args=[p.type for p in program.params], ) - otf_workflow = self.backend.executor + otf_workflow = self.backend.backend assert hasattr(otf_workflow, "translation") otf_workflow_translation = ( otf_workflow.translation.step 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 9934ffdb4a..30bd0401fd 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/backend.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/backend.py @@ -14,7 +14,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 import backend as next_backend, common, config from gt4py.next.otf import artifacts from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations from gt4py.next.program_processors.runners.dace.workflow import ( @@ -25,15 +25,20 @@ @dataclasses.dataclass(frozen=True) -class DaCeBackend(backend.Backend[Any]): - """DaCe backend with support for injecting an external workspace at load time.""" +class DaCeLoadingStep: + """ + Loading step that injects an external workspace into the loaded program. + + The workspace is owned by the caller, not by the toolchain; it is installed + onto the program wrapper before its first call, so that it is used when the + SDFG argument vector is constructed. + """ external_workspace: gtx_wfdcommon.ExternalWorkspace | None = None - def load_artifact(self, artifact: artifacts.CompilationArtifact) -> artifacts.ExecutableProgram: - program = super().load_artifact(artifact) + def __call__(self, artifact: artifacts.CompilationArtifact) -> artifacts.ExecutableProgram: + program = artifacts.load_artifact(artifact) assert isinstance(program, gtx_wfddecoration.DaCeDecoratedProgram) - # Inject the backend-level workspace so it is used when arguments are constructed. program.set_external_workspace(self.external_workspace or {}) return program @@ -48,7 +53,7 @@ def make_dace_backend( use_metrics: bool = True, use_zero_origin: bool = False, use_max_domain_range_on_unstructured_shift: bool | None = None, -) -> backend.Backend: +) -> next_backend.Toolchain: """Customize the dace backend with the given configuration parameters. Args: @@ -137,14 +142,14 @@ def make_dace_backend( use_max_domain_range_on_unstructured_shift=use_max_domain_range_on_unstructured_shift, ) - return DaCeBackend( + return next_backend.Toolchain( name=f"run_dace_{name_device}{'_opt' if auto_optimize else ''}", - executor=gtx_wfdfactory.make_dace_compile_workflow( + backend=gtx_wfdfactory.make_dace_compile_workflow( device_type=device_type, cached_translation=True, translation=translation ), allocator=allocator, - transforms=backend.DEFAULT_TRANSFORMS, - external_workspace=external_workspace, + frontend=next_backend.DEFAULT_TRANSFORMS, + loading=DaCeLoadingStep(external_workspace), ) 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 776d40df61..dd4c204c1f 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py @@ -109,7 +109,7 @@ def make_dace_compile_workflow( translation_step: stages.TranslationStep if cached_translation: translation_step = workflow.CachedStep[ - stages.CompilableProgramDef, artifacts.ProgramSource, str + stages.CompilableProgram, artifacts.ProgramSource, str ].persistent( bare_translation, input_fingerprinter=fingerprinting.strict_fingerprinter, 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 e45f3840b3..5f99d2c714 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py @@ -342,11 +342,11 @@ def make_sdfg_call_sync(sdfg: dace.SDFG, gpu: bool) -> None: @dataclasses.dataclass(frozen=True) class DaCeTranslator( workflow.ChainableWorkflowMixin[ - stages.CompilableProgramDef, + stages.CompilableProgram, artifacts.ProgramSource[artifacts.SDFGCodeSpec], ], workflow.ReplaceEnabledWorkflowMixin[ - stages.CompilableProgramDef, + stages.CompilableProgram, artifacts.ProgramSource[artifacts.SDFGCodeSpec], ], ): @@ -440,10 +440,10 @@ def _generate_sdfg_without_configuring_dace( return sdfg def __call__( - self, inp: stages.CompilableProgramDef + self, inp: stages.CompilableProgram ) -> artifacts.ProgramSource[artifacts.SDFGCodeSpec]: """Generate DaCe SDFG file from the GTIR definition.""" - program: itir.Program = inp.data + program: itir.Program = inp.definition assert isinstance(program, itir.Program) sdfg = self.generate_sdfg( diff --git a/src/gt4py/next/program_processors/runners/double_roundtrip.py b/src/gt4py/next/program_processors/runners/double_roundtrip.py index 3e96b30c60..ef7681f685 100644 --- a/src/gt4py/next/program_processors/runners/double_roundtrip.py +++ b/src/gt4py/next/program_processors/runners/double_roundtrip.py @@ -12,9 +12,9 @@ from gt4py.next.program_processors.runners import roundtrip -backend = next_backend.Backend( +backend = next_backend.Toolchain( name="double_roundtrip", - transforms=next_backend.DEFAULT_TRANSFORMS, - executor=roundtrip.Roundtrip(dispatch_backend=roundtrip.default, use_embedded=False), + frontend=next_backend.DEFAULT_TRANSFORMS, + backend=roundtrip.Roundtrip(dispatch_backend=roundtrip.default, use_embedded=False), allocator=roundtrip.default.allocator, ) diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py index 8d9a4fe435..b8c55a1f25 100644 --- a/src/gt4py/next/program_processors/runners/gtfn.py +++ b/src/gt4py/next/program_processors/runners/gtfn.py @@ -15,7 +15,7 @@ import gt4py._core.definitions as core_defs import gt4py.next.custom_layout_allocators as next_allocators from gt4py._core import filecache -from gt4py.next import backend, common, config, field_utils, fingerprinting +from gt4py.next import backend as next_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 artifacts, recipes, stages, workflow @@ -174,7 +174,7 @@ def make_gtfn_compile_workflow( translation_step: stages.TranslationStep if cached_translation: translation_step = workflow.CachedStep[ - stages.CompilableProgramDef, artifacts.ProgramSource, str + stages.CompilableProgram, artifacts.ProgramSource, str ].persistent( bare_translation, input_fingerprinter=fingerprinting.strict_fingerprinter, @@ -209,21 +209,21 @@ def make_gtfn_backend( gpu: bool = False, name_postfix: str = "", translation: gtfn_module.GTFNTranslationStep | None = None, - executor: workflow.Workflow[stages.CompilableProgramDef, artifacts.CompilationArtifact] + backend: workflow.Workflow[stages.CompilableProgram, artifacts.CompilationArtifact] | None = None, -) -> backend.Backend: +) -> next_backend.Toolchain: """ - Build a GTFN backend for the given device. + Build a GTFN toolchain for the given device. Args: gpu: Target the GPU instead of the CPU. name_postfix: Appended to the backend name, which must stay unique. translation: A pre-built translation step, forwarded to `make_gtfn_compile_workflow`. - executor: A pre-built compile workflow, replacing the default one. + backend: A pre-built compile pipeline, replacing the default one. Returns: - The configured backend. + The configured toolchain. """ allocator: next_allocators.FieldBufferAllocatorProtocol device_type: core_defs.DeviceType @@ -236,16 +236,16 @@ def make_gtfn_backend( device_type = core_defs.DeviceType.CPU name_device = "cpu" - if executor is None: - executor = make_gtfn_compile_workflow( + if backend is None: + backend = make_gtfn_compile_workflow( device_type=device_type, cached_translation=True, translation=translation ) - return backend.Backend( + return next_backend.Toolchain( name=f"run_gtfn_{name_device}{name_postfix}", - executor=executor, + backend=backend, allocator=allocator, - transforms=backend.DEFAULT_TRANSFORMS, + frontend=next_backend.DEFAULT_TRANSFORMS, ) diff --git a/src/gt4py/next/program_processors/runners/roundtrip.py b/src/gt4py/next/program_processors/runners/roundtrip.py index 09f173d3f9..1a4ed0e02b 100644 --- a/src/gt4py/next/program_processors/runners/roundtrip.py +++ b/src/gt4py/next/program_processors/runners/roundtrip.py @@ -221,7 +221,7 @@ class RoundtripArtifact: source_code: str entry_point_name: str column_axis: common.Dimension | None - dispatch_backend: next_backend.Backend | None + dispatch_backend: next_backend.Toolchain | None debug: bool def load(self) -> artifacts.ExecutableProgram: @@ -253,17 +253,17 @@ def decorated_fencil( @dataclasses.dataclass(frozen=True) -class Roundtrip(workflow.Workflow[stages.CompilableProgramDef, RoundtripArtifact]): +class Roundtrip(workflow.Workflow[stages.CompilableProgram, RoundtripArtifact]): debug: Optional[bool] = None use_embedded: bool = True - dispatch_backend: Optional[next_backend.Backend] = None + dispatch_backend: Optional[next_backend.Toolchain] = 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: stages.CompilableProgramDef) -> RoundtripArtifact: + def __call__(self, inp: stages.CompilableProgram) -> RoundtripArtifact: debug = config.DEBUG if self.debug is None else self.debug source_code, entry_point_name = _generate_source( - inp.data, + inp.definition, offset_provider=inp.args.offset_provider, debug=debug, use_embedded=self.use_embedded, @@ -280,41 +280,41 @@ def __call__(self, inp: stages.CompilableProgramDef) -> RoundtripArtifact: # TODO(tehrengruber): introduce factory -default = next_backend.Backend( +default = next_backend.Toolchain( name="roundtrip", - executor=Roundtrip( + backend=Roundtrip( transforms=functools.partial( itir_transforms.apply_common_transforms, extract_temporaries=False, ) ), allocator=next_allocators.StandardCPUFieldBufferAllocator(), - transforms=next_backend.DEFAULT_TRANSFORMS, + frontend=next_backend.DEFAULT_TRANSFORMS, ) -with_temporaries = next_backend.Backend( +with_temporaries = next_backend.Toolchain( name="roundtrip_with_temporaries", - executor=Roundtrip( + backend=Roundtrip( transforms=functools.partial( itir_transforms.apply_common_transforms, extract_temporaries=True, ) ), allocator=next_allocators.StandardCPUFieldBufferAllocator(), - transforms=next_backend.DEFAULT_TRANSFORMS, + frontend=next_backend.DEFAULT_TRANSFORMS, ) -no_transforms = next_backend.Backend( +no_transforms = next_backend.Toolchain( name="roundtrip", - executor=Roundtrip(transforms=lambda o, *, offset_provider: o), + backend=Roundtrip(transforms=lambda o, *, offset_provider: o), allocator=next_allocators.StandardCPUFieldBufferAllocator(), - transforms=next_backend.DEFAULT_TRANSFORMS, + frontend=next_backend.DEFAULT_TRANSFORMS, ) -gtir = next_backend.Backend( +gtir = next_backend.Toolchain( name="roundtrip_gtir", - executor=Roundtrip(transforms=itir_transforms.apply_fieldview_transforms), # type: ignore[arg-type] # don't understand why mypy complains + backend=Roundtrip(transforms=itir_transforms.apply_fieldview_transforms), # type: ignore[arg-type] # don't understand why mypy complains allocator=next_allocators.StandardCPUFieldBufferAllocator(), - transforms=next_backend.Transforms( + frontend=next_backend.Transforms( past_to_itir=past_to_itir.past_to_gtir_factory(), foast_to_itir=foast_to_gtir.adapted_foast_to_gtir_factory(cached=True), field_view_op_to_prog=foast_to_past.operator_to_program_factory( diff --git a/src/gt4py/next/typing.py b/src/gt4py/next/typing.py index 1c75458862..9fccb71595 100644 --- a/src/gt4py/next/typing.py +++ b/src/gt4py/next/typing.py @@ -25,7 +25,7 @@ CompiledProgramsKey: TypeAlias = Annotated[compiled_program.CompiledProgramsKey, _ONLY_FOR_TYPING] -Backend: TypeAlias = Annotated[backend.Backend, _ONLY_FOR_TYPING] +Toolchain: TypeAlias = Annotated[backend.Toolchain, _ONLY_FOR_TYPING] Allocator: TypeAlias = Annotated[constructors.Allocator, _ONLY_FOR_TYPING] @@ -34,10 +34,9 @@ __all__ = [ "Allocator", - "Backend", "FieldOperator", "OffsetProvider", "Program", - # from _core.definitions for convenience - "Scalar", + "Scalar", # from _core.definitions for convenience + "Toolchain", ] diff --git a/tests/next_tests/benchmarks/benchmark_program_call.py b/tests/next_tests/benchmarks/benchmark_program_call.py index 166509e8f3..ce0d05229b 100644 --- a/tests/next_tests/benchmarks/benchmark_program_call.py +++ b/tests/next_tests/benchmarks/benchmark_program_call.py @@ -47,7 +47,7 @@ @pytest.mark.parametrize("backend", BACKENDS, ids=lambda b: b.name) def benchmark_const_no_args_program( - benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Backend + benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Toolchain ): @gtx.field_operator def const() -> gtx.Field[Dims[IDim], gtx.float64]: @@ -69,7 +69,7 @@ def const_no_args(out: gtx.Field[Dims[IDim], gtx.float64]): @pytest.mark.parametrize("backend", BACKENDS, ids=lambda b: b.name) def benchmark_copy_01_arg_program( - benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Backend + benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Toolchain ): @gtx.field_operator def identity_fop( @@ -96,7 +96,7 @@ def copy_01_arg( @pytest.mark.parametrize("backend", BACKENDS, ids=lambda b: b.name) def benchmark_horizontal_copy_01_arg_program( - benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Backend + benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Toolchain ): @gtx.field_operator def identity_01_fop( @@ -143,7 +143,7 @@ def horizontal_copy_01_arg_program( @pytest.mark.parametrize("backend", BACKENDS, ids=lambda b: b.name) def benchmark_horizontal_copy_05_arg_program( - benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Backend + benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Toolchain ): @gtx.field_operator def identity_05_fop( @@ -202,7 +202,7 @@ def horizontal_copy_05_arg_program( @pytest.mark.parametrize("backend", BACKENDS, ids=lambda b: b.name) def benchmark_horizontal_copy_25_arg_program( - benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Backend + benchmark: ptb_fixture.BenchmarkFixture, backend: gtx_typing.Toolchain ): @gtx.field_operator def identity_25_fop( diff --git a/tests/next_tests/definitions.py b/tests/next_tests/definitions.py index e23d7d9319..8b84d35160 100644 --- a/tests/next_tests/definitions.py +++ b/tests/next_tests/definitions.py @@ -55,7 +55,7 @@ class ProgramBackendId(_PythonObjectIdMixin, str, enum.Enum): class EmbeddedDummyBackend: name: str allocator: constructors.Allocator - executor: Final = None + backend: Final = None numpy_execution = EmbeddedDummyBackend("EmbeddedNumPy", nd_array_field.np) diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index 08fb817856..9fcd817e36 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -546,7 +546,7 @@ def cartesian_case_no_backend(): @pytest.fixture def cartesian_case( - exec_alloc_descriptor: test_definitions.EmbeddedDummyBackend | next_backend.Backend, + exec_alloc_descriptor: test_definitions.EmbeddedDummyBackend | next_backend.Toolchain, ): return Case.from_cartesian_grid_descriptor( simple_cartesian_grid(), @@ -567,7 +567,7 @@ def unstructured_case_no_backend(mesh_descriptor: MeshDescriptor): @pytest.fixture def unstructured_case( mesh_descriptor: MeshDescriptor, - exec_alloc_descriptor: test_definitions.EmbeddedDummyBackend | next_backend.Backend, + exec_alloc_descriptor: test_definitions.EmbeddedDummyBackend | next_backend.Toolchain, ): return Case.from_mesh_descriptor( mesh_descriptor, @@ -713,7 +713,7 @@ def get_default_data( class Case: """Parametrizable components for single feature integration tests.""" - backend: Optional[next_backend.Backend] + backend: Optional[next_backend.Toolchain] offset_provider: dict[str, common.Connectivity | gtx.Dimension] default_sizes: dict[gtx.Dimension, int] grid_type: common.GridType @@ -727,7 +727,7 @@ def as_field(self): def from_cartesian_grid_descriptor( cls, grid_descriptor: CartesianGridDescriptor, - backend: Optional[next_backend.Backend] = None, + backend: Optional[next_backend.Toolchain] = None, allocator: Optional[next_allocators.FieldBufferAllocatorFactoryProtocol] = None, ) -> Case: return cls( @@ -749,7 +749,7 @@ def from_cartesian_grid_descriptor( def from_mesh_descriptor( cls, mesh_descriptor: MeshDescriptor, - backend: Optional[next_backend.Backend] = None, + backend: Optional[next_backend.Toolchain] = None, allocator: Optional[next_allocators.FieldBufferAllocatorFactoryProtocol] = None, ) -> Case: return cls( diff --git a/tests/next_tests/integration_tests/cases_utils.py b/tests/next_tests/integration_tests/cases_utils.py index b2be1e6217..e461490528 100644 --- a/tests/next_tests/integration_tests/cases_utils.py +++ b/tests/next_tests/integration_tests/cases_utils.py @@ -59,7 +59,7 @@ def _no_backend_allocator(*args: Any, **kwargs: Any) -> None: raise ValueError("No backend selected! Backend selection is mandatory in tests.") -class NoBackend(next_backend.Backend): +class NoBackend(next_backend.Toolchain): """Temporary default backend to not accidentally test the wrong backend.""" def __call__(self, program, *args, **kwargs) -> None: @@ -76,12 +76,12 @@ def __gt_allocator__( no_backend = NoBackend( name="no_backend", - executor=_no_backend_executor, + backend=_no_backend_executor, allocator=_no_backend_allocator, # TODO(tehrengruber): We don't want any transformations, but since `decorator.FieldOperator` # and `decorator.Program` unconditionally do linting on construction we need the # transformations. When this is up to the backend we can remove this again. - transforms=next_backend.DEFAULT_TRANSFORMS, + frontend=next_backend.DEFAULT_TRANSFORMS, ) diff --git a/tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py b/tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py index e613e3db92..741d8efff9 100644 --- a/tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py +++ b/tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py @@ -119,10 +119,12 @@ def uncached_dace_cpu(): The cache stores the optimized SDFG, so a second run of these tests would replay it and never call the transformation. """ - executor = gtx_dace.run_dace_cpu.executor + compile_pipeline = gtx_dace.run_dace_cpu.backend return dataclasses.replace( gtx_dace.run_dace_cpu, - executor=dataclasses.replace(executor, translation=executor.translation.step), + backend=dataclasses.replace( + compile_pipeline, translation=compile_pipeline.translation.step + ), ) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py index 9bd7e43273..1f1ee9c740 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py @@ -39,7 +39,7 @@ # becomes easier. -def make_builtin_field_operator(builtin_name: str, backend: Optional[next_backend.Backend]): +def make_builtin_field_operator(builtin_name: str, backend: Optional[next_backend.Toolchain]): # TODO(tehrengruber): creating a field operator programmatically should be # easier than what we need to do here. # construct annotation dictionary containing the input argument and return diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py index f13d2acb8c..616e79588d 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_temporaries_with_sizes.py @@ -31,10 +31,10 @@ # see https://docs.pytest.org/en/latest/how-to/fixtures.html#override-a-fixture-on-a-test-module-level @pytest.fixture def exec_alloc_descriptor(): - return backend.Backend( + return backend.Toolchain( name="run_gtfn_with_temporaries_and_sizes", - transforms=backend.DEFAULT_TRANSFORMS, - executor=gtfn.make_gtfn_compile_workflow( + frontend=backend.DEFAULT_TRANSFORMS, + backend=gtfn.make_gtfn_compile_workflow( translation=gtfn.gtfn_module.GTFNTranslationStep( symbolic_domain_sizes={ "Cell": "num_cells", diff --git a/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py b/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py index 9f51689d77..c7ef03a844 100644 --- a/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py +++ b/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_hooks.py @@ -126,7 +126,7 @@ def custom_compiled_program_callback( @pytest.mark.parametrize("backend", BACKENDS, ids=lambda b: getattr(b, "name", str(b))) -def test_program_call_hooks(backend: gtx_typing.Backend): +def test_program_call_hooks(backend: gtx_typing.Toolchain): size = 10 a_field = gtx.full([(IDim, size)], 1, dtype=gtx.float64) b_field = gtx.full([(IDim, size)], 1, dtype=gtx.float64) @@ -208,11 +208,11 @@ def test_program_call_hooks(backend: gtx_typing.Backend): @pytest.mark.parametrize( "backend", [b for b in BACKENDS if b is not None], ids=lambda b: getattr(b, "name", str(b)) ) -def test_compile_variant_hook(backend: gtx_typing.Backend): +def test_compile_variant_hook(backend: gtx_typing.Toolchain): def custom_compile_variant_hook( program_pool: "CompiledProgramsPool", key: gtx_typing.CompiledProgramsKey, - backend: gtx_typing.Backend, + backend: gtx_typing.Toolchain, argument_descriptors: dict[type, dict[str, Any]], offset_provider: common.OffsetProviderType | common.OffsetProvider, ) -> None: diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_ffront_fvm_nabla.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_ffront_fvm_nabla.py index 2b4847d3c3..760c2b01e3 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_ffront_fvm_nabla.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_ffront_fvm_nabla.py @@ -83,7 +83,7 @@ def test_ffront_compute_zavgS(exec_alloc_descriptor): zavgS = gtx.zeros({Edge: setup.edges_size}, allocator=exec_alloc_descriptor.allocator) compute_zavgS.with_backend( - None if exec_alloc_descriptor.executor is None else exec_alloc_descriptor + None if exec_alloc_descriptor.backend is None else exec_alloc_descriptor )( setup.input_field, setup.S_fields[0], @@ -106,7 +106,7 @@ def test_ffront_nabla(exec_alloc_descriptor): pnabla_MXX = gtx.zeros({Vertex: setup.nodes_size}, allocator=exec_alloc_descriptor.allocator) pnabla_MYY = gtx.zeros({Vertex: setup.nodes_size}, allocator=exec_alloc_descriptor.allocator) - pnabla.with_backend(None if exec_alloc_descriptor.executor is None else exec_alloc_descriptor)( + pnabla.with_backend(None if exec_alloc_descriptor.backend is None else exec_alloc_descriptor)( setup.input_field, setup.S_fields, setup.sign_field, diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py index f9b2060621..283f134876 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py @@ -100,7 +100,7 @@ def fen_solve_tridiag2(i_size, j_size, k_size, a, b, c, d, x): def test_tridiag(fencil, tridiag_reference, program_processor): program_processor, validate = program_processor - if isinstance(program_processor, backend.Backend) and "dace" in program_processor.name: + if isinstance(program_processor, backend.Toolchain) and "dace" in program_processor.name: pytest.xfail("Dace ITIR backend doesn't support the IR format used in this test.") a, b, c, d, x = tridiag_reference diff --git a/tests/next_tests/unit_tests/conftest.py b/tests/next_tests/unit_tests/conftest.py index d7b5d100e3..673d7ac56e 100644 --- a/tests/next_tests/unit_tests/conftest.py +++ b/tests/next_tests/unit_tests/conftest.py @@ -23,7 +23,7 @@ import next_tests -ProgramProcessor: TypeAlias = backend.Backend | program_formatter.ProgramFormatter +ProgramProcessor: TypeAlias = backend.Toolchain | program_formatter.ProgramFormatter def _program_processor(request) -> tuple[ProgramProcessor, bool]: 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 abf6ad73dd..bcc39f4b7d 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 @@ -90,8 +90,8 @@ def _verify_program_has_expected_true_value(program: itir.Program): def test_inlining_of_scalars_works(testee_prog): - input_pair = workflow.ConcreteArtifact( - data=testee_prog.definition_stage, + input_pair = workflow.ProgramWithArgs( + definition=testee_prog.definition_stage, args=arguments.CompileTimeArgs( args=list(testee_prog.past_stage.past_node.type.definition.pos_or_kw_args.values()), kwargs={}, @@ -103,7 +103,7 @@ def test_inlining_of_scalars_works(testee_prog): ), ) - transformed = backend.DEFAULT_TRANSFORMS(input_pair).data + transformed = backend.DEFAULT_TRANSFORMS(input_pair).definition _verify_program_has_expected_true_value(transformed) @@ -123,14 +123,14 @@ class _NoOpArtifact: def load(self): return lambda *args, **kwargs: None - def pirate(program: workflow.ConcreteArtifact): + def pirate(program: workflow.ProgramWithArgs): # Replaces the gtfn otf_workflow: steals the compilable program, then # returns a dummy artifact whose materialization is a no-op callable. nonlocal hijacked_program hijacked_program = program return _NoOpArtifact() - hacked_gtfn_backend = gtfn.make_gtfn_backend(name_postfix="_custom", executor=pirate) + hacked_gtfn_backend = gtfn.make_gtfn_backend(name_postfix="_custom", backend=pirate) testee = testee_prog.with_backend(hacked_gtfn_backend).compile(cond=[True], offset_provider={}) testee( @@ -139,7 +139,7 @@ def pirate(program: workflow.ConcreteArtifact): offset_provider={}, ) - _verify_program_has_expected_true_value(hijacked_program.data) + _verify_program_has_expected_true_value(hijacked_program.definition) def test_different_static_args_work_after_backend_change(testee_prog): @@ -298,8 +298,8 @@ 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 = workflow.ConcreteArtifact( - data=testee_prog.definition_stage, + input_pair = workflow.ProgramWithArgs( + definition=testee_prog.definition_stage, args=arguments.CompileTimeArgs( args=list(testee_prog.past_stage.past_node.type.definition.pos_or_kw_args.values()), kwargs={}, @@ -311,7 +311,7 @@ def test_inlining_of_static_domain_works(testee_prog, uids: utils.IDGeneratorPoo ), ) - transformed = backend.DEFAULT_TRANSFORMS(input_pair).data + transformed = backend.DEFAULT_TRANSFORMS(input_pair).definition _verify_program_has_expected_domain(transformed, domain, uids) diff --git a/tests/next_tests/unit_tests/otf_tests/test_runners.py b/tests/next_tests/unit_tests/otf_tests/test_runners.py index dfe0d5b5f6..5597a5b27c 100644 --- a/tests/next_tests/unit_tests/otf_tests/test_runners.py +++ b/tests/next_tests/unit_tests/otf_tests/test_runners.py @@ -100,11 +100,11 @@ def test_process_runner_falls_back_on_non_offloadable_task(process_runner): def test_make_compilation_task_decomposes_standard_backend(): - backend = next_backend.Backend( + backend = next_backend.Toolchain( name="test_backend", - executor=lambda compilable: _NoOpArtifact(), + backend=lambda compilable: _NoOpArtifact(), allocator=None, - transforms=lambda inp: inp, + frontend=lambda inp: inp, ) task = compilation_tasks.make_compilation_task( @@ -112,7 +112,7 @@ def test_make_compilation_task_decomposes_standard_backend(): ) assert task.no_offload_reason is None - assert task.executor is backend.executor + assert task.executor is backend.backend assert callable(task.compile().load()) @@ -129,11 +129,11 @@ def compile(self, program, compile_time_args): return self._wrapped.compile(program, compile_time_args=compile_time_args) backend = _WrapperBackend( - next_backend.Backend( + next_backend.Toolchain( name="test_backend", - executor=lambda compilable: _NoOpArtifact(), + backend=lambda compilable: _NoOpArtifact(), allocator=None, - transforms=lambda inp: inp, + frontend=lambda inp: inp, ) ) @@ -154,11 +154,11 @@ def test_offloaded_task_ships_connectivities_as_file_refs(): compile_time_args = dataclasses.replace( arguments.CompileTimeArgs.empty(), offset_provider={"V2E": conn} ) - backend = next_backend.Backend( + backend = next_backend.Toolchain( name="test_backend", - executor=lambda compilable: _NoOpArtifact(), + backend=lambda compilable: _NoOpArtifact(), allocator=None, - transforms=lambda inp: inp, + frontend=lambda inp: inp, ) task = compilation_tasks.make_compilation_task( 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 2f05f8debd..9a7445179d 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 @@ -76,8 +76,8 @@ def program_example(): def test_codegen(program_example): fencil, parameters = program_example module = gtfn_module.translate_program_cpu( - stages.CompilableProgramDef( - data=fencil, + stages.CompilableProgram( + definition=fencil, args=arguments.CompileTimeArgs.from_concrete(*parameters, **{"offset_provider": {}}), ) ) @@ -88,8 +88,8 @@ def test_codegen(program_example): def test_hash_and_diskcache(program_example, tmp_path): fencil, parameters = program_example - compilable_program = stages.CompilableProgramDef( - data=fencil, + compilable_program = stages.CompilableProgram( + definition=fencil, args=arguments.CompileTimeArgs.from_concrete(*parameters, **{"offset_provider": {}}), ) hash = fingerprinting.strict_fingerprinter(compilable_program) @@ -110,7 +110,7 @@ def test_hash_and_diskcache(program_example, tmp_path): # hash is different if program changes altered_program_id = copy.deepcopy(compilable_program) - altered_program_id.data.id = "example2" + altered_program_id.definition.id = "example2" assert fingerprinting.strict_fingerprinter( compilable_program ) != fingerprinting.strict_fingerprinter(altered_program_id) @@ -130,8 +130,8 @@ def test_hash_and_diskcache(program_example, tmp_path): def test_gtfn_file_cache(program_example): fencil, parameters = program_example - compilable_program = stages.CompilableProgramDef( - data=fencil, + compilable_program = stages.CompilableProgram( + definition=fencil, args=arguments.CompileTimeArgs.from_concrete(*parameters, **{"offset_provider": {}}), ) cached_gtfn_translation_step = gtfn.make_gtfn_compile_workflow( diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_backend.py index 488364b360..36ab6a9d37 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 @@ -194,7 +194,7 @@ def test_make_backend_accepts_external_workspace_with_external_mode(): external_workspace={core_defs.DeviceType.CPU: workspace}, ) - assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace + assert backend.loading.external_workspace[core_defs.DeviceType.CPU] is workspace def test_make_backend_infers_external_mode_when_workspace_is_provided(): @@ -208,10 +208,10 @@ def test_make_backend_infers_external_mode_when_workspace_is_provided(): ) assert ( - backend.executor.translation.step.auto_optimize_args["transient_memory_mode"] + backend.backend.translation.step.auto_optimize_args["transient_memory_mode"] == gtx_transformations.TransientMemoryMode.EXTERNAL ) - assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace + assert backend.loading.external_workspace[core_defs.DeviceType.CPU] is workspace def test_make_backend_warns_external_workspace_without_external_mode(): @@ -230,10 +230,10 @@ def test_make_backend_warns_external_workspace_without_external_mode(): # Explicit mode stays as requested by the caller; backend only warns. assert ( - backend.executor.translation.step.auto_optimize_args["transient_memory_mode"] + backend.backend.translation.step.auto_optimize_args["transient_memory_mode"] == gtx_transformations.TransientMemoryMode.POOL ) - assert backend.external_workspace[core_defs.DeviceType.CPU] is workspace + assert backend.loading.external_workspace[core_defs.DeviceType.CPU] is workspace def _parse_generated_code_from_sdfg(sdfg: dace.SDFG, gpu_api_prefix: str) -> str: @@ -318,9 +318,9 @@ def testee(a: cases.IField, b: cases.IField, out: cases.IField): out = cases.allocate(test_case, testee, "out")() captured_sdfg: dace.SDFG | None = None - translation_step = custom_backend.executor.translation.step + translation_step = custom_backend.backend.translation.step - def mocked_translator(inp: stages.CompilableProgramDef) -> dace.SDFG: + def mocked_translator(inp: stages.CompilableProgram) -> dace.SDFG: nonlocal captured_sdfg result = translation_step(inp) captured_sdfg = dace.SDFG.from_json(result.source_code) @@ -328,8 +328,8 @@ def mocked_translator(inp: stages.CompilableProgramDef) -> dace.SDFG: custom_backend = dataclasses.replace( custom_backend, - executor=dataclasses.replace( - custom_backend.executor, + backend=dataclasses.replace( + custom_backend.backend, translation=mocked_translator, ), ) @@ -368,7 +368,7 @@ def no_op_top_level_map_processing(*, sdfg: dace.SDFG, **kwargs) -> dace.SDFG: assert all( tdesc.lifetime == dace.AllocationLifetime.External for _, tdesc in transient_arrays ) - # load_artifact injected the backend-level workspace onto the program wrapper. + # The loading step injected the toolchain-level workspace onto the program wrapper. assert ( decorated_program._fun.external_workspace[device_type] is external_workspace[device_type] 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 87e212afc9..7fa373f665 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 @@ -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_workflow.ConcreteArtifact: +def _make_simple_field_operator_compilable_program() -> otf_workflow.ProgramWithArgs: """Return a compilable program wrapping a minimal GTIR field operator.""" ir = itir.Program( id="simple_field_operator", @@ -443,8 +443,8 @@ def _make_simple_field_operator_compilable_program() -> otf_workflow.ConcreteArt ), ], ) - return otf_workflow.ConcreteArtifact( - data=ir, + return otf_workflow.ProgramWithArgs( + definition=ir, args=otf_arguments.CompileTimeArgs( args=tuple(param.type for param in ir.params), kwargs={}, diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py index 5252f45565..c849d91245 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py @@ -34,15 +34,15 @@ def test_make_gtfn_backend_trait_device(): gpu_version = gtfn.make_gtfn_backend(gpu=True) assert cpu_version.name == "run_gtfn_cpu" - assert isinstance(cpu_version.executor.translation, workflow.CachedStep) - assert cpu_version.executor.translation.step.device_type is core_defs.DeviceType.CPU + assert isinstance(cpu_version.backend.translation, workflow.CachedStep) + assert cpu_version.backend.translation.step.device_type is core_defs.DeviceType.CPU assert gpu_version.name == "run_gtfn_gpu" - assert isinstance(gpu_version.executor.translation, workflow.CachedStep) - assert gpu_version.executor.translation.step.device_type is core_defs.DeviceType.CUDA + assert isinstance(gpu_version.backend.translation, workflow.CachedStep) + assert gpu_version.backend.translation.step.device_type is core_defs.DeviceType.CUDA # The compilation step now also carries device_type so it can stamp the artifact. - assert cpu_version.executor.compilation.device_type is core_defs.DeviceType.CPU - assert gpu_version.executor.compilation.device_type is core_defs.DeviceType.CUDA + assert cpu_version.backend.compilation.device_type is core_defs.DeviceType.CPU + assert gpu_version.backend.compilation.device_type is core_defs.DeviceType.CUDA assert custom_layout_allocators.is_field_allocator_for( cpu_version.allocator, core_defs.DeviceType.CPU @@ -58,9 +58,9 @@ def test_make_gtfn_backend_build_cache_config(monkeypatch): monkeypatch.setattr(config, "BUILD_CACHE_LIFETIME", config.BuildCacheLifetime.PERSISTENT) persistent_version = gtfn.make_gtfn_backend() - assert session_version.executor.compilation.cache_lifetime is config.BuildCacheLifetime.SESSION + assert session_version.backend.compilation.cache_lifetime is config.BuildCacheLifetime.SESSION assert ( - persistent_version.executor.compilation.cache_lifetime + persistent_version.backend.compilation.cache_lifetime is config.BuildCacheLifetime.PERSISTENT ) @@ -72,11 +72,11 @@ def test_make_gtfn_backend_build_type_config(monkeypatch): min_size_version = gtfn.make_gtfn_backend() assert ( - release_version.executor.compilation.builder_factory.cmake_build_type + release_version.backend.compilation.builder_factory.cmake_build_type is config.CMakeBuildType.RELEASE ) assert ( - min_size_version.executor.compilation.builder_factory.cmake_build_type + min_size_version.backend.compilation.builder_factory.cmake_build_type is config.CMakeBuildType.MIN_SIZE_REL ) @@ -94,8 +94,8 @@ def test_cmake_build_type_changes_build_folder(monkeypatch, tmp_path): monkeypatch.setattr(config, "CMAKE_BUILD_TYPE", config.CMakeBuildType.DEBUG) debug_version = gtfn.make_gtfn_backend() - release_compiler = release_version.executor.compilation - debug_compiler = debug_version.executor.compilation + release_compiler = release_version.backend.compilation + debug_compiler = debug_version.backend.compilation build_context_ids: list[str] = [] diff --git a/typing_tests/test_next.yaml b/typing_tests/test_next.yaml index 5db040f8df..a743f0401f 100644 --- a/typing_tests/test_next.yaml +++ b/typing_tests/test_next.yaml @@ -211,7 +211,7 @@ def foo(a: gtx.Field[gtx.Dims[KDim], gtx.int32]) -> gtx.Field[gtx.Dims[KDim], gtx.int32]: return a - def foo_with_backend(backend: gtx.typing.Backend | None) -> None: + def foo_with_backend(backend: gtx.typing.Toolchain | None) -> None: _ = foo.with_backend(backend) - case: with_backend_scanop @@ -227,7 +227,7 @@ ) -> float: return carry + val - def somescanop_with_backend(backend: gtx.typing.Backend | None) -> None: + def somescanop_with_backend(backend: gtx.typing.Toolchain | None) -> None: _ = somescanop.with_backend(backend) - case: with_backend_program @@ -251,7 +251,7 @@ ) -> None: foo(a, out=b, domain={KDim: (start, end)}) - def foo_prg_with_backend(backend: gtx.typing.Backend | None) -> None: + def foo_prg_with_backend(backend: gtx.typing.Toolchain | None) -> None: _ = foo_prg.with_backend(backend) - case: astype_return_type