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

mplang/__init__.py

@@ -12,35 +12,409 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Multi-Party Programming Language for Secure Computation.
+"""MPLang2: Next generation EDSL for multi-party computation.
 
-This module provides access to both v1 and v2 APIs:
+This is the temporary home for the new EDSL implementation during migration.
+Once migration is complete, this will replace the original mplang package.
 
-    # Use v1 (current default, will be deprecated)
-    import mplang as mp
-    # or explicitly:
-    import mplang.v1 as mp
+Public API is designed to be compatible with mplang v1 where possible.
+"""
 
-    # Use v2 (recommended for new code)
-    import mplang.v2 as mp
+from __future__ import annotations
 
-v1 API will be deprecated in a future release. Please migrate to v2.
-"""
+from collections.abc import Callable
+from typing import Any
 
 # Version is managed by hatch-vcs and available after package installation
 try:
-    from importlib.metadata import PackageNotFoundError, version
+    from importlib.metadata import version
 
     __version__ = version("mplang")
-except PackageNotFoundError:
+except Exception:
     # Fallback for development/editable installs when package is not installed
     __version__ = "0.0.0-dev"
 
+from mplang import dialects
+from mplang.backends.simp_driver.ops import DRIVER_HANDLERS
+from mplang.backends.simp_worker import SimpWorker
+from mplang.backends.simp_worker.mem import LocalMesh
+from mplang.backends.simp_worker.ops import WORKER_HANDLERS
+from mplang.dialects.simp import make_driver, make_simulator
+from mplang.edsl import (
+    Graph,
+    GraphPrinter,
+    Object,
+    Operation,
+    Primitive,
+    TracedFunction,
+    Tracer,
+    Value,
+    find_context,
+    find_context_with_state,
+    find_interpreter,
+    format_graph,
+    get_current_context,
+    get_default_context,
+    is_tracing,
+    jit,
+    pop_context,
+    primitive,
+    push_context,
+    register_default_context_factory,
+    set_root_context,
+    trace,
+)
+from mplang.edsl.registry import get_profiler
+
+# Type system
+from mplang.edsl.typing import (
+    MPType,
+    ScalarType,
+    SSType,
+    TableType,
+    TensorType,
+    VectorType,
+)
+
+# =============================================================================
+# Device API (compatible with mplang v1)
+# =============================================================================
+from mplang.libs.device import (
+    ClusterSpec,
+    Device,
+    Node,
+    device,
+    get_dev_attr,
+    is_device_obj,
+    jax_fn,
+    put,
+    set_dev_attr,
+)
+from mplang.libs.device import fetch as device_fetch
+from mplang.runtime.interpreter import Interpreter
+
+# =============================================================================
+# Context Management API (JAX-like pattern)
+# =============================================================================
+
+
+def _get_context(context: Interpreter | None) -> Interpreter:
+    """Get context from parameter or context stack."""
+    if context is not None:
+        return context
+    ctx = get_current_context()
+    if ctx is None:
+        raise RuntimeError(
+            "No context available. Either pass context explicitly or use "
+            "set_context()/push_context() to set a default context."
+        )
+    if not isinstance(ctx, Interpreter):
+        raise RuntimeError(
+            f"Current context is not an Interpreter: {type(ctx).__name__}. "
+            "Use mp.set_context(interpreter) to set the execution context."
+        )
+    return ctx
+
+
+# =============================================================================
+# Meta-APIs (compile, evaluate, fetch)
+# =============================================================================
+
+
+def evaluate(
+    fn: Callable[..., Any] | TracedFunction,
+    *args: Any,
+    context: Interpreter | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Evaluate a function or traced function.
+
+    Args:
+        fn: The function or TracedFunction to evaluate.
+        *args: Positional arguments to pass to the function.
+        context: Optional interpreter context. If None, uses current context.
+        **kwargs: Keyword arguments to pass to the function.
+
+    Returns:
+        The result of the function evaluation.
+
+    Example:
+        >>> with mp.make_simulator(3) as sim:
+        ...     result = mp.evaluate(traced)  # uses sim from context
+        >>> # Or explicitly:
+        >>> result = mp.evaluate(traced, context=sim)
+    """
+    from mplang.edsl.tracer import TracedFunction
+    from mplang.runtime.interpreter import InterpObject
+
+    interp = _get_context(context)
+
+    def unwrap_if_interp(val: Any) -> Any:
+        """Unwrap InterpObject to runtime value at execution boundary."""
+        if isinstance(val, InterpObject):
+            return val.runtime_obj
+        return val
+
+    with interp:
+        if isinstance(fn, TracedFunction):
+            inputs = fn.prepare_inputs(*args, **kwargs)
+            inputs = [unwrap_if_interp(v) for v in inputs]
+            raw_result = interp.evaluate_graph(fn.graph, inputs)
+            wrapped = [
+                InterpObject(v, fn.graph.outputs[i].type, interp)
+                for i, v in enumerate(raw_result)
+            ]
+            return fn.reconstruct_outputs(wrapped)
+
+        return fn(*args, **kwargs)
+
+
+def fetch(
+    result: Any,
+    *,
+    follow_device: bool = True,
+    context: Interpreter | None = None,
+) -> Any:
+    """Fetch results from interpreter context to Python.
+
+    This is a meta-function that operates at execution boundaries, not a traced
+    dialect operation. It brings data from the distributed/MPC runtime back to
+    the Python host.
+
+    Behavior in different contexts:
+    - **Tracing (compile)**: Returns the input unchanged (identity). The graph
+      outputs are determined by the function's return statement, not fetch calls.
+    - **Execution (evaluate)**: Actually fetches data from workers/parties.
+
+    Design Note (A vs B tradeoff):
+        Two designs were considered for fetch behavior during tracing:
+
+        - **Design A (chosen)**: fetch = identity during tracing. Graph outputs
+          are determined solely by the return statement. This is simpler and
+          avoids ambiguity when fetch and return reference different values.
+
+        - **Design B (alternative)**: fetch marks output points in the graph.
+          This would allow fetch(a), fetch(b), return b to output both a and b.
+          However, it complicates the semantics and requires tracking fetch
+          points separately from return values.
+
+        Design A was chosen for simplicity. If a value needs to be an output,
+        it should be returned. fetch's role is purely for execution-time I/O.
+
+    Args:
+        result: Object(s) to fetch. Can be a single InterpObject, DriverVar,
+            or nested structure containing them.
+        follow_device: If True and object has device attribute, dispatch to
+            device.fetch which fetches from the correct rank based on device.
+            If False, fetch from all parties.
+        context: Interpreter context. If None, uses current context.
+
+    Returns:
+        Fetched Python values. For device objects with follow_device=True,
+        returns single value from the device's rank(s). Otherwise returns
+        list of values (one per party) or single value for world_size=1.
+        During tracing, returns the input unchanged.
+    """
+    from jax.tree_util import tree_map
+
+    from mplang.backends.simp_driver.values import DriverVar
+    from mplang.edsl.context import is_tracing
+    from mplang.runtime.interpreter import InterpObject
+    from mplang.runtime.value import WrapValue
+
+    # Check if we are in tracing context - if so, return identity
+    if is_tracing():
+        # Design A: fetch = identity during tracing
+        # Graph outputs are determined by return statement, not fetch calls
+        return result
+
+    # Execution context - actually fetch data
+    interp = _get_context(context)
+
+    def _fetch_single(var: Any) -> Any:
+        """Fetch a single value from InterpObject."""
+        # InterpObject (from mp.evaluate) - extract runtime_obj
+        if isinstance(var, InterpObject):
+            if follow_device and is_device_obj(var):
+                return device_fetch(var)
+            var = var.runtime_obj  # extract and continue processing
+
+        # DriverVar (simp dialect) - remote fetch from workers
+        if isinstance(var, DriverVar):
+            from mplang.backends.simp_driver.state import SimpDriver
+
+            simp_state = interp.get_dialect_state("simp")
+            assert isinstance(simp_state, SimpDriver), "DriverVar requires simp state"
+
+            resolved: list[Any] = []
+            for rank, uri in enumerate(var.values):
+                if uri is None:
+                    resolved.append(None)
+                else:
+                    fetched = simp_state.fetch(rank, uri).result()
+                    if isinstance(fetched, WrapValue):
+                        fetched = fetched.data
+                    resolved.append(fetched)
+
+            return resolved[0] if len(resolved) == 1 else resolved
+
+        # WrapValue (TensorValue, TableValue, etc.) - unwrap
+        if isinstance(var, WrapValue):
+            return var.data
+
+        # Plain values pass through
+        return var
+
+    with interp:
+        return tree_map(_fetch_single, result)
+
+
+# Alias for compatibility
+def function(fn: Callable[..., Any] | None = None) -> Callable[..., Any]:
+    """Decorator defining a Multi-Party Function (MP Program).
+
+    This decorator "lifts" a local function into a distributed program by
+    automatically wrapping it in a `simp.pcall_static` that targets ALL available
+    parties in the current context.
+
+    Semantics: f(args) -> pcall(ALL, f, args)
+
+    Args:
+        fn: The function to decorate.
+
+    Returns:
+        A wrapper function that, when called, executes the original function
+        on all workers.
+    """
+    import functools
+
+    from mplang.dialects import simp
+
+    if fn is None:
+        return function
+
+    def has_simp_state(ctx: Any) -> bool:
+        if hasattr(ctx, "get_dialect_state"):
+            state = ctx.get_dialect_state("simp")
+            return state is not None and hasattr(state, "world_size")
+        return False
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        from mplang.edsl.context import find_context
+
+        # Find context with simp dialect state
+        ctx = find_context(has_simp_state)
+        if ctx is None:
+            raise RuntimeError(
+                "mp.function requires a context with world_size information "
+                "(e.g. SimpSimulator or Driver initialized)."
+            )
+
+        # ctx found by predicate so we know it has get_dialect_state
+        simp_state = ctx.get_dialect_state("simp")  # type: ignore[attr-defined]
+        world_size = simp_state.world_size  # type: ignore
+
+        all_parties = tuple(range(world_size))
+        return simp.pcall_static(all_parties, fn, *args, **kwargs)
+
+    return wrapper
+
+
+def compile(
+    fn: Callable[..., Any],
+    *args: Any,
+    context: Interpreter | None = None,
+    **kwargs: Any,
+) -> TracedFunction:
+    """Compile a function to get its IR without executing it.
+
+    Args:
+        fn: The function to compile.
+        *args: Example arguments for tracing.
+        context: Optional interpreter context. If None, uses current context.
+        **kwargs: Example keyword arguments for tracing.
+
+    Returns:
+        TracedFunction with the compiled graph.
+
+    Example:
+        >>> with mp.make_simulator(3) as sim:
+        ...     traced = mp.compile(job)  # uses sim from context
+    """
+    # If a context is explicitly provided, push it before tracing
+    # so that _resolve_cluster() can find it.
+    if context is not None:
+        with context:
+            return trace(fn, *args, **kwargs)
+
+    # Otherwise, rely on the caller having pushed an interpreter context.
+    # _resolve_cluster() will traverse the stack to find the interpreter.
+    return trace(fn, *args, **kwargs)
+
+
 # =============================================================================
