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

mplang/v2/backends/simp_worker/mem.py

@@ -0,0 +1,102 @@
+# 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.
+
+"""Simp Worker memory IPC runtime (LocalMesh, ThreadCommunicator)."""
+
+from __future__ import annotations
+
+import concurrent.futures
+import threading
+from typing import Any
+
+
+class ThreadCommunicator:
+    """Thread-based communicator for in-memory communication.
+
+    Args:
+        rank: This communicator's rank.
+        world_size: Total number of parties.
+        use_serde: If True, serialize/deserialize data through serde on send.
+    """
+
+    def __init__(self, rank: int, world_size: int, *, use_serde: bool = False):
+        self.rank = rank
+        self.world_size = world_size
+        self.use_serde = use_serde
+        self.peers: list[ThreadCommunicator] = []
+        # Mailbox keyed by (from_rank, tag): each key has exactly one message
+        self._mailbox: dict[tuple[int, str], Any] = {}
+        self._cond = threading.Condition()
+        self._sent_events: dict[str, threading.Event] = {}
+        self._shutdown = False
+
+    def set_peers(self, peers: list[ThreadCommunicator]) -> None:
+        assert len(peers) == self.world_size
+        self.peers = peers
+
+    def shutdown(self) -> None:
+        with self._cond:
+            self._shutdown = True
+            self._cond.notify_all()
+
+    def send(self, to: int, key: str, data: Any) -> None:
+        assert 0 <= to < self.world_size
+        if self.use_serde:
+            from mplang.v2.edsl import serde
+
+            data = serde.loads(serde.dumps(data))
+        self.peers[to]._on_receive(self.rank, key, data)
+
+    def recv(self, frm: int, key: str) -> Any:
+        mailbox_key = (frm, key)
+        with self._cond:
+            while mailbox_key not in self._mailbox and not self._shutdown:
+                self._cond.wait()
+            if self._shutdown:
+                raise RuntimeError("Communicator shut down")
+            return self._mailbox.pop(mailbox_key)
+
+    def _on_receive(self, frm: int, key: str, data: Any) -> None:
+        mailbox_key = (frm, key)
+        with self._cond:
+            if mailbox_key in self._mailbox:
+                raise RuntimeError(
+                    f"Mailbox overflow: key {mailbox_key} already exists"
+                )
+            self._mailbox[mailbox_key] = data
+            self._cond.notify_all()
+
+
+class LocalMesh:
+    """Communication mesh for local SIMP simulation.
+
+    Creates a set of ThreadCommunicators and a ThreadPoolExecutor for
+    worker-side execution.
+    """
+
+    def __init__(self, world_size: int, *, use_serde: bool = False):
+        self.world_size = world_size
+        self.use_serde = use_serde
+        self.comms = [
+            ThreadCommunicator(i, world_size, use_serde=use_serde)
+            for i in range(world_size)
+        ]
+        for comm in self.comms:
+            comm.set_peers(self.comms)
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=world_size)
+
+    def shutdown(self, wait: bool = True) -> None:
+        for comm in self.comms:
+            comm.shutdown()
+        self.executor.shutdown(wait=wait)

mplang/v2/backends/simp_worker/ops.py

