mplang-nightly 0.1.dev286__py3-none-any.whl → 0.1.dev288__py3-none-any.whl

mplang/backends/simp_worker/collective_algorithms.py

@@ -0,0 +1,228 @@
+# Copyright 2026 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.
+
+"""Collective communication algorithms (communicator-only).
+
+This module contains *pure* collective algorithms implemented only in terms of
+(a) a communicator and (b) an explicit participant set.
+
+It intentionally does NOT depend on:
+- Interpreter execution IDs / graph keys
+- SimpWorker current_parties
+- Operation objects
+
+Callers are expected to provide a collision-free key prefix for each collective
+instance.
+"""
+
+from __future__ import annotations
+
+import operator
+from collections.abc import Callable, Sequence
+from typing import Any, Protocol
+
+
+class Communicator(Protocol):
+    """Minimal communicator interface required by the algorithms."""
+
+    rank: int
+    world_size: int
+
+    def send(
+        self, to: int, key: str, data: Any, *, is_raw_bytes: bool = False
+    ) -> None: ...
+
+    def recv(self, frm: int, key: str) -> Any: ...
+
+
+def normalize_participants(
+    comm: Communicator, participants: Sequence[int]
+) -> tuple[int, ...]:
+    ps = tuple(sorted({int(r) for r in participants}))
+    if not ps:
+        raise ValueError("participants must be non-empty")
+    if any(r < 0 or r >= comm.world_size for r in ps):
+        raise ValueError(
+            f"participants out of range: {ps}, world_size={comm.world_size}"
+        )
+    if comm.rank not in ps:
+        raise ValueError(f"rank {comm.rank} is not in participants {ps}")
+    return ps
+
+
+def barrier(
+    comm: Communicator, *, participants: Sequence[int], key_prefix: str
+) -> None:
+    """Barrier using root gather + root release."""
+
+    ps = normalize_participants(comm, participants)
+    root = ps[0]
+
+    arrive_key = f"{key_prefix}_arrive"
+    release_key = f"{key_prefix}_release"
+
+    if comm.rank != root:
+        comm.send(root, arrive_key, True)
+        comm.recv(root, release_key)
+        return
+
+    for r in ps:
+        if r == root:
+            continue
+        _ = comm.recv(r, arrive_key)
+
+    for r in ps:
+        if r == root:
+            continue
+        comm.send(r, release_key, True)
+
+
+def broadcast(
+    comm: Communicator,
+    value: Any,
+    *,
+    root: int,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> Any:
+    """Broadcast a value from root to all participants."""
+
+    ps = normalize_participants(comm, participants)
+    if root not in ps:
+        raise ValueError(f"root {root} must be in participants {ps}")
+
+    bcast_key = f"{key_prefix}_bcast"
+
+    if comm.rank == root:
+        for r in ps:
+            if r == root:
+                continue
+            comm.send(r, bcast_key, value)
+        return value
+
+    return comm.recv(root, bcast_key)
+
+
+def allgather(
+    comm: Communicator, value: Any, *, participants: Sequence[int], key_prefix: str
+) -> list[Any]:
+    """Allgather implemented as gather-to-root then root broadcast."""
+
+    ps = normalize_participants(comm, participants)
+    root = ps[0]
+
+    gather_key = f"{key_prefix}_gather"
+    bcast_key = f"{key_prefix}_bcast"
+
+    if comm.rank != root:
+        comm.send(root, gather_key, value)
+        gathered = comm.recv(root, bcast_key)
+        if not isinstance(gathered, list):
+            raise TypeError(f"expected list from root broadcast, got {type(gathered)}")
+        return gathered
+
+    values_by_rank: dict[int, Any] = {root: value}
+    for r in ps:
+        if r == root:
+            continue
+        values_by_rank[r] = comm.recv(r, gather_key)
+
+    gathered = [values_by_rank[r] for r in ps]
+
+    for r in ps:
+        if r == root:
+            continue
+        comm.send(r, bcast_key, gathered)
+
+    return gathered
+
+
+def allreduce_bool_and(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> bool:
+    return _allreduce_bool(
+        comm,
+        value,
+        participants=participants,
+        key_prefix=key_prefix,
+        combine=operator.and_,
+    )
+
+
+def allreduce_bool_or(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> bool:
+    return _allreduce_bool(
+        comm,
+        value,
+        participants=participants,
+        key_prefix=key_prefix,
+        combine=operator.or_,
+    )
+
+
+def allreduce_bool_xor(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+) -> bool:
+    return _allreduce_bool(
+        comm,
+        value,
+        participants=participants,
+        key_prefix=key_prefix,
+        combine=operator.xor,
+    )
+
+
+def _allreduce_bool(
+    comm: Communicator,
+    value: bool,
+    *,
+    participants: Sequence[int],
+    key_prefix: str,
+    combine: Callable[[bool, bool], bool],
+) -> bool:
+    ps = normalize_participants(comm, participants)
+    root = ps[0]
+
+    gather_key = f"{key_prefix}_gather"
+    bcast_key = f"{key_prefix}_bcast"
+
+    if comm.rank != root:
+        comm.send(root, gather_key, bool(value))
+        return bool(comm.recv(root, bcast_key))
+
+    acc = bool(value)
+    for r in ps:
+        if r == root:
+            continue
+        acc = combine(acc, bool(comm.recv(r, gather_key)))
+
+    for r in ps:
+        if r == root:
+            continue
+        comm.send(r, bcast_key, acc)
+
+    return acc

mplang/backends/simp_worker/collectives.py

@@ -0,0 +1,275 @@
+# Copyright 2026 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-side collectives (wrapper layer).
+
+This module is the *context-aware wrapper* on top of
+`mplang.backends.simp_worker.collective_algorithms`.
+
+Responsibilities here:
+- Resolve "participants" from (explicit arg / op.attrs["parties"] /
+  worker.current_parties / world).
+- Build collision-free `key_prefix` using interpreter execution IDs.
+
+The underlying algorithms only depend on the communicator interface.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any, Protocol
+
+from mplang.backends.simp_worker import collective_algorithms as algo
+from mplang.edsl.graph import Operation
+
+
+class _ExecContext(Protocol):
+    def current_op_exec_id(self) -> int: ...
+
+    def current_graph_exec_key(self) -> str: ...
+
+
+class _Worker(Protocol):
+    rank: int
+    world_size: int
+    communicator: algo.Communicator
+    current_parties: tuple[int, ...] | None
+
+
+def resolve_participants(
+    worker: _Worker,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+) -> Sequence[int]:
+    """Resolve participant ranks.
+
+    Priority:
+    1) explicit participants argument
+    2) op.attrs["parties"] if present
+    3) worker.current_parties if set (pcall_static dynamic scope)
+    4) all ranks [0, world_size)
+
+    Note:
+        Normalization/validation (sorting, emptiness, range checks, rank
+        inclusion) is intentionally delegated to the lower-level algorithms in
+        `collective_algorithms`.
+    """
+
+    if participants is not None:
+        return participants
+
+    if op is not None:
+        parties = op.attrs.get("parties")
+        if parties is not None:
+            if not isinstance(parties, Sequence):
+                raise TypeError(
+                    "op.attrs['parties'] must be a sequence of rank integers"
+                )
+            return tuple(int(r) for r in parties)
+
+    if worker.current_parties is not None:
+        return worker.current_parties
+
+    return tuple(range(worker.world_size))
+
+
+def _collective_prefix(
+    interpreter: _ExecContext, *, op: Operation | None, name: str
+) -> str:
+    exec_id = interpreter.current_op_exec_id()
+    graph_key = interpreter.current_graph_exec_key()
+    op_name = op.name if op is not None else "_"
+    return f"coll_{graph_key}_{op_name}_{exec_id}_{name}"
+
+
+def barrier(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "barrier",
+) -> None:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    algo.barrier(worker.communicator, participants=ps, key_prefix=prefix)
+
+
+def broadcast(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: Any,
+    *,
+    root: int,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "broadcast",
+) -> Any:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.broadcast(
+        worker.communicator,
+        value,
+        root=int(root),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def allgather_obj(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: Any,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allgather_obj",
+) -> list[Any]:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allgather(
+        worker.communicator, value, participants=ps, key_prefix=prefix
+    )
+
+
+def allgather_bool(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allgather_bool",
+) -> list[bool]:
+    gathered = allgather_obj(
+        interpreter,
+        worker,
+        bool(value),
+        op=op,
+        participants=participants,
+        name=name,
+    )
+    return [bool(v) for v in gathered]
+
+
+def allreduce_bool_and(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allreduce_bool_and",
+) -> bool:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allreduce_bool_and(
+        worker.communicator,
+        bool(value),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def allreduce_bool_or(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allreduce_bool_or",
+) -> bool:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allreduce_bool_or(
+        worker.communicator,
+        bool(value),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def allreduce_bool_xor(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    value: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "allreduce_bool_xor",
+) -> bool:
+    ps = resolve_participants(worker, op=op, participants=participants)
+    prefix = _collective_prefix(interpreter, op=op, name=name)
+    return algo.allreduce_bool_xor(
+        worker.communicator,
+        bool(value),
+        participants=ps,
+        key_prefix=prefix,
+    )
+
+
+def verify_uniform_predicate(
+    interpreter: _ExecContext,
+    worker: _Worker,
+    pred: bool,
+    *,
+    op: Operation | None = None,
+    participants: Sequence[int] | None = None,
+    name: str = "uniform_predicate",
+) -> bool:
+    """Verify that `pred` is uniform across participants.
+
+    Uses AND/OR all-reduce to detect mismatch. If mismatch is detected, runs an
+    allgather to provide a helpful error message. All participants execute the
+    same comm steps to avoid deadlocks.
+    """
+
+    ps = resolve_participants(worker, op=op, participants=participants)
+
+    all_and = allreduce_bool_and(
+        interpreter,
+        worker,
+        bool(pred),
+        op=op,
+        participants=ps,
+        name=f"{name}_and",
+    )
+    all_or = allreduce_bool_or(
+        interpreter,
+        worker,
+        bool(pred),
+        op=op,
+        participants=ps,
+        name=f"{name}_or",
+    )
+
+    if all_and != all_or:
+        gathered = allgather_bool(
+            interpreter,
+            worker,
+            bool(pred),
+            op=op,
+            participants=ps,
+            name=f"{name}_gather",
+        )
+        ps_norm = algo.normalize_participants(worker.communicator, ps)
+        dist = dict(zip(ps_norm, gathered, strict=True))
+        raise RuntimeError(
+            "simp.uniform_cond predicate is not uniform across participants: "
+            f"participants={ps_norm}, values={dist}"
+        )
+
+    return bool(pred)

mplang/backends/simp_worker/ops.py

@@ -22,6 +22,7 @@ from __future__ import annotations
 
 from typing import Any
 
+from mplang.backends.simp_worker.collectives import verify_uniform_predicate
 from mplang.dialects import simp
 from mplang.edsl.graph import Operation
 from mplang.runtime.interpreter import Interpreter
@@ -117,12 +118,14 @@ def _uniform_cond_worker_impl(
     """Worker implementation of simp.uniform_cond."""
     from mplang.backends.tensor_impl import TensorValue
 
-    if op.attrs.get("verify_uniform", True):
-        pass  # TODO: Implement AllReduce verification
+    worker = _ensure_worker_context(interpreter, "uniform_cond_impl")
 
     if isinstance(pred, TensorValue):
         pred = bool(pred.unwrap())
 
+    if op.attrs.get("verify_uniform", True):
+        pred = verify_uniform_predicate(interpreter, worker, bool(pred), op=op)
+
     if pred:
         result = interpreter.evaluate_graph(op.regions[0], list(args))
     else:

mplang/backends/table_impl.py

@@ -346,7 +346,7 @@ class ParquetReader(pa.RecordBatchReader):
             pass
         if batches:
             return pa.Table.from_batches(batches)
-        return pa.Table.from_batches([])
+        return pa.Table.from_batches([], schema=self._schema)
 
     def read_next_batch(self) -> pa.RecordBatch:
         batch = next(self._iter)

{mplang_nightly-0.1.dev286.dist-info → mplang_nightly-0.1.dev288.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev286
+Version: 0.1.dev288
 Summary: Multi-Party Programming Language
 Author-email: SecretFlow Team <secretflow-contact@service.alipay.com>
 License:                                  Apache License

{mplang_nightly-0.1.dev286.dist-info → mplang_nightly-0.1.dev288.dist-info}/RECORD

@@ -13,7 +13,7 @@ mplang/backends/simp_design.md,sha256=u4YeDLKk5avsueXhzPtt4OUBToBmurOVK_BWTKTJ7w
 mplang/backends/spu_impl.py,sha256=mkyNCwHC4N7zZ23mCRdOqJWbKv8dxH-Tqi14cTyhGSA,9526
 mplang/backends/spu_state.py,sha256=tgwOiWN6zfP-fBx2rN2cNnC-Y9GrSArqU-ZBT3NdXls,6557
 mplang/backends/store_impl.py,sha256=BdrCfNhFbcjswIvfxD2XKGYeBfBrmEQA3pj2-EDzFcI,2213
-mplang/backends/table_impl.py,sha256=SFbZJpZS1SdRxxlYfQjc69qcy-2hVfVAYEapNRaTTiI,27617
+mplang/backends/table_impl.py,sha256=v9a_44QhvJf7plxj3yFbphf-Be07a6dNhwXBiiTbff8,27638
 mplang/backends/tee_impl.py,sha256=5yed94Lr55hgSDT6j3FjdTboqzhl4HpJ9lGk1YGZBXQ,7026
 mplang/backends/tensor_impl.py,sha256=iIOsvSXiRsQWbC8BBfPUpOUhAXJW8VY11l0BxsMVYPo,18963
 mplang/backends/simp_driver/__init__.py,sha256=pQkHGR8HagnPB4OqagPZY-Ul19y6HjIJ5dY6QXZrBH8,1243
@@ -23,9 +23,11 @@ mplang/backends/simp_driver/ops.py,sha256=WYObWDRCsiXH0UBWZX5vD5W98ZPkd88U_qBV8S
 mplang/backends/simp_driver/state.py,sha256=dNmYMFN2D2BBdgs6C0YLaHrfaBRMgs05UNxMWw6tZIs,1713
 mplang/backends/simp_driver/values.py,sha256=Lz1utNSIzH-dCzZAEjU6JRcxPsfKGfUJrYl6gIuMOGw,1509
 mplang/backends/simp_worker/__init__.py,sha256=gdrSY1-MDkupCoJ8xwwH7em7fgVWv3J4gBJ45uHdzgg,961
+mplang/backends/simp_worker/collective_algorithms.py,sha256=tjqQrXRpV71g7uKBGlITLM29VjsnYfBr7MBfx1uKzvM,5707
+mplang/backends/simp_worker/collectives.py,sha256=Jzo-eU4QXM8zXfbUdb0xvghh6LWDSfUsxglSYi3oGY8,7688
 mplang/backends/simp_worker/http.py,sha256=UiqXKcH5DZG_eTXcW7qDSjIsBWtA37tkqBEowJJ0hQE,13186
 mplang/backends/simp_worker/mem.py,sha256=hLJftzSMQIw64vjBdrSrFReFzWS6rlzf-0Q7SWkfhak,3605
-mplang/backends/simp_worker/ops.py,sha256=ntxfkD4e6Il4w7FshK1ODcUCUPMlipt33pDY_x5iC0U,5661
+mplang/backends/simp_worker/ops.py,sha256=qMfXH9lUc9nq9paWa-53XgDZQcL29nDF6Qt9XwKouQ0,5834
 mplang/backends/simp_worker/state.py,sha256=nIu0ybvdYqRqp0TkoSneUF2u31evDHucCRduVBaDals,1445
 mplang/dialects/__init__.py,sha256=CYMmkeQVU0Znr9n3_5clZKb16u7acJ5jl5Zjbx4Tn1U,1478
 mplang/dialects/bfv.py,sha256=m5YfobFCBqn0lg2zBM9RNs2AC7i4PUQH2qXjHLHwSy4,22332
@@ -99,8 +101,8 @@ mplang/tool/program.py,sha256=W3H8bpPirnoJ4ZrmyPYuMCPadJis20o__n_1MKqCsWU,11058
 mplang/utils/__init__.py,sha256=Hwrwti2nfPxWUXV8DN6T1QaqXH_Jsd27k8UMSdBGUns,1073
 mplang/utils/func_utils.py,sha256=aZ-X43w8JKJgiF-IUMS0G7QqrNeoTM5ZPzRNd-tKxpw,5180
 mplang/utils/logging.py,sha256=9dMhwprVbx1WMGJrgoQbWmV50vyYuLU4NSPnetcl1Go,7237
-mplang_nightly-0.1.dev286.dist-info/METADATA,sha256=jhn3cljBrdSLTIGeIfxXbCGXrjXj9Jw7Oq83HGpo9y0,16783
-mplang_nightly-0.1.dev286.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mplang_nightly-0.1.dev286.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev286.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev286.dist-info/RECORD,,
+mplang_nightly-0.1.dev288.dist-info/METADATA,sha256=hzQlzwB9anuop3UrpF_QX9zxGqVGwd0_XOZJhPk4bcM,16783
+mplang_nightly-0.1.dev288.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mplang_nightly-0.1.dev288.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev288.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev288.dist-info/RECORD,,