-# Default: Re-export v1 API for backward compatibility
+# Public API
 # =============================================================================
-# This will be changed to v2 in a future release
+__all__ = [  # noqa: RUF022
+    # Version
+    "__version__",
+    # Device API
+    "ClusterSpec",
+    "Device",
+    "Node",
+    "device",
+    "get_dev_attr",
+    "is_device_obj",
+    "jax_fn",
+    "put",
+    "set_dev_attr",
+    # Core EDSL
+    "Graph",
+    "GraphPrinter",
+    "Object",
+    "Operation",
+    "Primitive",
+    "TracedFunction",
+    "Tracer",
+    "Value",
+    "compile",
+    "evaluate",
+    "fetch",
+    "find_context",
+    "find_context_with_state",
+    "find_interpreter",
+    "format_graph",
+    "function",
+    "get_current_context",
+    "get_default_context",
+    "is_tracing",
+    "jit",
+    "pop_context",
+    "primitive",
+    "push_context",
+    "set_root_context",
+    "trace",
+    # Type system
+    "MPType",
+    "ScalarType",
+    "SSType",
+    "TableType",
+    "TensorType",
+    "VectorType",
+    # Backend / Runtime
+    "DRIVER_HANDLERS",
+    "Interpreter",
+    "LocalMesh",
+    "SimpWorker",
+    "WORKER_HANDLERS",
+    "make_driver",
+    "make_simulator",
+    # Dialects
+    "dialects",
+    "register_default_context_factory",
+    "get_profiler",
+]
 
