refactor[next]: Split compilation into a actual compilation and loading of the compilation artifact by havogt · Pull Request #2587 · GridTools/gt4py

havogt · 2026-04-27T20:12:45Z

This is a preparation to allow compilation in separate processes for faster jit compilation.

The new workflow: CompilationStep produces a CompilationArtifact which is picklable and has a load() method which returns the original ExecutableProgram.

…finalize-split

Copilot

Pull request overview

Prepares the “next” compilation pipeline for out-of-process / multi-process JIT by splitting “compile” into producing a (mostly) picklable compilation artifact and a separate load() step that materializes the in-process executable callable.

Changes:

Introduce stages.CompilationArtifact protocol and update OTF compilation contracts to return artifacts with load().
Refactor CPP (nanobind/GTFN) and DaCe compilation steps to return artifact dataclasses instead of live callables; move “decoration” into artifact.load().
Update Backend.compile() and tests to materialize executables via artifact.load() and add pickle round-trip contract tests.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py	Updates assertions to reflect `compilation.device_type` living on the compilation step.
tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py	Adds pickle round-trip contract test for DaCe artifact (ensuring live handles don’t serialize).
tests/next_tests/unit_tests/otf_tests/test_compiled_program.py	Adapts a hijacking test to return a dummy artifact with `load()` instead of a raw callable.
tests/next_tests/unit_tests/otf_tests/compilation_tests/test_compiler.py	Adds pickle round-trip contract test for CPP artifact.
tests/next_tests/integration_tests/feature_tests/otf_tests/test_nanobind_build.py	Updates integration flow to call `.load()` on the returned compilation artifact.
src/gt4py/next/program_processors/runners/roundtrip.py	Wraps roundtrip’s in-memory callable in a `RoundtripArtifact` implementing `load()`.
src/gt4py/next/program_processors/runners/gtfn.py	Introduces `GTFNCompilationArtifact`/`GTFNCompiler` to stamp device type and apply argument conversion in `load()`.
src/gt4py/next/program_processors/runners/dace/workflow/factory.py	Removes separate decoration stage wiring (now handled by artifact `load()`).
src/gt4py/next/program_processors/runners/dace/workflow/decoration.py	Adjusts imports/types to avoid cycles after moving decoration into artifact `load()`.
src/gt4py/next/program_processors/runners/dace/workflow/compilation.py	Adds `DaCeCompilationArtifact` and makes DaCe compilation return an on-disk artifact with `load()`.
src/gt4py/next/otf/stages.py	Adds `CompilationArtifact` protocol.
src/gt4py/next/otf/recipes.py	Updates `OTFCompileWorkflow` typing to end at `CompilationArtifact` (removes decoration step).
src/gt4py/next/otf/definitions.py	Updates `CompilationStep` protocol to return `CompilationArtifact`.
src/gt4py/next/otf/compilation/compiler.py	Refactors CPP compilation into `CPPCompiler` + `CPPCompilationArtifact(load)`.
src/gt4py/next/backend.py	Changes backend executor to return artifacts; `compile()` now calls `artifact.load()`.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

havogt · 2026-04-28T06:17:47Z

+            with dace.config.set_temporary("compiler", "use_cache", value=True):
+                sdfg_program = sdfg.compile(validate=False)


No, should be fine because we only reach load() because we went through a build step first, therefore load() is read only.

Copilot

Pull request overview

This PR refactors the “next” OTF compilation workflow so that compilation steps produce a picklable CompilationArtifact with a load() method, enabling future compilation in separate processes while keeping load/materialization process-local.

Changes:

Introduces stages.CompilationArtifact and updates CompilationStep, OTFCompileWorkflow, and Backend.compile() to compile → return artifact → load() into an ExecutableProgram.
Adds/updates concrete artifacts and compilers (CPP/GTFN/DaCe/Roundtrip) so backend-specific loading/wrapping happens in artifact.load() rather than a separate decoration stage.
Updates tests and adds minimal “pickle round-trip” contract tests for the new artifact types.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/next_tests/unit_tests/program_processor_tests/runners_tests/test_gtfn.py	Updates assertions to reflect device type living on the compilation step/artifact.
tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_compilation.py	Adds pickle round-trip contract test for `DaCeCompilationArtifact` (ensures live handle is dropped).
tests/next_tests/unit_tests/otf_tests/test_compiled_program.py	Adjusts test to use a dummy artifact (`load()` returns a no-op callable).
tests/next_tests/unit_tests/otf_tests/compilation_tests/test_compiler.py	Adds pickle round-trip contract test for `CPPCompilationArtifact`.
tests/next_tests/integration_tests/feature_tests/otf_tests/test_nanobind_build.py	Updates integration tests to build → artifact → `.load()` and passes `device_type`.
src/gt4py/next/program_processors/runners/roundtrip.py	Splits roundtrip into source generation + module load, introduces `RoundtripArtifact.load()`.
src/gt4py/next/program_processors/runners/gtfn.py	Introduces `GTFNCompilationArtifact`/`GTFNCompiler` so wrapping happens on `load()`.
src/gt4py/next/program_processors/runners/dace/workflow/factory.py	Removes decoration stage wiring; compilation now returns an artifact.
src/gt4py/next/program_processors/runners/dace/workflow/decoration.py	Adjusts imports/types to avoid cycles after moving decoration to artifact load.
src/gt4py/next/program_processors/runners/dace/workflow/compilation.py	Introduces `DaCeCompilationArtifact` and updates compiler to return it (with optional process-local live-program cache).
src/gt4py/next/otf/stages.py	Adds `CompilationArtifact` protocol.
src/gt4py/next/otf/recipes.py	Updates `OTFCompileWorkflow` typing to end at `CompilationArtifact` (removes decoration stage).
src/gt4py/next/otf/definitions.py	Updates `CompilationStep` protocol to return `CompilationArtifact`.
src/gt4py/next/otf/compilation/compiler.py	Refactors CPP compiler to return `CPPCompilationArtifact` and adds `load()`.
src/gt4py/next/backend.py	Updates `Backend.compile()` to call `artifact.load()` before returning an executable.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+    Protocol. Implementations are frozen dataclasses, picklable, and have no
+    live process-bound state — that is reconstructed by :meth:`load`,
+    which returns a directly-callable :class:`ExecutableProgram` taking
+    gt4py-shaped arguments.


+    def _load_compiled_program(self) -> CompiledDaceProgram:
+        sdfg = dace.SDFG.from_file(str(self.sdfg_dump))
+        sdfg.build_folder = str(self.build_folder)
+
+        with gtx_wfdcommon.dace_context(device_type=self.device_type):
+            with dace.config.set_temporary("compiler", "use_cache", value=True):
+                sdfg_program = sdfg.compile(validate=False)


+    source_code: str
+    entry_point_name: str
+    column_axis: common.Dimension | None
+    dispatch_backend: next_backend.Backend | None
+    debug: bool


havogt added 22 commits April 27, 2026 12:59

refactor[next]: otf split build and load

6a35549

refactorings to combine load+decoration

d251057

BuildArtifact can materialize() itself

b8efa44

remove OTFCompileWorkflow

915ce27

cleanup docstrings

c51ac55

refactor: remove lazy import

9b842e2

apply pre-commit

c022bde

update roundtrip

a107dd6

Merge remote-tracking branch 'upstream/main' into worktree-otf-build-…

fa18067

…finalize-split

separate gtfn from generic compiler

efecdf5

don't serialize/deserialize dace in the same process

68ead69

don't search for sdfg in materialize

4702660

restore generic compiler

77b6235

cleanup

3869d67

more cleanup

9017b33

add tests

5bffec0

cleanup

7f219d7

avoid decoration needs picklable

69c5ae8

cleanup

e05a336

cleanup

7e9fa08

Build->Compile, materialize->load

8f35cb7

cleanup

29508bf

havogt requested a review from Copilot April 27, 2026 20:14

Copilot AI reviewed Apr 27, 2026

View reviewed changes

havogt and others added 3 commits April 28, 2026 08:23

use tmp_path fixture

9f4e776

refactor roundtrip to resepect picklability

9c9234d

Merge branch 'main' into worktree-otf-build-finalize-split

528dcd5

havogt requested a review from Copilot April 28, 2026 07:33

Copilot started reviewing on behalf of havogt April 28, 2026 07:34 View session

Copilot AI reviewed Apr 28, 2026

View reviewed changes

sdfg as part of artifact (because of upcoming change)

edd9e93

msimberg mentioned this pull request May 4, 2026

Add persistent gt4py cache for CSCS CI C2SM/icon4py#1218

Draft

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor[next]: Split compilation into a actual compilation and loading of the compilation artifact#2587

refactor[next]: Split compilation into a actual compilation and loading of the compilation artifact#2587
havogt wants to merge 26 commits intoGridTools:mainfrom
havogt:worktree-otf-build-finalize-split

havogt commented Apr 27, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

havogt Apr 28, 2026

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		with dace.config.set_temporary("compiler", "use_cache", value=True):
		sdfg_program = sdfg.compile(validate=False)

Conversation

havogt commented Apr 27, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

havogt Apr 28, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants