logic-loom 0.3.0__py3-none-any.whl

logic_loom/__init__.py

@@ -0,0 +1,67 @@
+"""Logic-Loom: a compiler that understands mathematics.
+
+Instead of peephole-optimizing instructions, Logic-Loom reasons about
+the *algebra* of an expression.  It uses equality saturation over an
+e-graph to discover the whole space of equivalent forms, then extracts
+the cheapest one under a configurable, hardware-aware cost model.
+
+    >>> from logic_loom import optimize
+    >>> print(optimize("a*b + a*c"))
+    a * b + a * c  =>  a * (b + c)
+      cost 5.4 -> 3.3  (1.64x)
+"""
+
+from .analysis import Analysis, analyze, reachable_rules
+from .codegen import free_vars, to_code, to_llvm
+from .compiler import Result, build_egraph, optimize
+from .cost import (
+    DEFAULT_MODEL,
+    PROFILES,
+    CostModel,
+    expr_cost,
+    extract,
+    get_profile,
+)
+from .effects import is_effect_safe, tainted_classes
+from .egraph import EGraph
+from .expr import Expr, evaluate
+from .parser import parse
+from .rules import ALL_RULES, DEFAULT_RULES, EXTENDED_RULES, Rule, rule
+from .saturate import BackoffScheduler, SaturationReport, saturate
+from .viz import to_dot
+
+__version__ = "0.3.0"
+
+__all__ = [
+    "optimize",
+    "build_egraph",
+    "Result",
+    "Expr",
+    "evaluate",
+    "parse",
+    "EGraph",
+    "Rule",
+    "rule",
+    "DEFAULT_RULES",
+    "EXTENDED_RULES",
+    "ALL_RULES",
+    "saturate",
+    "SaturationReport",
+    "BackoffScheduler",
+    "extract",
+    "expr_cost",
+    "CostModel",
+    "PROFILES",
+    "DEFAULT_MODEL",
+    "get_profile",
+    "to_code",
+    "to_llvm",
+    "free_vars",
+    "to_dot",
+    "analyze",
+    "Analysis",
+    "reachable_rules",
+    "tainted_classes",
+    "is_effect_safe",
+    "__version__",
+]

logic_loom/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

logic_loom/analysis.py