-# Make v1 and v2 subpackages directly accessible
-from mplang import v1, v2  # noqa: F401
-from mplang.v1 import *  # noqa: F403
+# Register Interpreter as default context factory
+register_default_context_factory(Interpreter)

mplang/{v2/backends → backends}/__init__.py

@@ -41,13 +41,15 @@ def load_builtins() -> None:
     """Load all built-in backend implementations."""
     # Core backends that are expected to be present
     builtin_backends = [
-        "mplang.v2.backends.spu_impl",
-        "mplang.v2.backends.tensor_impl",
-        "mplang.v2.backends.table_impl",
-        "mplang.v2.backends.crypto_impl",
-        "mplang.v2.backends.tee_impl",
-        "mplang.v2.backends.bfv_impl",
-        "mplang.v2.backends.store_impl",
+        "mplang.backends.spu_impl",
+        "mplang.backends.tensor_impl",
+        "mplang.backends.table_impl",
+        "mplang.backends.crypto_impl",
+        "mplang.backends.tee_impl",
+        "mplang.backends.bfv_impl",
+        "mplang.backends.store_impl",
+        "mplang.backends.func_impl",
+        "mplang.backends.field_impl",
     ]
 
     for module_name in builtin_backends:

mplang/{v2/backends → backends}/bfv_impl.py

