diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 5d6d17cf0..433b5aee7 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -116,6 +116,14 @@ def inverse(self, value): return jnp.negative(value) +class SumInverseJax(ObjectInterpretation): + @implements(Sum.inverse) + def inverse(self, value): + if not _jax_args((value,)): + return fwd() + return jnp.negative(value) + + class ProductPlusJax(ObjectInterpretation): @implements(Product.plus) def plus(self, *args): @@ -801,6 +809,7 @@ def einsum( EvaluateIntp.extend( SumPlusJax(), + SumInverseJax(), ProductPlusJax(), MinPlusJax(), MaxPlusJax(), diff --git a/effectful/ops/monoid.py b/effectful/ops/monoid.py index fec2cdce8..f2cfee75b 100644 --- a/effectful/ops/monoid.py +++ b/effectful/ops/monoid.py @@ -5,14 +5,29 @@ import operator import typing from collections import UserDict, defaultdict -from collections.abc import Callable, Generator, Iterable, Mapping, Sequence, Sized +from collections.abc import ( + Callable, + Generator, + Iterable, + Mapping, + Sequence, + Sized, +) from dataclasses import dataclass from graphlib import TopologicalSorter -from typing import Annotated, Any +from typing import Annotated, Any, Literal import effectful.ops.syntax from effectful.internals.runtime import interpreter -from effectful.ops.semantics import coproduct, evaluate, fvsof, fwd, handler, typeof +from effectful.ops.semantics import ( + coproduct, + evaluate, + fvsof, + fwd, + handler, + sizeof, + typeof, +) from effectful.ops.syntax import ( ObjectInterpretation, Scoped, @@ -84,7 +99,7 @@ def inner_stream( ) -def inner_streams_first(streams: dict[Operation, Expr]) -> Iterable[Operation]: +def inner_streams_first(streams: Streams) -> Iterable[Operation]: """Iterable over streams where dependent streams precede their dependencies.""" stream_vars = set(streams.keys()) @@ -198,12 +213,28 @@ def __init__(self, name: str, identity: T, zero: T): CartesianProduct: MonoidWithZero[Sequence[Mapping]] = MonoidWithZero( name="CartesianProduct", identity=[{}], zero=[] ) -Union: Monoid[Sequence[Mapping]] = Monoid(name="Union", identity=[]) +Union: Monoid[Iterable] = Monoid(name="Union", identity=[]) +Intersection: MonoidWithZero[Iterable] = MonoidWithZero( + name="Intersection", + identity=Operation.define(Iterable, name="universal")(), + zero=[], +) And = MonoidWithZero(name="And", identity=True, zero=False) Or = Monoid(name="Or", identity=False) -def _conjuncts(mask) -> Sequence[Term]: +@Operation.define +def as_iterable(value: Iterable) -> Iterable: + raise NotHandled + + +def _unwrap_as_iterable[T](arg: Iterable[T]) -> Iterable[T]: + if isinstance(arg, Term) and arg.op is as_iterable: + return arg.args[0] # type: ignore + return arg + + +def _conjuncts(mask) -> Sequence[Expr]: """Return the conjuncts of an ``And`` mask as a flat tuple.""" match mask: case Term(And.plus, elems, {}): @@ -227,7 +258,7 @@ def __call__(self, t: T) -> bool: return t in self.elems -is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product, And, Or}) +is_commutative = _ExtensiblePredicate({Max, Min, Sum, Product, And, Or, Union}) is_idempotent = _ExtensiblePredicate({Max, Min, And, Or}) @@ -273,6 +304,7 @@ def of(self, t: S) -> S | None: (Product, Sum), (CartesianProduct, Union), (And, Or), + (Intersection, Union), ) @@ -341,12 +373,21 @@ def _(self, monoid, value, cond): is_equality = _ExtensiblePredicate({_NumberTerm.__eq__}) -def solve_group_equality(equality: Term[bool], index: int) -> Term[bool]: +def solve_group_equality( + equality: Term[bool], + index: int, + *, + side: typing.Literal["left", "right"] | None = None, +) -> Term[bool]: """Isolate one argument of a group ``plus`` in an equality. Given ``G.plus(a0, ..., an) == b``, isolate the argument at ``index``. - Negative indices follow normal Python indexing. The group expression may - occur on either side of the equality. + Negative indices follow normal Python indexing. ``side`` selects which + side to isolate when both sides are group expressions; when omitted, the + leftmost side that is a group expression is selected. + + The result preserves operand order and is therefore valid for + noncommutative groups. """ if not ( isinstance(equality, Term) @@ -357,6 +398,8 @@ def solve_group_equality(equality: Term[bool], index: int) -> Term[bool]: raise TypeError("expected an equality Term") if not isinstance(index, int): raise TypeError("argument index must be an integer") + if side not in (None, "left", "right"): + raise ValueError("side must be 'left' or 'right'") left, right = equality.args @@ -366,13 +409,23 @@ def group_plus(value): monoid = value.op.__self__ return monoid if isinstance(monoid, Group) else None - monoid = group_plus(left) - if monoid is not None: - plus_term, other = left, right - elif (monoid := group_plus(right)) is not None: - plus_term, other = right, left + left_group = group_plus(left) + right_group = group_plus(right) + if ( + left_group is not None + and right_group is not None + and left_group is not right_group + ): + raise TypeError("equality sides must use the same Group.plus") + + if side == "left" or (side is None and left_group is not None): + monoid, plus_term, other = left_group, left, right + elif side == "right" or (side is None and right_group is not None): + monoid, plus_term, other = right_group, right, left else: raise TypeError("expected an equality containing a Group.plus Term") + if monoid is None: + raise TypeError(f"expected {side} side to be a Group.plus Term") assert isinstance(plus_term, Term) try: @@ -390,95 +443,226 @@ def group_plus(value): return equality.op(target, isolated) -class ReduceEqualityMaskRange(ObjectInterpretation): - """M.reduce(M.mask(v, And.plus(i = x, *m)), {i: range(N)} ∪ S) ≡ - M.mask(M.reduce(M.mask(v, *m), {i: [x]} ∪ S), And.plus(0 <= x, x < N)) +@dataclass(frozen=True) +class _EqualitySolution: + stream_op: Operation + rhs: Expr + + +def _solve_stream_equality( + equality: Expr[bool], streams: Streams +) -> tuple[_EqualitySolution, ...]: + """Return all group-law solutions for bare stream operands in an equality.""" + if not ( + isinstance(equality, Term) + and is_equality(equality.op) + and len(equality.args) == 2 + and not equality.kwargs + ): + return () + + def stream_symbol(value): + if ( + isinstance(value, Term) + and not value.args + and not value.kwargs + and value.op in streams + ): + return value.op + return None + + left, right = equality.args + solutions: list[_EqualitySolution] = [] + + # Direct equalities do not need a group. Preserve the existing behavior of + # leaving equalities between two bare stream symbols for a separate rule. + for stream_term, expr in ((left, right), (right, left)): + stream_op = stream_symbol(stream_term) + if ( + stream_op is not None + and stream_symbol(expr) is None + and stream_op not in fvsof(expr) + ): + solutions.append(_EqualitySolution(stream_op, expr)) + + def group_of(value): + if not (isinstance(value, Term) and _is_monoid_plus(value.op)): + return None + monoid = value.op.__self__ + return monoid if isinstance(monoid, Group) else None + + left_group, right_group = group_of(left), group_of(right) + if ( + left_group is not None + and right_group is not None + and left_group is not right_group + ): + return tuple(solutions) + + sides: tuple[tuple[typing.Literal["left", "right"], Expr], ...] = ( + ("left", left), + ("right", right), + ) + for side, value in sides: + if group_of(value) is None: + continue + assert isinstance(value, Term) + for index, operand in enumerate(value.args): + stream_op = stream_symbol(operand) + if stream_op is None: + continue + solved = solve_group_equality(equality, index, side=side) + rhs = solved.args[1] + if stream_op not in fvsof(rhs): + solutions.append(_EqualitySolution(stream_op, rhs)) + + return tuple(solutions) - The equality constraint ``i = x`` on a range-stream reduce is discharged by - a gather (the stream becomes the singleton ``[x]``) guarded by a bounds - check. - When the reduce body is a ``plus`` of the same monoid, the rule distributes - the reduce over the plus -- but only when doing so exposes an eliminable - equality mask in some summand. This is a *targeted* split (it leaves - ``ReduceSplit`` conservative): summands whose reduced index appears in an - equality become gathers, while the rest stay as ordinary masked reduces. +class ReduceEqualityMask(ObjectInterpretation): + """Eliminate a batch of equalities by intersecting streams with singletons. + + For either orientation of an equality between a reduced variable and an + expression that is not another reduced-variable symbol:: + + M.reduce(M.mask(v, And.plus(i == x, *m)), {i: I} | S) + == M.reduce(M.mask(v, And.plus(*m)), + {i: Intersection.plus(I, [x])} | S) + + The expression ``x`` may also be isolated from an equality between two + ``Group.plus`` expressions. Group inverses are applied in order, so this is + valid for noncommutative groups. ``x`` may depend on streams in ``S``. + Keeping those streams in the outer bundle makes the intersection pointwise, + so this rewrite does not require idempotence or source-variable liveness + checks. A separate :class:`ReduceIntersectionSingletonRange` rule lowers + intersections with simple ranges to singleton gathers guarded by bounds + masks. """ @staticmethod - def _match_eq(cond, streams): - """If ``cond`` is ``stream_op == key`` (either order) where ``stream_op`` - is a ``range(0, N)`` stream and ``key`` is stream-independent, return - ``(stream_op, key)``; otherwise ``None``.""" - - def test(op, stream_op, mask_key): - return ( - is_equality(op) - and stream_op in streams - and _is_simple_range(streams[stream_op]) - and not (fvsof(mask_key) & set(streams)) - ) + def _rhs_cost(expr) -> tuple[float, ...]: + return (len(fvsof(expr)), sizeof(expr)) - match cond: - case Term(op, (Term(stream_op, (), {}), mask_key), {}) if test( - op, stream_op, mask_key - ): - return (stream_op, mask_key) - case Term(op, (mask_key, Term(stream_op, (), {})), {}) if test( - op, stream_op, mask_key - ): - return (stream_op, mask_key) - case _: - return None + @staticmethod + def _can_restrict(stream) -> bool: + """Whether an equality may add the first restriction to ``stream``. + + A stream already represented by an intersection, or by the singleton + produced when lowering one, has already had a substitution variable + selected. Leaving further equalities in the mask allows singleton + substitution to turn them into residual conditions instead of building + nested intersections. + """ + stream = _unwrap_as_iterable(stream) + if isinstance(stream, Term) and stream.op is Intersection.plus: + return False + return not (isinstance(stream, Sequence) and len(stream) == 1) - def _eliminate(self, monoid, value, mask, streams): - """Discharge one eliminable equality constraint via a gather, or return - ``None`` if no constraint is eliminable.""" + @classmethod + def _eliminate(cls, monoid, value, mask, streams): conds = _conjuncts(mask) - for i, cond in enumerate(conds): - matched = self._match_eq(cond, streams) - if matched is None: + matched_conds = [ + (i, solution.stream_op, solution.rhs) + for i, cond in enumerate(conds) + for solution in _solve_stream_equality(cond, streams) + if cls._can_restrict(streams[solution.stream_op]) + ] + ranked_conds = sorted(matched_conds, key=lambda c: (*cls._rhs_cost(c[2]), c[0])) + + # Build a compatible batch greedily. A conjunct and lhs may each occur + # only once in the batch, and none of the batch's lhs variables may + # occur in any of its rhs expressions. Other equalities remain in the + # mask; once a selected stream becomes an intersection or singleton, + # they are retained until substitution turns them into residual masks. + selected = [] + selected_indices = set() + selected_lhs = set() + selected_rhs_fvs = set() + for candidate in ranked_conds: + index, stream_op, expr = candidate + expr_fvs = fvsof(expr) + if ( + index in selected_indices + or stream_op in selected_lhs + or stream_op in selected_rhs_fvs + or selected_lhs & expr_fvs + ): continue - stream_op, mask_key = matched - stream = streams[stream_op] - return monoid.reduce( - monoid.mask( - monoid.mask( - value, - And.plus(stream.start <= mask_key, mask_key < stream.stop), - ), - And.plus(*(c for (j, c) in enumerate(conds) if i != j)), - ), - {stream_op: (mask_key,)} - | {k: v for (k, v) in streams.items() if k != stream_op}, - ) - return None + selected.append(candidate) + selected_indices.add(index) + selected_lhs.add(stream_op) + selected_rhs_fvs.update(expr_fvs) - def _summand_eliminable(self, monoid, summand, streams): - return ( - isinstance(summand, Term) - and _is_monoid_mask(summand.op) - and summand.op.__self__ == monoid - and self._eliminate(monoid, summand.args[0], summand.args[1], streams) - is not None + if not selected: + return None + + new_mask = monoid.mask( + value, + And.plus(*(c for i, c in enumerate(conds) if i not in selected_indices)), ) + replacements = { + stream_op: Intersection.plus(streams[stream_op], [expr]) + for _, stream_op, expr in selected + } + new_streams = { + stream_op: replacements.get(stream_op, stream) + for stream_op, stream in streams.items() + } + return monoid.reduce(new_mask, new_streams) @implements(Monoid.reduce) def _(self, monoid, body, streams): - if not isinstance(body, Term): + if not ( + isinstance(body, Term) + and _is_monoid_mask(body.op) + and body.op.__self__ == monoid + ): return fwd() - # single mask body: discharge an equality constraint directly - if _is_monoid_mask(body.op) and body.op.__self__ == monoid: - result = self._eliminate(monoid, body.args[0], body.args[1], streams) - return result if result is not None else fwd() + result = self._eliminate(monoid, body.args[0], body.args[1], streams) + return result if result is not None else fwd() + + +class ReduceIntersectionSingletonRange(ObjectInterpretation): + """Lower the intersection of a simple range and a singleton stream. + + M.reduce(v, {i: Intersection.plus(range(N), [x])} | S) + == M.reduce(M.mask(v, 0 <= x < N), {i: (x,)} | S) - # plus body: distribute the reduce only when it exposes an eliminable - # equality mask in some summand - if _is_monoid_plus(body.op) and body.op.__self__ == monoid: - if any(self._summand_eliminable(monoid, s, streams) for s in body.args): - return monoid.plus(*(monoid.reduce(s, streams) for s in body.args)) + The singleton expression may depend on streams in ``S``; the resulting + dependent singleton is eliminated later by :class:`EliminateSingletonStreams`. + """ + @staticmethod + def _match(stream): + match stream: + case Term(op, (lhs, rhs), {}) if op is Intersection.plus: + pass + case _: + return None + + lhs, rhs = _unwrap_as_iterable(lhs), _unwrap_as_iterable(rhs) + if _is_simple_range(lhs) and isinstance(rhs, list) and len(rhs) == 1: + return lhs, rhs[0] + if _is_simple_range(rhs) and isinstance(lhs, list) and len(lhs) == 1: + return rhs, lhs[0] + return None + + @implements(Monoid.reduce) + def reduce(self, monoid, body, streams): + for stream_op, stream in streams.items(): + matched = self._match(stream) + if matched is None: + continue + range_stream, value = matched + return monoid.reduce( + monoid.mask( + body, + And.plus(range_stream.start <= value, value < range_stream.stop), + ), + {stream_op: (value,)} + | {k: v for k, v in streams.items() if k is not stream_op}, + ) return fwd() @@ -733,7 +917,7 @@ def reduce(self, monoid, body, streams): @Operation.define def choose_contraction(factors: Sequence[Any], streams: Streams) -> Operation: - """Used by `ReduceFactorization` to choose a contraction when there is + """Used by `Factor` to choose a contraction when there is ambiguity. Takes the factors and streams that are eligible for contraction (innermost and non-universal). @@ -756,8 +940,24 @@ def choose_contraction(factors: Sequence[Any], streams: Streams) -> Operation: class Factor(ObjectInterpretation): + def mask_plus( + self, + outer_monoid: Monoid, + inner_monoid: Monoid, + *factor_conds: tuple[Literal["factor", "mask"], Expr], + ) -> Expr: + """Turn a flat list of factors and masks into a masked plus.""" + factors: list[Expr] = [] + conds: list[Expr] = [] + for k, f in factor_conds: + (factors if k == "factor" else conds).append(f) + + term = inner_monoid.plus(*factors) + term = term if not conds else outer_monoid.mask(term, And.plus(*conds)) + return term + @implements(Monoid.reduce) - def reduce(self, monoid, body, streams): + def reduce(self, monoid: Monoid, body, streams: Streams): """reduce(⊗(F_v ∪ F_rest), {v} ∪ S) = reduce(⊗F_rest ⊗ reduce(⊗F_v, {v}), S) where F_v = factors mentioning v, F_rest = the others. Fires only when @@ -781,10 +981,11 @@ def reduce(self, monoid, body, streams): return fwd() # Optionally peel an outer mask of the reduce monoid. - cond = None plus_term = body + conds: Sequence[Term] = () if _is_monoid_mask(body.op) and body.op.__self__ is monoid: plus_term, cond = body.args + conds = _conjuncts(cond) if not ( isinstance(plus_term, Term) @@ -795,8 +996,9 @@ def reduce(self, monoid, body, streams): inner = plus_term.op.__self__ stream_keys = set(streams) - cond_fvs = fvsof(cond) if cond is not None else set() - factors = [(a, fvsof(a) | cond_fvs) for a in plus_term.args] + factors: list[tuple[Literal["factor", "mask"], Expr]] = [ + ("factor", a) for a in plus_term.args + ] + [("mask", c) for c in conds] # candidates: innermost-eligible (no remaining stream depends on v), # non-universal (some factor doesn't mention v) @@ -804,9 +1006,7 @@ def reduce(self, monoid, body, streams): for k, v in streams.items(): if any(k in fvsof(vv) for kk, vv in streams.items() if k is not kk): continue - if len({i for i, (_, fvs) in enumerate(factors) if k in fvs}) == len( - factors - ): + if all(k in fvsof(factor) for (_, factor) in factors): continue # v is universal: leave it in the outer core eligible[k] = v @@ -818,22 +1018,20 @@ def reduce(self, monoid, body, streams): inner_stream = choose_contraction(plus_term.args, eligible) inner_factor_ids = frozenset( - i for i, (_, fvs) in enumerate(factors) if inner_stream in fvs + i for i, (_, factor) in enumerate(factors) if inner_stream in fvsof(factor) ) - inner_factors = [factors[i][0] for i in sorted(inner_factor_ids)] + inner_factors = [factors[i] for i in sorted(inner_factor_ids)] inner_stream_keys = {inner_stream} inner_deps = set().union( - *(factors[i][1] for i in inner_factor_ids), + *(fvsof(factors[i][1]) for i in inner_factor_ids), fvsof(streams[inner_stream]) & stream_keys, ) - outer_factors = [ - a for i, (a, _) in enumerate(factors) if i not in inner_factor_ids - ] + outer_factors = [a for i, a in enumerate(factors) if i not in inner_factor_ids] outer_stream_keys = stream_keys - inner_stream_keys outer_factor_deps = set().union( - *(vars for i, (_, vars) in enumerate(factors) if i not in inner_factor_ids) + *(fvsof(f) for i, (_, f) in enumerate(factors) if i not in inner_factor_ids) ) # find all streams that are used in the inner factors/streams and are @@ -854,12 +1052,12 @@ def reduce(self, monoid, body, streams): outer_stream_keys -= {s} inner_streams = {k: v for (k, v) in streams.items() if k in inner_stream_keys} - inner_red = monoid.reduce(inner.plus(*inner_factors), inner_streams) + inner_red = monoid.reduce( + self.mask_plus(monoid, inner, *inner_factors), inner_streams + ) rest_streams = {k: s for k, s in streams.items() if k in outer_stream_keys} - new_body = inner.plus(*outer_factors, inner_red) - if cond is not None: - new_body = monoid.mask(new_body, cond) + new_body = self.mask_plus(monoid, inner, *outer_factors, ("factor", inner_red)) return monoid.reduce(new_body, rest_streams) if rest_streams else new_body @@ -1271,6 +1469,16 @@ def inverse(self, value): return -value +class SumInverse(ObjectInterpretation): + """Scalar implementation of :meth:`Sum.inverse`.""" + + @implements(Sum.inverse) + def inverse(self, value): + if isinstance(value, Term) or not isinstance(value, int | float): + return fwd() + return -value + + class MinPlus(ObjectInterpretation): """Scalar implementation of :data:`Min`.""" @@ -1360,11 +1568,8 @@ class CartesianProductPlus(ObjectInterpretation): @implements(CartesianProduct.plus) def plus(self, *args): - if not args: - return fwd() - if any(isinstance(x, Term) for x in args): - return fwd() - if not all(isinstance(x, Iterable) for x in args): + args = tuple(_unwrap_as_iterable(arg) for arg in args) + if not args or any(isinstance(x, Term) for x in args): return fwd() return [_disjoint_merge(*vals) for vals in itertools.product(*args)] @@ -1372,13 +1577,46 @@ def plus(self, *args): class UnionPlus(ObjectInterpretation): @implements(Union.plus) def plus(self, *args): - if not args: + args = tuple(_unwrap_as_iterable(arg) for arg in args) + if not args or any(isinstance(x, Term) for x in args): + return fwd() + return list(itertools.chain(*args)) + + +class IntersectionPlus(ObjectInterpretation): + """Pure-Python implementation of :data:`Intersection`. + + The result preserves the order and multiplicity of the first iterable and + retains each of its elements that occurs in every subsequent iterable. + """ + + @implements(Intersection.plus) + def plus(self, *args): + args = tuple(_unwrap_as_iterable(arg) for arg in args) + if ( + not args + or any(isinstance(arg, Term) for arg in args) + or any(fvsof(arg) for arg in args) + ): return fwd() - if any(isinstance(x, Term) for x in args): + + first, *rest = args + rest = [tuple(values) for values in rest] + return [value for value in first if all(value in values for values in rest)] + + +class PlusCastIterable(ObjectInterpretation): + """Cast heterogeneous iterable arguments to a common iterable type.""" + + @implements(Monoid.plus) + def _plus(self, monoid, *args): + if monoid not in (Intersection, Union): return fwd() - if not all(isinstance(x, Iterable) for x in args): + + if not args or all(isinstance(a, Term) and a.op == as_iterable for a in args): return fwd() - return list(itertools.chain(*args)) + + return monoid.plus(*(as_iterable(arg) for arg in args)) is_scalar = _ExtensiblePredicate({Min, Max, Sum, Product, And, Or}) @@ -1842,6 +2080,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReducePartial(), DeltaConcrete(), SumPlus(), + SumInverse(), MinPlus(), MaxPlus(), ProductPlus(), @@ -1849,7 +2088,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: ArgMaxPlus(), CartesianProductPlus(), UnionPlus(), - ReduceEqualityMaskRange(), + IntersectionPlus(), ReduceWhereToMasks(), ) @@ -1882,6 +2121,8 @@ def extend(self, *intps: Interpretation) -> typing.Self: ReduceFusion(), ReduceUnion(), ReduceSplit(), + ReduceEqualityMask(), + ReduceIntersectionSingletonRange(), Factor(), ReduceDistributeCartesianProduct(), ReduceWeightedStream(), @@ -1897,6 +2138,7 @@ def extend(self, *intps: Interpretation) -> typing.Self: PlusConsecutiveDups(), PlusOrder(), PlusCastFloat(), + PlusCastIterable(), MaskFusion(), MaskBool(), WhereHoist(), diff --git a/effectful/ops/semantics.py b/effectful/ops/semantics.py index be3686dd6..1af1ec580 100644 --- a/effectful/ops/semantics.py +++ b/effectful/ops/semantics.py @@ -427,3 +427,37 @@ def fvsof[S](term: Expr[S]) -> collections.abc.Set[Operation]: """ result = evaluate(term, intp=_FVSOF_INTP) return result.fvs if isinstance(result, _FvsAnalysis) else frozenset() + + +class _SizeofAnalysisValue: + value: int + + def __init__(self, value: int | typing.Self): + self.value = value.value if isinstance(value, _SizeofAnalysisValue) else value + + def __add__(self, other): + if isinstance(other, _SizeofAnalysisValue): + return _SizeofAnalysisValue(self.value + other.value) + return NotImplemented + + __radd__ = __add__ + + +@functools.cache +def _sizeof_intp(): + def _apply(_, *args, **kwargs): + s_args = [ + x if isinstance(x, _SizeofAnalysisValue) else _SizeofAnalysisValue(1) + for x in (*args, *kwargs.values()) + ] + size = sum(s_args, start=_SizeofAnalysisValue(1)) + return size + + return {apply: _apply} + + +def sizeof(term: Expr) -> int: + result = evaluate(term, intp=_sizeof_intp()) + if not isinstance(result, _SizeofAnalysisValue): + return 1 + return result.value diff --git a/tests/_monoid_helpers.py b/tests/_monoid_helpers.py index 25490411e..ba194cea2 100644 --- a/tests/_monoid_helpers.py +++ b/tests/_monoid_helpers.py @@ -168,11 +168,11 @@ def __init__(self): def eq(self, a: Any, b: Any) -> bool: raise NotImplementedError + @staticmethod @abstractmethod def strategy( - self, arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", + ret: Literal["scalar", "stream"] | type = "scalar", ) -> SearchStrategy: raise NotImplementedError @@ -180,7 +180,7 @@ def _fresh_op( self, name: str, arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", + ret: Literal["scalar", "stream"] | type = "scalar", ) -> Operation: """Build a fresh, unhandled Operation whose parameter and return annotations are derived from this backend. @@ -190,7 +190,7 @@ def _fresh_op( each of type ``scalar_typ``. """ scalar = self.scalar_typ - out = self.stream_typ if ret == "stream" else scalar + out = self.stream_typ if ret == "stream" else scalar if ret == "scalar" else ret params = ", ".join(f"_a{i}" for i in range(len(arg_types))) ns: dict[str, Any] = {"NotHandled": NotHandled} exec(f"def _fn({params}):\n raise NotHandled\n", ns) @@ -328,28 +328,28 @@ class IntBackend(Backend): lambda x: [0, x, x + 1], ] + @staticmethod def strategy( - self, arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", + ret: Literal["scalar", "stream"] | type = "scalar", ) -> SearchStrategy: match arg_types, ret: - case (), "scalar": + case (), ("scalar" | builtins.int): return st.integers(min_value=-100, max_value=100).map(deffn) case (), "stream": scalars = st.integers(min_value=-100, max_value=100) return st.lists(scalars, max_size=2).map(deffn) case (builtins.int,), "scalar": - return st.sampled_from(self._unary_num_fns) + return st.sampled_from(IntBackend._unary_num_fns) case (builtins.int, builtins.int), "scalar": - return st.sampled_from(self._binary_num_fns) + return st.sampled_from(IntBackend._binary_num_fns) case (builtins.int, builtins.int, builtins.int), "scalar": return st.tuples( - st.sampled_from(self._binary_num_fns), - st.sampled_from(self._binary_num_fns), + st.sampled_from(IntBackend._binary_num_fns), + st.sampled_from(IntBackend._binary_num_fns), ).map(lambda fg: lambda a, b, c: fg[0](a, fg[1](b, c))) case (builtins.int,), "stream": - return st.sampled_from(self._unary_list_fns) + return st.sampled_from(IntBackend._unary_list_fns) raise NotImplementedError( f"No int strategy for op with return {ret!r} and {arg_types} args" ) @@ -385,12 +385,14 @@ class JaxBackend(Backend): lambda a, b: a * b, ] + @staticmethod def strategy( - self, arg_types: tuple[type, ...] = (), - ret: Literal["scalar", "stream"] = "scalar", + ret: Literal["scalar", "stream"] | type = "scalar", ) -> st.SearchStrategy[Callable]: match arg_types, ret: + case (), builtins.int: + return IntBackend.strategy((), builtins.int) case (), "scalar": return ( st.lists( @@ -412,16 +414,16 @@ def strategy( .map(deffn) ) case (jax.Array,), "scalar": - return st.sampled_from(self._unary_jax_scalar_fns) + return st.sampled_from(JaxBackend._unary_jax_scalar_fns) case (jax.Array, jax.Array), "scalar": - return st.sampled_from(self._binary_jax_scalar_fns) + return st.sampled_from(JaxBackend._binary_jax_scalar_fns) case (jax.Array, jax.Array, jax.Array), "scalar": return st.tuples( - st.sampled_from(self._binary_jax_scalar_fns), - st.sampled_from(self._binary_jax_scalar_fns), + st.sampled_from(JaxBackend._binary_jax_scalar_fns), + st.sampled_from(JaxBackend._binary_jax_scalar_fns), ).map(lambda fg: lambda a, b, c: fg[0](a, fg[1](b, c))) case (jax.Array,), "stream": - return st.sampled_from(self._unary_jax_stream_fns) + return st.sampled_from(JaxBackend._unary_jax_stream_fns) raise NotImplementedError( f"No jax strategy for op with return {ret!r} and {arg_types} args" diff --git a/tests/test_ops_monoid.py b/tests/test_ops_monoid.py index b57226a93..21fb8bc7c 100644 --- a/tests/test_ops_monoid.py +++ b/tests/test_ops_monoid.py @@ -18,6 +18,8 @@ EvaluateIntp, Factor, Group, + Intersection, + IntersectionPlus, InverseInverse, InversePlus, Max, @@ -28,6 +30,7 @@ NormalizeIntp, Or, PlusAssoc, + PlusCastIterable, PlusConsecutiveDups, PlusDistr, PlusEmpty, @@ -39,8 +42,9 @@ ReduceDisjunctiveDisequalityMask, ReduceDistributeCartesianProduct, ReduceEmpty, - ReduceEqualityMaskRange, + ReduceEqualityMask, ReduceFusion, + ReduceIntersectionSingletonRange, ReduceMaskHoist, ReducePartial, ReduceSplit, @@ -52,8 +56,10 @@ Sum, Union, WhereHoist, + as_iterable, distributes_over, is_commutative, + is_idempotent, solve_group_equality, ) from effectful.ops.semantics import coproduct, evaluate, fvsof, handler @@ -741,82 +747,156 @@ def test_reduce_mask_hoist_dependent_noop(monoid): @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_range_simple(backend: Backend, monoid): - """Best case: a single equality on the reduced range stream becomes a - singleton-stream gather guarded by the corresponding bounds check. - """ - a, c = backend.define_vars("a", "c", ret="scalar") - f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + """An equality restricts its stream to an intersection with a singleton.""" + a, c = backend.define_vars("a", "c", ret=int) - lhs = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(3)}) + lhs = monoid.reduce(monoid.mask(a(), a() == c()), {a: range(3)}) rhs = monoid.reduce( - monoid.mask( - monoid.mask(f(a()), And.plus(0 <= c(), c() < 3)), - And.plus(), - ), - {a: (c(),)}, + monoid.mask(a(), And.plus()), + {a: Intersection.plus(as_iterable(range(3)), as_iterable([c()]))}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) @pytest.mark.parametrize("monoid", ALL_MONOIDS) def test_reduce_equality_mask_range_residual_conjuncts(backend: Backend, monoid): - """Non-equality conjuncts are preserved inside the gathered singleton - reduce; only the equality on the reduced range stream is discharged. - """ - a, c, d, e = backend.define_vars("a", "c", "d", "e", ret="scalar") - f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + """Only the equality is moved into the stream-domain intersection.""" + a, c, d, e = backend.define_vars("a", "c", "d", "e", ret=int) lhs = monoid.reduce( - monoid.mask(f(a()), And.plus(d() < e(), a() == c(), c() < e())), + monoid.mask(a(), And.plus(d() < e(), a() == c(), c() < e())), {a: range(4)}, ) rhs = monoid.reduce( - monoid.mask( - monoid.mask(f(a()), And.plus(0 <= c(), c() < 4)), - And.plus(d() < e(), c() < e()), - ), - {a: (c(),)}, + monoid.mask(a(), And.plus(d() < e(), c() < e())), + {a: Intersection.plus(as_iterable(range(4)), as_iterable([c()]))}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) -@pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_equality_mask_range_noncanonical_range_noop(backend: Backend, monoid): - """The rule only handles ``range(0, N, 1)`` streams; for other ranges it - should leave the term unchanged. - """ - a, c = backend.define_vars("a", "c", ret="scalar") +def test_reduce_equality_mask_dependent_intersection(): + """The singleton expression may depend on another retained stream.""" + backend = IntBackend() + x, y, c, d = backend.define_vars("x", "y", "c", "d", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = Min.reduce( + Min.mask(g(x()), And.plus(x() == f(y()), c() < d())), + {x: X(), y: Y()}, + ) + rhs = Min.reduce( + Min.mask(g(x()), And.plus(c() < d())), + { + x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), + y: Y(), + }, + ) + + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) + ) - term = monoid.reduce(monoid.mask(f(a()), a() == c()), {a: range(1, 4)}) - backend.check_rewrite(lhs=term, rhs=term, rule=ReduceEqualityMaskRange()) + +def test_reduce_equality_mask_dependent_intersection_nonidempotent(): + """Pointwise intersection does not require an idempotent outer monoid.""" + backend = IntBackend() + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = Sum.reduce(Sum.mask(g(x()), x() == f(y())), {x: X(), y: Y()}) + rhs = Sum.reduce( + Sum.mask(g(x()), And.plus()), + {x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), y: Y()}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) + ) + + +def test_reduce_equality_mask_dependent_intersection_retains_source_uses(): + """Source variables remain available to the body after the rewrite.""" + backend = IntBackend() + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + f = backend.define_vars("f", arg_types=(backend.scalar_typ,), ret="scalar") + g = backend.define_vars( + "g", arg_types=(backend.scalar_typ, backend.scalar_typ), ret="scalar" + ) + + lhs = Min.reduce(Min.mask(g(x(), y()), x() == f(y())), {x: X(), y: Y()}) + rhs = Min.reduce( + Min.mask(g(x(), y()), And.plus()), + {x: Intersection.plus(as_iterable(X()), as_iterable([f(y())])), y: Y()}, + ) + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) + ) + + +def test_reduce_equality_mask_image_domain_symbol_side_noop(): + """Bare stream equalities require a separate domain-intersection rule.""" + backend = IntBackend() + x, y = backend.define_vars("x", "y", ret="scalar") + X, Y = backend.define_vars("X", "Y", ret="stream") + g = backend.define_vars("g", arg_types=(backend.scalar_typ,), ret="scalar") + + lhs = Min.reduce(Min.mask(g(x()), x() == y()), {x: X(), y: Y()}) + backend.check_rewrite( + lhs=lhs, rhs=lhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) + ) @pytest.mark.parametrize("monoid", ALL_MONOIDS) -def test_reduce_equality_mask_plus(backend: Backend, monoid): - """ReduceEqualityMaskRange distributes over a plus body, discharging an - equality on the reduced stream in one summand via a singleton-stream gather - while leaving the other summand as an ordinary masked reduce. - """ - a, c = backend.define_vars("a", "c", ret="scalar") - f, g = backend.define_vars("f", "g", arg_types=(backend.scalar_typ,), ret="scalar") +def test_reduce_equality_mask_noncanonical_range(backend: Backend, monoid): + """Equality elimination itself is independent of the domain representation.""" + a, c = backend.define_vars("a", "c", ret=int) - body = monoid.plus( - monoid.mask(f(a()), a() == c()), # eliminable: a == c over range - monoid.mask(g(a()), c() == 0), # not eliminable (no reduced-stream eq) + lhs = monoid.reduce(monoid.mask(a(), a() == c()), {a: range(1, 4)}) + rhs = monoid.reduce( + monoid.mask(a(), And.plus()), + {a: Intersection.plus(as_iterable(range(1, 4)), as_iterable([c()]))}, ) - lhs = monoid.reduce(body, {a: range(3)}) - rhs = monoid.plus( - monoid.reduce( - monoid.mask( - monoid.mask(f(a()), And.plus(0 <= c(), c() < 3)), - And.plus(), - ), - {a: (c(),)}, - ), - monoid.reduce(monoid.mask(g(a()), c() == 0), {a: range(3)}), + backend.check_rewrite( + lhs=lhs, rhs=rhs, rule=coproduct(ReduceEqualityMask(), PlusCastIterable()) + ) + + +def test_reduce_equality_mask_repeated_constraints_eliminate_intersections(): + """Repeated equalities for one stream still normalize to supported domains.""" + x, a, b = IntBackend().define_vars("x", "a", "b", ret=int) + lhs = Sum.reduce( + Sum.mask(x(), And.plus(x() == a(), x() == b())), + {x: range(3)}, + ) + + with handler(NormalizeIntp), handler(EvaluateIntp): + result = evaluate(lhs) + + assert Intersection.plus not in fvsof(result) + + +@pytest.mark.parametrize("monoid", ALL_MONOIDS) +def test_reduce_intersection_singleton_range(backend: Backend, monoid): + """A simple-range intersection lowers to a bounds-guarded singleton.""" + x, y = backend.define_vars("x", "y", ret=int) + + lhs = monoid.reduce( + x(), {x: Intersection.plus(as_iterable(range(3)), as_iterable([y()]))} + ) + rhs = monoid.reduce(monoid.mask(x(), And.plus(0 <= y(), y() < 3)), {x: (y(),)}) + backend.check_rewrite( + lhs=lhs, + rhs=rhs, + rule=coproduct(ReduceIntersectionSingletonRange(), PlusCastIterable()), ) - backend.check_rewrite(lhs=lhs, rhs=rhs, rule=ReduceEqualityMaskRange()) def test_reduce_independent_1(backend: Backend): @@ -928,6 +1008,38 @@ def test_reduce_lift_shared(outer, inner, backend: Backend): backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) +def test_reduce_factor_mask(backend: Backend): + a, b, c, d = backend.define_vars("a", "b", "c", "d", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + + lhs = Sum.reduce( + Sum.mask(Product.plus(a(), b()), And.plus(a() == c(), b() == d())), + {a: A(), b: B()}, + ) + rhs = Product.plus( + Sum.reduce(Sum.mask(Product.plus(a()), And.plus(a() == c())), {a: A()}), + Sum.reduce(Sum.mask(Product.plus(b()), And.plus(b() == d())), {b: B()}), + ) + backend.check_rewrite(lhs=lhs, rhs=rhs, rule=Factor()) + + +def test_reduce_factor_mask_2(backend: Backend): + a, b = backend.define_vars("a", "b", ret="scalar") + A, B = backend.define_vars("A", "B", ret="stream") + + lhs = Sum.reduce( + Sum.mask(Product.plus(a(), b()), And.plus(a() == b())), + {a: A(), b: B()}, + ) + lhs = Sum.reduce( + Product.plus( + b(), Sum.reduce(Sum.mask(Product.plus(a()), And.plus(a() == b())), {a: A()}) + ), + {b: B()}, + ) + backend.check_rewrite(lhs=lhs, rhs=lhs, rule=Factor()) + + @pytest.mark.parametrize("outer,inner", MONOID_PAIRS) def test_reduce_lift_shared_deps(outer, inner, backend: Backend): """A shared stream is lifted together with its dependencies: both ``c`` @@ -1245,6 +1357,35 @@ def _f(v: int) -> float: assert math.isclose(result, 10.0) +def test_intersection_plus_preserves_left_order_and_multiplicity(): + with handler(IntersectionPlus()): + result = Intersection.plus( + as_iterable([3, 1, 1, 2, 4]), + as_iterable([1, 2, 3]), + as_iterable([1, 3]), + ) + + assert result == [3, 1, 1] + + +def test_intersection_is_not_registered_commutative_or_idempotent(): + assert not is_commutative(Intersection) + assert not is_idempotent(Intersection) + + with handler(IntersectionPlus()): + assert Intersection.plus([1, 1], [1]) == [1, 1] + assert Intersection.plus([1], [1, 1]) == [1] + + +def test_equality_intersection_preserves_stream_multiplicity(): + x = Operation.define(int, name="x") + + with handler(NormalizeIntp), handler(EvaluateIntp): + result = evaluate(Sum.reduce(Sum.mask(x(), x() == 1), {x: [1, 1, 2]})) + + assert result == 2 + + # --------------------------------------------------------------------------- # CartesianProduct.plus (pure-Python ``CartesianProductPlus`` implementation) # ---------------------------------------------------------------------------