@@ -0,0 +1,167 @@
+# 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.
+
+"""Simp Worker ops (WORKER_HANDLERS).
+
+This module contains all simp operation implementations for the Worker Interpreter.
+These implementations execute locally on a single party.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from mplang.v2.dialects import simp
+from mplang.v2.edsl.graph import Operation
+from mplang.v2.runtime.interpreter import Interpreter
+
+
+def _ensure_worker_context(interpreter: Any, op_name: str) -> Any:
+    """Validate that interpreter has a Worker context."""
+    state = interpreter.get_dialect_state("simp")
+    if state is None or not hasattr(state, "communicator"):
+        raise RuntimeError(f"{op_name} requires simp Worker state (with communicator)")
+    return state
+
+
+def _pcall_static_worker_impl(
+    interpreter: Interpreter, op: Operation, *args: Any
+) -> Any:
+    """Worker implementation of pcall_static."""
+    worker = _ensure_worker_context(interpreter, "pcall_static_impl")
+
+    parties = op.attrs.get("parties")
+    if parties is None:
+        raise ValueError("pcall_static requires 'parties' attribute")
+
+    if worker.rank in parties:
+        fn_graph = op.regions[0]
+        prev_parties = worker.current_parties
+        worker.current_parties = parties
+
+        try:
+            result = interpreter.evaluate_graph(fn_graph, list(args))
+            # Return single value for single output (interpreter expects this)
+            return result[0] if len(op.outputs) == 1 else result
+        finally:
+            worker.current_parties = prev_parties
+    else:
+        # No data for this rank
+        return None if len(op.outputs) == 1 else [None] * len(op.outputs)
+
+
+def _pcall_dynamic_worker_impl(
+    interpreter: Interpreter, op: Operation, *args: Any
+) -> Any:
+    """Worker implementation of pcall_dynamic."""
+    fn_graph = op.regions[0]
+    result = interpreter.evaluate_graph(fn_graph, list(args))
+    return result[0] if len(op.outputs) == 1 else result
+
+
+def _shuffle_static_worker_impl(
+    interpreter: Interpreter, op: Operation, *args: Any
+) -> Any:
+    """Worker implementation of shuffle_static."""
+    worker = _ensure_worker_context(interpreter, "shuffle_static_impl")
+
+    routing = op.attrs.get("routing")
+    if routing is None:
+        return args[0]
+
+    comm = worker.communicator
+    my_rank = worker.rank
+    data = args[0]
+
+    for tgt, src in routing.items():
+        if src == my_rank and tgt != my_rank:
+            key = f"shuffle_{op.name}_{tgt}"
+            comm.send(tgt, key, data)
+
+    if my_rank in routing:
+        src = routing[my_rank]
+        if src == my_rank:
+            return data
+        key = f"shuffle_{op.name}_{my_rank}"
+        return comm.recv(src, key)
+    else:
+        return None
+
+
+def _converge_worker_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
+    """Worker implementation of simp.converge."""
+    for arg in args:
+        if arg is not None:
+            return arg
+    return None
+
+
+def _uniform_cond_worker_impl(
+    interpreter: Interpreter, op: Operation, pred: Any, *args: Any
+) -> Any:
+    """Worker implementation of simp.uniform_cond."""
+    from mplang.v2.backends.tensor_impl import TensorValue
+
+    if op.attrs.get("verify_uniform", True):
+        pass  # TODO: Implement AllReduce verification
+
+    if isinstance(pred, TensorValue):
+        pred = bool(pred.unwrap())
+
+    if pred:
+        result = interpreter.evaluate_graph(op.regions[0], list(args))
+    else:
+        result = interpreter.evaluate_graph(op.regions[1], list(args))
+    return result[0] if len(op.outputs) == 1 else result
+
+
+def _while_loop_worker_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
+    """Worker implementation of simp.while_loop."""
+    from mplang.v2.backends.tensor_impl import TensorValue
+
+    cond_graph = op.regions[0]
+    body_graph = op.regions[1]
+
+    num_state = len(op.outputs)
+    current_state = list(args[:num_state])
+    captures = list(args[num_state:])
+
+    while True:
+        region_inputs = current_state + captures
+
+        cond_res = interpreter.evaluate_graph(cond_graph, region_inputs)
+        # cond_res is a list, extract the single boolean
+        cond_val = cond_res[0] if cond_res else False
+
+        if isinstance(cond_val, TensorValue):
+            cond_val = bool(cond_val.unwrap())
+
+        if not cond_val:
+            break
+
+        body_res = interpreter.evaluate_graph(body_graph, region_inputs)
+        current_state = body_res  # body_res is always a list now
+
+    # Return single value for single output
+    return current_state[0] if len(current_state) == 1 else current_state
+
+
+WORKER_HANDLERS = {
+    simp.pcall_static_p.name: _pcall_static_worker_impl,
+    simp.pcall_dynamic_p.name: _pcall_dynamic_worker_impl,
+    simp.shuffle_static_p.name: _shuffle_static_worker_impl,
+    simp.converge_p.name: _converge_worker_impl,
+    simp.uniform_cond_p.name: _uniform_cond_worker_impl,
+    simp.while_loop_p.name: _while_loop_worker_impl,
+}

mplang/v2/backends/simp_worker/state.py

@@ -0,0 +1,49 @@
+# 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.
+
+"""Simp Worker state (SimpWorker)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import mplang.v2.backends.field_impl  # noqa: F401
+import mplang.v2.backends.tensor_impl  # noqa: F401
+from mplang.v2.runtime.dialect_state import DialectState
+from mplang.v2.runtime.object_store import ObjectStore
+
+
+class SimpWorker(DialectState):
+    """Worker state for SIMP execution.
+
+    This state provides capabilities (Store, Communicator) to the Interpreter.
+    Attached to Worker Interpreters.
+    """
+
+    dialect_name: str = "simp"
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        communicator: Any,
+        store: ObjectStore,
+        spu_endpoints: dict[int, str] | None = None,
+    ):
+        self.rank = rank
+        self.world_size = world_size
+        self.communicator = communicator
+        self.store = store
+        self.spu_endpoints = spu_endpoints
+        self.current_parties: tuple[int, ...] | None = None