@@ -29,12 +29,12 @@ import numpy as np
 import tenseal as ts
 import tenseal.sealapi as sealapi
 
-from mplang.v2.backends.tensor_impl import TensorValue
-from mplang.v2.dialects import bfv
-from mplang.v2.edsl import serde
-from mplang.v2.edsl.graph import Operation
-from mplang.v2.runtime.interpreter import Interpreter
-from mplang.v2.runtime.value import Value, WrapValue
+from mplang.backends.tensor_impl import TensorValue
+from mplang.dialects import bfv
+from mplang.edsl import serde
+from mplang.edsl.graph import Operation
+from mplang.runtime.interpreter import Interpreter
+from mplang.runtime.value import Value, WrapValue
 
 # =============================================================================
 # Helper for SEAL serialization

mplang/{v2/backends → backends}/crypto_impl.py

@@ -26,12 +26,12 @@ import coincurve
 import numpy as np
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 
-from mplang.v2.backends.tensor_impl import TensorValue
-from mplang.v2.dialects import crypto
-from mplang.v2.edsl import serde
-from mplang.v2.edsl.graph import Operation
-from mplang.v2.runtime.interpreter import Interpreter
-from mplang.v2.runtime.value import Value, WrapValue
+from mplang.backends.tensor_impl import TensorValue
+from mplang.dialects import crypto
+from mplang.edsl import serde
+from mplang.edsl.graph import Operation
+from mplang.runtime.interpreter import Interpreter
+from mplang.runtime.value import Value, WrapValue
 
 # =============================================================================
 # BytesValue - Wrapper for raw bytes (keys, hashes, ciphertexts)

