mplang-nightly 0.1.dev268__py3-none-any.whl → 0.1.dev270__py3-none-any.whl

mplang/v1/core/expr/walk.py

@@ -1,387 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Pure functional traversal utilities for MPLang expression graphs.
-
-This module provides semantic-agnostic walkers over expression graphs. It exposes
-both dataflow-only traversal (deps edges) and structural traversal (deps +
-contained regions like function bodies, then/else, loop body). All traversals are
-implemented iteratively to avoid Python recursion limits.
-
-Notes
-- These walkers never evaluate expressions nor decide runtime branches. For
-  execution order, use an evaluator/driver that consults semantic rules.
-- Topological order is produced w.r.t. deps edges, not region containment.
-"""
-
-from __future__ import annotations
-
-from collections import deque
-from collections.abc import Callable, Iterable, Iterator, Sequence
-from typing import cast
-
-from mplang.v1.core.expr.ast import (
-    AccessExpr,
-    CallExpr,
-    CondExpr,
-    ConvExpr,
-    EvalExpr,
-    Expr,
-    FuncDefExpr,
-    ShflExpr,
-    ShflSExpr,
-    TupleExpr,
-    VariableExpr,
-    WhileExpr,
-)
-
-Node = Expr
-GetDeps = Callable[[Node], Iterable[Node]]
-YieldCond = Callable[[Node], bool]
-
-
-def _identity_key(n: Node) -> int:
-    """Identity-based key for hashing nodes that may not be hashable."""
-    return id(n)
-
-
-# ---------------------------- default dependency getters ----------------------------
-
-
-def dataflow_deps(node: Node) -> Iterable[Node]:
-    """Default dataflow dependencies for core Expr nodes.
-
-    This includes only inputs that must be computed to evaluate the node itself.
-    Contained regions (function/branch/loop bodies) are NOT traversed here.
-    """
-    if isinstance(node, EvalExpr):
-        return node.args
-    if isinstance(node, TupleExpr):
-        return node.args
-    if isinstance(node, CondExpr):
-        # Pure dataflow: pred and actual args to branch functions
-        return [node.pred, *node.args]
-    if isinstance(node, WhileExpr):
-        # Initial state args required; bodies are regions, not dataflow deps
-        return list(node.args)
-    if isinstance(node, ConvExpr):
-        return node.vars
-    if isinstance(node, ShflSExpr):
-        return [node.src_val]
-    if isinstance(node, ShflExpr):
-        return [node.src, node.index]
-    if isinstance(node, AccessExpr):
-        return [node.src]
-    if isinstance(node, VariableExpr):
-        return []
-    if isinstance(node, FuncDefExpr):
-        # Definition is not a value-producing node in dataflow; no deps
-        return []
-    if isinstance(node, CallExpr):
-        # Arguments are dataflow deps; function body is a region
-        return node.args
-    # Fallback: try best-effort empty
-    return []
-
-
-def _structural_region_roots(node: Node) -> Iterable[Node]:
-    """Roots of contained regions for structural traversal.
-
-    - Cond: then_fn.body, else_fn.body
-    - While: cond_fn.body, body_fn.body
-    - Call: fn.body
-    - FuncDef: body
-    """
-    if isinstance(node, CondExpr):
-        return [node.then_fn.body, node.else_fn.body]
-    if isinstance(node, WhileExpr):
-        return [node.cond_fn.body, node.body_fn.body]
-    if isinstance(node, CallExpr):
-        return [node.fn.body]
-    if isinstance(node, FuncDefExpr):
-        return [node.body]
-    return []
-
-
-# ---------------------------------- core walkers ----------------------------------
-
-
-def walk(
-    roots: Node | Sequence[Node],
-    *,
-    get_deps: GetDeps,
-    traversal: str = "dfs_post_iter",
-    yield_condition: YieldCond | None = None,
-    detect_cycles: bool = True,
-) -> Iterator[Node]:
-    """Generic pure structural walker.
-
-    Args:
-        roots: Single root or a sequence of roots to start from.
-        get_deps: Function mapping node -> iterable of dependency nodes.
-        traversal: One of {'dfs_pre_iter','dfs_post_iter','bfs','topo'}.
-        yield_condition: Optional predicate to filter yielded nodes.
-        detect_cycles: If True, raises ValueError on cycles (for DFS/topo).
-
-    Yields:
-        Nodes in the chosen traversal order, filtered by yield_condition.
-    """
-    start: list[Node]
-    if isinstance(roots, (list, tuple)):
-        start = list(roots)
-    else:
-        start = [cast(Node, roots)]
-
-    if traversal == "bfs":
-        yield from _bfs(start, get_deps, yield_condition)
-        return
-    if traversal == "dfs_pre_iter":
-        yield from _dfs_pre_iter(start, get_deps, yield_condition, detect_cycles)
-        return
-    if traversal == "dfs_post_iter":
-        yield from _dfs_post_iter(start, get_deps, yield_condition, detect_cycles)
-        return
-    if traversal == "topo":
-        yield from _topo_kahn(start, get_deps, yield_condition, detect_cycles)
-        return
-
-    raise ValueError(f"Invalid traversal type: {traversal}")
-
-
-def walk_dataflow(
-    roots: Node | Sequence[Node],
-    *,
-    traversal: str = "dfs_post_iter",
-    yield_condition: YieldCond | None = None,
-    detect_cycles: bool = True,
-) -> Iterator[Node]:
-    """Walk using default dataflow dependencies for Expr nodes."""
-    return walk(
-        roots,
-        get_deps=dataflow_deps,
-        traversal=traversal,
-        yield_condition=yield_condition,
-        detect_cycles=detect_cycles,
-    )
-
-
-def walk_structural(
-    roots: Node | Sequence[Node],
-    *,
-    traversal: str = "dfs_post_iter",
-    yield_condition: YieldCond | None = None,
-    detect_cycles: bool = True,
-) -> Iterator[Node]:
-    """Walk including region containment (function bodies, branches, loop bodies).
-
-    This augments dataflow dependencies with region roots once, so structure is
-    fully traversed without runtime branch choices or loop iteration expansion.
-    """
-
-    def deps_plus_regions(n: Node) -> Iterable[Node]:
-        yield from dataflow_deps(n)
-        yield from _structural_region_roots(n)
-
-    return walk(
-        roots,
-        get_deps=deps_plus_regions,
-        traversal=traversal,
-        yield_condition=yield_condition,
-        detect_cycles=detect_cycles,
-    )
-
-
-# -------------------------------- traversal engines --------------------------------
-
-
-def _maybe_yield(n: Node, pred: YieldCond | None) -> Iterator[Node]:
-    if pred is None or pred(n):
-        yield n
-
-
-def _bfs(
-    roots: Sequence[Node],
-    get_deps: GetDeps,
-    yield_condition: YieldCond | None,
-) -> Iterator[Node]:
-    seen: set[int] = set()
-    q: deque[Node] = deque(roots)
-    while q:
-        n = q.popleft()
-        k = _identity_key(n)
-        if k in seen:
-            continue
-        seen.add(k)
-        yield from _maybe_yield(n, yield_condition)
-        for d in get_deps(n):
-            q.append(d)
-
-
-def _dfs_pre_iter(
-    roots: Sequence[Node],
-    get_deps: GetDeps,
-    yield_condition: YieldCond | None,
-    detect_cycles: bool,
-) -> Iterator[Node]:
-    seen: set[int] = set()
-    onstack: set[int] = set()
-    stack: list[tuple[Node, int]] = []  # (node, next_child_index)
-
-    for root in roots:
-        kroot = _identity_key(root)
-        if kroot in seen:
-            continue
-        stack.append((root, 0))
-        onstack.add(kroot)
-
-        while stack:
-            node, idx = stack[-1]
-            k = _identity_key(node)
-            if k not in seen:
-                # Pre-order yield on first encounter
-                seen.add(k)
-                yield from _maybe_yield(node, yield_condition)
-
-            deps = list(get_deps(node))
-            if idx < len(deps):
-                child = deps[idx]
-                stack[-1] = (node, idx + 1)
-                kc = _identity_key(child)
-                if kc in onstack:
-                    if detect_cycles:
-                        raise ValueError("Cycle detected during dfs_pre_iter walk")
-                    # skip on cycles if not detecting
-                elif kc not in seen:
-                    stack.append((child, 0))
-                    onstack.add(kc)
-                else:
-                    # already processed child
-                    pass
-            else:
-                stack.pop()
-                onstack.discard(k)
-
-
-def _dfs_post_iter(
-    roots: Sequence[Node],
-    get_deps: GetDeps,
-    yield_condition: YieldCond | None,
-    detect_cycles: bool,
-) -> Iterator[Node]:
-    seen: set[int] = set()
-    onstack: set[int] = set()
-    done: set[int] = set()
-    stack: list[tuple[Node, int]] = []  # (node, next_child_index)
-
-    for root in roots:
-        kroot = _identity_key(root)
-        if kroot in done:
-            continue
-        stack.append((root, 0))
-        onstack.add(kroot)
-
-        while stack:
-            node, idx = stack[-1]
-            k = _identity_key(node)
-            deps = list(get_deps(node))
-            if k not in seen:
-                seen.add(k)
-
-            if idx < len(deps):
-                child = deps[idx]
-                stack[-1] = (node, idx + 1)
-                kc = _identity_key(child)
-                if kc in done:
-                    continue
-                if kc in onstack:
-                    if detect_cycles:
-                        raise ValueError("Cycle detected during dfs_post_iter walk")
-                    # else ignore to avoid infinite loop
-                    continue
-                stack.append((child, 0))
-                onstack.add(kc)
-            else:
-                # all children processed
-                stack.pop()
-                onstack.discard(k)
-                if k not in done:
-                    done.add(k)
-                    yield from _maybe_yield(node, yield_condition)
-
-
-def _collect_closure(roots: Sequence[Node], get_deps: GetDeps) -> list[Node]:
-    """Collect reachable nodes and adjacency (by identity) for topo sort.
-
-    Returns (nodes_list, parents_map) where parents_map[v] is the set of keys of
-    dependency nodes for v (edges dep -> v).
-    """
-    nodes: list[Node] = []
-    seen: set[int] = set()
-
-    q: deque[Node] = deque(roots)
-    while q:
-        n = q.popleft()
-        kn = _identity_key(n)
-        if kn in seen:
-            continue
-        seen.add(kn)
-        nodes.append(n)
-        for d in get_deps(n):
-            q.append(d)
-    return nodes
-
-
-def _topo_kahn(
-    roots: Sequence[Node],
-    get_deps: GetDeps,
-    yield_condition: YieldCond | None,
-    detect_cycles: bool,
-) -> Iterator[Node]:
-    # Build closure and in-degree from deps edges (dep -> node)
-    nodes = _collect_closure(roots, get_deps)
-
-    # reverse map: parent -> set(children)
-    children: dict[int, set[int]] = {}
-    indeg: dict[int, int] = {}
-    for n in nodes:
-        kn = _identity_key(n)
-        indeg.setdefault(kn, 0)
-        for d in get_deps(n):
-            kd = _identity_key(d)
-            children.setdefault(kd, set()).add(kn)
-            indeg[kn] = indeg.get(kn, 0) + 1
-            indeg.setdefault(kd, indeg.get(kd, 0))
-
-    # queue of zero in-degree nodes (ready after all deps)
-    q: deque[int] = deque(k for k, v in indeg.items() if v == 0)
-    key2node: dict[int, Node] = {_identity_key(n): n for n in nodes}
-    produced = 0
-    seen_keys: set[int] = set()
-
-    while q:
-        k = q.popleft()
-        if k in seen_keys:
-            continue
-        seen_keys.add(k)
-        n = key2node[k]
-        produced += 1
-        yield from _maybe_yield(n, yield_condition)
-        for c in children.get(k, set()):
-            indeg[c] -= 1
-            if indeg[c] == 0:
-                q.append(c)
-
-    if detect_cycles and produced < len(indeg):
-        raise ValueError("Cycle detected during topo walk")

mplang/v1/core/interp.py

@@ -1,160 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Interpreter context and InterpVar implementation.
-
-This module provides the interpreter context for eager evaluation and InterpVar
-which references computed values in an interpreter.
-"""
-
-from __future__ import annotations
-
-from abc import abstractmethod
-from collections.abc import Sequence
-from typing import Any, cast
-
-from mplang.v1.core.cluster import ClusterSpec
-from mplang.v1.core.expr.ast import Expr, VariableExpr
-from mplang.v1.core.mpobject import MPContext, MPObject
-from mplang.v1.core.mptype import MPType, TensorLike
-from mplang.v1.core.tracer import TracedFunction
-from mplang.v1.utils.func_utils import var_demorph, var_morph
-
-
-# TODO(jint): Should we use inheritance or composition here?
-class InterpContext(MPContext):
-    """Context for eager evaluation using an interpreter.
-
-    InterpContext executes computations immediately and stores results
-    in an underlying interpreter.
-    """
-
-    def __init__(
-        self,
-        cluster_spec: ClusterSpec,
-    ):
-        super().__init__(cluster_spec)
-
-    @abstractmethod
-    def evaluate(self, expr: Expr, bindings: dict[str, MPObject]) -> Sequence[MPObject]:
-        """Evaluate an expression in this context.
-
-        Args:
-            expr: The expression to evaluate.
-            bindings: A dictionary of variable bindings.
-
-        Returns:
-            The result of the evaluation as an MPObject.
-        """
-        raise NotImplementedError("Should be overridden in subclasses.")
-
-    @abstractmethod
-    def fetch(self, obj: MPObject) -> list[TensorLike]:
-        """Fetch the value of an MPObject from this InterpContext to the current Python interpreter.
-
-        The MPObject must have been created by this InterpContext. If the object
-        was not produced by this context, a ValueError will be raised.
-
-        Args:
-            obj: The MPObject to fetch. Must be produced by this InterpContext.
-
-        Returns:
-            A list of tensor-like values with length equal to psize(). For each party i,
-            if the i-th bit of obj.pmask is 0 (indicating party i does not hold this value),
-            the i-th element in the returned list will be None. Otherwise, it contains
-            the actual tensor value held by party i.
-
-        Raises:
-            ValueError: If obj was not produced by this InterpContext.
-        """
-        raise NotImplementedError("Should be overridden in subclasses.")
-
-
-class InterpVar(MPObject):
-    """A variable that references a value in an interpreter.
-
-    InterpVar represents a value that has been computed and exists
-    in the interpreter's variable store.
-    """
-
-    def __init__(self, ctx: InterpContext, mptype: MPType):
-        self._ctx = ctx
-        self._mptype = mptype
-
-    @property
-    def ctx(self) -> MPContext:
-        """The context this variable belongs to."""
-        return self._ctx
-
-    @property
-    def mptype(self) -> MPType:
-        """The type of this variable."""
-        # TODO: fetch type from the Interpreter and cache it.
-        return self._mptype
-
-    def __repr__(self) -> str:
-        return f"InterpVar(mptype={self.mptype})"
-
-
-def apply(ctx: InterpContext, fn: TracedFunction, *args: Any, **kwargs: Any) -> Any:
-    is_mpobj = lambda x: isinstance(x, MPObject)
-    in_args, in_imms, in_struct = var_morph((args, kwargs), is_mpobj)
-
-    # All variables must be in the same context as the function.
-    if not all(isinstance(var, InterpVar) and var.ctx is ctx for var in in_args):
-        raise ValueError("All input variables must be InterpVars in the same context.")
-
-    # Check if the function signature matches the input types.
-    if fn.in_struct != in_struct:
-        raise ValueError(f"Input structure mismatch: {fn.in_struct} != {in_struct}")
-    if fn.in_imms != in_imms:
-        # Should trigger re-trace in JAX
-        raise ValueError(f"Input immutables mismatch: {fn.in_imms} != {in_imms}")
-    if len(fn.in_vars) != len(in_args):
-        raise ValueError(f"Input types mismatch: {fn.in_vars} != {in_args}")
-    # check parameter type match
-    for param, arg in zip(fn.in_vars, in_args, strict=False):
-        if param.mptype != arg.mptype:
-            raise ValueError(
-                f"Input variable type mismatch: {param.mptype} != {arg.mptype}"
-            )
-
-    # Prepare for the captured variables, which should also be in the same context.
-    for captured, _traced in fn.capture_map.items():
-        if not isinstance(captured, InterpVar) or captured.ctx is not ctx:
-            raise ValueError(
-                f"Capture {captured} must be in this({ctx}) context, got {captured.ctx}."
-            )
-
-    arg_binding: dict[str, MPObject] = {
-        cast(VariableExpr, var.expr).name: obj
-        for var, obj in zip(fn.in_vars, in_args, strict=False)
-    }
-    capture_binding = {
-        cast(VariableExpr, var.expr).name: captured
-        for captured, var in fn.capture_map.items()
-    }
-
-    if len(fn.out_vars) == 0:
-        out_vars: list[MPObject] = []
-    else:
-        func_expr = fn.make_expr()
-        assert func_expr is not None, "Function expression should not be None."
-        out_vars = list(
-            ctx.evaluate(func_expr.body, {**arg_binding, **capture_binding})
-        )
-
-    assert isinstance(out_vars, list), f"Expected list, got {type(out_vars)}"
-    return var_demorph(out_vars, fn.out_imms, fn.out_struct)