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

mplang/v1/core/__init__.py

@@ -0,0 +1,157 @@
+# 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.
+
+"""
+Core components for multi-party computation.
+
+This package provides the fundamental building blocks for multi-party computation,
+including type systems, tracing mechanisms, and interpreter contexts.
+"""
+
+# Core type system
+# Communication interfaces & core symbols
+from mplang.v1.core.cluster import ClusterSpec, Device, Node, RuntimeInfo
+from mplang.v1.core.comm import (
+    CollectiveMixin,
+    CommunicatorBase,
+    ICollective,
+    ICommunicator,
+)
+from mplang.v1.core.context_mgr import cur_ctx, set_ctx, with_ctx
+from mplang.v1.core.dtypes import (
+    BINARY,
+    BOOL,
+    COMPLEX64,
+    COMPLEX128,
+    DATE,
+    DECIMAL,
+    FLOAT16,
+    FLOAT32,
+    FLOAT64,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    INTERVAL,
+    JSON,
+    STRING,
+    TIME,
+    TIMESTAMP,
+    UINT8,
+    UINT16,
+    UINT32,
+    UINT64,
+    UUID,
+    DType,
+)
+from mplang.v1.core.interp import InterpContext, InterpVar
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.mpir import IrReader, IrWriter
+from mplang.v1.core.mpobject import MPContext, MPObject
+from mplang.v1.core.mptype import MPType, Rank, Shape
+from mplang.v1.core.pfunc import PFunction, get_fn_name
+
+# Import primitive-dependent symbols at the end to avoid cycles when frontend ops
+# import from the core facade during package initialization.
+from mplang.v1.core.primitive import (
+    builtin_function,
+    function,
+    pconv,
+    peval,
+    pmask,
+    pshfl,
+    pshfl_s,
+    psize,
+    uniform_cond,
+    while_loop,
+)
+from mplang.v1.core.table import TableLike, TableType
+from mplang.v1.core.tensor import ScalarType, TensorLike, TensorType
+from mplang.v1.core.tracer import (
+    TraceContext,
+    TracedFunction,
+    TraceVar,
+    VarNamer,
+    trace,
+)
+
+__all__ = [
+    "BINARY",
+    "BOOL",
+    "COMPLEX64",
+    "COMPLEX128",
+    "DATE",
+    "DECIMAL",
+    "FLOAT16",
+    "FLOAT32",
+    "FLOAT64",
+    "INT8",
+    "INT16",
+    "INT32",
+    "INT64",
+    "INTERVAL",
+    "JSON",
+    "STRING",
+    "TIME",
+    "TIMESTAMP",
+    "UINT8",
+    "UINT16",
+    "UINT32",
+    "UINT64",
+    "UUID",
+    "ClusterSpec",
+    "CollectiveMixin",
+    "CommunicatorBase",
+    "DType",
+    "Device",
+    "ICollective",
+    "ICommunicator",
+    "InterpContext",
+    "InterpVar",
+    "IrReader",
+    "IrWriter",
+    "MPContext",
+    "MPObject",
+    "MPType",
+    "Mask",
+    "Node",
+    "PFunction",
+    "Rank",
+    "RuntimeInfo",
+    "ScalarType",
+    "Shape",
+    "TableLike",
+    "TableType",
+    "TensorLike",
+    "TensorType",
+    "TraceContext",
+    "TraceVar",
+    "TracedFunction",
+    "VarNamer",
+    "builtin_function",
+    "cur_ctx",
+    "function",
+    "get_fn_name",
+    "pconv",
+    "peval",
+    "pmask",
+    "pshfl",
+    "pshfl_s",
+    "psize",
+    "set_ctx",
+    "trace",
+    "uniform_cond",
+    "while_loop",
+    "with_ctx",
+]

mplang/{core → v1/core}/cluster.py

@@ -20,6 +20,7 @@ MPLang cluster configuration.
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from functools import cached_property
 from typing import Any
 
 
@@ -132,11 +133,11 @@ class ClusterSpec:
                         "which is not defined in nodes"
                     )
 
-        # ensure local devices have exactly one member
+        # ensure ppu devices have exactly one member
         for device in self.devices.values():