@@ -0,0 +1,94 @@
+"""Static analysis to tame the search before saturation begins.
+
+Equality saturation explores blindly; a little static reasoning about the
+input lets us avoid work that provably cannot help.
+
+Two analyses live here:
+
+1. **Reachable-rule pruning.** A rule can only ever fire if the operators
+   in its left-hand side are present -- and operators only appear if some
+   *other* fireable rule introduces them.  Computing this set as a
+   fixed-point lets us drop rules that can never match.  This is sound:
+   the pruned rules would have contributed nothing, so the result is
+   identical, only reached faster.
+
+2. **Complexity estimate.** Counting associative/commutative operators
+   predicts how badly the e-graph might blow up, which we use to size the
+   resource limits automatically.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Set
+
+from .expr import Expr
+
+
+def operators(e: Expr, acc: Set[str] | None = None) -> Set[str]:
+    """The set of operator/function heads appearing in ``e``."""
+    if acc is None:
+        acc = set()
+    if e.kind == "app":
+        acc.add(e.name)
+        for a in e.args:
+            operators(a, acc)
+    return acc
+
+
+def _size(e: Expr) -> int:
+    return 1 + sum(_size(a) for a in e.args)
+
+
+def _count_heads(e: Expr, heads) -> int:
+    n = 1 if (e.kind == "app" and e.name in heads) else 0
+    return n + sum(_count_heads(a, heads) for a in e.args)
+
+
+@dataclass
+class Analysis:
+    size: int
+    variables: int
+    ac_ops: int                 # number of + and * nodes (the explosion driver)
+    node_limit: int
+    match_limit: int
+
+    def summary(self) -> str:
+        return (f"size={self.size} vars={self.variables} ac_ops={self.ac_ops} "
+                f"-> node_limit={self.node_limit} match_limit={self.match_limit}")
+
+
+def reachable_rules(expr: Expr, rules: List) -> List:
+    """Keep only rules whose operators can become reachable from ``expr``."""
+    reachable = operators(expr)
+    rule_ops = [(r, operators(r.lhs), operators(r.rhs)) for r in rules]
+    changed = True
+    while changed:
+        changed = False
+        for _r, lhs_ops, rhs_ops in rule_ops:
+            if lhs_ops <= reachable and not rhs_ops <= reachable:
+                reachable |= rhs_ops
+                changed = True
+    return [r for r, lhs_ops, _ in rule_ops if lhs_ops <= reachable]
+
+
+def analyze(expr: Expr) -> Analysis:
+    size = _size(expr)
+    variables = len({v for v in _vars(expr)})
+    ac = _count_heads(expr, {"+", "*"})
+    # Small inputs can saturate fully; large AC-heavy ones need tighter
+    # budgets so we stop early with the best form found instead of thrashing.
+    node_limit = min(20_000, max(2_000, 400 * (ac + 1)))
+    match_limit = 2_000 if ac <= 3 else max(300, 2_000 // ac)
+    return Analysis(size=size, variables=variables, ac_ops=ac,
+                    node_limit=node_limit, match_limit=match_limit)
+
+
+def _vars(e: Expr, acc: Set[str] | None = None) -> Set[str]:
+    if acc is None:
+        acc = set()
+    if e.kind == "var":
+        acc.add(e.name)
+    for a in e.args:
+        _vars(a, acc)
+    return acc

logic_loom/cli.py

@@ -0,0 +1,88 @@
+"""Command-line interface:  python -m logic_loom "a*b + a*c"."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from . import __version__
+from .codegen import to_code, to_llvm
+from .compiler import build_egraph, optimize
+from .cost import PROFILES
+from .rules import ALL_RULES, DEFAULT_RULES
+from .viz import to_dot
+
+
+def main(argv=None) -> int:
+    p = argparse.ArgumentParser(
+        prog="logic_loom",
+        description="A compiler that understands mathematics. "
+        "Give it an expression; it returns the cheapest equivalent form.",
+    )
+    p.add_argument("expression", nargs="*", help="math expression, e.g. 'a*b + a*c'")
+    p.add_argument("-v", "--verbose", action="store_true",
+                   help="show saturation statistics")
+    p.add_argument("--extended", action="store_true",
+                   help="enable transcendental-function rules (exp/log/sqrt/trig)")
+    p.add_argument("--profile", choices=list(PROFILES), default="default",
+                   help="cost profile to optimize for (default: default)")
+    p.add_argument("--target", choices=["c", "rust", "js", "llvm"],
+                   help="also emit the optimized form as source code / IR")
+    p.add_argument("--impure", default="",
+                   help="comma-separated names of side-effecting functions")
+    p.add_argument("--explain", action="store_true",
+                   help="report the domain assumptions the result relies on")
+    p.add_argument("--dot", action="store_true",
+                   help="print the saturated e-graph as Graphviz DOT")
+    p.add_argument("--max-iters", type=int, default=30)
+    p.add_argument("--node-limit", type=int, default=None)
+    p.add_argument("--version", action="version", version=f"logic-loom {__version__}")
+    args = p.parse_args(argv)
+
+    rules = ALL_RULES if args.extended else DEFAULT_RULES
+    impure = {s.strip() for s in args.impure.split(",") if s.strip()}
+    sources = [" ".join(args.expression)] if args.expression else _read_stdin()
+    if not sources:
+        p.print_help()
+        return 1
+
+    for src in sources:
+        src = src.strip()
+        if not src:
+            continue
+
+        if args.dot:
+            eg, root, _ = build_egraph(
+                src, rules=rules, max_iters=args.max_iters,
+                node_limit=args.node_limit, impure=impure)
+            print(to_dot(eg, root))
+            continue
+
+        r = optimize(src, rules=rules, profile=args.profile,
+                     impure=impure, max_iters=args.max_iters,
+                     node_limit=args.node_limit)
+        print(r)
+        if args.target == "llvm":
+            print(to_llvm(r.optimized))
+        elif args.target:
+            print(f"  {args.target}: {to_code(r.optimized, args.target)}")
+        if args.explain and r.assumptions:
+            print(f"  assumes (for soundness): {'; '.join(r.assumptions)}")
+        if args.verbose:
+            rep = r.report
+            print(f"  [{rep.stop_reason}] profile={r.model.name} "
+                  f"iterations={rep.iterations} e-nodes={rep.nodes} "
+                  f"e-classes={rep.classes}")
+            if rep.banned:
+                print(f"  throttled rules: {', '.join(rep.banned)}")
+    return 0
+
+
+def _read_stdin():
+    if sys.stdin.isatty():
+        return []
+    return sys.stdin.read().splitlines()
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

