diff --git a/effectful/handlers/jax/_handlers.py b/effectful/handlers/jax/_handlers.py index 9c933af43..3ce8392db 100644 --- a/effectful/handlers/jax/_handlers.py +++ b/effectful/handlers/jax/_handlers.py @@ -14,7 +14,6 @@ from effectful.ops.semantics import apply, evaluate, fvsof, typeof from effectful.ops.syntax import ( ConstructorOperation, - PureInterpretation, Scoped, _BaseTerm, _CustomSingleDispatchCallable, @@ -23,7 +22,7 @@ defop, syntactic_eq, ) -from effectful.ops.types import Expr, NotHandled, Operation, Term +from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term # + An element of an array index expression. IndexElement = None | int | slice | Sequence[int] | EllipsisType | jax.Array @@ -43,7 +42,7 @@ def is_eager_array(x): @functools.cache -def _sizesof_intp() -> tuple[PureInterpretation, Operation]: +def _sizesof_intp() -> tuple[Interpretation, Operation]: """Construct the singleton interpretation used by ``sizesof``.""" from effectful.internals.product_n import argsof, productN @@ -90,17 +89,15 @@ def _getitem(arr, index): return functools.reduce(_merge, itertools.chain(arg_sizes, sizes), {}) return ( - PureInterpretation( - productN( - { - _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, - _getitem_term: { - apply: _retain, - jax_getitem: _retain_getitem, - ConstructorOperation.__apply__: apply.__default_rule__, - }, - } - ) + productN( + { + _sizes: {apply: _apply_sizes, jax_getitem: _getitem}, + _getitem_term: { + apply: _retain, + jax_getitem: _retain_getitem, + ConstructorOperation.__apply__: apply.__default_rule__, + }, + } ), _sizes, ) diff --git a/effectful/handlers/torch.py b/effectful/handlers/torch.py index 62fa8742e..ffd65b31a 100644 --- a/effectful/handlers/torch.py +++ b/effectful/handlers/torch.py @@ -16,14 +16,13 @@ from effectful.ops.semantics import apply, evaluate, fvsof, handler, typeof from effectful.ops.syntax import ( ConstructorOperation, - PureInterpretation, Scoped, _BaseTerm, defdata, defop, syntactic_eq, ) -from effectful.ops.types import Expr, NotHandled, Operation, Term +from effectful.ops.types import Expr, Interpretation, NotHandled, Operation, Term # + An element of a tensor index expression. IndexElement = None | int | slice | Sequence[int] | EllipsisType | torch.Tensor @@ -43,7 +42,7 @@ def _getitem_ellipsis_and_none( @functools.cache -def _sizesof_intp() -> tuple[PureInterpretation, Operation]: +def _sizesof_intp() -> tuple[Interpretation, Operation]: """Construct the singleton interpretation used by ``sizesof``.""" from effectful.internals.product_n import argsof, productN @@ -94,17 +93,15 @@ def _getitem(x, key): return functools.reduce(_merge, itertools.chain(arg_sizes, index_sizes), {}) return ( - PureInterpretation( - productN( - { - sizes: {apply: _apply_sizes, torch_getitem: _getitem}, - getitem_term: { - apply: _retain, - torch_getitem: _retain_getitem, - ConstructorOperation.__apply__: apply.__default_rule__, - }, - } - ) + productN( + { + sizes: {apply: _apply_sizes, torch_getitem: _getitem}, + getitem_term: { + apply: _retain, + torch_getitem: _retain_getitem, + ConstructorOperation.__apply__: apply.__default_rule__, + }, + } ), sizes, ) diff --git a/effectful/internals/runtime.py b/effectful/internals/runtime.py index 9fd07dade..0af21be85 100644 --- a/effectful/internals/runtime.py +++ b/effectful/internals/runtime.py @@ -1,42 +1,103 @@ import contextlib -import dataclasses +import contextvars import functools import inspect import typing -from collections.abc import Callable, Mapping, MutableMapping -from threading import local +from collections.abc import Callable, Mapping +from effectful.internals.weak import ( + AutoIdKeyDictionary, + WeakIdKeyDictionary, + weak_memoize, +) from effectful.ops.types import Interpretation, Operation +type CacheEntry = AutoIdKeyDictionary[Interpretation, typing.Any] +type EvalCache = AutoIdKeyDictionary[typing.Any, CacheEntry] -@dataclasses.dataclass -class Runtime[S, T](local): - interpretation: "Interpretation[S, T]" - cache: MutableMapping[int, typing.Any] | None +EVAL_CACHE: contextvars.ContextVar[EvalCache | None] = contextvars.ContextVar( + "EVAL_CACHE", default=None +) -@functools.lru_cache(maxsize=1) -def get_runtime() -> Runtime: - return Runtime(interpretation={}, cache=None) +INTERPRETATION: contextvars.ContextVar[Interpretation] = contextvars.ContextVar( + "INTERPRETATION", default=typing.cast(Interpretation, {}) +) -def get_interpretation(): - return get_runtime().interpretation +get_interpretation = INTERPRETATION.get @contextlib.contextmanager def interpreter(intp: "Interpretation"): - r = get_runtime() - old_intp, old_cache = r.interpretation, r.cache + token = INTERPRETATION.set(intp) try: - old_intp, r.interpretation = r.interpretation, intp - old_cache, r.cache = ( - r.cache, - old_cache if old_intp is intp and old_cache is not None else {}, - ) yield intp finally: - r.interpretation, r.cache = old_intp, old_cache + INTERPRETATION.reset(token) + + +@contextlib.contextmanager +def cache(store: EvalCache | None = None): + """Memoize evaluation under any interpretation for the duration of this block. + + Installs ``store``, or a fresh cache if none is given, and yields it so that a + later block can reuse it:: + + with cache() as store: + ... + with cache(store): + ... + + :func:`effectful.ops.semantics.evaluate` installs one for the duration of a + call when none is active, so a lone call is memoized internally whether or not + a scope is open. Holding a scope is what shares that work *between* calls. + """ + store = AutoIdKeyDictionary() if store is None else store + token = EVAL_CACHE.set(store) + try: + yield store + finally: + EVAL_CACHE.reset(token) + + +def cache_get( + store: EvalCache, + expr: typing.Any, + intp: "Interpretation", + default: typing.Any = None, +) -> typing.Any: + """Look ``expr`` up under ``intp``, returning ``default`` if it is not cached.""" + inner = store.get(expr) + return default if inner is None else inner.get(intp, default) + + +def cache_put( + store: EvalCache, expr: typing.Any, intp: "Interpretation", value: typing.Any +) -> None: + """Record that ``expr`` evaluates to ``value`` under ``intp``.""" + inner = store.get(expr) + if inner is None: + # Same flavour as the outer store, so the inner map accepts the plain + # ``dict`` interpretations that ``coproduct`` builds. + inner = type(store)() + store[expr] = inner + inner[intp] = value + + +def copy_cache_entries(src, dst) -> None: + """Copy everything cached for ``src`` onto ``dst``. + + :func:`effectful.ops.syntax._build_term` computes a node's type analysis on a + throwaway term and then needs it attributed to the term it actually returns. + """ + store = EVAL_CACHE.get() + if store is None: + return + inner = store.get(src) + if inner is not None: + for intp, value in inner.items(): + cache_put(store, dst, intp, value) @Operation.define @@ -44,6 +105,7 @@ def _get_args() -> tuple[tuple, Mapping]: return ((), {}) +@weak_memoize(cache=WeakIdKeyDictionary()) def _restore_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]: sig = inspect.signature(fn) if not sig.parameters: @@ -57,6 +119,7 @@ def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T: return _cont_wrapper +@weak_memoize(cache=WeakIdKeyDictionary()) def _save_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]: from effectful.ops.semantics import handler @@ -72,6 +135,24 @@ def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T: return _cont_wrapper +@weak_memoize(cache=WeakIdKeyDictionary()) +def _save_then_restore_args[**P, T](fn: Callable[P, T]) -> Callable[P, T]: + # should be equivalent to _restore_args(_save_args(fn)), just fused + from effectful.ops.semantics import handler + + sig = inspect.signature(fn) + if not sig.parameters: + return fn + + @functools.wraps(fn) + def _cont_wrapper(*a: P.args, **k: P.kwargs) -> T: + a, k = (a, k) if a or k else _get_args() + with handler({_get_args: lambda: (a, k)}): + return fn(*a, **k) + + return _cont_wrapper + + def _set_prompt[**P, T]( prompt: Operation[P, T], cont: Callable[P, T], body: Callable[P, T] ) -> Callable[P, T]: @@ -79,7 +160,7 @@ def _set_prompt[**P, T]( @functools.wraps(body) def bound_body(*a: P.args, **k: P.kwargs) -> T: - next_cont = get_interpretation().get(prompt, prompt.__default_rule__) + next_cont = INTERPRETATION.get().get(prompt, prompt.__default_rule__) with handler({prompt: handler({prompt: next_cont})(cont)}): return body(*a, **k) diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index 4fd5c481c..6fac67e35 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -5,15 +5,12 @@ import operator import types import typing -import weakref -from collections.abc import Callable -from typing import Any +from effectful.internals.runtime import cache from effectful.ops.syntax import ( ConstructorOperation, DataclassConstructorOperation, ObjectInterpretation, - PureInterpretation, _BaseTerm, _CustomSingleDispatchCallable, defop, @@ -31,7 +28,7 @@ @defop -def fwd(*args, **kwargs) -> Any: +def fwd(*args, **kwargs) -> typing.Any: """Forward execution to the next most enclosing handler. :func:`fwd` should only be called in the context of a handler. @@ -92,8 +89,8 @@ def coproduct(intp: Interpretation, intp2: Interpretation) -> Interpretation: """ from effectful.internals.runtime import ( _get_args, - _restore_args, _save_args, + _save_then_restore_args, _set_prompt, ) @@ -105,7 +102,7 @@ def coproduct(intp: Interpretation, intp2: Interpretation) -> Interpretation: # calling fwd in the right handler should dispatch to the left handler i1 = intp.get(op) res[op] = ( - _set_prompt(fwd, _restore_args(_save_args(i1)), _save_args(i2)) + _set_prompt(fwd, _save_then_restore_args(i1), _save_args(i2)) if i1 is not None else _save_args(i2) ) @@ -130,9 +127,14 @@ def as_tuple(*args) -> tuple: return tuple(args) +_MISSING: typing.Any = object() + + @_CustomSingleDispatchCallable def evaluate[T]( - __dispatch: Callable[[type], Callable[..., Expr[T]]], + __dispatch: collections.abc.Callable[ + [type], collections.abc.Callable[..., Expr[T]] + ], expr: Expr[T], *, intp: Interpretation | None = None, @@ -155,18 +157,30 @@ def evaluate[T]( 6 """ - from effectful.internals.runtime import get_runtime, interpreter - - with interpreter(intp if intp is not None else get_runtime().interpretation): - cache = get_runtime().cache - assert cache is not None, "Cache should be initialized by interpreter" - key = id(expr) - if key in cache: - ref, result = cache[key] - if ref is expr: - return result + from effectful.internals.runtime import ( + EVAL_CACHE, + cache_get, + cache_put, + get_interpretation, + interpreter, + ) + + with interpreter(intp if intp is not None else get_interpretation()) as current: + store = EVAL_CACHE.get() + if store is None: + # No cache installed. Open one for the duration of this call and start + # over. Without it, an expression that reaches a subexpression along + # several paths re-evaluates it once per path, which is exponential in + # the depth of a DAG. Only the outermost call takes this branch, so the + # extra re-entry is paid once. + with cache(): + return evaluate(expr, intp=current) + + result = cache_get(store, expr, current, _MISSING) + if result is not _MISSING: + return result result = __dispatch(type(expr))(expr) - cache[key] = (expr, result) + cache_put(store, expr, current, result) return result @@ -190,44 +204,11 @@ def _evaluate_dataclass[T](expr: T, **kwargs) -> T: ) -_EVALUATION_CACHE_ATTR = "__effectful_evaluation_cache__" - - -def _term_cache( - expr: Term, -) -> weakref.WeakKeyDictionary[PureInterpretation, Any] | None: - """Return the cache owned by ``expr``, or ``None`` if it cannot store one.""" - try: - return getattr(expr, _EVALUATION_CACHE_ATTR) - except AttributeError: - cache: weakref.WeakKeyDictionary[PureInterpretation, Any] = ( - weakref.WeakKeyDictionary() - ) - try: - setattr(expr, _EVALUATION_CACHE_ATTR, cache) - except (AttributeError, TypeError): - return None - return cache - - @evaluate.register(Term) def _evaluate_term(expr: Term, **kwargs): - from effectful.internals.runtime import get_interpretation - - current_intp = get_interpretation() - if isinstance(current_intp, PureInterpretation): - cache = _term_cache(expr) - if cache is not None and current_intp in cache: - return cache[current_intp] - else: - cache = None - args = tuple(evaluate(arg) for arg in expr.args) kwargs = {k: evaluate(v) for k, v in expr.kwargs.items()} - result = expr.op(*args, **kwargs) - if cache is not None: - cache[current_intp] = result - return result + return expr.op(*args, **kwargs) @evaluate.register(Operation) @@ -338,7 +319,7 @@ def _dataclass_constructor_apply(self, op, *args, **kwargs): return Box(op.__type_rule__(*args, **kwargs)) -_TYPEOF_INTP = PureInterpretation(_TypeofIntp()) +_TYPEOF_INTP = _TypeofIntp() def _typeof(term: Expr): @@ -377,7 +358,7 @@ def typeof[T](term: Expr[T]) -> type[T]: @functools.cache -def _fvsof_intp() -> tuple[PureInterpretation, Operation]: +def _fvsof_intp() -> tuple[Interpretation, Operation]: """Construct the singleton interpretation used by ``fvsof``.""" from effectful.internals.product_n import argsof, productN @@ -432,19 +413,17 @@ def _apply_fvs(op, *args, **kwargs): _fvsof_binders = defop(object, name="fvsof_binders") return ( - PureInterpretation( - productN( - { - _fvsof_fvs: { - apply: _apply_fvs, - ConstructorOperation.__apply__: _apply_passthrough_fvs, - }, - _fvsof_binders: { - apply: _apply_binders, - ConstructorOperation.__apply__: _apply_collection_binders, - }, - } - ) + productN( + { + _fvsof_fvs: { + apply: _apply_fvs, + ConstructorOperation.__apply__: _apply_passthrough_fvs, + }, + _fvsof_binders: { + apply: _apply_binders, + ConstructorOperation.__apply__: _apply_collection_binders, + }, + } ), _fvsof_fvs, ) diff --git a/effectful/ops/syntax.py b/effectful/ops/syntax.py index d73546028..5b372f0f1 100644 --- a/effectful/ops/syntax.py +++ b/effectful/ops/syntax.py @@ -11,7 +11,6 @@ from effectful.ops.types import ( Annotation, Expr, - Interpretation, NotHandled, Operation, Term, @@ -445,12 +444,17 @@ def _build_term[T]( type from the types of its arguments and dispatches on that type to pick a constructor. """ - from effectful.ops.semantics import _simple_type, _typeof + from effectful.internals.runtime import copy_cache_entries + from effectful.ops.semantics import typeof - typed_args = tuple(_typeof(arg) for arg in args) - typed_kwargs = {k: _typeof(v) for k, v in kwargs.items()} - dispatch_type = _simple_type(op.__type_rule__(*typed_args, **typed_kwargs)) - return __dispatch(dispatch_type)(dispatch_type, op, *args, **kwargs) + # Compute the type on a throwaway node so that the analysis is cached against + # something, then move that cache onto the node actually returned: a parent's + # later typeof on this child is then a hit rather than a fresh traversal. + raw_term: Expr[T] = _BaseTerm(op, *args, **kwargs) + dispatch_type: type = typeof(raw_term) + result = __dispatch(dispatch_type)(dispatch_type, op, *args, **kwargs) + copy_cache_entries(raw_term, result) + return result @_CustomSingleDispatchCallable @@ -980,25 +984,6 @@ def __getitem__(self, item: Operation[..., T]) -> Callable[..., V]: return self.implementations[item].__get__(self, type(self)) -class PureInterpretation[T, V](Mapping[Operation[..., T], Callable[..., V]]): - """Mark an interpretation as pure so its evaluation results can be cached.""" - - def __init__(self, intp: Interpretation[T, V]): - self.intp = intp - - def __iter__(self): - return iter(self.intp) - - def __len__(self): - return len(self.intp) - - def __getitem__(self, item: Operation[..., T]) -> Callable[..., V]: - return self.intp[item] - - __hash__ = object.__hash__ - __eq__ = object.__eq__ - - class _ImplementedOperation[**P, **Q, T, V]: impl: Callable[Q, V] | None op: Operation[P, T] diff --git a/effectful/ops/types.py b/effectful/ops/types.py index 688141229..6d9b99c19 100644 --- a/effectful/ops/types.py +++ b/effectful/ops/types.py @@ -80,19 +80,18 @@ class Operation[**Q, V]: def __init__(self, default: Callable[Q, V], name: str | None = None): functools.update_wrapper(self, default) - # update_wrapper copies the wrapped callable's __dict__. Do not retain a - # signature cached by another Operation, since `default` may now be a - # bound version of that operation. - self.__dict__.pop("_signature", None) + # update_wrapper copies the wrapped callable's __dict__, + # clear any cached_property values that may have been copied + for klass in type(self).__mro__: + for var, val in vars(klass).items(): + if isinstance(val, functools.cached_property): + self.__dict__.pop(var, None) + self.__default__ = default self.__name__ = name or default.__name__ - @property - def __signature__(self): - return self._signature - @functools.cached_property - def _signature(self): + def __signature__(self): # Resolve forward references (e.g. -> "MyClass") using the # default function's __globals__. This handles module-level # forward refs; local forward refs will raise NameError. @@ -411,6 +410,12 @@ def __type_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> type[V]: subst_type = substitute(return_anno, unify(self.__signature__, bound_sig)) return typing.cast(type[V], subst_type) + @functools.cached_property + def _signature_with_scopes(self): + from effectful.ops.syntax import Scoped + + return Scoped.infer_annotations(self.__signature__) + @typing.final def __fvs_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> inspect.BoundArguments: """Returns the sets of variables that appear free in each argument and @@ -425,7 +430,7 @@ def __fvs_rule__(self, *args: Q.args, **kwargs: Q.kwargs) -> inspect.BoundArgume """ from effectful.ops.syntax import Scoped - sig = Scoped.infer_annotations(self.__signature__) + sig = self._signature_with_scopes bound_sig = sig.bind(*args, **kwargs) bound_sig.apply_defaults() @@ -504,8 +509,14 @@ def _instance_op(instance, *args, **kwargs): else: return self + @functools.cached_property + def _default_rule_with_args(self): + from effectful.internals.runtime import _restore_args + + return _restore_args(self.__default_rule__) + def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: - from effectful.internals.runtime import _restore_args, get_interpretation + from effectful.internals.runtime import get_interpretation from effectful.ops.semantics import fwd, handler intp = get_interpretation() @@ -514,9 +525,7 @@ def __call__(self, *args: Q.args, **kwargs: Q.kwargs) -> V: if self_handler is not None: # ensure that fwd is bound to the default rule. if this handler has # a bound fwd, it will override this binding - fwd_intp = typing.cast( - Interpretation, {fwd: _restore_args(self.__default_rule__)} - ) + fwd_intp = typing.cast(Interpretation, {fwd: self._default_rule_with_args}) with handler(fwd_intp): return self_handler(*args, **kwargs) elif args and isinstance(args[0], Operation) and self is args[0].__apply__: diff --git a/tests/test_ops_semantics.py b/tests/test_ops_semantics.py index 04e595f9e..f950949cf 100644 --- a/tests/test_ops_semantics.py +++ b/tests/test_ops_semantics.py @@ -1,6 +1,5 @@ import contextlib import dataclasses -import functools import itertools import logging from collections.abc import Callable, Mapping @@ -19,7 +18,6 @@ ) from effectful.ops.syntax import ( ObjectInterpretation, - PureInterpretation, Scoped, deffn, defop, @@ -473,7 +471,7 @@ def Nested(*args, **kwargs): def test_memoized_interpretation(): - from effectful.internals.runtime import interpreter + from effectful.internals.runtime import cache, interpreter @defop def node(x: object) -> object: @@ -490,38 +488,40 @@ def _(self, op, *args, **kwargs): self.calls += 1 return (op.__name__, args, kwargs) - intp_impl = Intp() - intp = PureInterpretation(intp_impl) + intp = Intp() expected = ("node", (("node", (1,), {}),), {}) - assert interpreter(intp)(evaluate)(term) == expected - assert intp_impl.calls == 2 + # ``evaluate`` installs a cache for the duration of a call when none is + # active, so results are shared between separate calls only while a scope + # holds one open. + with cache(): + assert interpreter(intp)(evaluate)(term) == expected + assert intp.calls == 2 - # The root cache is checked before its children are traversed, including - # when evaluation is expressed directly through a handler. - with interpreter(intp): - assert evaluate(term) == expected - assert intp_impl.calls == 2 + # The root cache is checked before its children are traversed, including + # when evaluation is expressed directly through a handler. + with interpreter(intp): + assert evaluate(term) == expected + assert intp.calls == 2 - # Child results are cached independently and can be reused directly. - assert interpreter(intp)(evaluate)(term.args[0]) == expected[1][0] - assert intp_impl.calls == 2 + # Child results are cached independently and can be reused directly. + assert interpreter(intp)(evaluate)(term.args[0]) == expected[1][0] + assert intp.calls == 2 - # A composition has a distinct identity even when its added handler is not - # used while evaluating this term. - combined_intp = coproduct(intp, {plus_1: lambda x: x}) - assert interpreter(combined_intp)(evaluate)(term) == expected - assert intp_impl.calls == 4 + # A composition has a distinct identity even when its added handler is not + # used while evaluating this term. + combined_intp = coproduct(intp, {plus_1: lambda x: x}) + assert interpreter(combined_intp)(evaluate)(term) == expected + assert intp.calls == 4 - # A separate pure interpretation has its own cache namespace. - other_intp_impl = Intp() - other_intp = PureInterpretation(other_intp_impl) - assert interpreter(other_intp)(evaluate)(term) == expected - assert other_intp_impl.calls == 2 + # A separate interpretation has its own cache namespace. + other_intp = Intp() + assert interpreter(other_intp)(evaluate)(term) == expected + assert other_intp.calls == 2 def test_memoized_interpretation_does_not_cache_failures(): - from effectful.internals.runtime import interpreter + from effectful.internals.runtime import cache, interpreter @defop def node() -> object: @@ -540,15 +540,17 @@ def _(self, op, *args, **kwargs): raise ValueError("failed analysis") return "success" - intp_impl = Intp() - intp = PureInterpretation(intp_impl) - with pytest.raises(ValueError, match="failed analysis"): - interpreter(intp)(evaluate)(term) + intp = Intp() + with cache(): + with pytest.raises(ValueError, match="failed analysis"): + interpreter(intp)(evaluate)(term) - assert interpreter(intp)(evaluate)(term) == "success" - assert intp_impl.calls == 2 - assert interpreter(intp)(evaluate)(term) == "success" - assert intp_impl.calls == 2 + # The failure left nothing behind, so the retry recomputes; the result of + # that retry is what gets cached and reused. + assert interpreter(intp)(evaluate)(term) == "success" + assert intp.calls == 2 + assert interpreter(intp)(evaluate)(term) == "success" + assert intp.calls == 2 @pytest.mark.parametrize( @@ -773,6 +775,8 @@ def test_defdata_large(benchmark): """Test defdata with large nested operations that form a binary tree of arbitrary size.""" import random + from effectful.internals.runtime import cache + @defop def f[T, A, B]( v: Annotated[Operation[[], int], Scoped[A]], @@ -807,7 +811,14 @@ def build_tree(depth: int) -> Any: return f(defop(int), left, right) # Test a very large tree (depth 8 = 255 leaf nodes) - benchmark(functools.partial(build_tree, 7)) + def run(): + # A scope per iteration rather than one around ``benchmark``: each round + # builds fresh objects, so a shared cache would only accumulate entries + # that can never be hit. + with cache(): + return build_tree(7) + + benchmark(run) def test_evaluate_deep(): @@ -899,8 +910,16 @@ def get_mixed() -> Literal[1, "a"]: typeof(get_mixed()) +@pytest.mark.timeout(20) def test_evaluate_dag_no_exponential_blowup(): - """A DAG of nested tuples sharing the same Term is O(n), not O(2^n).""" + """A DAG of nested tuples sharing the same Term is O(n), not O(2^n). + + Bounded by a timeout because the regression this guards against does not + fail, it hangs: losing memoization turns the depth-20 DAG below into 2**20 + evaluations. + """ + from effectful.internals.runtime import cache + call_count = 0 @defop @@ -919,11 +938,15 @@ def counted_handler(): for _ in range(depth): node = (node, node) - call_count = 0 - with handler({counted: counted_handler}): - result = evaluate(node) + # One cache scope spanning both the evaluation and the term construction + # below. Without it each would install its own, so nothing computed by the + # first would be available to the second. + with cache(): + call_count = 0 + with handler({counted: counted_handler}): + result = evaluate(node) - deffn(node, counted)(0) + deffn(node, counted)(0) # The handler should only be called once (the shared Term) assert call_count == 1 diff --git a/tests/test_ops_syntax.py b/tests/test_ops_syntax.py index a5fdb749c..38d1ab307 100644 --- a/tests/test_ops_syntax.py +++ b/tests/test_ops_syntax.py @@ -1418,6 +1418,7 @@ def test_defop_forward_ref_mutual_recursion(): def test_bench_term_construction(benchmark): """Benchmark polymorphic type checking during term construction.""" + from effectful.internals.runtime import cache @defop def _benchmark_identity[T](value: T) -> T: @@ -1455,7 +1456,14 @@ def _make_benchmark_term(size: int) -> Term[int]: assert isinstance(value, Term) return value - result = benchmark(_make_benchmark_term, 25) + def run(): + # A scope per iteration rather than one around ``benchmark``: each round + # builds fresh objects, so a shared cache would only accumulate entries + # that can never be hit. + with cache(): + return _make_benchmark_term(25) + + result = benchmark(run) assert isinstance(result, Term) @@ -1471,6 +1479,7 @@ def test_bench_nested_binder_construction(benchmark): reaches this path -- it is fast even when nested construction is exponential in depth. """ + from effectful.internals.runtime import cache @defop def _benchmark_let[S, T, A]( @@ -1497,5 +1506,186 @@ def _make_nested_term(depth: int) -> Term[int]: assert isinstance(body, Term) return body - result = benchmark(_make_nested_term, 10) + def run(): + with cache(): + return _make_nested_term(10) + + result = benchmark(run) + assert isinstance(result, Term) + + +# --------------------------------------------------------------------------- +# Operation calls whose arguments are large dataclasses. +# +# ``_build_term`` computes a node's dispatch type with ``typeof``, which is a +# full recursive traversal of every argument, so an operation call costs O(size +# of its arguments). These model a robotics workload -- a component tree built +# up one part at a time -- where that traversal dominated the run time. + + +@dataclasses.dataclass(frozen=True) +class _Pose: + """Leaf payload, mirroring a rigid-body transform.""" + + x: float = 0.0 + y: float = 0.0 + z: float = 0.0 + + +@dataclasses.dataclass(frozen=True) +class _Part: + """One child of an assembly, holding an unsolved pose.""" + + name: str + pose: _Pose + + +@dataclasses.dataclass(frozen=True) +class _Assembly: + """Container that accumulates children, like a robot component.""" + + parts: tuple[_Part, ...] = () + + def add(self, part: "_Part") -> "_Assembly": + return _Assembly(self.parts + (part,)) + + +@defop +def _link(parent: _Assembly, child: _Part, delta: _Pose) -> None: + """Unhandled, so calling it builds a term via ``defdata`` -> ``_build_term``.""" + raise NotHandled + + +def _free_pose() -> _Pose: + """A free variable of type ``_Pose``, standing for a component's unsolved pose.""" + return defop(_Pose)() + + +def _make_assembly(k: int) -> _Assembly: + """An assembly of ``k`` parts, each with a distinct free pose.""" + a = _Assembly() + for i in range(k): + a = a.add(_Part(f"p{i}", _free_pose())) + return a + + +def test_typeof_dataclass_is_cached_within_a_scope(monkeypatch): + """Repeating ``typeof`` on the same dataclass does no work inside one scope. + + Counted rather than timed: the property is that the second traversal is + free, which a call count states exactly where a wall-clock threshold would + only approximate it and would be flaky besides. + + Before evaluation was memoized this was the central problem -- ``Term`` + arguments were cached but plain dataclasses were re-walked on every call, so + an operation taking a large dataclass paid for the whole traversal every + time it was called. + """ + from effectful.internals.runtime import cache + from effectful.ops import semantics + + traversals = 0 + original = semantics._evaluate_dataclass + + def counting(expr, **kwargs): + nonlocal traversals + traversals += 1 + return original(expr, **kwargs) + + monkeypatch.setattr(semantics, "_evaluate_dataclass", counting) + + assembly = _make_assembly(8) + with cache(): + typeof(assembly) + after_first = traversals + typeof(assembly) + after_second = traversals + + # The first call walks the container and every part inside it. + assert after_first > 8 + # The second visits nothing. + assert after_second == after_first + + +def test_bench_dataclass_argument_op_call(benchmark): + """Benchmark one operation call whose argument is a large dataclass. + + Expected to stay O(k) in the number of children: the dispatch type of the + node genuinely depends on every one of them, so no cache can remove the + traversal of an argument being seen for the first time. Caching addresses + repeat visits, which is what the two tests below exercise. + """ + from effectful.internals.runtime import cache + + assembly = _make_assembly(160) + child, delta = _Part("t", _free_pose()), _Pose(1.0) + + def run(): + # A fresh scope per round, so this stays the cost of a *cold* call + # rather than a cache hit on the previous round. + with cache(): + return _link(assembly, child, delta) + + result = benchmark(run) + assert isinstance(result, Term) + + +def test_bench_dataclass_growing_container(benchmark): + """Benchmark N operation calls against a container that grows by one per call. + + This is the shape that motivated memoizing evaluation: a model assembled + component by component, where every call re-analyses the whole container. + Holding one :func:`cache` scope across the loop turns each element into a + cache hit rather than a fresh traversal, which is worth roughly 4x here. + + It is a constant-factor win and not an asymptotic one, so this benchmark is + expected to remain quadratic in ``n``. Call ``i`` still visits all ``i`` + elements: ``_Assembly.add`` allocates a new instance wrapping a new tuple + every call, and a new container has no cache entry of its own. Making this + linear needs the container's own analysis to be reusable -- structural + sharing so a new tuple shares a cached prefix, or a container-type rule that + does not inspect every element. + + ``test_bench_dataclass_fixed_container`` is the control: same call count, + argument size held constant. + """ + from effectful.internals.runtime import cache + + n = 100 + + def run(): + a, delta = _Assembly(), _Pose(1.0) + term = None + with cache(): + for i in range(n): + child = _Part(f"c{i}", _free_pose()) + term = _link(a, child, delta) + a = a.add(child) + return term + + result = benchmark(run) + assert isinstance(result, Term) + + +def test_bench_dataclass_fixed_container(benchmark): + """Control for :func:`test_bench_dataclass_growing_container`. + + The same number of operation calls, but the container argument stays one + element wide, so this is linear in ``n``. The gap between the two is the + cost attributable to argument size rather than to call count. + """ + from effectful.internals.runtime import cache + + n = 100 + + def run(): + hub = _Assembly((_Part("hub", _free_pose()),)) + delta = _Pose(1.0) + term = None + with cache(): + for i in range(n): + term = _link(hub, _Part(f"c{i}", _free_pose()), delta) + return term + + result = benchmark(run) assert isinstance(result, Term)