-            if device.kind.lower() == "local" and len(device.members) != 1:
+            if device.kind.lower() == "ppu" and len(device.members) != 1:
                 raise ValueError(
-                    f"Local device '{device.name}' must have exactly one member"
+                    f"PPU device '{device.name}' must have exactly one member"
                 )
 
     def get_node(self, name: str) -> Node:
@@ -169,6 +170,16 @@ class ClusterSpec:
             },
         }
 
+    @cached_property
+    def endpoints(self) -> list[str]:
+        eps: list[str] = []
+        for n in sorted(
+            self.nodes.values(),
+            key=lambda x: x.rank,  # type: ignore[attr-defined]
+        ):
+            eps.append(n.endpoint)
+        return eps
+
     @classmethod
     def from_dict(cls, config: dict[str, Any]) -> ClusterSpec:
         """Parses a raw config dictionary and returns a validated ClusterSpec."""
@@ -237,12 +248,13 @@ class ClusterSpec:
         world_size: int,
         *,
         endpoints: list[str] | None = None,
+        spu_world_size: int | None = None,
         spu_protocol: str = "SEMI2K",
         spu_field: str = "FM128",
         runtime_version: str = "simulated",
         runtime_platform: str = "simulated",
         op_bindings: list[dict[str, str]] | None = None,
-        enable_local_device: bool = True,
+        enable_ppu_device: bool = True,
         enable_spu_device: bool = True,
     ) -> ClusterSpec:
         """Convenience constructor used heavily in tests.
@@ -263,8 +275,8 @@ class ClusterSpec:
         op_bindings:
             Optional list of length ``world_size`` supplying per-node op_bindings
             override dicts (defaults to empty dicts).
-        enable_local_device:
-            If True (default), create one ``local_{rank}`` device per node.
+        enable_ppu_device:
+            If True (default), create one ``P{rank}`` PPU device per node.
         enable_spu_device:
             If True (default) create a shared SPU device named ``SP0``.
         """
@@ -282,9 +294,9 @@ class ClusterSpec:
                 f"{len(op_bindings)} != {world_size}"
             )
 
-        if not enable_local_device and not enable_spu_device:
+        if not enable_ppu_device and not enable_spu_device:
             raise ValueError(
-                "At least one of enable_local_device or enable_spu_device must be True"
+                "At least one of enable_ppu_device or enable_spu_device must be True"
             )
 
         nodes: dict[str, Node] = {}
@@ -303,21 +315,25 @@ class ClusterSpec:
             )
 
         devices: dict[str, Device] = {}
-        # Optional per-node local devices
-        if enable_local_device:
+        # Optional per-node PPU devices
+        if enable_ppu_device:
             for i in range(world_size):