logic_loom/codegen.py

@@ -0,0 +1,189 @@
+"""Emit an optimized expression as source code in a real language.
+
+The whole point of optimizing an expression is to *run* it. This module
+turns an :class:`Expr` tree into a snippet of C, Rust or JavaScript, so
+the optimized form can be pasted straight into a program.
+
+    >>> from logic_loom import optimize, to_code
+    >>> e = optimize("a*x*x + b*x + c").optimized
+    >>> print(to_code(e, "c"))
+    x * (a * x + b) + c
+"""
+
+from __future__ import annotations
+
+from .expr import Expr
+
+# Per-language rendering of operators and functions.
+_LANGS = {
+    "c": {
+        "pow": lambda a, b: f"pow({a}, {b})",
+        "funcs": {"sin": "sin", "cos": "cos", "tan": "tan",
+                  "exp": "exp", "log": "log", "sqrt": "sqrt"},
+    },
+    "rust": {
+        "pow": lambda a, b: f"({a}).powf({b})",
+        "funcs": {"sin": "{0}.sin()", "cos": "{0}.cos()", "tan": "{0}.tan()",
+                  "exp": "{0}.exp()", "log": "{0}.ln()", "sqrt": "{0}.sqrt()"},
+    },
+    "js": {
+        "pow": lambda a, b: f"Math.pow({a}, {b})",
+        "funcs": {"sin": "Math.sin", "cos": "Math.cos", "tan": "Math.tan",
+                  "exp": "Math.exp", "log": "Math.log", "sqrt": "Math.sqrt"},
+    },
+}
+
+_INFIX = {"+": "+", "-": "-", "*": "*", "/": "/"}
+_PREC = {"+": 1, "-": 1, "*": 2, "/": 2, "neg": 3, "^": 4}
+
+
+def to_code(e: Expr, lang: str = "c") -> str:
+    """Render ``e`` as an expression in ``lang`` (``"c"``, ``"rust"`` or ``"js"``)."""
+    lang = lang.lower()
+    if lang not in _LANGS:
+        raise ValueError(f"unsupported language {lang!r}; choose from {list(_LANGS)}")
+    return _emit(e, _LANGS[lang], 0)
+
+
+def _emit(e: Expr, spec, parent_prec: int) -> str:
+    if e.kind == "num":
+        v = e.value
+        if isinstance(v, float) and v.is_integer():
+            v = int(v)
+        return str(v)
+    if e.kind == "var":
+        return e.name
+    if e.kind == "patvar":
+        raise ValueError("cannot generate code from a pattern variable")
+
+    op = e.name
+    if op == "neg":
+        return f"-{_emit(e.args[0], spec, _PREC['neg'])}"
+
+    if op == "^":
+        a = _emit(e.args[0], spec, 0)
+        b = _emit(e.args[1], spec, 0)
+        return spec["pow"](a, b)
+
+    if op in _INFIX:
+        prec = _PREC[op]
+        left = _emit(e.args[0], spec, prec)
+        right = _emit(e.args[1], spec, prec + 1)
+        s = f"{left} {_INFIX[op]} {right}"
+        return f"({s})" if prec < parent_prec else s
+
+    # function call
+    funcs = spec["funcs"]
+    if op in funcs:
+        args = [_emit(a, spec, 0) for a in e.args]
+        tmpl = funcs[op]
+        if "{0}" in tmpl:                     # method style (Rust)
+            return tmpl.format(*args)
+        return f"{tmpl}({', '.join(args)})"   # call style (C / JS)
+
+    # unknown function: emit verbatim
+    args = ", ".join(_emit(a, spec, 0) for a in e.args)
+    return f"{op}({args})"
+
+
+# --------------------------------------------------------------------- #
+# LLVM IR transpiler
+# --------------------------------------------------------------------- #
+# This lets Logic-Loom plug into a real toolchain: emit the optimized
+# expression as an LLVM IR function that clang/opt can compile, inline,
+# and vectorize alongside the rest of a C/C++/Rust program.
+_LLVM_BINOP = {"+": "fadd", "-": "fsub", "*": "fmul", "/": "fdiv"}
+_LLVM_INTRINSIC = {
+    "sin": "@llvm.sin.f64", "cos": "@llvm.cos.f64",
+    "exp": "@llvm.exp.f64", "log": "@llvm.log.f64",
+    "sqrt": "@llvm.sqrt.f64", "^": "@llvm.pow.f64",
+}
+
+
+def _llvm_const(v) -> str:
+    return f"{float(v):e}"
+
+
+def free_vars(e: Expr):
+    """Sorted list of variable names in ``e`` (the function parameters)."""
+    seen = set()
+
+    def walk(node):
+        if node.kind == "var":
+            seen.add(node.name)
+        for a in node.args:
+            walk(a)
+
+    walk(e)
+    return sorted(seen)
+
+
+def to_llvm(e: Expr, name: str = "f") -> str:
+    """Render ``e`` as an LLVM IR function ``double @name(double, ...)``."""
+    params = free_vars(e)
+    body: list[str] = []
+    counter = [0]
+    used: set[str] = set()
+
+    def fresh() -> str:
+        counter[0] += 1
+        return f"%t{counter[0]}"
+
+    def emit(node: Expr) -> str:
+        if node.kind == "num":
+            return _llvm_const(node.value)
+        if node.kind == "var":
+            return f"%{node.name}"
+        if node.kind == "patvar":
+            raise ValueError("cannot generate IR from a pattern variable")
+
+        op = node.name
+        if op == "neg":
+            x = emit(node.args[0])
+            r = fresh()
+            body.append(f"  {r} = fneg double {x}")
+            return r
+        if op in _LLVM_BINOP:
+            a = emit(node.args[0])
+            b = emit(node.args[1])
+            r = fresh()
+            body.append(f"  {r} = {_LLVM_BINOP[op]} double {a}, {b}")
+            return r
+        if op in _LLVM_INTRINSIC:
+            args = [emit(a) for a in node.args]
+            fn = _LLVM_INTRINSIC[op]
+            used.add(op)
+            r = fresh()
+            joined = ", ".join(f"double {a}" for a in args)
+            body.append(f"  {r} = call double {fn}({joined})")
+            return r
+        # external function fallback
+        args = [emit(a) for a in node.args]
+        used.add(op)
+        r = fresh()
+        joined = ", ".join(f"double {a}" for a in args)
+        body.append(f"  {r} = call double @{op}({joined})")
+        return r
+
+    ret = emit(e)
+    sig = ", ".join(f"double %{p}" for p in params)
+
+    decls = []
+    for op in sorted(used):
+        if op in _LLVM_INTRINSIC:
+            fn = _LLVM_INTRINSIC[op]
+            arity = 2 if op == "^" else 1
+            decls.append(f"declare double {fn}({', '.join(['double'] * arity)})")
+        else:
+            decls.append(f"declare double @{op}(double)")
+
+    lines = []
+    lines.extend(decls)
+    if decls:
+        lines.append("")
+    lines.append(f"define double @{name}({sig}) {{")
+    lines.append("entry:")
+    lines.extend(body)
+    lines.append(f"  ret double {ret}")
+    lines.append("}")
+    return "\n".join(lines)

