diff --git a/docs/development/ADRs/next/0027-Toolchain-Naming-and-Pipeline-Simplification.md b/docs/development/ADRs/next/0027-Toolchain-Naming-and-Pipeline-Simplification.md index abecc4aa57..14a219c205 100644 --- a/docs/development/ADRs/next/0027-Toolchain-Naming-and-Pipeline-Simplification.md +++ b/docs/development/ADRs/next/0027-Toolchain-Naming-and-Pipeline-Simplification.md @@ -88,7 +88,7 @@ Named pipelines become frozen dataclasses with an explicit, fully typed 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. +`CachedStep` is kept, minus the mixin-provided `.replace` / `.chain`. What ADR 0011's decisions become: @@ -114,16 +114,16 @@ 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. +The decisions in this ADR landed as a stack of PRs: the naming decisions first +(#2741), then stage observability (#2742), then the pipeline simplification +together with the deferred `OTFCompileWorkflow` → `CompilePipeline` rename +(this stack's final PR). ## 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`. + and mutating frozen stages) was replaced by a dace-owned translate-only + toolchain 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 diff --git a/docs/user/next/advanced/HackTheToolchain.md b/docs/user/next/advanced/HackTheToolchain.md index 2dc40b7386..37d3fc232f 100644 --- a/docs/user/next/advanced/HackTheToolchain.md +++ b/docs/user/next/advanced/HackTheToolchain.md @@ -4,8 +4,6 @@ import typing from gt4py import next as gtx from gt4py.next.otf import workflow -from gt4py.next.ffront import field_operator_ast as foast, stages as ff_stages -from gt4py import eve ``` @@ -13,37 +11,24 @@ from gt4py import eve ## Replace Steps -```python -cached_lowering_toolchain = gtx.backend.DEFAULT_TRANSFORMS.replace( - past_to_itir=gtx.ffront.past_to_itir.past_to_gtir_factory(cached=False) -) -``` - -## Skip Steps / Change Order +Pipelines are frozen dataclasses whose fields are the steps, so a variant is built with `dataclasses.replace`. ```python -DUMMY_FOP = workflow.ProgramWithArgs( - definition=ff_stages.DSLFieldOperatorDef(definition=None), args=None +cached_lowering_toolchain = dataclasses.replace( + gtx.backend.DEFAULT_TRANSFORMS, + past_to_itir=gtx.ffront.past_to_itir.past_to_gtir_factory(cached=False), ) ``` -```python -gtx.backend.DEFAULT_TRANSFORMS.step_order(DUMMY_FOP) -``` - -```python -@dataclasses.dataclass(frozen=True) -class SkipLinting(gtx.backend.Transforms): - def step_order(self, inp): - order = super().step_order(inp) - if "past_lint" in order: - order.remove("past_lint") # not running "past_lint" - return order +## Skip Steps +Which steps run is decided by `Transforms.__call__` from the type of the input definition; that selection is not configurable. Behavior is customized by replacing a step, so skipping one means replacing it with an identity step. -same_steps = dataclasses.asdict(gtx.backend.DEFAULT_TRANSFORMS) -skip_linting_transforms = SkipLinting(**same_steps) -skip_linting_transforms.step_order(DUMMY_FOP) +```python +skip_linting_transforms = dataclasses.replace( + gtx.backend.DEFAULT_TRANSFORMS, + past_lint=lambda past_def: past_def, # identity step: linting skipped +) ``` ## Alternative Factory @@ -56,18 +41,20 @@ class Cpp2BindingsGen: ... class PureCpp2WorkflowFactory(gtx.program_processors.runners.gtfn.GTFNCompileWorkflowFactory): - translation: workflow.Workflow[ + translation: workflow.Step[ gtx.otf.stages.CompilableProgram, gtx.otf.artifacts.ProgramSource ] = MyCodeGen() - bindings: workflow.Workflow[ - gtx.otf.artifacts.ProgramSource, gtx.otf.artifacts.ExtensionSource - ] = Cpp2BindingsGen() + bindings: workflow.Step[gtx.otf.artifacts.ProgramSource, gtx.otf.artifacts.ExtensionSource] = ( + Cpp2BindingsGen() + ) PureCpp2WorkflowFactory(cmake_build_type=gtx.config.CMAKE_BUILD_TYPE.DEBUG) ``` -## Invent new Workflow Types +## Invent new Pipeline Types + +A pipeline is just a frozen dataclass of steps with an explicit, fully typed `__call__`. Nothing else is needed, so a non-linear shape is written the same way as a linear one. ```mermaid graph LR @@ -86,15 +73,11 @@ OUT_T = typing.TypeVar("OUT_T") @dataclasses.dataclass(frozen=True) -class FullyModularDiamond( - workflow.ChainableWorkflowMixin[IN_T, OUT_T], - workflow.ReplaceEnabledWorkflowMixin[IN_T, OUT_T], - typing.Protocol[IN_T, OUT_T, A_T, B_T, X_T, Y_T], -): - split: workflow.Workflow[IN_T, tuple[A_T, X_T]] - track_a: workflow.Workflow[A_T, B_T] - track_x: workflow.Workflow[X_T, Y_T] - combine: workflow.Workflow[tuple[B_T, Y_T], OUT_T] +class Diamond(typing.Generic[IN_T, OUT_T, A_T, B_T, X_T, Y_T]): + split: workflow.Step[IN_T, tuple[A_T, X_T]] + track_a: workflow.Step[A_T, B_T] + track_x: workflow.Step[X_T, Y_T] + combine: workflow.Step[tuple[B_T, Y_T], OUT_T] def __call__(self, inp: IN_T) -> OUT_T: a, x = self.split(inp) @@ -103,20 +86,10 @@ class FullyModularDiamond( return self.combine((b, y)) -@dataclasses.dataclass(frozen=True) -class PartiallyModularDiamond( - workflow.ChainableWorkflowMixin[IN_T, OUT_T], - workflow.ReplaceEnabledWorkflowMixin[IN_T, OUT_T], - typing.Protocol[IN_T, OUT_T, A_T, B_T, X_T, Y_T], -): - track_a: workflow.Workflow[A_T, B_T] - track_x: workflow.Workflow[X_T, Y_T] - - def split(inp: IN_T) -> tuple[A_T, X_T]: ... - - def combine(b: B_T, y: Y_T) -> OUT_T: ... - - def __call__(inp: IN_T) -> OUT_T: - a, x = self.split(inp) - return self.combine(b=self.track_a(a), y=self.track_x(x)) +Diamond( + split=lambda inp: (inp, inp), + track_a=lambda a: a + 1, + track_x=lambda x: x * 2, + combine=lambda by: by[0] + by[1], +)(3) ``` diff --git a/docs/user/next/advanced/WorkflowPatterns.md b/docs/user/next/advanced/WorkflowPatterns.md index 66910324d7..dff6b93e9d 100644 --- a/docs/user/next/advanced/WorkflowPatterns.md +++ b/docs/user/next/advanced/WorkflowPatterns.md @@ -15,29 +15,26 @@ jupyter: ```python editable=true slideshow={"slide_type": ""} import dataclasses -import re import factory import gt4py.next as gtx - -import devtools ``` -# How to read (toolchain) workflows +# How to read (toolchain) pipelines -## Basic workflow (single step) +## Basic step ```mermaid graph LR -StageA -->|basic workflow| StageB +StageA -->|step| StageB ``` Where "Stage" describes any data structure, and where `StageA` contains all the input data and `StageB` contains all the output data. @@ -60,7 +57,9 @@ simple_add_1(1) -This is already a (single step) workflow. We can build a more complex one by chaining it multiple times. +This is already a step: the type `gtx.otf.workflow.Step[S, T]` is nothing but an alias of `Callable[[S], T]`, so every single-argument callable is a step and there is no base class, decorator or wrapper to apply. + +Composing steps is therefore plain Python: ```mermaid graph LR @@ -71,49 +70,28 @@ inp(A: int) -->|simple_add_1| b(A + 1) -->|simple_add_1| c(A + 2) -->|simple_add ```python editable=true slideshow={"slide_type": ""} -manual_add_3 = ( - gtx.otf.workflow.StepSequence.start(simple_add_1).chain(simple_add_1).chain(simple_add_1) -) - -manual_add_3(1) -``` - - - -### Simplest Composable Step - -All we have to do for chaining to work out of the box is add the `make_step` decorator! +add_3: gtx.otf.workflow.Step[int, int] = lambda inp: simple_add_1(simple_add_1(simple_add_1(inp))) - - -```python editable=true slideshow={"slide_type": ""} -@gtx.otf.workflow.make_step -def chainable_add_1(inp: int) -> int: - return inp + 1 -``` - -```python editable=true slideshow={"slide_type": ""} -add_3 = chainable_add_1.chain(chainable_add_1).chain(chainable_add_1) add_3(1) ``` ### Example in the Wild ```python -gtx.ffront.func_to_past.func_to_past.steps.inner[0]?? +gtx.ffront.func_to_past.func_to_past?? ``` ### Step with Parameters -Sometimes we want to allow for different configurations of a step. +Sometimes we want to allow for different configurations of a step. A frozen dataclass with a `__call__` gives us a configurable, immutable step. ```python editable=true slideshow={"slide_type": ""} @dataclasses.dataclass(frozen=True) -class MathOp(gtx.otf.workflow.ChainableWorkflowMixin[int, int]): +class MathOp: op: str rhs: int = 0 @@ -127,7 +105,14 @@ class MathOp(gtx.otf.workflow.ChainableWorkflowMixin[int, int]): return lhs * rhs -add_3_times_2 = MathOp("add", 3).chain(MathOp("mul", 2)) +add_3_step = MathOp("add", 3) +times_2_step = MathOp("mul", 2) + + +def add_3_times_2(inp: int) -> int: + return times_2_step(add_3_step(inp)) + + add_3_times_2(1) ``` @@ -142,11 +127,10 @@ gtx.program_processors.runners.roundtrip.Roundtrip?? ### Wrapper Steps Sometimes we want to make a step behave slightly differently without modifying the step itself. In this case we can wrap it into a wrapper step. These behave a little bit like (limited) decorators. -Below we will go through the existing wrapper steps, which you might encounter. #### Caching / memoizing -For example we might want to cach the output (memoize) for which we need to add a way of hashing the input: +For example we might want to cache the output (memoize) for which we need to add a way of hashing the input: ```mermaid graph LR @@ -157,24 +141,22 @@ inp(A: int) --> ha{{"input_fingerprinter(A)"}} --> h("hash(A)") --> ck{{"check c ck -->|hit| hit("in cache") --> out ``` -For this we can use the `CachedStep`, you will see something like below +For this we can use the `CachedStep`, the one wrapper step the toolchain still has. You will see something like below ```python editable=true slideshow={"slide_type": ""} -def debug_print(inp: int) -> int: +def debug_print_and_calc(inp: int) -> int: print("cache miss!") - return inp + return add_3_times_2(inp) # NOTE: `CachedStep` keys the cache on a fingerprint of the step itself plus the # `input_fingerprinter` of the input. `.in_memory` pairs a `dict` cache with the # lenient fingerprinter, which hashes the step structurally -- so non-importable # callables (lambdas, closures) are fine for this single-process cache. -debug_print_step = gtx.otf.workflow.make_step(debug_print) - cached_calc = gtx.otf.workflow.CachedStep.in_memory( - step=debug_print_step.chain(add_3_times_2), + step=debug_print_and_calc, input_fingerprinter=lambda i: str(i), # using ints as their own hash ) @@ -186,28 +168,28 @@ cached_calc(1) ### Example in the Wild ```python -gtx.backend.DEFAULT_PROG_TRANSFORMS.past_lint?? +gtx.backend.DEFAULT_TRANSFORMS.past_lint?? ``` -Though we execute the workflow three times we only get the debug print once, it worked! Btw, hashing is rarely that easy in the wild... +Though we execute the step three times we only get the debug print once, it worked! Btw, hashing is rarely that easy in the wild... -Let's say we want to make our calculation workflow compatible with string input. We can add a conversion step (which only works with strings). +Let's say we want to make our calculation compatible with string input. We can add a conversion step (which only works with strings). ```python editable=true slideshow={"slide_type": ""} -# A plain conversion step turning a string into an int, chained into the -# workflow below and reused by `StrToIntFactory(cached=True)`. +# A plain conversion step turning a string into an int, composed into the +# pipeline below and reused by `StrToIntFactory(cached=True)`. def to_int(inp: str) -> int: assert isinstance(inp, str), "Can not work with 'int'!" # yes, this is horribly contrived return int(inp) -to_int_step = gtx.otf.workflow.make_step(to_int) +def str_calc(inp: str) -> int: + return add_3_times_2(to_int(inp)) -str_calc = to_int_step.chain(add_3_times_2) str_calc("1") ``` @@ -222,8 +204,8 @@ If a step can be useful with different combinations of parameters and wrappers, ```python editable=true slideshow={"slide_type": ""} @dataclasses.dataclass(frozen=True) -class AnyStrToInt(gtx.otf.workflow.ChainableWorkflowMixin[str | int, int]): - inner_step: gtx.otf.workflow.Workflow[str, int] = to_int +class AnyStrToInt: + inner_step: gtx.otf.workflow.Step[str, int] = to_int def __call__(self, inp: str | int) -> int: return self.inner_step(inp) @@ -254,158 +236,47 @@ uncached.inner_step ### Example in the Wild ```python -gtx.ffront.past_passes.linters.LinterFactory?? -``` - - - -## Composition 1: Chaining - -So far we have only seen compsition of workflows by chaining. Any sequence of steps can be represented as a chain. Chains can be built of smaller chains, so a Workflow could be composed and then reused in a bigger workflow. - -However, chains are of limited use in the real world, because it's a pain to access a specific step. This we might want to do in order to: - -- run that step in isolation for debugging or other purposes -- build a new chain with a step swapped out (workflows are immutable). - -Imagine swapping out `sub_third` in `complicated_workflow` below (without copy pasting code): - -```python -complicated_workflow = ( - start_step.chain(first_sub_first.chain(first_sub_second).chain(first_sub_third)) - .chain(second_sub_first.chain(second_sub_second)) - .chain(last) -) -``` - -```mermaid -graph TD -c{{complicated_workflow}} --> 0 --> s{{start_step}} -c --> 1 -->|0| a1{{first_sub_first}} -1 -->|1| a2{{first_sub_second}} -1 -->|2| a3{{first_sub_third}} -c --> 2 -->|0| b1{{second_sub_first}} -2 -->|1| b2{{second_sub_second}} -c --> 3 -->|0| l{{last}} +gtx.ffront.past_passes.linters.linter_factory?? ``` - - - - -## Composition 2: Sequence of Named Steps - -Let's say we want a string processing workflow where the intermediate stages are also of value on their own. We would want to access individual steps, specifically each step as it was configured for this workflow (with parameters, caching, etc identical). - -For this we can use `NamedStepSequence`, giving each step a name, by which we can access it later. For this we have to create a dataclass and derive from `NamedStepSequence`. Each step is then a field of the dataclass, type hinted as a `Workflow`. The resulting workflow will run the steps in order of their apperance in the class body. - -To use the same "complicated workflow" example from above: - -```python -@dataclasses.dataclass(frozen=True) -class FirstSub(gtx.otf.workflow.NamedStepSequence[B, E]): - first: Workflow[B, C] - second: Workflow[C, D] - third: Workflow[D, E] + +## Named pipelines -@dataclasses.dataclass(frozen=True) -class SecondSub(gtx.otf.workflow.NamedStepSequence[E, G]): - first: Workflow[E, F] - second: Workflow[F, G] +Real toolchain pipelines are frozen dataclasses whose fields are the named steps and whose `__call__` spells the composition out explicitly. There is no combinator machinery left: reading `__call__` tells you exactly which steps run and in which order, and the whole composition is statically typed. +There are two of them: -@dataclasses.dataclass(frozen=True) -class ComplicatedWorkflow(gtx.otf.workflow.NamedStepSequence[A, F]): - start_step: Workflow[A, B] - first_sub: Workflow[B, E] - second_sub: Workflow[E, G] - last: Workflow[G, F] - - -complicated_workflow = ComplicatedWorkflow( - start_step=start_step, - first_sub=FirstSub(first=first_sub_first, second=first_sub_second, third=first_sub_third), - second_sub=SecondSub(first=second_sub_first, second=second_sub_second), - last=last, -) -``` +- `gtx.backend.Transforms` — the frontend pipeline, from a program definition in any stage to a `CompilableProgram`. Which of its steps run depends on the type of the input definition (a DSL program definition starts one step earlier than a PAST one), which is why its `__call__` is a `match` rather than a straight line. +- `gtx.backend.CompilePipeline` — the compiled backends' `translation` / `bindings` / `compilation` pipeline. -```mermaid -graph TD - -w{{complicated_workflow: ComplicatedWorkflow}} -->|".start_step"| a{{start_step}} -w -->|".first_sub.first"| b{{first_sub_first}} -w -->|".first_sub.second"| c{{first_sub_second}} -w -->|".first_sub.third"| d{{first_sub_third}} -w -->|".second_sub.first"| e{{second_sub_first}} -w -->|".second_sub_second"| f{{second_sub_second}} -w -->|".last"| g{{last}} -``` +Naming the steps buys two things: they can be run in isolation for debugging, and a variant pipeline is one `dataclasses.replace` away, without touching the code that uses it. ```python editable=true slideshow={"slide_type": ""} -## Here we define how the steps are composed -@dataclasses.dataclass(frozen=True) -class StrProcess(gtx.otf.workflow.NamedStepSequence): - hexify_colors: gtx.otf.workflow.Workflow[str, str] - replace_tabs: gtx.otf.workflow.Workflow[str, str] - - -## Here we define the steps themselves -@dataclasses.dataclass(frozen=True) -class HexifyColors(gtx.otf.workflow.ChainableWorkflowMixin): - color_scheme: dict[str, str] = dataclasses.field( - default_factory=lambda: {"blue": "#0000ff", "green": "#00ff00", "red": "#ff0000"} - ) - - def __call__(self, inp: str) -> str: - result = inp - for color, hexcode in self.color_scheme.items(): - result = result.replace(color, hexcode) - return result - - -def spaces_to_tabs(inp: str) -> str: - return re.sub(r" ", r"\t", inp) +## The steps are plain attributes, so they can be run on their own ... +transforms = gtx.backend.DEFAULT_TRANSFORMS +transforms.past_lint ``` - - -Note that with all this there comes an extra feature: We can easily create variants with different steps, without having to change the code that will use the composed workflow. Even if the calling code calls steps in isolation! - - - ```python editable=true slideshow={"slide_type": ""} -CUSTOM_COLORS = {"blue": "#55aaff", "green": "#00ff00", "red": "#ff0000"} - -proc = StrProcess( - hexify_colors=HexifyColors(color_scheme=CUSTOM_COLORS), replace_tabs=spaces_to_tabs +## ... and a variant is built at composition time, never with per-call flags. +dataclasses.replace( + transforms, past_to_itir=gtx.ffront.past_to_itir.past_to_gtir_factory(cached=False) ) - -proc(""" -p { - background-color: blue; - color: red; -} -""") -``` - -```python editable=true slideshow={"slide_type": ""} -proc.hexify_colors("blue") ``` -`NamedStepSequence`s still work with wrapper steps, parameters and chaining. They can also be nested. So for a complex workflow there would be innumerous possible variants. Therefore expect to often see them paired with factories. +Both pipelines announce every step they run on the `gtx.otf.workflow.stage_hook` event hook, as `(name, artifact)` where `name` is the field name of the step. That is the sanctioned way to observe intermediate stages; setting `GT4PY_DUMP_STAGES=` registers a subscriber that writes each stage artifact to disk. ### Example in the Wild ```python editable=true slideshow={"slide_type": ""} -gtx.backend.DEFAULT_PROG_TRANSFORMS?? +gtx.backend.DEFAULT_TRANSFORMS?? ``` ```python diff --git a/src/gt4py/next/AGENTS.md b/src/gt4py/next/AGENTS.md index 106cc82c3c..70986482f2 100644 --- a/src/gt4py/next/AGENTS.md +++ b/src/gt4py/next/AGENTS.md @@ -29,11 +29,11 @@ field operators / programs (ffront) - `embedded/` — field-level embedded execution: field operators evaluated directly on NumPy / CuPy / JAX-backed `Field`s, without lowering to the IR. Distinct from `iterator/embedded.py`, which runs one level lower. -- `otf/` — on-the-fly compilation toolchain (workflow steps, caching, - argument descriptors). +- `otf/` — on-the-fly compilation toolchain (pipeline steps: `Step`, + `CachedStep`; caching; argument descriptors). - `backend.py` — the toolchain root object `Toolchain` (formerly `Backend`): - `frontend` (definition transforms) + `backend` (compile pipeline) + - allocator. + `frontend` (the `Transforms` definition pipeline) + `backend` (the + `CompilePipeline` translate / bind / 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 3f53fb7354..3261447300 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 Any, Generic, NoReturn from gt4py._core import definitions as core_defs from gt4py.next import custom_layout_allocators as next_allocators @@ -24,7 +24,7 @@ ) from gt4py.next.ffront.past_passes import linters as past_linters from gt4py.next.iterator import ir as itir -from gt4py.next.otf import arguments, artifacts, recipes, stages, toolchain, workflow +from gt4py.next.otf import arguments, artifacts, stages, workflow def jit_to_aot_args( @@ -33,113 +33,206 @@ def jit_to_aot_args( return arguments.CompileTimeArgs.from_concrete(*inp.args, **inp.kwargs) -def adapted_jit_to_aot_args_factory() -> workflow.Workflow[ - stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.JITArgs], - stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.CompileTimeArgs], -]: - """Wrap `jit_to_aot` into a workflow adapter to fit into backend transform workflows.""" - return toolchain.ArgsOnlyAdapter(jit_to_aot_args) - - @dataclasses.dataclass(frozen=True) -class Transforms( - workflow.MultiWorkflow[ - stages.ConcreteProgramDef[stages.IRDefinitionT, stages.ArgsDefinitionT], - stages.CompilableProgram, - ] -): +class Transforms: """ - Modular workflow for transformations with access to intermediates. - - The set and order of transformation steps depends on the input type. - Thus this workflow can be applied to DSL field operator and program definitions, - as well as their AST representations. Even to Iterator IR programs, although in that - case it will be a no-op. - - The input to the workflow as well as each step must be a `CompilableProgram`. The arguments - inside the `CompilableProgram` passed to the whole workflow may be concrete (`JITArgs`) - or compile-time (`CompileTimeArgs`). The individual steps (apart from `.aotify_args`) - require compile-time arguments. Some of the steps can work with an empty `CompileTimeArgs` instance. + Pipeline of the definition transformations, with access to intermediates. + + The set of transformation steps that run depends on the type of the input + definition: DSL field operator and program definitions, their AST + representations, and Iterator IR programs (for which the pipeline is a + no-op) are all supported. + + The input is a `ProgramWithArgs` pair whose arguments may be concrete + (`JITArgs`) or compile-time (`CompileTimeArgs`); `aotify_args` turns the + former into the latter up front. Of the remaining steps only + `field_view_op_to_prog`, `field_view_prog_args_transform` and + `past_to_itir` consume the arguments at all; the others transform the bare + definition. `stage_hook` is emitted after each executed step. + Customization is composition-time: build a variant with + `dataclasses.replace(transforms, past_lint=...)`. """ - aotify_args: workflow.Workflow[ - stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.JITArgs], - stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.CompileTimeArgs], - ] = dataclasses.field(default_factory=adapted_jit_to_aot_args_factory) + aotify_args: workflow.Step[arguments.JITArgs, arguments.CompileTimeArgs] = jit_to_aot_args - func_to_foast: workflow.Workflow[ - ffront_stages.ConcreteDSLFieldOperatorDef, ffront_stages.ConcreteFOASTOperatorDef - ] = dataclasses.field(default_factory=func_to_foast.adapted_func_to_foast_factory) + func_to_foast: workflow.Step[ + ffront_stages.DSLFieldOperatorDef, ffront_stages.FOASTOperatorDef + ] = dataclasses.field(default_factory=func_to_foast.func_to_foast_factory) - func_to_past: workflow.Workflow[ - ffront_stages.ConcreteDSLProgramDef, ffront_stages.ConcretePASTProgramDef - ] = dataclasses.field(default_factory=func_to_past.adapted_func_to_past_factory) + func_to_past: workflow.Step[ffront_stages.DSLProgramDef, ffront_stages.PASTProgramDef] = ( + dataclasses.field(default_factory=func_to_past.func_to_past_factory) + ) - foast_to_itir: workflow.Workflow[ - ffront_stages.ConcreteFOASTOperatorDef, itir.FunctionDefinition - ] = dataclasses.field(default_factory=foast_to_gtir.adapted_foast_to_gtir_factory) + # Not part of the pipeline: `__call__` never runs this step. It is kept as a + # configuration point for downstream code (and for `roundtrip`, its only + # in-repo writer) that needs to pin the FOAST -> ITIR lowering of a toolchain. + foast_to_itir: workflow.Step[ffront_stages.FOASTOperatorDef, itir.FunctionDefinition] = ( + dataclasses.field(default_factory=foast_to_gtir.foast_to_gtir_factory) + ) - field_view_op_to_prog: workflow.Workflow[ + field_view_op_to_prog: workflow.Step[ ffront_stages.ConcreteFOASTOperatorDef, ffront_stages.ConcretePASTProgramDef ] = dataclasses.field(default_factory=foast_to_past.operator_to_program_factory) - past_lint: workflow.Workflow[ - ffront_stages.ConcretePASTProgramDef, ffront_stages.ConcretePASTProgramDef - ] = dataclasses.field(default_factory=past_linters.adapted_linter_factory) + past_lint: workflow.Step[ffront_stages.PASTProgramDef, ffront_stages.PASTProgramDef] = ( + dataclasses.field(default_factory=past_linters.linter_factory) + ) - field_view_prog_args_transform: workflow.Workflow[ + field_view_prog_args_transform: workflow.Step[ ffront_stages.ConcretePASTProgramDef, ffront_stages.ConcretePASTProgramDef ] = dataclasses.field(default_factory=past_process_args.transform_program_args_factory) - past_to_itir: workflow.Workflow[ - ffront_stages.ConcretePASTProgramDef, stages.CompilableProgram - ] = dataclasses.field(default_factory=past_to_itir.past_to_gtir_factory) + past_to_itir: workflow.Step[ffront_stages.ConcretePASTProgramDef, stages.CompilableProgram] = ( + dataclasses.field(default_factory=past_to_itir.past_to_gtir_factory) + ) + + def __call__( + self, inp: stages.ConcreteProgramDef[stages.IRDefinitionT, stages.ArgsDefinitionT] + ) -> stages.CompilableProgram: + """ + Transform any supported program definition into a `CompilableProgram`. + + Which steps run is selected from the type of the input definition; + `stage_hook` is emitted after each of them, carrying the very object + handed on to the next step. + + Args: + inp: A program definition in any supported stage, paired with the + arguments it is compiled for. Concrete (`JITArgs`) arguments + are turned into compile-time ones by `aotify_args` first. - def step_order(self, inp: stages.ConcreteProgramDef) -> list[str]: - steps: list[str] = [] + Returns: + The Iterator IR program paired with its compile-time arguments. An + input that already holds an `itir.Program` is passed through + unchanged. + + Raises: + ValueError: If the input definition is not one of the supported + stages. Raised before any step runs, so nothing is emitted. + """ + # The type guard runs before anything else, so an unsupported definition + # is rejected before any step runs and before any `stage_hook` fires. + if not isinstance( + inp.definition, + ( + ffront_stages.DSLFieldOperatorDef, + ffront_stages.FOASTOperatorDef, + ffront_stages.DSLProgramDef, + ffront_stages.PASTProgramDef, + itir.Program, + ), + ): + raise ValueError("Unexpected input.") + + pair: stages.ConcreteProgramDef[stages.IRDefinitionT, arguments.CompileTimeArgs] if isinstance(inp.args, arguments.JITArgs): - steps.append("aotify_args") - match inp.definition: - case ffront_stages.DSLFieldOperatorDef(): - steps.extend( - [ - "func_to_foast", - "field_view_op_to_prog", - "past_lint", - "field_view_prog_args_transform", - "past_to_itir", - ] - ) - case ffront_stages.FOASTOperatorDef(): - steps.extend( - [ - "field_view_op_to_prog", - "past_lint", - "field_view_prog_args_transform", - "past_to_itir", - ] + pair = workflow.ProgramWithArgs(inp.definition, self.aotify_args(inp.args)) + workflow.stage_hook("aotify_args", pair) + else: + pair = inp + + match pair.definition: + case ffront_stages.DSLFieldOperatorDef() as dsl_operator: + foast_pair = workflow.ProgramWithArgs(self.func_to_foast(dsl_operator), pair.args) + workflow.stage_hook("func_to_foast", foast_pair) + return self._transform_foast_operator(foast_pair) + case ffront_stages.FOASTOperatorDef() as foast_operator: + return self._transform_foast_operator( + workflow.ProgramWithArgs(foast_operator, pair.args) ) - case ffront_stages.DSLProgramDef(): - steps.extend( - [ - "func_to_past", - "past_lint", - "field_view_prog_args_transform", - "past_to_itir", - ] + case ffront_stages.DSLProgramDef() as dsl_program: + past_pair = workflow.ProgramWithArgs(self.func_to_past(dsl_program), pair.args) + workflow.stage_hook("func_to_past", past_pair) + return self._transform_past_program(past_pair) + case ffront_stages.PASTProgramDef() as past_program: + return self._transform_past_program( + workflow.ProgramWithArgs(past_program, pair.args) ) - case ffront_stages.PASTProgramDef(): - steps.extend(["past_lint", "field_view_prog_args_transform", "past_to_itir"]) case itir.Program(): - pass + # Nothing left to transform. This is the same object as `inp` + # unless `aotify_args` had to rebuild the pair above. + return pair case _: - raise ValueError("Unexpected input.") - return steps + # Unreachable while the guard above lists exactly these stages; + # it fails loudly if a stage is added to only one of the two. + raise AssertionError( + f"Unhandled definition type '{type(pair.definition).__name__}'." + ) + + def _transform_foast_operator( + self, operator: ffront_stages.ConcreteFOASTOperatorDef + ) -> stages.CompilableProgram: + past_pair = self.field_view_op_to_prog(operator) + workflow.stage_hook("field_view_op_to_prog", past_pair) + return self._transform_past_program(past_pair) + + def _transform_past_program( + self, program: ffront_stages.ConcretePASTProgramDef + ) -> stages.CompilableProgram: + linted_pair = workflow.ProgramWithArgs(self.past_lint(program.definition), program.args) + workflow.stage_hook("past_lint", linted_pair) + args_pair = self.field_view_prog_args_transform(linted_pair) + workflow.stage_hook("field_view_prog_args_transform", args_pair) + compilable = self.past_to_itir(args_pair) + workflow.stage_hook("past_to_itir", compilable) + return compilable + + def step_order(self, inp: Any) -> NoReturn: + """ + Tombstone of the removed input-dependent step-order hook. + + Raises: + TypeError: Always. Overriding this method used to customize which + steps run; it is now dead code, so it must fail loudly rather + than silently restore the default behaviour. + """ + raise TypeError( + "'Transforms.step_order' was removed: the steps are now selected in" + " 'Transforms.__call__'. Build a variant with 'dataclasses.replace'" + " instead of overriding the step order (see ADR 0027)." + ) DEFAULT_TRANSFORMS: Transforms = Transforms() +@dataclasses.dataclass(frozen=True) +class CompilePipeline: + """ + The standard three-step compile pipeline of the compiled backends. + + Turns a `CompilableProgram` into a loadable compilation artifact through + source-code translation, bindings generation and compilation, emitting + `stage_hook` after each step. Customization is composition-time: build a + variant with `dataclasses.replace(pipeline, translation=...)`. + """ + + translation: stages.TranslationStep + bindings: workflow.Step[artifacts.ProgramSource, artifacts.ExtensionSource] + compilation: workflow.Step[artifacts.ExtensionSource, artifacts.CompilationArtifact] + + def __call__(self, program: stages.CompilableProgram) -> artifacts.CompilationArtifact: + """ + Run translation, bindings generation and compilation, in that order. + + `stage_hook` is emitted after each of the three steps. + + Args: + program: The Iterator IR program paired with the compile-time + arguments it is compiled for. + + Returns: + The loadable artifact produced by the compilation step. + """ + source = self.translation(program) + workflow.stage_hook("translation", source) + extension = self.bindings(source) + workflow.stage_hook("bindings", extension) + artifact = self.compilation(extension) + workflow.stage_hook("compilation", artifact) + return artifact + + @dataclasses.dataclass(frozen=True) class Toolchain(Generic[core_defs.DeviceTypeT]): """ @@ -152,9 +245,9 @@ class Toolchain(Generic[core_defs.DeviceTypeT]): """ name: str - backend: workflow.Workflow[stages.CompilableProgram, artifacts.CompilationArtifact] + backend: workflow.Step[stages.CompilableProgram, artifacts.CompilationArtifact] allocator: next_allocators.FieldBufferAllocatorProtocol[core_defs.DeviceTypeT] - frontend: workflow.Workflow[stages.ConcreteProgramDef, stages.CompilableProgram] + frontend: workflow.Step[stages.ConcreteProgramDef, stages.CompilableProgram] def compile( self, program: stages.IRDefinitionT, compile_time_args: arguments.CompileTimeArgs @@ -187,12 +280,12 @@ def translate( Raises: NotImplementedError: If this toolchain's `backend` is not the - standard `OTFCompileWorkflow` pipeline shape. + standard `CompilePipeline` pipeline shape. """ - if not isinstance(self.backend, recipes.OTFCompileWorkflow): + if not isinstance(self.backend, CompilePipeline): raise NotImplementedError( f"Toolchain '{self.name}' does not support partial runs: 'translate'" - " requires the standard 'OTFCompileWorkflow' compile pipeline" + " requires the standard 'CompilePipeline' compile pipeline" " ('translation' / 'bindings' / 'compilation' steps), but this" f" toolchain's backend is a '{type(self.backend).__name__}'." " Monolithic backends execute in a single step and produce no" diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 81ff528aec..69d36b76cc 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -264,8 +264,7 @@ def __gt_type__(self) -> ts_ffront.ProgramType: # TODO(ricoh): linting should become optional, up to the backend. def __post_init__(self) -> None: - no_args_past = workflow.ProgramWithArgs(self.past_stage, arguments.CompileTimeArgs.empty()) - _ = self._frontend_transforms.past_lint(no_args_past).definition + _ = self._frontend_transforms.past_lint(self.past_stage) @property def __name__(self) -> str: @@ -286,11 +285,7 @@ def definition(self) -> types.FunctionType: @functools.cached_property def past_stage(self) -> ffront_stages.PASTProgramDef: - # backwards compatibility for backends that do not support the full toolchain - no_args_def = workflow.ProgramWithArgs( - self.definition_stage, arguments.CompileTimeArgs.empty() - ) - return self._frontend_transforms.func_to_past(no_args_def).definition + return self._frontend_transforms.func_to_past(self.definition_stage) @property def _frontend_transforms(self) -> next_backend.Transforms: @@ -607,11 +602,7 @@ def __post_init__(self) -> None: @functools.cached_property def foast_stage(self) -> ffront_stages.FOASTOperatorDef: - return self._frontend_transforms.func_to_foast( - workflow.ProgramWithArgs( - definition=self.definition_stage, args=arguments.CompileTimeArgs.empty() - ) - ).definition + return self._frontend_transforms.func_to_foast(self.definition_stage) @property def __name__(self) -> str: diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 10bc754526..94edf1b13b 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -24,11 +24,11 @@ type_specifications as ts_ffront, ) from gt4py.next.ffront.foast_passes import utils as foast_utils -from gt4py.next.ffront.stages import ConcreteFOASTOperatorDef, FOASTOperatorDef +from gt4py.next.ffront.stages import FOASTOperatorDef from gt4py.next.iterator import ir as itir from gt4py.next.iterator.ir_utils import ir_makers as im from gt4py.next.iterator.transforms import constant_folding -from gt4py.next.otf import arguments, toolchain, workflow +from gt4py.next.otf import arguments, workflow from gt4py.next.type_system import type_info, type_specifications as ts, type_translation as tt @@ -43,21 +43,13 @@ def foast_to_gtir(inp: ffront_stages.FOASTOperatorDef) -> itir.FunctionDefinitio def foast_to_gtir_factory( cached: bool = True, -) -> workflow.Workflow[FOASTOperatorDef, itir.FunctionDefinition]: - """Wrap `foast_to_gtir` into a chainable and, optionally, cached workflow step.""" - wf = foast_to_gtir +) -> workflow.Step[FOASTOperatorDef, itir.FunctionDefinition]: + """Return the `foast_to_gtir` step, in-memory cached unless `cached` is unset.""" if cached: - wf = workflow.CachedStep.in_memory( - step=wf, input_fingerprinter=ffront_stages.semantic_fingerprinter + return workflow.CachedStep.in_memory( + step=foast_to_gtir, input_fingerprinter=ffront_stages.semantic_fingerprinter ) - return wf - - -def adapted_foast_to_gtir_factory( - **kwargs: Any, -) -> workflow.Workflow[ConcreteFOASTOperatorDef, itir.FunctionDefinition]: - """Wrap the `foast_to_gtir` workflow step into an adapter to fit into backend transform workflows.""" - return toolchain.StripArgsAdapter(foast_to_gtir_factory(**kwargs)) + return foast_to_gtir def promote_to_list(node_type: ts.TypeSpec) -> Callable[[itir.Expr], itir.Expr]: diff --git a/src/gt4py/next/ffront/foast_to_past.py b/src/gt4py/next/ffront/foast_to_past.py index 488aa2a253..9ce2d7b479 100644 --- a/src/gt4py/next/ffront/foast_to_past.py +++ b/src/gt4py/next/ffront/foast_to_past.py @@ -34,7 +34,7 @@ class ItirShim: """ operator_def: ConcreteFOASTOperatorDef - foast_to_itir: workflow.Workflow[ConcreteFOASTOperatorDef, itir.FunctionDefinition] + foast_to_itir: workflow.Step[ffront_stages.FOASTOperatorDef, itir.FunctionDefinition] def __gt_closure_vars__(self) -> Optional[dict[str, Any]]: return self.operator_def.definition.closure_vars @@ -44,20 +44,20 @@ def __gt_type__(self) -> ts.CallableType: return self.operator_def.definition.foast_node.type def __gt_itir__(self) -> itir.FunctionDefinition: - return self.foast_to_itir(self.operator_def) + return self.foast_to_itir(self.operator_def.definition) # 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.operator_def) + return self.foast_to_itir(self.operator_def.definition) @dataclasses.dataclass(frozen=True) -class OperatorToProgram(workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePASTProgramDef]): +class OperatorToProgram: """ Generate a PAST program definition from a FOAST operator definition. - This workflow step must must be given a FOAST -> ITIR lowering step so that it can place + This step must be given a FOAST -> ITIR lowering step so that it can place valid `ItirShim` instances into the closure variables of the generated program. Example: @@ -69,7 +69,7 @@ class OperatorToProgram(workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePAST ... def copy(a: gtx.Field[[IDim], gtx.float32]) -> gtx.Field[[IDim], gtx.float32]: ... return a - >>> op_to_prog = OperatorToProgram(foast_to_gtir.adapted_foast_to_gtir_factory()) + >>> op_to_prog = OperatorToProgram(foast_to_gtir.foast_to_gtir_factory()) >>> compile_time_args = arguments.CompileTimeArgs( ... args=( @@ -93,7 +93,7 @@ class OperatorToProgram(workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePAST ... ) """ - foast_to_itir: workflow.Workflow[ConcreteFOASTOperatorDef, itir.FunctionDefinition] + foast_to_itir: workflow.Step[ffront_stages.FOASTOperatorDef, itir.FunctionDefinition] def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: # TODO(tehrengruber): implement mechanism to deduce default values @@ -186,13 +186,13 @@ def __call__(self, inp: ConcreteFOASTOperatorDef) -> ConcretePASTProgramDef: def operator_to_program_factory( foast_to_itir_step: Optional[ - workflow.Workflow[ConcreteFOASTOperatorDef, itir.FunctionDefinition] + workflow.Step[ffront_stages.FOASTOperatorDef, itir.FunctionDefinition] ] = None, cached: bool = True, -) -> workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePASTProgramDef]: +) -> workflow.Step[ConcreteFOASTOperatorDef, ConcretePASTProgramDef]: """Optionally wrap `OperatorToProgram` in a `CachedStep`.""" - wf: workflow.Workflow[ConcreteFOASTOperatorDef, ConcretePASTProgramDef] = OperatorToProgram( - foast_to_itir_step or foast_to_gtir.adapted_foast_to_gtir_factory() + wf: workflow.Step[ConcreteFOASTOperatorDef, ConcretePASTProgramDef] = OperatorToProgram( + foast_to_itir_step or foast_to_gtir.foast_to_gtir_factory() ) if cached: wf = workflow.CachedStep.in_memory( diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py index 831710a22f..48612c5fae 100644 --- a/src/gt4py/next/ffront/func_to_foast.py +++ b/src/gt4py/next/ffront/func_to_foast.py @@ -36,13 +36,8 @@ from gt4py.next.ffront.foast_passes.dead_closure_var_elimination import DeadClosureVarElimination from gt4py.next.ffront.foast_passes.iterable_unpack import UnpackedAssignPass from gt4py.next.ffront.foast_passes.type_deduction import FieldOperatorTypeDeduction -from gt4py.next.ffront.stages import ( - ConcreteDSLFieldOperatorDef, - ConcreteFOASTOperatorDef, - DSLFieldOperatorDef, - FOASTOperatorDef, -) -from gt4py.next.otf import toolchain, workflow +from gt4py.next.ffront.stages import DSLFieldOperatorDef, FOASTOperatorDef +from gt4py.next.otf import workflow from gt4py.next.type_system import type_info, type_specifications as ts, type_translation @@ -100,21 +95,13 @@ def func_to_foast(inp: DSLFieldOperatorDef) -> FOASTOperatorDef: def func_to_foast_factory( cached: bool = True, -) -> workflow.Workflow[DSLFieldOperatorDef, FOASTOperatorDef]: - """Wrap `func_to_foast` in a chainable and optionally cached workflow step.""" - wf = workflow.make_step(func_to_foast) +) -> workflow.Step[DSLFieldOperatorDef, FOASTOperatorDef]: + """Return the `func_to_foast` step, in-memory cached unless `cached` is unset.""" if cached: - wf = workflow.CachedStep.in_memory( - step=wf, input_fingerprinter=ffront_stages.semantic_fingerprinter + return workflow.CachedStep.in_memory( + step=func_to_foast, input_fingerprinter=ffront_stages.semantic_fingerprinter ) - return wf - - -def adapted_func_to_foast_factory( - **kwargs: Any, -) -> workflow.Workflow[ConcreteDSLFieldOperatorDef, ConcreteFOASTOperatorDef]: - """Wrap the `func_to_foast step in an adapter to fit into transform toolchains.`""" - return toolchain.DataOnlyAdapter(func_to_foast_factory(**kwargs)) + return func_to_foast class FieldOperatorParser(DialectParser[foast.FunctionDefinition]): diff --git a/src/gt4py/next/ffront/func_to_past.py b/src/gt4py/next/ffront/func_to_past.py index 292a56767b..cad0b5fb9a 100644 --- a/src/gt4py/next/ffront/func_to_past.py +++ b/src/gt4py/next/ffront/func_to_past.py @@ -27,13 +27,8 @@ from gt4py.next.ffront.dialect_parser import DialectParser from gt4py.next.ffront.past_passes.closure_var_type_deduction import ClosureVarTypeDeduction from gt4py.next.ffront.past_passes.type_deduction import ProgramTypeDeduction -from gt4py.next.ffront.stages import ( - ConcreteDSLProgramDef, - ConcretePASTProgramDef, - DSLProgramDef, - PASTProgramDef, -) -from gt4py.next.otf import toolchain, workflow +from gt4py.next.ffront.stages import DSLProgramDef, PASTProgramDef +from gt4py.next.otf import workflow from gt4py.next.type_system import type_specifications as ts, type_translation @@ -72,29 +67,19 @@ def func_to_past(inp: DSLProgramDef) -> PASTProgramDef: ) -def func_to_past_factory(cached: bool = True) -> workflow.Workflow[DSLProgramDef, PASTProgramDef]: +def func_to_past_factory(cached: bool = True) -> workflow.Step[DSLProgramDef, PASTProgramDef]: """ - Wrap `func_to_past` in a chainable and optionally cached workflow step. + Return the `func_to_past` step, in-memory cached unless `cached` is unset. - Caching is switched off by default, because whether recompiling is necessary can only be known after - the closure variables have been collected (which is done in this step). In special cases where it can - be guaranteed that the closure variables do not change, switching caching on should be safe. + Caching is only safe as long as the closure variables of the program definition do not + change: whether recompiling is necessary can be known only after the closure variables + have been collected, which is what this step does. """ - wf = workflow.make_step(func_to_past) if cached: - wf = workflow.CachedStep.in_memory( - wf, input_fingerprinter=ffront_stages.semantic_fingerprinter + return workflow.CachedStep.in_memory( + func_to_past, input_fingerprinter=ffront_stages.semantic_fingerprinter ) - return wf - - -def adapted_func_to_past_factory( - **kwargs: Any, -) -> workflow.Workflow[ConcreteDSLProgramDef, ConcretePASTProgramDef]: - """ - Wrap an adapter around the DSL definition -> PAST definition step to fit into transform toolchains. - """ - return toolchain.DataOnlyAdapter(func_to_past_factory(**kwargs)) + return func_to_past @dataclasses.dataclass(frozen=True, kw_only=True) diff --git a/src/gt4py/next/ffront/past_passes/linters.py b/src/gt4py/next/ffront/past_passes/linters.py index 15fe25e649..5167989b6b 100644 --- a/src/gt4py/next/ffront/past_passes/linters.py +++ b/src/gt4py/next/ffront/past_passes/linters.py @@ -6,11 +6,9 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -from typing import Any - from gt4py.next.ffront import gtcallable, stages as ffront_stages, transform_utils -from gt4py.next.ffront.stages import ConcretePASTProgramDef, PASTProgramDef -from gt4py.next.otf import toolchain, workflow +from gt4py.next.ffront.stages import PASTProgramDef +from gt4py.next.otf import workflow def lint_misnamed_functions( @@ -44,20 +42,15 @@ def lint_undefined_symbols( return inp -def linter_factory( - cached: bool = True, adapter: bool = True -) -> workflow.Workflow[PASTProgramDef, PASTProgramDef]: - wf: workflow.Workflow[PASTProgramDef, PASTProgramDef] = workflow.StepSequence( - (lint_misnamed_functions, lint_undefined_symbols) - ) - if cached: - wf = workflow.CachedStep.in_memory( - step=wf, input_fingerprinter=ffront_stages.semantic_fingerprinter - ) - return wf +def lint_program(inp: ffront_stages.PASTProgramDef) -> ffront_stages.PASTProgramDef: + """Run all PAST program linters (in order).""" + return lint_undefined_symbols(lint_misnamed_functions(inp)) -def adapted_linter_factory( - **kwargs: Any, -) -> workflow.Workflow[ConcretePASTProgramDef, ConcretePASTProgramDef]: - return toolchain.DataOnlyAdapter(linter_factory(**kwargs)) +def linter_factory(cached: bool = True) -> workflow.Step[PASTProgramDef, PASTProgramDef]: + """Return the PAST linting step, in-memory cached unless `cached` is unset.""" + if cached: + return workflow.CachedStep.in_memory( + step=lint_program, input_fingerprinter=ffront_stages.semantic_fingerprinter + ) + return lint_program diff --git a/src/gt4py/next/ffront/past_process_args.py b/src/gt4py/next/ffront/past_process_args.py index ced8cf7dd8..edbeea31d3 100644 --- a/src/gt4py/next/ffront/past_process_args.py +++ b/src/gt4py/next/ffront/past_process_args.py @@ -38,7 +38,7 @@ def transform_program_args( def transform_program_args_factory( cached: bool = True, -) -> workflow.Workflow[ffront_stages.ConcretePASTProgramDef, ffront_stages.ConcretePASTProgramDef]: +) -> workflow.Step[ffront_stages.ConcretePASTProgramDef, ffront_stages.ConcretePASTProgramDef]: wf = transform_program_args if cached: wf = workflow.CachedStep.in_memory( diff --git a/src/gt4py/next/ffront/past_to_itir.py b/src/gt4py/next/ffront/past_to_itir.py index 5182918488..a553b4d0f5 100644 --- a/src/gt4py/next/ffront/past_to_itir.py +++ b/src/gt4py/next/ffront/past_to_itir.py @@ -149,13 +149,13 @@ def past_to_gtir(inp: ConcretePASTProgramDef) -> stages.CompilableProgram: def past_to_gtir_factory( cached: bool = True, -) -> workflow.Workflow[ConcretePASTProgramDef, stages.CompilableProgram]: - wf = workflow.make_step(past_to_gtir) +) -> workflow.Step[ConcretePASTProgramDef, stages.CompilableProgram]: + """Return the `past_to_gtir` step, in-memory cached unless `cached` is unset.""" if cached: - wf = workflow.CachedStep.in_memory( - wf, input_fingerprinter=ffront_stages.semantic_fingerprinter + return workflow.CachedStep.in_memory( + past_to_gtir, input_fingerprinter=ffront_stages.semantic_fingerprinter ) - return wf + return past_to_gtir def _column_axis(all_closure_vars: dict[str, Any]) -> Optional[common.Dimension]: diff --git a/src/gt4py/next/instrumentation/stage_dump.py b/src/gt4py/next/instrumentation/stage_dump.py index 26e51b242e..9219866d7b 100644 --- a/src/gt4py/next/instrumentation/stage_dump.py +++ b/src/gt4py/next/instrumentation/stage_dump.py @@ -25,7 +25,7 @@ environment, before it receives the parent's configuration. A programmatic `config.DUMP_STAGES = ...; enable()` therefore only dumps the frontend stages; `enable` warns about it. -- Monolithic backends (e.g. `roundtrip`) have no `OTFCompileWorkflow` and +- Monolithic backends (e.g. `roundtrip`) have no `CompilePipeline` and execute in a single unnamed step, so only their frontend stages are dumped. """ diff --git a/src/gt4py/next/otf/artifacts.py b/src/gt4py/next/otf/artifacts.py index 0d13d0bebc..1f74915452 100644 --- a/src/gt4py/next/otf/artifacts.py +++ b/src/gt4py/next/otf/artifacts.py @@ -198,7 +198,7 @@ def build(self) -> None: ... @runtime_checkable class CompilationArtifact(Protocol): - """The output of an ``OTFCompileWorkflow``. + """The output of a ``CompilePipeline``. Each backend defines its own concrete artifact dataclass; all share this Protocol. Implementations are frozen dataclasses, picklable, and carry no diff --git a/src/gt4py/next/otf/compilation/compiler.py b/src/gt4py/next/otf/compilation/compiler.py index e97548e626..5cca2caf39 100644 --- a/src/gt4py/next/otf/compilation/compiler.py +++ b/src/gt4py/next/otf/compilation/compiler.py @@ -10,11 +10,11 @@ import dataclasses import pathlib -from typing import Protocol, TypeGuard, TypeVar +from typing import Generic, Protocol, TypeGuard, TypeVar from gt4py._core import definitions as core_defs, locking from gt4py.next import config, fingerprinting -from gt4py.next.otf import artifacts, workflow +from gt4py.next.otf import artifacts from gt4py.next.otf.compilation import build_data, cache, importer @@ -78,16 +78,7 @@ def load(self) -> artifacts.ExecutableProgram: @dataclasses.dataclass(frozen=True) -class CPPCompiler( - workflow.ChainableWorkflowMixin[ - artifacts.ExtensionSource[CPPLikeCodeSpecT, artifacts.PythonCodeSpec], - CPPCompilationArtifact, - ], - workflow.ReplaceEnabledWorkflowMixin[ - artifacts.ExtensionSource[CPPLikeCodeSpecT, artifacts.PythonCodeSpec], - CPPCompilationArtifact, - ], -): +class CPPCompiler(Generic[CPPLikeCodeSpecT]): """Drive a CPP-style build system into a ``CPPCompilationArtifact``. Backends override ``_make_artifact`` to use their own artifact subclass. diff --git a/src/gt4py/next/otf/recipes.py b/src/gt4py/next/otf/recipes.py deleted file mode 100644 index a515fd5a98..0000000000 --- a/src/gt4py/next/otf/recipes.py +++ /dev/null @@ -1,24 +0,0 @@ -# GT4Py - GridTools Framework -# -# Copyright (c) 2014-2024, ETH Zurich -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import dataclasses - -from gt4py.next.otf import artifacts, stages, workflow - - -@dataclasses.dataclass(frozen=True) -class OTFCompileWorkflow( - workflow.NamedStepSequence[stages.CompilableProgram, artifacts.CompilationArtifact] -): - """The typical compiled backend steps composed into a workflow.""" - - translation: stages.TranslationStep - bindings: workflow.Workflow[artifacts.ProgramSource, artifacts.ExtensionSource] - compilation: workflow.Workflow[artifacts.ExtensionSource, artifacts.CompilationArtifact] diff --git a/src/gt4py/next/otf/stages.py b/src/gt4py/next/otf/stages.py index e374929deb..3a4a23a769 100644 --- a/src/gt4py/next/otf/stages.py +++ b/src/gt4py/next/otf/stages.py @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import Protocol, TypeAlias, TypeVar +from typing import TypeAlias, TypeVar from gt4py.next.ffront import stages as ffront_stages from gt4py.next.iterator import ir as itir @@ -43,28 +43,13 @@ CompilableProgram: TypeAlias = ConcreteProgramDef[itir.Program, arguments.CompileTimeArgs] -class TranslationStep( - workflow.ReplaceEnabledWorkflowMixin[CompilableProgram, artifacts.ProgramSource[CodeSpecT]], - Protocol[CodeSpecT], -): - """Translate a GT4Py program to source code (ProgramCall -> ProgramSource).""" +#: Translate a compilable program to source code (CompilableProgram -> ProgramSource). +TranslationStep: TypeAlias = workflow.Step[CompilableProgram, artifacts.ProgramSource[CodeSpecT]] - ... - - -class CompilationStep( - workflow.Workflow[ - artifacts.ExtensionSource[CodeSpecT, TargetCodeSpecT], artifacts.CompilationArtifact - ], - Protocol[CodeSpecT, TargetCodeSpecT], -): - """Run the build system and produce an ``artifacts.CompilationArtifact``. - - Each backend defines its own concrete artifact dataclass (frozen, - picklable, with a ``load`` method); they all satisfy the - ``artifacts.CompilationArtifact`` Protocol structurally. - """ - - def __call__( - self, source: artifacts.ExtensionSource[CodeSpecT, TargetCodeSpecT] - ) -> artifacts.CompilationArtifact: ... +#: Run the build system and produce a loadable `artifacts.CompilationArtifact`. +#: Each backend defines its own concrete artifact dataclass (frozen, picklable, +#: with a `load` method); they all satisfy the `artifacts.CompilationArtifact` +#: protocol structurally. +CompilationStep: TypeAlias = workflow.Step[ + artifacts.ExtensionSource[CodeSpecT, TargetCodeSpecT], artifacts.CompilationArtifact +] diff --git a/src/gt4py/next/otf/toolchain.py b/src/gt4py/next/otf/toolchain.py deleted file mode 100644 index 195f555885..0000000000 --- a/src/gt4py/next/otf/toolchain.py +++ /dev/null @@ -1,60 +0,0 @@ -# GT4Py - GridTools Framework -# -# Copyright (c) 2014-2024, ETH Zurich -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -from __future__ import annotations - -import dataclasses -import typing -from typing import Generic - -from gt4py.next.otf import workflow - - -S = typing.TypeVar("S") -T = typing.TypeVar("T") -DefT = typing.TypeVar("DefT") -ArgsT = typing.TypeVar("ArgsT") - - -@dataclasses.dataclass(frozen=True) -class DataOnlyAdapter( - workflow.ChainableWorkflowMixin, - workflow.ReplaceEnabledWorkflowMixin, - workflow.Workflow[workflow.ProgramWithArgs[S, ArgsT], workflow.ProgramWithArgs[T, ArgsT]], - Generic[ArgsT, S, T], -): - step: workflow.Workflow[S, T] - - def __call__(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.ProgramWithArgs[DefT, S], workflow.ProgramWithArgs[DefT, T]], - Generic[DefT, S, T], -): - step: workflow.Workflow[S, T] - - 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.ProgramWithArgs[S, ArgsT], T], - Generic[ArgsT, S, T], -): - step: workflow.Workflow[S, T] - - 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 5b9e2c8225..241a80d7a8 100644 --- a/src/gt4py/next/otf/workflow.py +++ b/src/gt4py/next/otf/workflow.py @@ -8,14 +8,10 @@ from __future__ import annotations -import abc import dataclasses import functools import pathlib -import typing -from typing import Any, Callable, Generic, Protocol, TypeVar - -from typing_extensions import Self +from typing import Any, Callable, Generic, TypeAlias, TypeVar from gt4py._core import filecache from gt4py.eve.extended_typing import OpaqueMutableMapping @@ -24,18 +20,19 @@ StartT = TypeVar("StartT") -StartT_contra = TypeVar("StartT_contra", contravariant=True) EndT = TypeVar("EndT") -EndT_co = TypeVar("EndT_co", covariant=True) -NewEndT = TypeVar("NewEndT") -IntermediateT = TypeVar("IntermediateT") HashT = TypeVar("HashT") -DataT = TypeVar("DataT") -ArgT = TypeVar("ArgT") DefT = TypeVar("DefT") ArgsT = TypeVar("ArgsT") +#: A pipeline step: any callable taking a single input and returning the next +#: stage. This alias is the whole composition "framework": pipelines are plain +#: frozen dataclasses of steps with an explicit `__call__`, and customization +#: is composition-time via `dataclasses.replace`. +Step: TypeAlias = Callable[[StartT], EndT] + + @dataclasses.dataclass class ProgramWithArgs(Generic[DefT, ArgsT]): """Pair of a program definition in any stage with the arguments it is compiled for. @@ -55,7 +52,8 @@ def stage_hook(name: str, artifact: Any) -> None: """ Event hook emitted when a named pipeline step produces an artifact. - It is emitted by the named step pipelines after each executed step, and by + It is emitted by the named pipelines (`gt4py.next.backend.Transforms` and + `gt4py.next.backend.CompilePipeline`) after each executed step, and by `Toolchain.translate` for the translation step it runs directly. The step names of the standard pipelines (`func_to_past`, `past_to_itir`, `translation`, `bindings`, `compilation`, ...) are therefore observable. @@ -72,202 +70,8 @@ def stage_hook(name: str, artifact: Any) -> None: """ -def make_step(function: Workflow[StartT, EndT]) -> ChainableWorkflowMixin[StartT, EndT]: - """ - Wrap a function in the workflow step convenience wrapper. - - Examples: - --------- - >>> @make_step - ... def times_two(x: int) -> int: - ... return x * 2 - - >>> def stringify(x: int) -> str: - ... return str(x) - - >>> # create a workflow int -> int -> str - >>> times_two.chain(stringify)(3) - '6' - """ - return StepSequence.start(function) - - -@typing.runtime_checkable -class Workflow(Protocol[StartT_contra, EndT_co]): - """ - Workflow protocol. - - Anything that implements this interface can be a workflow of one or more steps. - - callable - - take a single input argument - """ - - def __call__(self, inp: StartT_contra) -> EndT_co: ... - - -class ReplaceEnabledWorkflowMixin(Workflow[StartT_contra, EndT_co], Protocol): - """ - Subworkflow replacement mixin. - - Any subclass MUST be a dataclass for `.replace` to work - """ - - def replace(self, **kwargs: Any) -> Self: - """ - Build a new instance with replaced substeps. - - Raises: - TypeError: If `self` is not a dataclass. - """ - if not dataclasses.is_dataclass(self): - raise TypeError(f"'{self.__class__}' is not a dataclass.") - assert not isinstance(self, type) - return dataclasses.replace(self, **kwargs) - - -class ChainableWorkflowMixin(Workflow[StartT, EndT_co], Protocol[StartT, EndT_co]): - def chain( - self, next_step: Workflow[EndT_co, NewEndT] - ) -> ChainableWorkflowMixin[StartT, NewEndT]: - return make_step(self).chain(next_step) - - -@dataclasses.dataclass(frozen=True) -class NamedStepSequence( - ChainableWorkflowMixin[StartT, EndT], - ReplaceEnabledWorkflowMixin[StartT, EndT], -): - """ - Workflow with linear succession of named steps. - - Examples: - --------- - >>> import dataclasses - - >>> def parse(x: str) -> int: - ... return int(x) - - >>> def plus_half(x: int) -> float: - ... return x + 0.5 - - >>> def stringify(x: float) -> str: - ... return str(x) - - >>> @dataclasses.dataclass(frozen=True) - ... class ParseOpPrint(NamedStepSequence[str, str]): - ... parse: Workflow[str, int] - ... op: Workflow[int, float] - ... print: Workflow[float, str] - - >>> pop = ParseOpPrint(parse=parse, op=plus_half, print=stringify) - - >>> pop.step_order - ['parse', 'op', 'print'] - - >>> pop(73) - '73.5' - - >>> def plus_tenth(x: int) -> float: - ... return x + 0.1 - - - >>> pop.replace(op=plus_tenth)(73) - '73.1' - """ - - def __call__(self, inp: StartT) -> EndT: - """Compose the steps in the order defined in the `.step_order` class attribute.""" - step_result: Any = inp - for step_name in self.step_order: - step_result = getattr(self, step_name)(step_result) - stage_hook(step_name, step_result) - return step_result - - @functools.cached_property - def step_order(self) -> list[str]: - """ - Read step order from class definition by default. - - Only attributes who are type hinted to be of a type that - conforms to the Workflow protocol are considered steps. - """ - step_names: list[str] = [] - annotations = typing.get_type_hints(self.__class__) - for field in dataclasses.fields(self): - field_type = annotations[field.name] - field_type = typing.get_origin(field_type) or field_type - if issubclass(field_type, Workflow): - step_names.append(field.name) - return step_names - - -@dataclasses.dataclass(frozen=True) -class MultiWorkflow( - ChainableWorkflowMixin[StartT, EndT], - ReplaceEnabledWorkflowMixin[StartT, EndT], -): - """A flexible workflow, where the sequence of steps depends on the input type.""" - - def __call__(self, inp: StartT) -> EndT: - step_result: Any = inp - for step_name in self.step_order(inp): - step_result = getattr(self, step_name)(step_result) - stage_hook(step_name, step_result) - return step_result - - @abc.abstractmethod - def step_order(self, inp: StartT) -> list[str]: - pass - - -@dataclasses.dataclass(frozen=True) -class StepSequence( - ChainableWorkflowMixin[StartT, EndT], -): - """ - Composable workflow of single input callables. - - Examples: - --------- - >>> def plus_one(x: int) -> int: - ... return x + 1 - - >>> def plus_half(x: int) -> float: - ... return x + 0.5 - - >>> def stringify(x: float) -> str: - ... return str(x) - - >>> StepSequence.start(plus_one).chain(plus_half).chain(stringify)(73) - '74.5' - - """ - - steps: tuple[Workflow[Any, Any], ...] - - def __call__(self, inp: StartT) -> EndT: - step_result: Any = inp - for step in self.steps: - step_result = step(step_result) - return step_result - - def chain(self, next_step: Workflow[EndT, NewEndT]) -> ChainableWorkflowMixin[StartT, NewEndT]: - return typing.cast( - ChainableWorkflowMixin[StartT, NewEndT], - self.__class__((*self.steps, next_step)), - ) - - @classmethod - def start(cls, first_step: Workflow[StartT, EndT]) -> ChainableWorkflowMixin[StartT, EndT]: - return cls((first_step,)) - - @dataclasses.dataclass(frozen=True) -class CachedStep( - ChainableWorkflowMixin[StartT, EndT], - ReplaceEnabledWorkflowMixin[StartT, EndT], - Generic[StartT, EndT, HashT], -): +class CachedStep(Generic[StartT, EndT, HashT]): """ Cached workflow of single input callables. @@ -312,7 +116,7 @@ class CachedStep( False """ - step: Workflow[StartT, EndT] + step: Step[StartT, EndT] input_fingerprinter: Callable[[StartT], HashT] = dataclasses.field( metadata=utils.gt4py_metadata(fingerprint=False) ) @@ -327,7 +131,7 @@ class CachedStep( @classmethod def in_memory( cls, - step: Workflow[StartT, EndT], + step: Step[StartT, EndT], *, input_fingerprinter: Callable[[StartT], HashT], cache: OpaqueMutableMapping[str, EndT] | None = None, @@ -347,7 +151,7 @@ def in_memory( @classmethod def persistent( cls, - step: Workflow[StartT, EndT], + step: Step[StartT, EndT], *, input_fingerprinter: Callable[[StartT], HashT], cache: OpaqueMutableMapping[str, EndT] | str | pathlib.Path, 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 9e4af6f7b4..ade014c4f1 100644 --- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py +++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py @@ -21,7 +21,7 @@ from gt4py.next.ffront import fbuiltins from gt4py.next.iterator import ir as itir from gt4py.next.iterator.transforms import pass_manager -from gt4py.next.otf import artifacts, stages, workflow +from gt4py.next.otf import artifacts, stages from gt4py.next.otf.binding import cpp_interface, interface from gt4py.next.program_processors.codegens.gtfn.codegen import GTFNCodegen, GTFNIMCodegen from gt4py.next.program_processors.codegens.gtfn.gtfn_ir_to_gtfn_im_ir import GTFN_IM_lowering @@ -37,16 +37,7 @@ def get_param_description(name: str, type_: Any) -> interface.Parameter: @dataclasses.dataclass(frozen=True) -class GTFNTranslationStep( - workflow.ReplaceEnabledWorkflowMixin[ - stages.CompilableProgram, - artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec], - ], - workflow.ChainableWorkflowMixin[ - stages.CompilableProgram, - artifacts.ProgramSource[artifacts.HeaderAndSourceCodeSpec], - ], -): +class GTFNTranslationStep: code_spec: Optional[artifacts.HeaderAndSourceCodeSpec] = None # TODO replace by more general mechanism, see https://github.com/GridTools/gt4py/issues/1135 enable_itir_transforms: bool = True diff --git a/src/gt4py/next/program_processors/runners/dace/program.py b/src/gt4py/next/program_processors/runners/dace/program.py index d529937685..f95558c814 100644 --- a/src/gt4py/next/program_processors/runners/dace/program.py +++ b/src/gt4py/next/program_processors/runners/dace/program.py @@ -18,7 +18,7 @@ from gt4py.next.ffront import decorator from gt4py.next.iterator import ir as itir, transforms as itir_transforms from gt4py.next.iterator.transforms import extractors as extractors -from gt4py.next.otf import arguments, recipes, workflow +from gt4py.next.otf import arguments, workflow from gt4py.next.program_processors.runners.dace import sdfg_args as gtx_dace_args from gt4py.next.program_processors.runners.dace.workflow import translation as gtx_dace_translation from gt4py.next.type_system import type_specifications as ts @@ -206,17 +206,17 @@ def _translation_only_toolchain(backend: gtx_backend.Toolchain) -> gtx_backend.T Raises: NotImplementedError: If `backend` is not shaped like the standard dace - toolchain, i.e. an 'OTFCompileWorkflow' whose translation step is a + toolchain, i.e. a 'CompilePipeline' whose translation step is a 'DaCeTranslator' (optionally wrapped in a 'CachedStep'). """ pipeline = backend.backend - if not isinstance(pipeline, recipes.OTFCompileWorkflow): + if not isinstance(pipeline, gtx_backend.CompilePipeline): raise NotImplementedError( f"Toolchain '{backend.name}' cannot be converted to an SDFG: SDFG" - " conversion requires the standard 'OTFCompileWorkflow' compile" + " conversion requires the standard 'CompilePipeline' compile" f" pipeline, but this toolchain's backend is a '{type(pipeline).__name__}'." ) - translation: workflow.Workflow[Any, Any] = pipeline.translation + translation: workflow.Step[Any, Any] = pipeline.translation if isinstance(translation, workflow.CachedStep): # The persistent translation cache is keyed on the untransformed program, # so it must not see the pre-transformed one handed to this variant. diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py b/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py index 4d825c0c9b..fa0a29f818 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/__init__.py @@ -10,7 +10,7 @@ The main module is `backend`, that exports the backends for CPU and GPU devices. The `backend` module uses `factory` to define a workflow that implements the -`OTFCompileWorkflow` recipe. The different stages are implemeted in separate modules: +`CompilePipeline` recipe. The different stages are implemeted in separate modules: - `translation` for lowering of GTIR to SDFG and applying SDFG transformations - `compilation` for compiling the SDFG into a program - `decoration` to parse the program arguments and pass them to the program call diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py index 434f1db277..6df4258051 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/compilation.py @@ -22,7 +22,7 @@ from gt4py._core import definitions as core_defs, locking from gt4py.next import common, config, fingerprinting -from gt4py.next.otf import artifacts, workflow +from gt4py.next.otf import artifacts from gt4py.next.otf.compilation import cache as gtx_cache from gt4py.next.program_processors.runners.dace.workflow import ( common as gtx_wfdcommon, @@ -195,16 +195,7 @@ def load(self) -> artifacts.ExecutableProgram: @dataclasses.dataclass(frozen=True) -class DaCeCompiler( - workflow.ChainableWorkflowMixin[ - artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec], - DaCeCompilationArtifact, - ], - workflow.ReplaceEnabledWorkflowMixin[ - artifacts.ExtensionSource[artifacts.SDFGCodeSpec, artifacts.PythonCodeSpec], - DaCeCompilationArtifact, - ], -): +class DaCeCompiler: """Run the DaCe build system and produce an on-disk ``DaCeCompilationArtifact``.""" bind_func_name: str 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 eeb5f8d7a0..31e42e8842 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/factory.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/factory.py @@ -14,8 +14,8 @@ import factory from gt4py._core import definitions as core_defs, filecache -from gt4py.next import config, fingerprinting -from gt4py.next.otf import recipes, workflow +from gt4py.next import backend as next_backend, config, fingerprinting +from gt4py.next.otf import workflow from gt4py.next.otf.compilation import cache from gt4py.next.program_processors.runners.dace.workflow import bindings as bindings_step from gt4py.next.program_processors.runners.dace.workflow.compilation import ( @@ -31,7 +31,7 @@ class DaCeWorkflowFactory(factory.Factory): class Meta: - model = recipes.OTFCompileWorkflow + model = next_backend.CompilePipeline class Params: auto_optimize: bool = False 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 1bc70de72d..5926307b58 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py @@ -18,7 +18,7 @@ from gt4py.next import common from gt4py.next.instrumentation import metrics from gt4py.next.iterator import ir as itir, transforms as itir_transforms -from gt4py.next.otf import artifacts, stages, workflow +from gt4py.next.otf import artifacts, stages from gt4py.next.otf.binding import interface from gt4py.next.program_processors.runners.dace import ( lowering as gtx_dace_lowering, @@ -341,16 +341,7 @@ def make_sdfg_call_sync(sdfg: dace.SDFG, gpu: bool) -> None: @dataclasses.dataclass(frozen=True) -class DaCeTranslator( - workflow.ChainableWorkflowMixin[ - stages.CompilableProgram, - artifacts.ProgramSource[artifacts.SDFGCodeSpec], - ], - workflow.ReplaceEnabledWorkflowMixin[ - stages.CompilableProgram, - artifacts.ProgramSource[artifacts.SDFGCodeSpec], - ], -): +class DaCeTranslator: device_type: core_defs.DeviceType auto_optimize: bool auto_optimize_args: dict[str, Any] | None diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py index 26a160a8ce..77325a447c 100644 --- a/src/gt4py/next/program_processors/runners/gtfn.py +++ b/src/gt4py/next/program_processors/runners/gtfn.py @@ -19,7 +19,7 @@ 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, workflow +from gt4py.next.otf import artifacts, workflow from gt4py.next.otf.binding import nanobind from gt4py.next.otf.compilation import cache, compiler from gt4py.next.otf.compilation.build_systems import compiledb @@ -130,7 +130,7 @@ class Meta: class GTFNCompileWorkflowFactory(factory.Factory): class Meta: - model = recipes.OTFCompileWorkflow + model = next_backend.CompilePipeline class Params: device_type: core_defs.DeviceType = core_defs.DeviceType.CPU @@ -162,7 +162,7 @@ class Params: ) translation = factory.LazyAttribute(lambda o: o.bare_translation) - bindings: workflow.Workflow[artifacts.ProgramSource, artifacts.ExtensionSource] = ( + bindings: workflow.Step[artifacts.ProgramSource, artifacts.ExtensionSource] = ( factory.LazyAttribute( # type: ignore[assignment] # factory-boy typing not precise enough lambda o: nanobind.ExtensionGenerator( unstructured_horizontal_has_unit_stride=o.unstructured_horizontal_has_unit_stride diff --git a/src/gt4py/next/program_processors/runners/roundtrip.py b/src/gt4py/next/program_processors/runners/roundtrip.py index 2950d4cd37..c393d51f2a 100644 --- a/src/gt4py/next/program_processors/runners/roundtrip.py +++ b/src/gt4py/next/program_processors/runners/roundtrip.py @@ -27,7 +27,7 @@ ) from gt4py.next.ffront import foast_to_gtir, foast_to_past, past_to_itir from gt4py.next.iterator import ir as itir, transforms as itir_transforms -from gt4py.next.otf import artifacts, stages, workflow +from gt4py.next.otf import artifacts, stages from gt4py.next.type_system import type_info, type_specifications as ts @@ -254,7 +254,7 @@ def decorated_fencil( @dataclasses.dataclass(frozen=True) -class Roundtrip(workflow.Workflow[stages.CompilableProgram, RoundtripArtifact]): +class Roundtrip: debug: Optional[bool] = None use_embedded: bool = True dispatch_backend: Optional[next_backend.Toolchain] = None @@ -317,10 +317,10 @@ def __call__(self, inp: stages.CompilableProgram) -> RoundtripArtifact: allocator=next_allocators.StandardCPUFieldBufferAllocator(), 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), + foast_to_itir=foast_to_gtir.foast_to_gtir_factory(cached=True), field_view_op_to_prog=foast_to_past.operator_to_program_factory( - foast_to_itir_step=foast_to_gtir.adapted_foast_to_gtir_factory() + foast_to_itir_step=foast_to_gtir.foast_to_gtir_factory() ), ), ) -foast_to_gtir_step = foast_to_gtir.adapted_foast_to_gtir_factory(cached=True) +foast_to_gtir_step = foast_to_gtir.foast_to_gtir_factory(cached=True) diff --git a/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_stage_dump.py b/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_stage_dump.py index 1a23bb9639..c022cc303b 100644 --- a/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_stage_dump.py +++ b/tests/next_tests/integration_tests/feature_tests/instrumentation_tests/test_stage_dump.py @@ -29,7 +29,7 @@ #: The step names the standard pipelines announce for a DSL program definition, -#: in pipeline order (`Transforms.step_order`, then the `OTFCompileWorkflow` fields). +#: in pipeline order (`Transforms.__call__`, then the `CompilePipeline` fields). EXPECTED_STAGES = [ "func_to_past", "past_lint", diff --git a/tests/next_tests/integration_tests/feature_tests/otf_tests/test_nanobind_build.py b/tests/next_tests/integration_tests/feature_tests/otf_tests/test_nanobind_build.py index 59bf9dff3d..787e9dc76f 100644 --- a/tests/next_tests/integration_tests/feature_tests/otf_tests/test_nanobind_build.py +++ b/tests/next_tests/integration_tests/feature_tests/otf_tests/test_nanobind_build.py @@ -12,7 +12,6 @@ from gt4py._core import definitions as core_defs from gt4py.next import config -from gt4py.next.otf import workflow from gt4py.next.otf.binding import nanobind from gt4py.next.otf.compilation import compiler from gt4py.next.otf.compilation.build_systems import cmake, compiledb @@ -24,15 +23,14 @@ def test_gtfn_cpp_with_cmake(program_source_with_name): example_program_source = program_source_with_name("gtfn_cpp_with_cmake") - build_the_program = workflow.make_step(nanobind.ExtensionGenerator()).chain( - compiler.CPPCompiler( - cache_lifetime=config.BuildCacheLifetime.SESSION, - builder_factory=cmake.CMakeFactory(), - device_type=core_defs.DeviceType.CPU, - fingerprint_builder_factory=False, - ) + generate_bindings = nanobind.ExtensionGenerator() + compile_extension = compiler.CPPCompiler( + cache_lifetime=config.BuildCacheLifetime.SESSION, + builder_factory=cmake.CMakeFactory(), + device_type=core_defs.DeviceType.CPU, + fingerprint_builder_factory=False, ) - compiled_program = build_the_program(example_program_source).load() + compiled_program = compile_extension(generate_bindings(example_program_source)).load() buf = (np.zeros(shape=(6, 5), dtype=np.float32), (0, 0)) tup = [ (np.zeros(shape=(6, 5), dtype=np.float32), (0, 0)), @@ -45,14 +43,13 @@ def test_gtfn_cpp_with_cmake(program_source_with_name): def test_gtfn_cpp_with_compiledb(program_source_with_name): example_program_source = program_source_with_name("gtfn_cpp_with_compiledb") - build_the_program = workflow.make_step(nanobind.ExtensionGenerator()).chain( - compiler.CPPCompiler( - cache_lifetime=config.BuildCacheLifetime.SESSION, - builder_factory=compiledb.CompiledbFactory(), - device_type=core_defs.DeviceType.CPU, - ) + generate_bindings = nanobind.ExtensionGenerator() + compile_extension = compiler.CPPCompiler( + cache_lifetime=config.BuildCacheLifetime.SESSION, + builder_factory=compiledb.CompiledbFactory(), + device_type=core_defs.DeviceType.CPU, ) - compiled_program = build_the_program(example_program_source).load() + compiled_program = compile_extension(generate_bindings(example_program_source)).load() buf = (np.zeros(shape=(6, 5), dtype=np.float32), (0, 0)) tup = [ (np.zeros(shape=(6, 5), dtype=np.float32), (0, 0)), diff --git a/tests/next_tests/unit_tests/otf_tests/test_workflow.py b/tests/next_tests/unit_tests/otf_tests/test_workflow.py index ebefb6a748..a7a9c29cc0 100644 --- a/tests/next_tests/unit_tests/otf_tests/test_workflow.py +++ b/tests/next_tests/unit_tests/otf_tests/test_workflow.py @@ -16,28 +16,6 @@ from gt4py.next.otf import workflow -@dataclasses.dataclass -class StageOne: - x: int - - -@dataclasses.dataclass -class StageTwo: - pre: StageOne - y: str - - -@dataclasses.dataclass(frozen=True) -class NamedStepsExample(workflow.NamedStepSequence[int, str]): - repeat: workflow.Workflow[int, list[int]] - strify: workflow.Workflow[list[int], str] - - -@dataclasses.dataclass(frozen=True) -class SingleStep(workflow.NamedStepSequence[int, StageTwo]): - step: workflow.Workflow[int, StageTwo] - - @dataclasses.dataclass(frozen=True) class _StepWithValue: v: int @@ -46,40 +24,6 @@ def __call__(self, x: int) -> int: return x + self.v -def step_one(inp: StageOne) -> StageTwo: - return StageTwo(inp, str(inp.x)) - - -def step_two(inp: StageTwo) -> str: - return inp.y - - -def test_single_step(): - step1: workflow.Workflow[StageOne, StageTwo] = workflow.make_step(step_one) - assert step1(StageOne(3)) == step_one(StageOne(3)) - - -def test_chain_step_sequence(): - wf: workflow.Workflow[StageOne, str] = workflow.StepSequence.start(step_one).chain(step_two) - inp = StageOne(5) - assert wf(inp) == step_two(step_one(inp)) - - -def test_named_steps(): - """Test composing named steps""" - - wf = NamedStepsExample(repeat=lambda inp: [inp] * 3, strify=lambda inp: str(inp)) - assert wf.repeat(4) == [4, 4, 4] - assert wf.strify([1, 2, 3]) == "[1, 2, 3]" - assert wf(4) == "[4, 4, 4]" - - -def test_chain_from_named(): - initial_workflow: workflow.Workflow[StageOne, StageTwo] = SingleStep(step=step_one) - full_workflow: workflow.Workflow[StageOne, str] = initial_workflow.chain(step_two) - assert full_workflow(StageOne(42)) == "42" - - def _append_one(inp: list[int]) -> list[int]: return [*inp, 1] @@ -94,13 +38,6 @@ def hashing(inp: list[int]) -> int: assert wf([3, 2, 1]) == [1, 2, 3, 1] -def test_replace(): - """Test replacing a named step.""" - wf = NamedStepsExample(repeat=lambda inp: [inp] * 3, strify=lambda inp: str(inp)) - wf_repl = wf.replace(repeat=lambda inp: [inp] * 4) - assert wf_repl(4) == "[4, 4, 4, 4]" - - def test_fingerprint_is_defined(): assert callable(fingerprinting.strict_fingerprinter) assert isinstance(fingerprinting.strict_fingerprinter("hello"), str) diff --git a/tests/next_tests/unit_tests/test_backend.py b/tests/next_tests/unit_tests/test_backend.py index 58337fb5d4..ae434da784 100644 --- a/tests/next_tests/unit_tests/test_backend.py +++ b/tests/next_tests/unit_tests/test_backend.py @@ -6,18 +6,20 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -"""Backend-free tests of the `Toolchain` partial runs.""" +"""Backend-free tests of the `Transforms` / `CompilePipeline` pipelines and the `Toolchain` partial runs.""" from __future__ import annotations +import dataclasses from typing import Any import pytest import gt4py.next as gtx from gt4py.next import backend as next_backend, custom_layout_allocators as next_allocators +from gt4py.next.ffront import stages as ffront_stages from gt4py.next.iterator import ir as itir -from gt4py.next.otf import arguments, artifacts, recipes, stages, workflow +from gt4py.next.otf import arguments, artifacts, stages, workflow from gt4py.next.otf.binding import interface from gt4py.next.type_system import type_specifications as ts @@ -70,7 +72,7 @@ def fake_translation(inp: stages.CompilableProgram) -> artifacts.ProgramSource: toolchain = next_backend.Toolchain( name="fake", - backend=recipes.OTFCompileWorkflow( + backend=next_backend.CompilePipeline( translation=fake_translation, bindings=_unreachable_step, compilation=_unreachable_step, @@ -96,7 +98,7 @@ def stage_callback(name: str, artifact: Any) -> None: toolchain = next_backend.Toolchain( name="fake", - backend=recipes.OTFCompileWorkflow( + backend=next_backend.CompilePipeline( translation=lambda inp: SENTINEL_SOURCE, bindings=_unreachable_step, compilation=_unreachable_step, @@ -132,10 +134,175 @@ def recording_frontend(inp: Any) -> Any: frontend=recording_frontend, ) - with pytest.raises(NotImplementedError, match="OTFCompileWorkflow") as exc_info: + with pytest.raises(NotImplementedError, match="CompilePipeline") as exc_info: toolchain.translate(copy_prog.definition_stage, compile_time_args) message = str(exc_info.value) assert "monolithic" in message # names the toolchain assert "function" in message # names the offending backend type assert frontend_calls == [] # fails fast, before running the frontend + + +@pytest.fixture +def emitted_stages(): + """Record every `stage_hook` event raised while the fixture is active.""" + emitted: list[tuple[str, Any]] = [] + + def stage_callback(name: str, artifact: Any) -> None: + emitted.append((name, artifact)) + + workflow.stage_hook.register(stage_callback) + try: + yield emitted + finally: + workflow.stage_hook.remove(stage_callback) + + +def _empty_itir_program() -> itir.Program: + return itir.Program(id="noop", function_definitions=[], params=[], declarations=[], body=[]) + + +def test_compile_pipeline_runs_steps_in_order_and_emits_hooks(emitted_stages): + called: list[str] = [] + + def translation(inp: Any) -> str: + called.append("translation") + return "source" + + def bindings(inp: str) -> str: + called.append("bindings") + return "extension" + + def compilation(inp: str) -> str: + called.append("compilation") + return "artifact" + + pipeline = next_backend.CompilePipeline( + translation=translation, bindings=bindings, compilation=compilation + ) + + result = pipeline(_empty_itir_program()) + + assert called == ["translation", "bindings", "compilation"] + assert result == "artifact" + assert emitted_stages == [ + ("translation", "source"), + ("bindings", "extension"), + ("compilation", "artifact"), + ] + + +def test_transforms_emits_stages_for_dsl_program(compile_time_args, emitted_stages): + result = next_backend.DEFAULT_TRANSFORMS( + workflow.ProgramWithArgs(copy_prog.definition_stage, compile_time_args) + ) + + assert [name for name, _ in emitted_stages] == [ + "func_to_past", + "past_lint", + "field_view_prog_args_transform", + "past_to_itir", + ] + assert all(isinstance(artifact, workflow.ProgramWithArgs) for _, artifact in emitted_stages) + assert isinstance(result, workflow.ProgramWithArgs) + assert isinstance(result.definition, itir.Program) + + +def test_transforms_emits_stages_for_dsl_field_operator(compile_time_args, emitted_stages): + result = next_backend.DEFAULT_TRANSFORMS( + workflow.ProgramWithArgs(copy_op.definition_stage, compile_time_args) + ) + + assert [name for name, _ in emitted_stages] == [ + "func_to_foast", + "field_view_op_to_prog", + "past_lint", + "field_view_prog_args_transform", + "past_to_itir", + ] + assert isinstance(result.definition, itir.Program) + + +def test_transforms_emits_stages_for_foast_operator(compile_time_args, emitted_stages): + result = next_backend.DEFAULT_TRANSFORMS( + workflow.ProgramWithArgs(copy_op.foast_stage, compile_time_args) + ) + + assert [name for name, _ in emitted_stages] == [ + "field_view_op_to_prog", + "past_lint", + "field_view_prog_args_transform", + "past_to_itir", + ] + assert isinstance(result.definition, itir.Program) + + +def test_transforms_emits_stages_for_past_program(compile_time_args, emitted_stages): + result = next_backend.DEFAULT_TRANSFORMS( + workflow.ProgramWithArgs(copy_prog.past_stage, compile_time_args) + ) + + assert [name for name, _ in emitted_stages] == [ + "past_lint", + "field_view_prog_args_transform", + "past_to_itir", + ] + assert isinstance(result.definition, itir.Program) + + +def test_transforms_itir_program_is_passthrough(compile_time_args, emitted_stages): + pair = workflow.ProgramWithArgs(_empty_itir_program(), compile_time_args) + + assert next_backend.DEFAULT_TRANSFORMS(pair) is pair + assert emitted_stages == [] + + +def test_transforms_aotify_and_replace(compile_time_args, emitted_stages): + transforms = dataclasses.replace( + next_backend.DEFAULT_TRANSFORMS, aotify_args=lambda jit_args: compile_time_args + ) + pair = workflow.ProgramWithArgs(_empty_itir_program(), arguments.JITArgs(args=(), kwargs={})) + + result = transforms(pair) + + assert [name for name, _ in emitted_stages] == ["aotify_args"] + assert result.args is compile_time_args + assert result.definition is pair.definition + + +def test_transforms_rejects_unexpected_input(compile_time_args, emitted_stages): + with pytest.raises(ValueError, match="Unexpected input"): + next_backend.DEFAULT_TRANSFORMS(workflow.ProgramWithArgs(42, compile_time_args)) + + # The type guard runs before `aotify_args`, so nothing is emitted either. + with pytest.raises(ValueError, match="Unexpected input"): + next_backend.DEFAULT_TRANSFORMS( + workflow.ProgramWithArgs(42, arguments.JITArgs(args=(), kwargs={})) + ) + + assert emitted_stages == [] + + +def test_transforms_steps_are_plain_callables(): + """The pipeline fields hold bare callables, not combinator instances.""" + assert callable(next_backend.DEFAULT_TRANSFORMS.past_lint) + assert isinstance(next_backend.DEFAULT_TRANSFORMS.past_lint, workflow.CachedStep) + # A data-only step consumes the bare stage, not the `ProgramWithArgs` pair. + linted = next_backend.DEFAULT_TRANSFORMS.past_lint(copy_prog.past_stage) + assert isinstance(linted, ffront_stages.PASTProgramDef) + + +def test_transforms_step_order_fails_loudly(): + """Overriding the removed step-order hook must break, not silently no-op.""" + + @dataclasses.dataclass(frozen=True) + class SkipLinting(next_backend.Transforms): + def step_order(self, inp): + order = super().step_order(inp) + return [step for step in order if step != "past_lint"] + + with pytest.raises(TypeError, match="'Transforms.step_order' was removed"): + next_backend.DEFAULT_TRANSFORMS.step_order(None) + + with pytest.raises(TypeError, match="'Transforms.step_order' was removed"): + SkipLinting().step_order(None)