mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev271__py3-none-any.whl

mplang/v1/core/expr/__init__.py

@@ -1,80 +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.
-
-"""
-Expression system for multi-party computation graph construction.
-
-This package provides a modern, extensible expression-based architecture for building
-multi-party computation graphs using the visitor pattern.
-"""
-
-# Core expression types
-from mplang.v1.core.expr.ast import (
-    AccessExpr,
-    CallExpr,
-    CondExpr,
-    ConvExpr,
-    EvalExpr,
-    Expr,
-    FuncDefExpr,
-    ShflExpr,
-    ShflSExpr,
-    TupleExpr,
-    VariableExpr,
-    WhileExpr,
-)
-
-# Built-in evaluator engines
-from mplang.v1.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.v1.core.expr.printer import Printer
-from mplang.v1.core.expr.transformer import ExprTransformer
-
-# Utility functions
-from mplang.v1.core.expr.utils import (
-    deduce_mask,
-    ensure_scalar,
-    ensure_tensorlist_equal,
-    type_equal,
-)
-
-# Visitor pattern interface
-from mplang.v1.core.expr.visitor import ExprVisitor
-from mplang.v1.core.expr.walk import walk, walk_dataflow, walk_structural
-
-__all__ = [
-    "AccessExpr",
-    "CallExpr",
-    "CondExpr",
-    "ConvExpr",
-    "EvalExpr",
-    "Expr",
-    "ExprTransformer",
-    "ExprVisitor",
-    "FuncDefExpr",
-    "IEvaluator",
-    "Printer",
-    "ShflExpr",
-    "ShflSExpr",
-    "TupleExpr",
-    "VariableExpr",
-    "WhileExpr",
-    "create_evaluator",
-    "deduce_mask",
-    "ensure_scalar",
-    "ensure_tensorlist_equal",
-    "type_equal",
-    "walk",
-    "walk_dataflow",
-    "walk_structural",
-]

mplang/v1/core/expr/ast.py