logic_loom/compiler.py

@@ -0,0 +1,135 @@
+"""High-level API: turn a math string into its cheapest equivalent form."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Iterable, List, Optional
+
+from .analysis import analyze, reachable_rules
+from .cost import DEFAULT_MODEL, CostModel, expr_cost, extract, get_profile
+from .egraph import EGraph
+from .expr import Expr
+from .parser import parse
+from .rules import DEFAULT_RULES, Rule
+from .saturate import BackoffScheduler, SaturationReport, saturate
+
+
+@dataclass
+class Result:
+    source: str
+    original: Expr
+    optimized: Expr
+    original_cost: float
+    optimized_cost: float
+    report: SaturationReport
+    model: CostModel = DEFAULT_MODEL
+    assumptions: List[str] = field(default_factory=list)
+
+    @property
+    def improved(self) -> bool:
+        return self.optimized_cost < self.original_cost - 1e-9
+
+    @property
+    def speedup(self) -> float:
+        if self.optimized_cost <= 0:
+            return float("inf")
+        return self.original_cost / self.optimized_cost
+
+    def __str__(self) -> str:
+        arrow = "=>" if self.improved else "=="
+        return (
+            f"{self.original}  {arrow}  {self.optimized}\n"
+            f"  cost {self.original_cost:.1f} -> {self.optimized_cost:.1f}"
+            f"  ({self.speedup:.2f}x)"
+        )
+
+
+def _resolve_model(model, profile) -> CostModel:
+    if profile is not None:
+        return get_profile(profile)
+    if model is not None:
+        return model
+    return DEFAULT_MODEL
+
+
+def _prepare(source, rules, auto, node_limit):
+    """Parse and, if auto, statically prune rules and size the limits."""
+    original = parse(source)
+    rules = list(rules) if rules is not None else list(DEFAULT_RULES)
+    if auto:
+        rules = reachable_rules(original, rules)
+        an = analyze(original)
+        nl = node_limit if node_limit is not None else an.node_limit
+        scheduler = BackoffScheduler(match_limit=an.match_limit)
+    else:
+        nl = node_limit if node_limit is not None else 5_000
+        scheduler = None
+    return original, rules, nl, scheduler
+
+
+def build_egraph(
+    source: str,
+    *,
+    rules: Optional[List[Rule]] = None,
+    auto: bool = True,
+    max_iters: int = 30,
+    node_limit: Optional[int] = None,
+    impure: Optional[Iterable[str]] = None,
+):
+    """Parse and saturate ``source``; return ``(egraph, root_id, report)``."""
+    original, rules, nl, scheduler = _prepare(source, rules, auto, node_limit)
+    eg = EGraph()
+    root = eg.add_expr(original)
+    report = saturate(eg, rules, max_iters=max_iters, node_limit=nl,
+                      scheduler=scheduler, impure=set(impure or ()))
+    return eg, root, report
+
+
+def optimize(
+    source: str,
+    *,
+    rules: Optional[List[Rule]] = None,
+    model: Optional[CostModel] = None,
+    profile: Optional[str] = None,
+    impure: Optional[Iterable[str]] = None,
+    auto: bool = True,
+    max_iters: int = 30,
+    node_limit: Optional[int] = None,
+) -> Result:
+    """Parse, saturate and extract the cheapest equivalent of ``source``.
+
+    Parameters
+    ----------
+    rules    : rewrite rules to use (defaults to ``DEFAULT_RULES``).
+    model    : a :class:`~logic_loom.cost.CostModel` for extraction.
+    profile  : name of a built-in cost profile ("x86", "arm", "gpu", ...);
+               overrides ``model`` when given.
+    impure   : names of side-effecting functions; rewrites that would
+               duplicate, drop, or reorder their calls are forbidden.
+    auto     : enable static rule pruning and automatic limit sizing.
+    """
+    cost_model = _resolve_model(model, profile)
+    original, used_rules, nl, scheduler = _prepare(source, rules, auto, node_limit)
+
+    eg = EGraph()
+    root = eg.add_expr(original)
+    report = saturate(eg, used_rules, max_iters=max_iters, node_limit=nl,
+                      scheduler=scheduler, impure=set(impure or ()))
+    optimized, opt_cost = extract(eg, root, cost_model)
+
+    by_name = {r.name: r for r in used_rules}
+    assumptions = sorted({
+        a for name in report.fired
+        for a in by_name.get(name, Rule(name, original, original)).assumes
+    })
+
+    return Result(
+        source=source,
+        original=original,
+        optimized=optimized,
+        original_cost=expr_cost(original, cost_model),
+        optimized_cost=opt_cost,
+        report=report,
+        model=cost_model,
+        assumptions=assumptions,
+    )