mplang/{v2/backends → backends}/field_impl.py

@@ -29,11 +29,11 @@ import threading
 import jax.numpy as jnp
 import numpy as np
 
-from mplang.v2.backends.tensor_impl import TensorValue, _unwrap, _wrap
-from mplang.v2.dialects import field
-from mplang.v2.edsl.graph import Operation
-from mplang.v2.kernels import py_kernels
-from mplang.v2.runtime.interpreter import Interpreter
+from mplang.backends.tensor_impl import TensorValue, _unwrap, _wrap
+from mplang.dialects import field
+from mplang.edsl.graph import Operation
+from mplang.kernels import py_kernels
+from mplang.runtime.interpreter import Interpreter
 
 # =============================================================================
 # Kernel Loading

mplang/{v2/backends → backends}/func_impl.py

@@ -23,10 +23,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, ClassVar
 
-from mplang.v2.dialects.func import call_p, func_def_p
-from mplang.v2.edsl import serde
-from mplang.v2.edsl.graph import Graph, Operation
-from mplang.v2.runtime.value import Value
+from mplang.dialects.func import call_p, func_def_p
+from mplang.edsl import serde
+from mplang.edsl.graph import Graph, Operation
+from mplang.runtime.value import Value
 
 if TYPE_CHECKING:
     from typing import Self

mplang/{v2/backends → backends}/phe_impl.py

@@ -22,9 +22,9 @@ from typing import Any, cast
 from lightphe import LightPHE
 from lightphe.models.Ciphertext import Ciphertext
 
-from mplang.v2.dialects import phe
-from mplang.v2.edsl.graph import Operation
-from mplang.v2.runtime.interpreter import Interpreter
+from mplang.dialects import phe
+from mplang.edsl.graph import Operation
+from mplang.runtime.interpreter import Interpreter
 
 
 class PHEContext:

mplang/{v2/backends → backends}/simp_design.md

@@ -67,7 +67,7 @@ class SimpDriver(DialectState, ABC):
     """Abstract interface for simp drivers."""
     dialect_name = "simp"
     world_size: int
-    
+
     @abstractmethod
     def submit(self, rank, graph, inputs, job_id=None) -> Future: ...
     @abstractmethod

mplang/{v2/backends → backends}/simp_driver/__init__.py

@@ -17,16 +17,16 @@
 Provides Driver-side state, values, and ops for the simp dialect.
 """
 
-from mplang.v2.backends.simp_driver.http import SimpHttpDriver, make_driver
-from mplang.v2.backends.simp_driver.mem import (
+from mplang.backends.simp_driver.http import SimpHttpDriver, make_driver
+from mplang.backends.simp_driver.mem import (
     LocalCluster,
     MemCluster,
     SimpMemDriver,
     make_simulator,
 )
-from mplang.v2.backends.simp_driver.ops import DRIVER_HANDLERS
-from mplang.v2.backends.simp_driver.state import SimpDriver
-from mplang.v2.backends.simp_driver.values import DriverVar
+from mplang.backends.simp_driver.ops import DRIVER_HANDLERS
+from mplang.backends.simp_driver.state import SimpDriver
+from mplang.backends.simp_driver.values import DriverVar
 
 __all__ = [
     "DRIVER_HANDLERS",

mplang/{v2/backends → backends}/simp_driver/http.py

@@ -23,16 +23,16 @@ from typing import TYPE_CHECKING, Any
 
 import httpx
 
-from mplang.v2.backends.simp_driver.state import SimpDriver
-from mplang.v2.edsl import serde
-from mplang.v2.runtime.interpreter import Interpreter
-from mplang.v2.runtime.object_store import ObjectStore
+from mplang.backends.simp_driver.state import SimpDriver
+from mplang.edsl import serde
+from mplang.runtime.interpreter import Interpreter
+from mplang.runtime.object_store import ObjectStore
 
 if TYPE_CHECKING:
     from concurrent.futures import Future
 
-    from mplang.v2.edsl.graph import Graph
-    from mplang.v2.libs.device import ClusterSpec
+    from mplang.edsl.graph import Graph
+    from mplang.libs.device import ClusterSpec
 
 
 class SimpHttpDriver(SimpDriver):
@@ -142,10 +142,10 @@ def make_driver(endpoints: list[str], *, cluster_spec: Any = None) -> Interprete
         >>> with interp:
         ...     result = my_func()
     """
