feat[next]: Add memory and disk-based caching to more workflow steps

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
   'cytoolz>=0.12.1',
   'deepdiff>=5.6.0',
   'devtools>=0.6',
+  'diskcache>=5.6.3',
   'factory-boy>=3.3.0',
   'frozendict>=2.3',
   'gridtools-cpp>=2.3.6,==2.*',

src/gt4py/next/config.py

@@ -73,6 +73,9 @@ def env_flag_to_bool(name: str, default: bool) -> bool:
 )
 
 
+GTFN_SOURCE_CACHE_DIR: str = os.environ.get(f"{_PREFIX}_GTFN_SOURCE_CACHE_DIR", "gtfn_cache")
+
+
 #: Whether generated code projects should be kept around between runs.
 #: - SESSION: generated code projects get destroyed when the interpreter shuts down
 #: - PERSISTENT: generated code projects are written to BUILD_CACHE_DIR and persist between runs

src/gt4py/next/ffront/func_to_past.py

@@ -64,7 +64,7 @@ def func_to_past(inp: DSL_PRG) -> PRG:
     )
 
 
-def func_to_past_factory(cached: bool = False) -> workflow.Workflow[DSL_PRG, PRG]:
+def func_to_past_factory(cached: bool = True) -> workflow.Workflow[DSL_PRG, PRG]:
     """
     Wrap `func_to_past` in a chainable and optionally cached workflow step.
 

src/gt4py/next/ffront/stages.py

@@ -100,6 +100,7 @@ def add_content_to_fingerprint(obj: Any, hasher: xtyping.HashlibAlgorithm) -> No
 
 @add_content_to_fingerprint.register(FieldOperatorDefinition)
 @add_content_to_fingerprint.register(FoastOperatorDefinition)
+@add_content_to_fingerprint.register(ProgramDefinition)
 @add_content_to_fingerprint.register(PastProgramDefinition)
 @add_content_to_fingerprint.register(toolchain.CompilableProgram)
 @add_content_to_fingerprint.register(arguments.CompileTimeArgs)
@@ -121,10 +122,14 @@ def add_func_to_fingerprint(obj: types.FunctionType, hasher: xtyping.HashlibAlgo
     for item in sourcedef:
         add_content_to_fingerprint(item, hasher)
 
+    closure_vars = source_utils.get_closure_vars_from_function(obj)
+    for item in sorted(closure_vars.items(), key=lambda x: x[0]):
+        add_content_to_fingerprint(item, hasher)
+
 
 @add_content_to_fingerprint.register
 def add_dict_to_fingerprint(obj: dict, hasher: xtyping.HashlibAlgorithm) -> None:
-    for key, value in obj.items():
+    for key, value in sorted(obj.items()):
         add_content_to_fingerprint(key, hasher)
         add_content_to_fingerprint(value, hasher)
 
@@ -148,4 +153,3 @@ def add_foast_located_node_to_fingerprint(
 ) -> None:
     add_content_to_fingerprint(obj.location, hasher)
     add_content_to_fingerprint(str(obj), hasher)
-    add_content_to_fingerprint(str(obj), hasher)

src/gt4py/next/iterator/ir.py

@@ -208,7 +208,9 @@ class FencilDefinition(Node, ValidatedSymbolTableTrait):
     closures: List[StencilClosure]
     implicit_domain: bool = False
 
-    _NODE_SYMBOLS_: ClassVar[List[Sym]] = [Sym(id=name) for name in BUILTINS]
+    _NODE_SYMBOLS_: ClassVar[List[Sym]] = [
+        Sym(id=name) for name in sorted(BUILTINS)
+    ]  # sorted for serialization stability
 
 
 class Stmt(Node): ...

src/gt4py/next/otf/workflow.py

@@ -12,6 +12,7 @@
 import dataclasses
 import functools
 import typing
+from collections.abc import MutableMapping
 from typing import Any, Callable, Generic, Protocol, TypeVar
 
 from typing_extensions import Self
@@ -253,16 +254,15 @@ class CachedStep(
 
     step: Workflow[StartT, EndT]
     hash_function: Callable[[StartT], HashT] = dataclasses.field(default=hash)  # type: ignore[assignment]
-
-    _cache: dict[HashT, EndT] = dataclasses.field(repr=False, init=False, default_factory=dict)
+    cache: MutableMapping[HashT, EndT] = dataclasses.field(repr=False, default_factory=dict)
 
     def __call__(self, inp: StartT) -> EndT:
         """Run the step only if the input is not cached, else return from cache."""
         hash_ = self.hash_function(inp)
         try:
-            result = self._cache[hash_]
+            result = self.cache[hash_]
         except KeyError:
-            result = self._cache[hash_] = self.step(inp)
+            result = self.cache[hash_] = self.step(inp)
         return result
 
 

src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py

@@ -213,6 +213,7 @@ def generate_stencil_source(
             generated_code = GTFNIMCodegen.apply(gtfn_im_ir)
         else:
             generated_code = GTFNCodegen.apply(gtfn_ir)
+
         return codegen.format_source("cpp", generated_code, style="LLVM")
 
     def __call__(

src/gt4py/next/program_processors/runners/gtfn.py

@@ -8,16 +8,19 @@
 
 import functools
 import warnings
-from typing import Any
+from typing import Any, Optional
 
+import diskcache
 import factory
 import numpy.typing as npt
 
 import gt4py._core.definitions as core_defs
 import gt4py.next.allocators as next_allocators
+from gt4py.eve import utils
 from gt4py.eve.utils import content_hash
 from gt4py.next import backend, common, config
-from gt4py.next.iterator import transforms
+from gt4py.next.common import Connectivity, Dimension
+from gt4py.next.iterator import ir as itir, transforms
 from gt4py.next.otf import arguments, recipes, stages, workflow
 from gt4py.next.otf.binding import nanobind
 from gt4py.next.otf.compilation import compiler
@@ -116,6 +119,37 @@ def compilation_hash(otf_closure: stages.CompilableProgram) -> int:
     )
 
 
+def generate_stencil_source_hash_function(inp: stages.CompilableProgram) -> str:
+    """
+    Generates a unique hash string for a stencil source program representing
+    the program, sorted offset_provider, and column_axis.
+    """
+    program: itir.FencilDefinition | itir.Program = inp.data
+    offset_provider: dict[str, Connectivity | Dimension] = inp.args.offset_provider
+    column_axis: Optional[common.Dimension] = inp.args.column_axis
+
+    program_hash = utils.content_hash(
+        (
+            program,
+            sorted(offset_provider.items(), key=lambda el: el[0]),
+            column_axis,
+        )
+    )
+
+    return program_hash
+
+
+class FileCache(diskcache.Cache):
+    """
+    This class extends `diskcache.Cache` to ensure the cache is closed upon deletion,
+    i.e. it ensures that any resources associated with the cache are properly
+    released when the instance is garbage collected.
+    """
+
+    def __del__(self) -> None:
+        self.close()
+
+
 class GTFNCompileWorkflowFactory(factory.Factory):
     class Meta:
         model = recipes.OTFCompileWorkflow
@@ -129,10 +163,23 @@ class Params:
             lambda o: compiledb.CompiledbFactory(cmake_build_type=o.cmake_build_type)
         )
 
-    translation = factory.SubFactory(
-        gtfn_module.GTFNTranslationStepFactory,
-        device_type=factory.SelfAttribute("..device_type"),
-    )
+        cached_translation = factory.Trait(
+            translation=factory.LazyAttribute(
+                lambda o: workflow.CachedStep(
+                    o.translation_,
+                    hash_function=generate_stencil_source_hash_function,
+                    cache=FileCache(str(config.BUILD_CACHE_DIR / config.GTFN_SOURCE_CACHE_DIR)),
+                )
+            ),
+        )
+
+        translation_ = factory.SubFactory(
+            gtfn_module.GTFNTranslationStepFactory,
+            device_type=factory.SelfAttribute("..device_type"),
+        )
+
+    translation = factory.LazyAttribute(lambda o: o.translation_)
+
     bindings: workflow.Workflow[stages.ProgramSource, stages.CompilableSource] = (
         nanobind.bind_source
     )
@@ -193,7 +240,7 @@ class Params:
     name_postfix="_imperative", otf_workflow__translation__use_imperative_backend=True
 )
 
-run_gtfn_cached = GTFNBackendFactory(cached=True)
+run_gtfn_cached = GTFNBackendFactory(cached=True, otf_workflow__cached_translation=True)
 
 run_gtfn_with_temporaries = GTFNBackendFactory(use_temporaries=True)
 

tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py

@@ -7,9 +7,12 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from functools import reduce
-
+from gt4py.next.otf import languages, stages, workflow
+from gt4py.next.otf.binding import interface
 import numpy as np
 import pytest
+import diskcache
+from gt4py.eve import SymbolName
 
 import gt4py.next as gtx
 from gt4py.next import (
@@ -30,7 +33,7 @@
 from gt4py.next.program_processors.runners import gtfn
 from gt4py.next.type_system import type_specifications as ts
 from gt4py.next import utils as gt_utils
-
+from gt4py.next import config
 from next_tests.integration_tests import cases
 from next_tests.integration_tests.cases import (
     C2E,

tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py

@@ -8,13 +8,29 @@
 
 import numpy as np
 import pytest
+import tempfile
+import pathlib
+import os
+import pickle
+import copy
+import diskcache
+
 
 import gt4py.next as gtx
+import gt4py.next.config
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.ir_utils import ir_makers as im
-from gt4py.next.otf import arguments, languages, stages
+from gt4py.next.otf import arguments, languages, stages, workflow, toolchain
 from gt4py.next.program_processors.codegens.gtfn import gtfn_module
+from gt4py.next.program_processors.runners import gtfn
 from gt4py.next.type_system import type_translation
+from next_tests.integration_tests import cases
+
+from next_tests.integration_tests.cases import cartesian_case
+
+from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
+    exec_alloc_descriptor,
+)
 
 
 @pytest.fixture
@@ -71,3 +87,89 @@ def test_codegen(fencil_example):
     assert module.entry_point.name == fencil.id
     assert any(d.name == "gridtools_cpu" for d in module.library_deps)
     assert module.language is languages.CPP
+
+
+def test_hash_and_diskcache(fencil_example):
+    fencil, parameters = fencil_example
+    compilable_program = stages.CompilableProgram(
+        data=fencil,
+        args=arguments.CompileTimeArgs.from_concrete_no_size(
+            *parameters, **{"offset_provider": {}}
+        ),
+    )
+
+    hash = gtfn.generate_stencil_source_hash_function(compilable_program)
+    path = str(gt4py.next.config.BUILD_CACHE_DIR / gt4py.next.config.GTFN_SOURCE_CACHE_DIR)
+    with diskcache.Cache(path) as cache:
+        cache[hash] = compilable_program
+
+    # check content of cash file
+    with diskcache.Cache(path) as reopened_cache:
+        assert hash in reopened_cache
+        compilable_program_from_cache = reopened_cache[hash]
+        assert compilable_program == compilable_program_from_cache
+        del reopened_cache[hash]  # delete data
+
+    # hash creation is deterministic
+    assert hash == gtfn.generate_stencil_source_hash_function(compilable_program)
+    assert hash == gtfn.generate_stencil_source_hash_function(compilable_program_from_cache)
+
+    # hash is different if program changes
+    altered_program = copy.deepcopy(compilable_program)
+    altered_program.data.id = "example2"
+    assert gtfn.generate_stencil_source_hash_function(
+        compilable_program
+    ) != gtfn.generate_stencil_source_hash_function(altered_program)
+
+
+def test_gtfn_file_cache(fencil_example):
+    fencil, parameters = fencil_example
+    compilable_program = stages.CompilableProgram(
+        data=fencil,
+        args=arguments.CompileTimeArgs.from_concrete_no_size(
+            *parameters, **{"offset_provider": {}}
+        ),
+    )
+    cached_gtfn_translation_step = gtfn.GTFNBackendFactory(
+        gpu=False, cached=True, otf_workflow__cached_translation=True
+    ).executor.step.translation
+
+    bare_gtfn_translation_step = gtfn.GTFNBackendFactory(
+        gpu=False, cached=True, otf_workflow__cached_translation=False
+    ).executor.step.translation
+
+    cached_gtfn_translation_step(
+        compilable_program
+    )  # run cached translation step once to populate cache
+    assert bare_gtfn_translation_step(compilable_program) == cached_gtfn_translation_step(
+        compilable_program
+    )
+
+    cache_key = gtfn.generate_stencil_source_hash_function(compilable_program)
+    assert cache_key in cached_gtfn_translation_step.cache
+    assert (
+        bare_gtfn_translation_step(compilable_program)
+        == cached_gtfn_translation_step.cache[cache_key]
+    )
+
+
+def test_gtfn_file_cache_whole_workflow(cartesian_case):
+    if cartesian_case.backend != gtfn.run_gtfn:
+        pytest.skip("Skipping backend.")
+    cartesian_case.backend = gtfn.GTFNBackendFactory(
+        gpu=False, cached=True, otf_workflow__cached_translation=True
+    )
+
+    @gtx.field_operator
+    def testee(a: cases.IJKField) -> cases.IJKField:
+        field_tuple = (a, a)
+        field_0 = field_tuple[0]
+        field_1 = field_tuple[1]
+        return field_0
+
+    # first call: this generates the cache file
+    cases.verify_with_default_data(cartesian_case, testee, ref=lambda a: a)
+    # clearing the OTFCompileWorkflow cache such that the OTFCompileWorkflow step is executed again
+    object.__setattr__(cartesian_case.backend.executor, "cache", {})
+    # second call: the cache file is used
+    cases.verify_with_default_data(cartesian_case, testee, ref=lambda a: a)