Skip to content

refactor[next]: replace workflow combinators with explicit typed pipelines #2743

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
egparedes wants to merge 1 commit into from
Open

refactor[next]: replace workflow combinators with explicit typed pipelines #2743

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions docs/development/ADRs/next/0027-Toolchain-Naming-and-Pipeline-Simplification.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand All @@ -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
Expand Down
85 changes: 29 additions & 56 deletions docs/user/next/advanced/HackTheToolchain.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,46 +4,31 @@ 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
```

<link href="https://fonts.googleapis.com/icon?family=Material+Icons" rel="stylesheet"><script src="https://spcl.github.io/dace/webclient2/dist/sdfv.js"></script>
<link href="https://spcl.github.io/dace/webclient2/sdfv.css" rel="stylesheet">

## 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
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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)
```
Loading