@@ -1,542 +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.
-
-"""
-Abstract Syntax Tree (AST) nodes for multi-party computation expressions.
-
-This module defines the AST nodes for representing multi-party computation expressions.
-Each node type represents a different kind of operation or construct in the multi-party
-computation language, following the visitor pattern for extensible processing.
-"""
-
-from __future__ import annotations
-
-import logging
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
-
-from mplang.v1.core.expr.utils import deduce_mask
-from mplang.v1.core.mask import Mask
-from mplang.v1.core.mptype import MPType, Rank
-from mplang.v1.core.pfunc import PFunction
-from mplang.v1.core.table import TableType
-from mplang.v1.core.tensor import TensorType
-
-if TYPE_CHECKING:
-    from mplang.v1.core.expr.visitor import ExprVisitor
-
-
-class Expr(ABC):
-    """Base class for all expression types in the multi-party computation graph.
-
-    This expression system is designed to be Multi-Input Multi-Output (MIMO),
-    meaning each expression node can conceptually have multiple outputs. This is
-    fundamental to supporting multi-output PFunctions and constructing complex
-    dataflow graphs efficiently.
-
-    Attributes:
-        mptypes (list[MPType]): The list of output types for this expression. This
-            is the core property that enables MIMO capabilities. It's computed
-            lazily and cached.
-        mptype (MPType): A convenience property for the common case of a single-output
-            expression. It raises a ValueError if the expression does not have
-            exactly one output, providing a useful runtime check.
-    """
-
-    def __init__(self) -> None:
-        self._mptypes: list[MPType] | None = None
-
-    @property
-    def num_outputs(self) -> int:
-        """Return the number of outputs this expression produces."""
-        return len(self.mptypes)
-
-    @property
-    def mptypes(self) -> list[MPType]:
-        if self._mptypes is None:
-            self._mptypes = self._compute_mptypes()
-        return self._mptypes
-
-    @property
-    def mptype(self) -> MPType:
-        """Convenience property for single-output expressions."""
-        types = self.mptypes
-        if len(types) != 1:
-            raise ValueError(f"Expression has {len(types)} outputs, expected 1")
-        return types[0]
-
-    @abstractmethod
-    def _compute_mptypes(self) -> list[MPType]:
-        """Computes the types of the expression's outputs."""
-
-    @abstractmethod
-    def accept(self, visitor: ExprVisitor) -> Any:
-        """Accept a visitor for the visitor pattern."""
-
-
-# ============================================================================
-# Concrete Expression Classes
-# ============================================================================
-
-
-class EvalExpr(Expr):
-    """Expression for multi-party function evaluation."""
-
-    def __init__(
-        self, pfunc: PFunction, args: list[Expr], rmask: Mask | int | None = None
-    ):
-        super().__init__()
-        # Type checking - basic validation that we have the right number of inputs
-        if len(args) != len(pfunc.ins_info):
-            raise ValueError(
-                f"Expected {len(pfunc.ins_info)} arguments, got {len(args)}"
-            )
-        rmask = Mask(rmask) if rmask is not None else None
-
-        self.pfunc = pfunc
-        self.args = args
-        self.rmask = rmask
-
-    def _compute_mptypes(self) -> list[MPType]:
-        """Compute output MPTypes based on PFunction and mask deduction logic.
-
-        The logic follows these steps:
-        1. Determine output TensorType (dtype + shape) from PFunction
-        2. If rmask is explicitly provided (caller has strong mask knowledge):
-            2.1 Deduce pmask from args (intersection of all arg pmasks)
-                2.1.1 If deduced pmask is not None (trace time known):
-                    - If rmask is subset of deduced pmask: use rmask
-                    - If rmask is not subset of deduced pmask: raise error
-               2.1.2 If deduced pmask is None (trace time unknown): force use rmask
-        3. If rmask is not provided (caller lets expr deduce it): use deduced pmask from args
-        """
-        # Deduce pmask from arguments (including None values - if any arg has None, result is None)
-        arg_pmasks = [arg.mptype.pmask for arg in self.args]
-        deduced_pmask = deduce_mask(*arg_pmasks)
-
-        # Determine effective output pmask
-        effective_pmask: Mask | None
-        if self.rmask is not None:
-            # rmask is explicitly provided - caller has strong mask knowledge
-            if deduced_pmask is not None:
-                # pmask is known at trace time - validate subset relationship
-                if not Mask(self.rmask).is_subset(deduced_pmask):
-                    raise ValueError(
-                        f"Specified rmask {self.rmask} is not a subset of deduced pmask {deduced_pmask}."
-                    )
-                effective_pmask = self.rmask
-            else:
-                # pmask is unknown at trace time - force use rmask
-                effective_pmask = self.rmask
-        else:
-            # rmask not provided - use deduced pmask from args
-            effective_pmask = deduced_pmask
-
-        # Create result MPTypes based on PFunction output info
-        result_types = []
-        for out_info in self.pfunc.outs_info:
-            if isinstance(out_info, TensorType):
-                # Tensor type
-                result_types.append(
-                    MPType.tensor(out_info.dtype, out_info.shape, effective_pmask)
-                )
-            elif isinstance(out_info, TableType):
-                # Table type
-                result_types.append(MPType.table(out_info, effective_pmask))
-            else:
-                raise TypeError(f"Unsupported output type: {type(out_info)}")
-        return result_types
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_eval(self)
-
-
-class TupleExpr(Expr):
-    """Expression for creating a tuple from multiple single-output expressions.
-
-    In a Multi-Input Multi-Output (MIMO) expression system, this primitive
-    creates a logical tuple from multiple single-output expressions. Unlike
-    the previous FlattenExpr, TupleExpr requires all input expressions to
-    have exactly one output each.
-
-    This expression acts as a "tuple construction" primitive. It takes a list
-    of single-output expressions and produces a new logical expression whose
-    outputs are the list of all input expression outputs.
-
-    For example, if expr1 has output [A] and expr2 has output [B],
-    TupleExpr([expr1, expr2]) will have outputs [A, B].
-
-    This is the opposite of AccessExpr, which extracts a single element
-    from a multi-output expression.
-    """
-
-    def __init__(self, args: list[Expr]):
-        super().__init__()
-        # Validate that all arguments are single-output expressions
-        for i, arg in enumerate(args):
-            if arg.num_outputs != 1:
-                raise ValueError(
-                    f"TupleExpr requires all arguments to be single-output expressions, "
-                    f"but argument {i} has {arg.num_outputs} outputs"
-                )
-        self.args = args
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # TupleExpr creates a tuple from single-output expressions
-        result_types = []
-        for arg in self.args:
-            result_types.append(
-                arg.mptype
-            )  # Use mptype since we validated single output
-        return result_types
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_tuple(self)
-
-
-class CondExpr(Expr):
-    """Expression for conditional execution.
-
-    Added fields:
-        verify_uniform: whether runtime should assert the predicate is uniform across parties.
-    """
-
-    def __init__(
-        self,
-        pred: Expr,
-        then_fn: FuncDefExpr,
-        else_fn: FuncDefExpr,
-        args: list[Expr],
-        verify_uniform: bool = False,
-    ):
-        super().__init__()
-        self.pred = pred
-        self.then_fn = then_fn
-        self.else_fn = else_fn
-        self.args = args
-        self.verify_uniform = verify_uniform
-
-    def _compute_mptypes(self) -> list[MPType]:
-        for t_type, e_type in zip(
-            self.then_fn.mptypes, self.else_fn.mptypes, strict=False
-        ):
-            if t_type != e_type:
-                raise TypeError(
-                    f"Then branch type {t_type} does not match else branch type {e_type}"
-                )
-        return self.then_fn.mptypes
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_cond(self)
-
-
-class WhileExpr(Expr):
-    """Expression for while loop."""
-
-    def __init__(
-        self,
-        cond_fn: FuncDefExpr,
-        body_fn: FuncDefExpr,
-        args: list[Expr],
-    ):
-        super().__init__()
-        if not args:
-            raise ValueError("WhileExpr requires at least one argument (init value)")
-        self.cond_fn = cond_fn
-        self.body_fn = body_fn
-        self.args = args
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # The result types of a while loop are the same as the body function's outputs.
-        # This supports multi-value loop-carried state (PyTree leaves) and ensures
-        # evaluator can determine how many values are produced by the loop.
-        return self.body_fn.mptypes
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_while(self)
-
-
-class ConvExpr(Expr):
-    """Expression for convergence of multiple variables."""
-
-    def __init__(self, vars: list[Expr]):
-        super().__init__()
-
-        # Validate all vars have identical out-length.
-        for v in vars:
-            if v.num_outputs != 1:
-                raise ValueError("All variables in ConvExpr must have the same arity.")
-
-        self.vars = vars
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # Collect the idx-th mptype from every var.
-        types = [v.mptype for v in self.vars]
-        # Validate dtype / shape consistency.
-        first = types[0]
-        for c in types[1:]:
-            if c.raw_type() != first.raw_type():
-                raise TypeError(f"Inconsistent type in pconv: {c} vs {first}")
-
-        # Deduce the pmask by intersecting all pmasks.
-        pmasks = [t.pmask for t in types]
-        dynamic_pmask = False
-        if any(pmask is None for pmask in pmasks):
-            logging.warning("pconv called with None pmask.")
-            dynamic_pmask = True
-
-        non_none_pmasks = [pmask for pmask in pmasks if pmask is not None]
-        for i, mask1 in enumerate(non_none_pmasks):
-            for mask2 in non_none_pmasks[i + 1 :]:
-                if not Mask(mask1).is_disjoint(mask2):
-                    raise ValueError(
-                        f"pconv called with non-disjoint pmasks: {pmasks}."
-                    )
-
-        # deduce output pmask.
-        if dynamic_pmask:
-            out_pmask = None
-        else:
-            valid_pmasks = [pmask for pmask in pmasks if pmask is not None]
-            if valid_pmasks:
-                out_pmask = Mask(valid_pmasks[0])
-                for mask in valid_pmasks[1:]:
-                    out_pmask = out_pmask.union(mask)
-            else:
-                out_pmask = None
-
-        return [MPType(first.raw_type(), out_pmask, first.attrs)]
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_conv(self)
-
-
-class ShflSExpr(Expr):
-    """Expression for static shuffle operation.
-
-    Redistributes data from source ranks to target ranks based on a specified
-    mapping. Each party in the output mask (`pmask`) receives data from a
-    corresponding source rank specified in `src_ranks`.
-
-    Rationale for Design (Pull vs. Push Model):
-        This operation uses a "pull" model, where each receiving party explicitly
-        states its data source (`src_ranks`). This contrasts with a "push" model,
-        where each sending party would specify a destination.
-
-        The pull model is chosen because it guarantees that every party in the
-        output `pmask` receives exactly one value, upholding the semantic
-        integrity of the computation graph.
-
-        A push model, on the other hand, would be semantically ambiguous. For
-        example, two different source parties could attempt to send data to the
-        same destination, or some parties might receive no data at all. This
-        would break the Single Instruction, Multiple Programs (SIMP) paradigm by
-        creating an unpredictable number of outputs at each party.
-
-        While the pull model might have performance implications if multiple
-        receivers pull from the same source (potentially creating a network
-        bottleneck at that source), this is a performance consideration rather
-        than a correctness issue. The chosen design prioritizes semantic
-        predictability and correctness.
-    """
-
-    def __init__(self, src_val: Expr, pmask: Mask, src_ranks: list[Rank]):
-        """Initialize static shuffle expression.
-
-        Args:
-            src_val (Expr): The input tensor to be shuffled.
-            pmask (Mask): The mask indicating which parties will hold the output.
-                         Only parties with non-zero bits in pmask will receive output.
-            src_ranks (list[Rank]): List of source ranks. The i-th output party
-                                   (i-th non-zero bit in pmask) receives data from
-                                   src_ranks[i].
-
-        Raises:
-            ValueError: If src_val has multiple outputs, if src_ranks length doesn't
-                       match pmask bit count, or if any rank in src_ranks is not
-                       present in src_val.pmask.
-
-        Example:
-            If pmask indicates parties [0, 2] should receive output and src_ranks = [1, 3], then:
-            - Party 0 receives data from rank 1
-            - Party 2 receives data from rank 3
-        """
-        super().__init__()
-        if src_val.num_outputs != 1:
-            raise ValueError(
-                f"ShflSExpr requires a single output source, got {src_val.num_outputs}"
-            )
-
-        # Assign values first before validation
-        self.src_val = src_val
-        self.pmask = pmask
-        self.src_ranks = src_ranks
-
-        # Now do validation using the assigned values
-        if len(self.src_ranks) != Mask(self.pmask).num_parties():
-            raise ValueError(
-                f"src_ranks length ({len(self.src_ranks)}) not match {self.pmask}"
-            )
-        for i, rank in enumerate(self.src_ranks):
-            src_pmask = self.src_val.mptype.pmask
-            if src_pmask is not None and rank not in Mask(src_pmask):
-                raise ValueError(
-                    f"Source rank {rank} at index {i} is not present in src {Mask(src_pmask)}"
-                )
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # The types are the same as the source value, but with a new pmask.
-        src_type = self.src_val.mptype
-        return [MPType(src_type._type, self.pmask, src_type.attrs)]
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_shfl_s(self)
-
-
-class ShflExpr(Expr):
-    """Expression for dynamic shuffle operation."""
-
-    def __init__(self, src: Expr, index: Expr):
-        super().__init__()
-        self.src = src
-        self.index = index
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # Dynamic shuffle is complex. The resulting pmask is often unknown
-        # at compile time. We'll assume the tensor types remain the same
-        # but the pmask becomes None (runtime-determined).
-        src_types = self.src.mptypes
-        result_types = []
-        for src_type in src_types:
-            result_types.append(
-                MPType.tensor(src_type.dtype, src_type.shape, None, **src_type.attrs)
-            )
-        return result_types
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_shfl(self)
-
-
-class AccessExpr(Expr):
-    """Expression for accessing a specific output of a multi-output expression.
-
-    As the counterpart to TupleExpr, AccessExpr is the "un-packing" or "selection"
-    primitive in the MIMO system. It takes a (potentially multi-output) expression
-    and an index, and produces a new single-output expression representing just
-    the selected output.
-
-    This is essential for routing specific outputs from a multi-output function
-    or a flattened stream to subsequent operations that expect single inputs.
-    """
-
-    def __init__(self, src: Expr, index: int):
-        super().__init__()
-        self.src = src
-        self.index = index
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # Access a specific output from the expression's output list
-        expr_types = self.src.mptypes
-        if self.index < 0 or self.index >= len(expr_types):
-            raise IndexError(
-                f"Index {self.index} out of range for expression with {len(expr_types)} outputs"
-            )
-        return [expr_types[self.index]]
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_access(self)
-
-
-class VariableExpr(Expr):
-    """Expression for variable reference/lookup."""
-
-    def __init__(self, name: str, mptype: MPType):
-        super().__init__()
-        self.name = name
-        self.mptype_value = mptype
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # Return the explicitly provided type for this variable.
-        return [self.mptype_value]
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_variable(self)
-
-
-class FuncDefExpr(Expr):
-    """Expression representing a function definition with parameters and body.
-
-    This class captures the essence of lambda abstraction in functional programming.
-    The body expression tree may contain free variables (VariableExpr nodes) that
-    reference parameter names. When the function is called, arguments are bound
-    to parameters positionally, resolving these free variables.
-
-    Example:
-        Consider a function that adds two variables:
-        ```
-        # Body expression tree contains free variables "x" and "y"
-        body = EvalExpr(
-            add_pfunc, [VariableExpr("x", int_type), VariableExpr("y", int_type)]
-        )
-
-        # Parameters define the binding order - note "y" comes before "x"
-        params = ["z", "y", "x"]  # extra parameter "z", different order
-
-        func_def = FuncDefExpr(params, body)
-
-        # When called with [expr0, expr1, expr2]:
-        # - "z" binds to expr0 (unused in body, but valid)
-        # - "y" binds to expr1 (resolves VariableExpr("y") in body)
-        # - "x" binds to expr2 (resolves VariableExpr("x") in body)
-        call = CallExpr(func_def, [expr0, expr1, expr2])
-        ```
-
-    Key insights:
-    - Free variables in the body are placeholders waiting for concrete expressions
-    - Parameters act as a "binding contract" - they define which arguments map to which variables
-    - Parameter order matters for positional binding, not alphabetical or usage order
-    - Parameters can include names not used in the body (dead parameters)
-    - All free variables in the body should have corresponding parameters for well-formed functions
-    """
-
-    def __init__(self, params: list[str], body: Expr):
-        super().__init__()
-        self.params = params
-        self.body = body
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # The types of a function are the types of its body.
-        return self.body.mptypes
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_func_def(self)
-
-
-class CallExpr(Expr):
-    """Expression for function call."""
-
-    def __init__(self, name: str, fn: FuncDefExpr, args: list[Expr]):
-        super().__init__()
-        self.name = name
-        self.fn = fn
-        self.args = args
-
-    def _compute_mptypes(self) -> list[MPType]:
-        # The result types are the types of the function's body, with parameter
-        # types substituted. For simplicity, we return the function's declared
-        # return types. A full implementation would require substitution logic.
-        return self.fn.mptypes
-
-    def accept(self, visitor: ExprVisitor) -> Any:
-        return visitor.visit_call(self)