-    from mplang.v2.backends.simp_driver.ops import DRIVER_HANDLERS
+    from mplang.backends.simp_driver.ops import DRIVER_HANDLERS
 
     if cluster_spec is None:
-        from mplang.v2.libs.device import ClusterSpec
+        from mplang.libs.device import ClusterSpec
 
         cluster_spec = ClusterSpec.simple(
             world_size=len(endpoints), endpoints=endpoints

mplang/{v2/backends → backends}/simp_driver/mem.py

@@ -22,17 +22,17 @@ import pathlib
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, cast
 
-from mplang.v2.backends.simp_driver.state import SimpDriver
-from mplang.v2.backends.simp_worker import WORKER_HANDLERS, SimpWorker
-from mplang.v2.backends.simp_worker.mem import LocalMesh
-from mplang.v2.runtime.interpreter import ExecutionTracer, Interpreter
-from mplang.v2.runtime.object_store import ObjectStore
+from mplang.backends.simp_driver.state import SimpDriver
+from mplang.backends.simp_worker import WORKER_HANDLERS, SimpWorker
+from mplang.backends.simp_worker.mem import LocalMesh
+from mplang.runtime.interpreter import ExecutionTracer, Interpreter
+from mplang.runtime.object_store import ObjectStore
 
 if TYPE_CHECKING:
     from concurrent.futures import Future
 
-    from mplang.v2.edsl.graph import Graph
-    from mplang.v2.libs.device import ClusterSpec
+    from mplang.edsl.graph import Graph
+    from mplang.libs.device import ClusterSpec
 
 
 class MemCluster:
@@ -245,10 +245,10 @@ def make_simulator(
         >>> with interp:
         ...     result = my_func()
     """
-    from mplang.v2.backends.simp_driver.ops import DRIVER_HANDLERS
+    from mplang.backends.simp_driver.ops import DRIVER_HANDLERS
 
     if cluster_spec is None:
-        from mplang.v2.libs.device import ClusterSpec
+        from mplang.libs.device import ClusterSpec
 
         cluster_spec = ClusterSpec.simple(world_size)
 

mplang/{v2/backends → backends}/simp_driver/ops.py

@@ -23,10 +23,10 @@ from __future__ import annotations
 
 from typing import Any
 
-from mplang.v2.backends.simp_driver.values import DriverVar
-from mplang.v2.dialects import simp
-from mplang.v2.edsl.graph import Graph, Operation
-from mplang.v2.edsl.typing import CustomType
+from mplang.backends.simp_driver.values import DriverVar
+from mplang.dialects import simp
+from mplang.edsl.graph import Graph, Operation
+from mplang.edsl.typing import CustomType
 
 
 def _get_driver_context(interpreter: Any) -> Any:

mplang/{v2/backends → backends}/simp_driver/state.py

@@ -19,12 +19,12 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any
 
-from mplang.v2.runtime.dialect_state import DialectState
+from mplang.runtime.dialect_state import DialectState
 
 if TYPE_CHECKING:
     from concurrent.futures import Future
 
-    from mplang.v2.edsl.graph import Graph
+    from mplang.edsl.graph import Graph
 
 
 class SimpDriver(DialectState, ABC):