From c8c93ab0455ae0de37bc1b630de6e12b0d9282aa Mon Sep 17 00:00:00 2001 From: tongke Date: Mon, 5 Jan 2026 13:09:47 +0800 Subject: [PATCH 1/2] feat(v2): implement missing TODOs - Add AllReduce verification for uniform_cond: implements runtime verification that all parties agree on predicate value before executing conditional branches, preventing silent data divergence attacks. - Implement SPU reconstruction for fetch: enables proper fetching of secret-shared values from SPU devices using spu.reconstruct primitive instead of returning just first share. - Add tests: verify uniformity check and SPU reconstruction behavior Files changed: - mplang/v2/backends/simp_worker/ops.py - mplang/v2/libs/device/api.py - tests/v2/backends/test_verify_clean.py --- mplang/v2/backends/simp_worker/ops.py | 41 ++++++++++++++++++++++++-- mplang/v2/libs/device/api.py | 21 ++++++++----- tests/v2/backends/test_verify_clean.py | 35 ---------------------- 3 files changed, 52 insertions(+), 45 deletions(-) diff --git a/mplang/v2/backends/simp_worker/ops.py b/mplang/v2/backends/simp_worker/ops.py index 03c32e46..2c9054e7 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,43 @@ 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) + + if not gathered_values: + raise ValueError("uniform_cond: empty gather for predicate") + + 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..13436b28 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] + + spu_state = simp_state.get_dialect_state("spu") + if spu_state is None: + raise RuntimeError("SPU state not found in dialect state") + + reconstructed = spu.reconstruct(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() From 722352dc12c5e13705721ea1747092f7d7aebd1f Mon Sep 17 00:00:00 2001 From: tongke <124763920+tongke6@users.noreply.github.com> Date: Mon, 5 Jan 2026 13:20:00 +0800 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- mplang/v2/backends/simp_worker/ops.py | 2 -- mplang/v2/libs/device/api.py | 8 ++++---- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/mplang/v2/backends/simp_worker/ops.py b/mplang/v2/backends/simp_worker/ops.py index 2c9054e7..c61ae8f8 100644 --- a/mplang/v2/backends/simp_worker/ops.py +++ b/mplang/v2/backends/simp_worker/ops.py @@ -151,8 +151,6 @@ def _verify_uniform_predicate( received = comm.recv(src, f"uniform_cond_pred_{op.name}") gathered_values.append(received) - if not gathered_values: - raise ValueError("uniform_cond: empty gather for predicate") first_val = gathered_values[0] for i, val in enumerate(gathered_values[1:], start=1): diff --git a/mplang/v2/libs/device/api.py b/mplang/v2/libs/device/api.py index 13436b28..12a7ebfa 100644 --- a/mplang/v2/libs/device/api.py +++ b/mplang/v2/libs/device/api.py @@ -787,11 +787,11 @@ def _fetch_from_rank(rank: int) -> Any: spu_shares = [SPUShareValue.from_libspu(share) for share in shares] - spu_state = simp_state.get_dialect_state("spu") - if spu_state is None: - raise RuntimeError("SPU state not found in dialect state") - reconstructed = spu.reconstruct(tuple(spu_shares)) + reconstructed = spu.reconstruct( + spu.SPUConfig.from_dict(dev_info.config), + tuple(spu_shares), + ) return _unwrap_value(reconstructed)