mplang/v2/backends/spu_impl.py

@@ -0,0 +1,275 @@
+# 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.
+
+"""SPU Runtime Implementation.
+
+Implements execution logic for SPU primitives using libspu.
+"""
+
+from __future__ import annotations
+
+import base64
+from typing import Any, ClassVar
+
+import numpy as np
+import spu.api as spu_api
+import spu.libspu as libspu
+
+from mplang.v2.backends.spu_state import SPUState
+from mplang.v2.backends.tensor_impl import TensorValue
+from mplang.v2.dialects import spu
+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 WrapValue
+
+# =============================================================================
+# SPU Share Wrapper
+# =============================================================================
+
+
+@serde.register_class
+class SPUShareValue(WrapValue[libspu.Share]):
+    """Wrapper for libspu.Share representing an SPU secret share.
+
+    This wraps the external libspu library's Share type to provide
+    proper serialization support via the Value base class.
+
+    In-memory, we hold the libspu.Share directly to avoid copying.
+    Serialization extracts meta/share_chunks when needed.
+    """
+
+    _serde_kind: ClassVar[str] = "spu_impl.SPUShareValue"
+
+    def _convert(self, data: Any) -> libspu.Share:
+        if isinstance(data, SPUShareValue):
+            return data.unwrap()
+        if isinstance(data, libspu.Share):
+            return data
+        raise TypeError(f"Expected libspu.Share, got {type(data)}")
+
+    @property
+    def libspu_share(self) -> libspu.Share:
+        """Get the underlying libspu.Share object."""
+        return self._data
+
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "meta": base64.b64encode(self._data.meta).decode("ascii"),
+            "share_chunks": [
+                base64.b64encode(chunk).decode("ascii")
+                for chunk in self._data.share_chunks
+            ],
+        }
+
+    @classmethod
+    def from_json(cls, data: dict[str, Any]) -> SPUShareValue:
+        share = libspu.Share()
+        share.meta = base64.b64decode(data["meta"])
+        share.share_chunks = [
+            base64.b64decode(chunk_b64) for chunk_b64 in data["share_chunks"]
+        ]
+        return cls(share)
+
+    @classmethod
+    def from_libspu(cls, share: libspu.Share) -> SPUShareValue:
+        """Create SPUShareValue from a libspu.Share (zero-copy)."""
+        return cls(share)
+
+
+# =============================================================================
+# SPU Config Helpers
+# =============================================================================
+
+
+def to_runtime_config(config: spu.SPUConfig) -> libspu.RuntimeConfig:
+    """Convert SPUConfig to libspu.RuntimeConfig.
+
+    This is a runtime-only function that maps the string-based configuration
+    to libspu enums. Should only be called in the backend implementation.
+    """
+    runtime_config = libspu.RuntimeConfig()
+    # ProtocolKind uses "SEMI2K" not "PROT_SEMI2K"
+    runtime_config.protocol = getattr(libspu.ProtocolKind, config.protocol)
+    runtime_config.field = getattr(libspu.FieldType, config.field)
+    runtime_config.fxp_fraction_bits = config.fxp_fraction_bits
+    return runtime_config
+
+
+@spu.makeshares_p.def_impl
+def makeshares_impl(
+    interpreter: Interpreter, op: Operation, data: TensorValue
+) -> tuple[SPUShareValue, ...]:
+    """Generate secret shares for data using spu.Io."""
+    count = op.attrs["count"]
+    config: spu.SPUConfig = op.attrs["config"]
+
+    # We create a standalone Io for share generation (no link needed for make_shares)
+    runtime_config = to_runtime_config(config)
+    io = spu_api.Io(count, runtime_config)
+
+    # Unwrap TensorValue
+    arr = data.unwrap()
+
+    # data is expected to be numpy array
+    arr = np.asarray(arr)
+
+    # Generate shares (VIS_SECRET)
+    libspu_shares = io.make_shares(arr, libspu.Visibility.VIS_SECRET)
+
+    # Wrap libspu.Share objects in SPUShareValue
+    return tuple(SPUShareValue.from_libspu(share) for share in libspu_shares)
+
+
+@spu.reconstruct_p.def_impl
+def reconstruct_impl(
+    interpreter: Interpreter, op: Operation, *shares: SPUShareValue
+) -> TensorValue:
+    """Reconstruct data from secret shares using spu.Io."""
+    count = len(shares)
+    config: spu.SPUConfig = op.attrs["config"]
+
+    runtime_config = to_runtime_config(config)
+    io = spu_api.Io(count, runtime_config)
+
+    # Unwrap SPUShareValue to libspu.Share
+    libspu_shares = [share.libspu_share for share in shares]
+
+    # Reconstruct
+    result = io.reconstruct(libspu_shares)
+
+    # Wrap result as TensorValue
+    return TensorValue.wrap(result)
+
+
+@spu.exec_p.def_impl
+def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
+    """Execute SPU kernel using spu.Runtime.
+
+    The SPU config must contain parties info to correctly map global rank
+    to local SPU rank and determine SPU world size.
+    """
+    from mplang.v2.backends.simp_worker.state import SimpWorker
+
+    # Get SPU config from attrs (passed through from run_jax)
+    config: spu.SPUConfig = op.attrs["config"]
+
+    # Get parties from SimpWorker state (injected by pcall_static_impl)
+    context = interpreter.get_dialect_state("simp")
+    if not isinstance(context, SimpWorker):
+        raise RuntimeError(f"spu.exec requires SimpWorker, got {type(context)}")
+
+    parties = context.current_parties
+    if parties is None:
+        raise RuntimeError(
+            "spu.exec requires 'current_parties' in SimpWorker state. "
+            "Ensure it is called within a pcall_static block."
+        )
+
+    global_rank = context.rank
+
+    if global_rank not in parties:
+        raise RuntimeError(
+            f"Global rank {global_rank} is not in current parties {parties}"
+        )
+
+    # Convert global rank to local SPU rank
+    local_rank = parties.index(global_rank)
+    spu_world_size = len(parties)
+
+    # Get SPU endpoints from interpreter (set by WorkerInterpreter for BRPC mode)
+    # spu_endpoints is a dict mapping global_rank -> brpc_endpoint
+    spu_endpoints_map: dict[int, str] | None = getattr(
+        interpreter, "spu_endpoints", None
+    )
+    if spu_endpoints_map is None:
+        # Try getting from SimpWorker context (context is already SimpWorker)
+        spu_endpoints_map = getattr(context, "spu_endpoints", None)
+
+    # Build ordered list of endpoints for SPU parties
+    spu_endpoints: list[str] | None = None
+    if spu_endpoints_map is not None:
+        spu_endpoints = []
+        for party_rank in parties:
+            if party_rank not in spu_endpoints_map:
+                raise RuntimeError(
+                    f"SPU endpoint not found for party {party_rank}. "
+                    f"Available: {list(spu_endpoints_map.keys())}"
+                )
+            spu_endpoints.append(spu_endpoints_map[party_rank])
+
+    # Get communicator for Channels mode (reuse existing communication)
+    # If no BRPC endpoints configured, use Channels mode
+    communicator = None
+    if spu_endpoints is None:
+        # Use worker's communicator for channel reuse
+        # (SimpWorker already imported at function start)
+        communicator = context.communicator
+
+    # Get or create SPUState for caching Runtime/Io
+    spu_state = interpreter.get_dialect_state(SPUState.dialect_name)
+    if not isinstance(spu_state, SPUState):
+        spu_state = SPUState()
+        interpreter.set_dialect_state(SPUState.dialect_name, spu_state)
+
+    runtime, io = spu_state.get_or_create(
+        local_rank,
+        spu_world_size,
+        config,
+        spu_endpoints,
+        communicator=communicator,
+        parties=list(parties),
+    )
+
+    executable_code = op.attrs["executable"]
+    input_names = op.attrs["input_names"]
+    output_names = op.attrs["output_names"]
+
+    # Create Executable
+    executable = libspu.Executable(
+        name="spu_kernel",
+        input_names=input_names,
+        output_names=output_names,
+        code=executable_code,
+    )
+
+    # Set inputs
+    for name, share in zip(input_names, args, strict=True):
+        # Handle SPUShareValue wrapper - unwrap to libspu.Share
+        if isinstance(share, SPUShareValue):
+            libspu_share = share.libspu_share
+        else:
+            # Handle public input (numpy array)
+            # Generate shares with VIS_PUBLIC
+            # make_shares expects numpy array
+            if not isinstance(share, (np.ndarray, np.generic, int, float)):
+                share = np.array(share)
+
+            shares = io.make_shares(share, libspu.Visibility.VIS_PUBLIC)
+            libspu_share = shares[local_rank]
+
+        runtime.set_var(name, libspu_share)
+
+    # Run
+    runtime.run(executable)
+
+    # Get outputs and wrap in SPUShareValue
+    results = []
+    for name in output_names:
+        libspu_share = runtime.get_var(name)
+        results.append(SPUShareValue.from_libspu(libspu_share))
+
+    if len(results) == 1:
+        return results[0]
+    return tuple(results)