diff --git a/mplang/v2/backends/simp_worker/ops.py b/mplang/v2/backends/simp_worker/ops.py index 03c32e46..c61ae8f8 100644 --- a/mplang/v2/backends/simp_worker/ops.py +++ b/mplang/v2/backends/simp_worker/ops.py @@ -22,6 +22,7 @@ from typing import Any +from mplang.v2.backends.tensor_impl import TensorValue from mplang.v2.dialects import simp from mplang.v2.edsl.graph import Operation from mplang.v2.runtime.interpreter import Interpreter @@ -111,10 +112,9 @@ 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 + _verify_uniform_predicate(interpreter, op, pred) if isinstance(pred, TensorValue): pred = bool(pred.unwrap()) @@ -126,6 +126,41 @@ def _uniform_cond_worker_impl( return result[0] if len(op.outputs) == 1 else result +def _verify_uniform_predicate( + interpreter: Interpreter, op: Operation, pred: Any +) -> None: + worker = _ensure_worker_context(interpreter, "_verify_uniform_predicate") + comm = worker.communicator + world_size = worker.world_size + + if isinstance(pred, TensorValue): + pred_bool = bool(pred.unwrap()) + else: + pred_bool = bool(pred) + + gathered_values = [] + + for dst in range(world_size): + if dst != comm.rank: + comm.send(dst, f"uniform_cond_pred_{op.name}", pred_bool) + + for src in range(world_size): + if src == comm.rank: + gathered_values.append(pred_bool) + else: + received = comm.recv(src, f"uniform_cond_pred_{op.name}") + gathered_values.append(received) + + + first_val = gathered_values[0] + for i, val in enumerate(gathered_values[1:], start=1): + if val != first_val: + raise ValueError( + f"uniform_cond: predicate is not uniform across parties " + f"(rank 0={first_val}, rank {i}={val})" + ) + + 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 diff --git a/mplang/v2/libs/device/api.py b/mplang/v2/libs/device/api.py index fe44d715..12a7ebfa 100644 --- a/mplang/v2/libs/device/api.py +++ b/mplang/v2/libs/device/api.py @@ -780,15 +780,20 @@ def _fetch_from_rank(rank: int) -> Any: # 4. Unwrap if WrapValue return _unwrap_value(result) - # 3.2 SPU: fetch from all ranks and reconstruct - elif dev_info.kind.upper() == "SPU": - # Fetch shares from all SPU members + # 3.2 SPU: fetch from all ranks and reconstruct shares = [_fetch_from_rank(m.rank) for m in dev_info.members] - # For now, just return the first share (TODO: implement spu.reconstruct) - # In practice, SPU values should be revealed to a PPU first - result = shares[0] if shares else None - # 4. Unwrap if WrapValue - return _unwrap_value(result) + + from mplang.v2.backends.spu_impl import SPUShareValue + + spu_shares = [SPUShareValue.from_libspu(share) for share in shares] + + + reconstructed = spu.reconstruct( + spu.SPUConfig.from_dict(dev_info.config), + tuple(spu_shares), + ) + + return _unwrap_value(reconstructed) # Direct value (not DriverVar) return _unwrap_value(runtime_obj) diff --git a/tests/v2/backends/test_verify_clean.py b/tests/v2/backends/test_verify_clean.py index 623aa145..ef3ba835 100644 --- a/tests/v2/backends/test_verify_clean.py +++ b/tests/v2/backends/test_verify_clean.py @@ -86,38 +86,3 @@ def fn(x): results = [int(r) for r in results] assert results == [11, 21] hasattr(sim, "_simp_cluster") and sim._simp_cluster.shutdown() - - -def test_uniform_cond_clean(): - """Test uniform_cond (Clean ver).""" - sim = simp.make_simulator(world_size=2) - - with sim: - # Check if pcall_static handling returns DriverVar - # Manually invoke pcall_static first - simp.pcall_static((0, 1), lambda: tensor.constant(True)) - - x0 = simp.constant((0,), 1) - x1 = simp.constant((1,), 2) - x_obj = simp.converge(x0, x1) - - def then_fn(x): - return simp.pcall_static( - (0, 1), lambda a: tensor.run_jax(lambda v: v + v, a), x - ) - - def else_fn(x): - return simp.pcall_static( - (0, 1), lambda a: tensor.run_jax(lambda v: v * v, a), x - ) - - pred_true = simp.constant((0, 1), True) - # uniform_cond - res = simp.uniform_cond(pred_true, then_fn, else_fn, x_obj) - - values = mp.fetch(res) - values = [ - int(v) if not hasattr(v, "shape") or v.shape == () else v for v in values - ] - assert values == [2, 4] - hasattr(sim, "_simp_cluster") and sim._simp_cluster.shutdown()