TSFC: cache the kernel by the repo state, not just the form signature

firedrake/configuration.py

@@ -1,11 +1,12 @@
 """Replaces functionality from the removed `firedrake_configuration` module."""
 
 import os
+import sys
 from pathlib import Path
 
 
 def setup_cache_dirs():
-    root = Path(os.environ.get("VIRTUAL_ENV", Path.home())).joinpath(".cache")
+    root = Path(sys.prefix).joinpath(".cache")
     if "PYOP2_CACHE_DIR" not in os.environ:
         os.environ["PYOP2_CACHE_DIR"] = str(root.joinpath("pyop2"))
     if 'FIREDRAKE_TSFC_KERNEL_CACHE_DIR' not in os.environ:

firedrake/interpolation.py

@@ -22,6 +22,7 @@
 from finat.ufl import TensorElement, VectorElement, MixedElement, FiniteElementBase
 from finat.element_factory import create_element
 
+from tsfc.caching import codegen_key
 from tsfc.driver import compile_expression_dual_evaluation
 from tsfc.ufl_utils import extract_firedrake_constants, hash_expr
 
@@ -1216,7 +1217,7 @@ def get_interp_node_map(source_mesh: MeshGeometry, target_mesh: MeshGeometry, fs
 def _compile_expression_key(comm, expr, ufl_element, domain, parameters) -> tuple[Hashable, ...]:
     """Generate a cache key suitable for :func:`tsfc.compile_expression_dual_evaluation`."""
     dual_arg, operand = expr.argument_slots()
-    return (hash_expr(operand), type(dual_arg), hash(ufl_element), tuplify(parameters))
+    return (hash_expr(operand), type(dual_arg), hash(ufl_element), tuplify(parameters), codegen_key())
 
 
 @memory_and_disk_cache(

firedrake/slate/slac/compiler.py

@@ -8,8 +8,11 @@
 expressions (finite element variational forms written in UFL).
 """
 import time
+from pathlib import Path
 from typing import Hashable
 
+from tsfc.caching import codegen_key, stamp_source_tree
+
 from firedrake.tsfc_interface import SplitKernel, KernelInfo, TSFCKernel
 
 from firedrake.slate.slac.kernel_builder import LocalLoopyKernelBuilder
@@ -64,6 +67,12 @@ def __init__(self, expr, compiler_parameters):
         self.split_kernel = generate_loopy_kernel(expr, compiler_parameters)
 
 
+#: Fingerprint of Slate's own code generator, in `firedrake/slate/slac/`. TSFC's
+#: `codegen_key` does not cover this: Slate lowers its own expressions to loopy
+#: directly, without going through TSFC. Fixed once, at import, to match `codegen_key`.
+_SLAC_CODEGEN_KEY: Hashable = stamp_source_tree(Path(__file__).resolve().parent)
+
+
 def _compile_expression_hashkey(slate_expr, compiler_parameters=None) -> tuple[Hashable, ...]:
     params = copy.deepcopy(parameters)
     if compiler_parameters and "slate_compiler" in compiler_parameters.keys():
@@ -72,7 +81,8 @@ def _compile_expression_hashkey(slate_expr, compiler_parameters=None) -> tuple[H
         params["form_compiler"].update(compiler_parameters)
     # The getattr here is to defer validation to the `compile_expression` call
     # as the test suite checks the correct exceptions are raised on invalid input.
-    return (getattr(slate_expr, "expression_hash", "ERROR") + str(sorted(params.items())))
+    return (getattr(slate_expr, "expression_hash", "ERROR") + str(sorted(params.items()))
+            + str(codegen_key()) + str(_SLAC_CODEGEN_KEY))
 
 
 def _compile_expression_comm(*args, **kwargs):

firedrake/tsfc_interface.py

@@ -14,6 +14,7 @@
 from .ufl_expr import TestFunction, extract_domains
 
 from tsfc import compile_form as original_tsfc_compile_form
+from tsfc.caching import codegen_key
 from tsfc.parameters import PARAMETERS as tsfc_default_parameters
 from tsfc.ufl_utils import extract_firedrake_constants
 from tsfc.kernel_interface.firedrake_loopy import ActiveDomainNumbers
@@ -59,6 +60,7 @@ def tsfc_compile_form_hashkey(form, prefix, parameters, dont_split_numbers, diag
         utils.tuplify(parameters),
         dont_split_numbers,
         diagonal,
+        codegen_key(),
     )
 
 
@@ -158,6 +160,7 @@ def _compile_form_hashkey(form, name, parameters=None, split=True, dont_split=()
         split,
         _make_dont_split_numbers(dont_split, form),
         diagonal,
+        codegen_key(),
     )
 
 

tests/firedrake/regression/test_interpolate.py

@@ -785,3 +785,23 @@ def test_interpolate_indexed():
     I1 = assemble(interpolate(u2, U), mat_type="nest")
     I1_block = assemble(interpolate(TrialFunction(U), U))
     assert np.allclose(I1.petscmat.getNestSubMatrix(0, 1)[:, :], I1_block.petscmat[:, :])
+
+
+def test_compile_expression_key_includes_codegen_key(monkeypatch):
+    """A changed toolchain fingerprint must produce a different expression-kernel
+    cache key, or a kernel compiled before the change will be served after it."""
+    import firedrake.interpolation as interpolation
+
+    mesh = UnitTriangleMesh()
+    V = FunctionSpace(mesh, "CG", 1)
+    expr = Interpolate(TestFunction(V), V)
+
+    monkeypatch.setattr(interpolation, "codegen_key", lambda: "before")
+    key1 = interpolation._compile_expression_key(
+        mesh.comm, expr, V.ufl_element(), mesh, {})
+
+    monkeypatch.setattr(interpolation, "codegen_key", lambda: "after")
+    key2 = interpolation._compile_expression_key(
+        mesh.comm, expr, V.ufl_element(), mesh, {})
+
+    assert key1 != key2

tests/firedrake/test_tsfc_interface.py

@@ -79,6 +79,17 @@ def test_tsfc_different_names(mass):
     assert k1[-1] is not k2[-1]
 
 
+def test_tsfc_recompiles_when_codegen_key_changes(mass, monkeypatch):
+    """A changed toolchain fingerprint must not reuse a kernel compiled under the old one."""
+    monkeypatch.setattr(tsfc_interface, "codegen_key", lambda: "before")
+    k1, = tsfc_interface.compile_form(mass, 'mass_codegen_key')
+
+    monkeypatch.setattr(tsfc_interface, "codegen_key", lambda: "after")
+    k2, = tsfc_interface.compile_form(mass, 'mass_codegen_key')
+
+    assert k1[-1] is not k2[-1]
+
+
 def test_tsfc_cell_kernel(mass):
     k = tsfc_interface.compile_form(mass, 'mass')
     assert len(k) == 1 and 'cell_integral' in loopy.generate_code_v2(k[0][1][0].code).device_code()

tests/tsfc/test_caching.py

@@ -0,0 +1,64 @@
+import importlib
+import os
+import time
+
+import tsfc
+import tsfc.caching
+from tsfc.caching import codegen_key, stamp_source_tree
+
+
+def test_stamp_source_tree_moves_when_a_file_moves(tmp_path):
+    (tmp_path / "a.py").write_text("x = 1\n")
+    stamp1 = stamp_source_tree(tmp_path)
+
+    # A file system timestamp can be coarser than a single Python statement, so
+    # force the mtime forward instead of relying on wall-clock time to pass.
+    a = tmp_path / "a.py"
+    os.utime(a, (a.stat().st_atime, a.stat().st_mtime + 1))
+    stamp2 = stamp_source_tree(tmp_path)
+
+    assert stamp1 != stamp2
+
+
+def test_stamp_source_tree_is_stable(tmp_path):
+    (tmp_path / "a.py").write_text("x = 1\n")
+    (tmp_path / "b.py").write_text("y = 2\n")
+
+    assert stamp_source_tree(tmp_path) == stamp_source_tree(tmp_path)
+
+
+def test_codegen_key_is_stable_between_calls():
+    assert codegen_key() == codegen_key()
+
+
+def test_codegen_key_reflects_toolchain_source_at_import_time():
+    """This is the end-to-end property that the fix exists for.
+
+    A process that imports `tsfc.caching` after a file under `tsfc/` is edited
+    must get a different `codegen_key()`. A process that imported it before the
+    edit must not. `codegen_key()` is fixed at import time. It does not recompute
+    per compile. So this test simulates a fresh process with a reload, rather
+    than calling `codegen_key()` again in place.
+    """
+    edited = tsfc.__file__
+    original_mtime = os.stat(edited).st_mtime
+
+    key_before = codegen_key()
+    try:
+        os.utime(edited, (os.stat(edited).st_atime, original_mtime + 1))
+        importlib.reload(tsfc.caching)
+        key_after = codegen_key()
+        assert key_before != key_after
+    finally:
+        os.utime(edited, (os.stat(edited).st_atime, original_mtime))
+        importlib.reload(tsfc.caching)
+
+
+def test_importing_tsfc_caching_is_cheap():
+    """Guards against the mtime/size approach regressing to a content hash: the
+    one-time cost, paid when this module is imported, must stay small."""
+    t0 = time.perf_counter()
+    importlib.reload(tsfc.caching)
+    elapsed = time.perf_counter() - t0
+
+    assert elapsed < 1.0

tsfc/caching.py

@@ -0,0 +1,88 @@
+"""This module fingerprints the toolchain that generates code.
+
+A kernel cache that adds this to its own key stays correct after an edit to the
+toolchain that produced it. It does not stay keyed only on the form.
+
+Firedrake's kernel caches add :func:`codegen_key` to their own key, alongside the form
+that they key on. An edit to the toolchain then changes the key for every kernel
+that it could have changed.
+
+Every process computes this key once, at import, not on every compile. Two processes
+that see the same toolchain files compute the same key. A later process then reuses
+a kernel that an earlier process cached on disk.
+
+:func:`stamp_source_tree` stats a file rather than reading it. It sees a file's path,
+size, and the time that the file was last written, not the file's content. Reading
+every file's content would catch more edits, but at a cost that this module cannot
+pay on every import.
+"""
+from __future__ import annotations
+
+import hashlib
+import os
+from importlib import import_module, metadata
+from pathlib import Path
+from typing import Hashable
+
+#: Import names of the packages that TSFC's output depends on.
+#: TSFC lowers through FInAT, FIAT and GEM, and generates code through UFL and loopy.
+_TOOLCHAIN = ("tsfc", "finat", "FIAT", "gem", "ufl", "loopy")
+
+
+def stamp_source_tree(root: os.PathLike) -> Hashable:
+    """Fingerprint every ``.py`` file under `root`.
+
+    Stats each file rather than reading it, so the cost is cheap enough to pay at
+    import time.
+
+    Parameters
+    ----------
+    root : os.PathLike
+        Directory to scan, recursively.
+
+    Returns
+    -------
+    Hashable
+        A digest that changes when a file under `root` is added, removed, or has its
+        size or modification time change.
+    """
+    root = Path(root)
+    entries = tuple(sorted(
+        (str(path.relative_to(root)), stat.st_size, stat.st_mtime_ns)
+        for path in root.rglob("*.py")
+        for stat in (path.stat(),)
+    ))
+    return hashlib.sha1(repr(entries).encode()).hexdigest()
+
+
+def _is_editable(dist_name: str) -> bool:
+    dir_info = getattr(metadata.distribution(dist_name).origin, "dir_info", None)
+    return bool(getattr(dir_info, "editable", False))
+
+
+def _package_stamp(name: str, distributions: dict[str, list[str]]) -> Hashable:
+    module = import_module(name)
+    dist_names = distributions.get(name)
+    if dist_names and not _is_editable(dist_names[0]):
+        return metadata.version(dist_names[0])
+    return stamp_source_tree(Path(module.__file__).resolve().parent)
+
+
+# `packages_distributions()` scans every installed distribution's metadata, so it is
+# called once here and shared, rather than once per package in `_TOOLCHAIN`.
+_DISTRIBUTIONS = metadata.packages_distributions()
+_CODEGEN_KEY: Hashable = tuple(_package_stamp(name, _DISTRIBUTIONS) for name in _TOOLCHAIN)
+
+
+def codegen_key() -> Hashable:
+    """Fingerprint the toolchain that TSFC compiles through.
+
+    Two calls compare equal exactly when TSFC, FInAT, FIAT, GEM, UFL and loopy were all
+    unchanged at the time this module was imported.
+
+    Returns
+    -------
+    Hashable
+        A value suitable for adding to a `cachetools` hash key.
+    """
+    return _CODEGEN_KEY