-                devices[f"local_{i}"] = Device(
-                    name=f"local_{i}",
-                    kind="local",
+                devices[f"P{i}"] = Device(
+                    name=f"P{i}",
+                    kind="ppu",
                     members=[nodes[f"node{i}"]],
                 )
 
         # Shared SPU device
         if enable_spu_device:
+            if spu_world_size is None:
+                spu_world_size = world_size
+            spu_members = [nodes[f"node{i}"] for i in range(spu_world_size)]
+
             devices["SP0"] = Device(
                 name="SP0",
                 kind="SPU",
-                members=list(nodes.values()),
+                members=spu_members,
                 config={
                     "protocol": spu_protocol,
                     "field": spu_field,

mplang/{core → v1/core}/comm.py

@@ -19,7 +19,7 @@ import threading
 from abc import ABC, abstractmethod
 from typing import Any
 
-from mplang.core.mask import Mask
+from mplang.v1.core.mask import Mask
 
 
 class ICommunicator(ABC):
@@ -48,6 +48,10 @@ class ICommunicator(ABC):
     def recv(self, frm: int, key: str) -> Any:
         """Receive data from peer with the given key"""
 
+    @abstractmethod
+    def onSent(self, frm: int, key: str, data: Any) -> None:
+        """Called when a key is sent to self"""
+
 
 class ICollective(ABC):
     """Interface for collective communication"""

mplang/{core → v1/core}/context_mgr.py

@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     # Imported only for typing to avoid import cycles at runtime.
-    from mplang.core.mpobject import MPContext
+    from mplang.v1.core.mpobject import MPContext
 
 # The global working context.
 _g_ctx: MPContext | None = None

mplang/{core/dtype.py → v1/core/dtypes.py}

@@ -21,7 +21,8 @@ import numpy as np
 
 try:
     # Check if JAX is available
-    import jax  # noqa: F401
+    import jax
+    import jax.numpy as jnp
 
     _JAX_AVAILABLE = True
 except ImportError:
@@ -140,6 +141,10 @@ class DType:
         """Convert from JAX dtype to custom DType."""
         if not _JAX_AVAILABLE:
             raise ImportError("JAX is not available")
+        # Special handling for PRNG KeyTy: <class jax._src.prng.KeyTy>
+        if jnp.issubdtype(jax_dtype, jax.dtypes.prng_key):
+            return cls.from_numpy(np.uint32)
+
         # JAX dtypes are essentially NumPy dtypes
         return cls.from_numpy(jax_dtype)
 
@@ -172,6 +177,13 @@ class DType:
             # TypeError if it's not a pandas dtype we can handle
             pass
 
+        try:
+            return cls._from_arrow_dtype(dtype_like)
+        except (ImportError, TypeError):
+            # ImportError if pyarrow is not installed
+            # TypeError if it's not a pyarrow dtype we can handle
+            pass
+
         if isinstance(dtype_like, type) and dtype_like in (bool, int, float, complex):
             return cls.from_python_type(dtype_like)
         elif hasattr(dtype_like, "dtype") and not isinstance(dtype_like, type):
@@ -220,6 +232,37 @@ class DType:
 
         raise TypeError(f"Unsupported pandas dtype: {dtype_like}")
 
+    @classmethod
+    def _from_arrow_dtype(cls, dtype_like: Any) -> DType:
+        try:
+            import pyarrow as pa
+        except ImportError:
+            raise ImportError("pyarrow not available") from None
+
+        if not isinstance(dtype_like, pa.DataType):
+            raise TypeError("Not a pyarrow dtype")
+
+        ARROW_DTYPE_MAPPING = {
+            pa.bool_(): BOOL,
+            pa.int8(): INT8,
+            pa.int16(): INT16,
+            pa.int32(): INT32,
+            pa.int64(): INT64,
+            pa.uint8(): UINT8,
+            pa.uint16(): UINT16,
+            pa.uint32(): UINT32,
+            pa.uint64(): UINT64,
+            pa.float16(): FLOAT16,
+            pa.float32(): FLOAT32,
+            pa.float64(): FLOAT64,
+            pa.string(): STRING,
+            pa.large_string(): STRING,
+        }
+        result = ARROW_DTYPE_MAPPING.get(dtype_like)
+        if result is not None:
+            return result
+        raise TypeError(f"Unsupported arrow dtype: {dtype_like}")
+
     def to_numpy(self) -> np.dtype:
         """Convert custom DType to NumPy dtype."""
         return np.dtype(self.name)
@@ -228,7 +271,6 @@ class DType:
         """Convert custom DType to JAX dtype."""
         if not _JAX_AVAILABLE:
             raise ImportError("JAX is not available")
-        import jax.numpy as jnp
 
         return jnp.dtype(self.name)
 

mplang/{core → v1/core}/expr/__init__.py

@@ -20,7 +20,7 @@ multi-party computation graphs using the visitor pattern.
 """
 
 # Core expression types
-from mplang.core.expr.ast import (
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -36,12 +36,12 @@ from mplang.core.expr.ast import (
 )
 
 # Built-in evaluator engines
-from mplang.core.expr.evaluator import IEvaluator, create_evaluator
-from mplang.core.expr.printer import Printer
-from mplang.core.expr.transformer import ExprTransformer
+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.core.expr.utils import (
+from mplang.v1.core.expr.utils import (
     deduce_mask,
     ensure_scalar,
     ensure_tensorlist_equal,
@@ -49,8 +49,8 @@ from mplang.core.expr.utils import (
 )
 
 # Visitor pattern interface
-from mplang.core.expr.visitor import ExprVisitor
-from mplang.core.expr.walk import walk, walk_dataflow, walk_structural
+from mplang.v1.core.expr.visitor import ExprVisitor
+from mplang.v1.core.expr.walk import walk, walk_dataflow, walk_structural
 
 __all__ = [
     "AccessExpr",

mplang/{core → v1/core}/expr/ast.py

@@ -26,15 +26,15 @@ import logging
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any
 
-from mplang.core.expr.utils import deduce_mask
-from mplang.core.mask import Mask
-from mplang.core.mptype import MPType, Rank
-from mplang.core.pfunc import PFunction
-from mplang.core.table import TableType
-from mplang.core.tensor import TensorType
+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.core.expr.visitor import ExprVisitor
+    from mplang.v1.core.expr.visitor import ExprVisitor
 
 
 class Expr(ABC):
@@ -286,8 +286,8 @@ class ConvExpr(Expr):
         # Validate dtype / shape consistency.
         first = types[0]
         for c in types[1:]:
-            if (c.dtype, c.shape) != (first.dtype, first.shape):
-                raise TypeError(f"Inconsistent dtype/shape in pconv: {c} vs {first}")
+            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]
@@ -316,7 +316,7 @@ class ConvExpr(Expr):
             else:
                 out_pmask = None
 
-        return [MPType.tensor(first.dtype, first.shape, out_pmask, **first.attrs)]
+        return [MPType(first.raw_type(), out_pmask, first.attrs)]
 
     def accept(self, visitor: ExprVisitor) -> Any:
         return visitor.visit_conv(self)
@@ -398,9 +398,7 @@ class ShflSExpr(Expr):
     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.tensor(src_type.dtype, src_type.shape, self.pmask, **src_type.attrs)
-        ]
+        return [MPType(src_type._type, self.pmask, src_type.attrs)]
 
     def accept(self, visitor: ExprVisitor) -> Any:
         return visitor.visit_shfl_s(self)
@@ -528,8 +526,9 @@ class FuncDefExpr(Expr):
 class CallExpr(Expr):
     """Expression for function call."""
 
-    def __init__(self, fn: FuncDefExpr, args: list[Expr]):
+    def __init__(self, name: str, fn: FuncDefExpr, args: list[Expr]):
         super().__init__()
+        self.name = name
         self.fn = fn
         self.args = args
 

mplang/{core → v1/core}/expr/evaluator.py

@@ -27,8 +27,8 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any, Protocol
 
-from mplang.core.comm import ICommunicator
-from mplang.core.expr.ast import (
+from mplang.v1.core.comm import ICommunicator
+from mplang.v1.core.expr.ast import (
     AccessExpr,
     CallExpr,
     CondExpr,
@@ -42,11 +42,12 @@ from mplang.core.expr.ast import (
     VariableExpr,
     WhileExpr,
 )
-from mplang.core.expr.visitor import ExprVisitor
-from mplang.core.expr.walk import walk_dataflow
-from mplang.core.mask import Mask
-from mplang.core.pfunc import PFunction
-from mplang.kernels.context import RuntimeContext
+from mplang.v1.core.expr.visitor import ExprVisitor
+from mplang.v1.core.expr.walk import walk_dataflow
+from mplang.v1.core.mask import Mask
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.kernels.context import RuntimeContext
+from mplang.v1.kernels.value import Value
 
 
 class IEvaluator(Protocol):
@@ -149,12 +150,12 @@ class EvalSemantic:
     def _as_optional_int(val: Any) -> int | None:
         """Convert a value to int if possible, preserving None.
 
-        Handles Python ints, numpy scalars with .item(), and None.
+        Handles Python ints, floats, numpy scalar types (e.g., np.int32, np.float64), and None.
+        Uses int(val) for conversion which works with numpy scalars via __int__().
         """
+        val = EvalSemantic._unwrap_value(val)
         if val is None:
             return None
-        if hasattr(val, "item"):
-            return int(val.item())
         return int(val)
 
     def _simple_allgather(self, value: Any) -> list[Any]:
@@ -167,6 +168,7 @@ class EvalSemantic:
         Returns a list of length world_size with entries ordered by rank.
         """
         ws = self.comm.world_size
+        value = self._unwrap_value(value)
         # Trivial fast-path
         if ws == 1:
             return [value]
@@ -185,7 +187,12 @@ class EvalSemantic:
 
     def _verify_uniform_predicate(self, pred: Any) -> None:
         # Runtime uniformity check (O(P^2) send/recv emulation).
-        vals = self._simple_allgather(bool(pred))
+        # Use Value.to_bool() if available, otherwise unwrap and convert
+        if isinstance(pred, Value):
+            pred_bool = pred.to_bool()
+        else:
+            pred_bool = bool(self._unwrap_value(pred))
+        vals = self._simple_allgather(pred_bool)
         if not vals:
             raise ValueError("uniform_cond: empty gather for predicate")
         first = vals[0]
@@ -209,13 +216,36 @@ class EvalSemantic:
         assert len(cond_result) == 1, (
             f"Condition function must return a single value, got {cond_result}"
         )
-        cond_value = cond_result[0]
-        if cond_value is None:
+        cond_val = cond_result[0]
+        if cond_val is None:
             raise RuntimeError(
                 "while_loop condition produced None on rank "
                 f"{self.rank}; ensure the predicate yields a boolean for every party."
             )
-        return cond_value
+        # Use Value.to_bool() if available for cleaner conversion
+        if isinstance(cond_val, Value):
+            return cond_val.to_bool()
+        return bool(self._unwrap_value(cond_val))
+
+    @staticmethod
+    def _unwrap_value(value: Any) -> Any:
+        """Convert Value payloads to numpy/python equivalents when possible."""
+        if value is None:
+            return None
+
+        if isinstance(value, Value):
+            # Try to_numpy first for broader compatibility
+            to_numpy = getattr(value, "to_numpy", None)
+            if callable(to_numpy):
+                arr = to_numpy()
+                import numpy as np
+
+                if isinstance(arr, np.ndarray):
+                    if arr.size == 1:
+                        return arr.item()
+                    return arr
+                return arr
+        return value
 
 
 class RecursiveEvaluator(EvalSemantic, ExprVisitor):
@@ -296,19 +326,25 @@ class RecursiveEvaluator(EvalSemantic, ExprVisitor):
           * Add optional static uniform inference (data provenance) to elide the
             runtime check when predicate uniformity is provable at trace time.
         """
-        pred = self._value(expr.pred)
-        if pred is None:
+        pred_val = self._value(expr.pred)
+        if pred_val is None:
             return [None] * len(expr.mptypes)
 
         if expr.verify_uniform:
-            self._verify_uniform_predicate(pred)
+            self._verify_uniform_predicate(pred_val)
+
+        # Convert to bool using Value.to_bool() if available
+        if isinstance(pred_val, Value):
+            pred = pred_val.to_bool()
+        else:
+            pred = bool(self._unwrap_value(pred_val))
 
         # Only evaluate selected branch locally
-        if pred:
-            then_call = CallExpr(expr.then_fn, expr.args)
+        if bool(pred):
+            then_call = CallExpr("then", expr.then_fn, expr.args)
             return self._values(then_call)
         else:
-            else_call = CallExpr(expr.else_fn, expr.args)
+            else_call = CallExpr("else", expr.else_fn, expr.args)
             return self._values(else_call)
 
     def visit_call(self, expr: CallExpr) -> Any:
@@ -435,15 +471,20 @@ class IterativeEvaluator(EvalSemantic):
                 res = self._iter_eval_graph(node.fn.body, {**env, **sub_env})
                 symbols[id(node)] = res
             elif isinstance(node, CondExpr):
-                pred_v = self._first(symbols[id(node.pred)])
+                pred_val = self._first(symbols[id(node.pred)])
                 arg_vals = [self._first(symbols[id(a)]) for a in node.args]
-                if pred_v is None:
+                if pred_val is None:
                     symbols[id(node)] = [None] * len(node.mptypes)
                 else:
                     # Optional uniform verification identical to recursive evaluator (DRY helper).
                     if node.verify_uniform:
-                        self._verify_uniform_predicate(pred_v)
-                    if bool(pred_v):
+                        self._verify_uniform_predicate(pred_val)
+                    # Convert to bool using Value.to_bool() if available
+                    if isinstance(pred_val, Value):
+                        pred = pred_val.to_bool()
+                    else:
+                        pred = bool(self._unwrap_value(pred_val))
+                    if pred:
                         sub_env = dict(zip(node.then_fn.params, arg_vals, strict=True))
                         res = self._iter_eval_graph(
                             node.then_fn.body, {**env, **sub_env}