mplang-nightly 0.1.dev267__tar.gz → 0.1.dev269__tar.gz

{mplang_nightly-0.1.dev267 → mplang_nightly-0.1.dev269}/PKG-INFO

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

mplang_nightly-0.1.dev269/mplang/py.typed

@@ -0,0 +1,13 @@
+# 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.

{mplang_nightly-0.1.dev267 → mplang_nightly-0.1.dev269}/mplang/v2/backends/simp_driver/http.py

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
     from concurrent.futures import Future
 
     from mplang.v2.edsl.graph import Graph
-    from mplang.v2.edsl.spec import ClusterSpec
+    from mplang.v2.libs.device import ClusterSpec
 
 
 class SimpHttpDriver(SimpDriver):

{mplang_nightly-0.1.dev267 → mplang_nightly-0.1.dev269}/mplang/v2/backends/simp_driver/mem.py

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
     from concurrent.futures import Future
 
     from mplang.v2.edsl.graph import Graph
-    from mplang.v2.edsl.spec import ClusterSpec
+    from mplang.v2.libs.device import ClusterSpec
 
 
 class MemCluster:

{mplang_nightly-0.1.dev267 → mplang_nightly-0.1.dev269}/mplang/v2/kernels/okvs_opt.cpp

@@ -36,6 +36,10 @@ extern "C" {
         uint64_t h1, h2, h3;
     };
 
+    // Declaration of the safe (robust) solver implemented in okvs.cpp
+    // Signature: solve_okvs(keys, values, output, n, m, seed_ptr)
+    void solve_okvs(uint64_t* keys, uint64_t* values, uint64_t* output, uint64_t n, uint64_t m, uint64_t* seed_ptr);
+
     // Stateless Bin Selection
     // Maps a key to a deterministic bin index [0, NUM_BINS).
     inline uint64_t get_bin_index(uint64_t key, __m128i seed) {
@@ -245,10 +249,21 @@ extern "C" {
             uint64_t valid_m = m_per_bin[b];
             
             if(!solve_bin(bin_keys[b], bin_vals[b], &P_vec[offset], valid_m, seed)) {
+                // On failure, log and fall back to the robust solver for this bin.
+                // The fallback is executed inside a critical section to avoid nested OpenMP
+                // regions and to serialize rare fallbacks.
                 #pragma omp critical
                 {
-                    fprintf(stderr, "[ERROR] Bin %lu failed OKVS peeling. Items: %lu / M: %lu (Ratio: %.2f). Try increasing expansion factor.\n", 
+                    fprintf(stderr, "[WARN] Bin %lu failed optimized peeling; falling back to safe solver. Items: %lu / M: %lu (Ratio: %.2f)\n",
                         b, bin_keys[b].size(), valid_m, (double)valid_m / bin_keys[b].size());
+
+                    // Prepare pointers for the safe solver
+                    uint64_t* keys_ptr = &bin_keys[b][0];
+                    uint64_t* vals_ptr = &bin_vals[b][0];
+                    uint64_t* out_ptr = output + (offset * 2ULL); // each 128-bit slot == 2 uint64_t
+
+                    // Call the safe solver implemented in okvs.cpp
+                    solve_okvs(keys_ptr, vals_ptr, out_ptr, bin_keys[b].size(), valid_m, seed_ptr);
                 }
             }
         }

{mplang_nightly-0.1.dev267 → mplang_nightly-0.1.dev269}/mplang/v2/libs/mpc/psi/okvs_gct.py

@@ -40,12 +40,16 @@ def get_okvs_expansion(n: int) -> float:
     - For N → ∞: Theoretical minimum is ε ≈ 0.23 (M = 1.23N)
     - For finite N: Larger ε needed due to variance in random hash collisions
 
-    Empirical safe thresholds (failure probability < 0.01%):
-    - N < 1,000:    ε = 4.5  (M = 5.5N)  - very small sets need extra wide margin
-                                           to handle worst-case hash collisions
-    - N < 10,000:   ε = 0.4  (M = 1.4N)
-    - N < 100,000:  ε = 0.3  (M = 1.3N)
-    - N ≥ 100,000:  ε = 0.35 (M = 1.35N) - large sets converge near theory
+    Empirical safe thresholds (failure probability < 0.001%):
+    - N ≤ 200:      ε = 24.0 (M = 25.0N)  - extremely small sets need very wide margin
+    - N < 1,000:    ε = 11.0 (M = 12.0N)  - small sets need extra wide safety margin
+    - N < 10,000:   ε = 0.6  (M = 1.6N)
+    - N < 100,000:  ε = 0.4  (M = 1.4N)
+    - N ≥ 100,000:  ε = 0.35 (M = 1.35N)  - large sets converge near theory
+
+    Note: These expansion factors account for the 128-byte alignment requirement
+    in the OKVS implementation. The factors are intentionally conservative to
+    ensure high success rates (>99.9%) for the probabilistic peeling algorithm.
 
     Args:
         n: Number of key-value pairs to encode
@@ -53,12 +57,14 @@ def get_okvs_expansion(n: int) -> float:
     Returns:
         Expansion factor ε such that M = (1+ε)*N is safe for peeling
     """
-    if n < 1000:
-        return 5.5  # Small scale: need very wide safety margin for stability
+    if n <= 200:
+        return 25.0  # Extremely small scale: need very wide margin for stability
+    elif n < 1000:
+        return 12.0  # Small scale: need wide safety margin for stability
     elif n <= 10000:
-        return 1.4  # Medium scale
+        return 1.6  # Medium scale
     elif n <= 100000:
-        return 1.3  # Large scale
+        return 1.4  # Large scale
     else:
         # Mega-Binning requires ~1.35 for stability with 1024 bins
         return 1.35

mplang_nightly-0.1.dev269/mplang/v2/libs/mpc/psi/rr22.py

@@ -0,0 +1,303 @@
+# 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.
+
+"""Private Set Intersection using VOLE and OKVS (RR22-Style).
+
+This module implements a high-performance PSI protocol based on the "Blazing Fast PSI"
+(RR22) paper. The protocol relies on Vector Oblivious Linear Evaluation (VOLE) and
+Oblivious Key-Value Stores (OKVS) to achieve efficient set intersection with linear
+communication O(N) and computation complexity.
+
+Protocol Overview:
+The core idea is to mask a "Polynomial" (encoded via OKVS) with VOLE-correlated randomness,
+such that the mask can only be removed (and the polynomial verified) if the parties share
+the same element.
+
+Phases:
+1.  **Correlated Randomness (VOLE)**:
+    Sender and Receiver establish a shared correlation:
+    W = V + U * Delta
+    - PSI Receiver holds `U` and `V` (these are generated by the OT "sender"
+      role in `silent_vole_random_u`).
+    - PSI Sender holds `W` and `Delta` (these are generated by the OT "receiver"
+      role).
+    - `U` is random; `Delta` is a fixed secret scalar held by the Sender.
+
+2.  **Encoding (OKVS)**:
+    The Receiver encodes its input set Y into an OKVS storage `P` such that
+    Decode(P, y) = H(y) for all y in Y. The function `H(y)` is implemented via
+    AES/Davies–Meyer expansion acting as a random oracle.
+
+3.  **Masking & Exchange**:
+    The Receiver masks the OKVS storage `P` with its VOLE share `U`:
+    Q = P + U
+    The masked storage `Q` is sent to the Sender (so Sender sees a masked OKVS).
+
+4.  **Decoding & Tag Generation (Sender)**:
+    The Sender holds `W` and `Delta` and computes the linear combination:
+    K = Q * Delta + W
+    Using W = V + U * Delta and Q = P + U, this simplifies to
+    K = P * Delta + V.
+    The Sender decodes `K` for each of its items x to obtain P(x)*Delta + V(x),
+    then subtracts H(x)*Delta (computed locally) to recover `V(x)`. The value
+    `Tag = V(x)` serves as the sender-side tag for item x.
+
+5.  **Verification (Receiver)**:
+    The Sender hashes and truncates tags and sends the truncated hashes to the
+    Receiver. The Receiver locally decodes `V(y)` from its OKVS, hashes it with
+    the same domain separation and truncation, and compares to the received
+    truncated hashes to determine membership in the intersection.
+"""
+
+from typing import Any, cast
+
+import jax.numpy as jnp
+
+import mplang.v2.edsl as el
+import mplang.v2.libs.mpc.ot.silent as silent_ot
+from mplang.v2.dialects import field, simp, tensor
+
+
+def psi_intersect(
+    sender: int,
+    receiver: int,
+    n: int,
+    sender_items: el.Object,
+    receiver_items: el.Object,
+) -> el.Object:
+    """Execute OKVS-based PSI Protocol (Original RR22 Logic).
+
+    This implementation follows the RR22 paper's role assignment where:
+    - PSI Sender holds Delta (and W).
+    - PSI Receiver holds U and V.
+
+    This enables the "One Decode" optimization on the Sender side and prevents
+    offline brute-force attacks by the Receiver (though Sender could brute-force).
+
+    Args:
+        sender: Rank of Sender.
+        receiver: Rank of Receiver.
+        n: Number of items.
+        sender_items: Object located at Sender.
+        receiver_items: Object located at Receiver.
+
+    Returns:
+        Intersection mask (0/1) located at Receiver.
+    """
+
+    # Validation
+    if sender == receiver:
+        raise ValueError(
+            f"Sender ({sender}) and Receiver ({receiver}) must be different."
+        )
+    if n <= 0:
+        raise ValueError(f"Input size n must be positive, got {n}.")
+
+    # =========================================================================
+    # Phase 1. Parameter Setup
+    # =========================================================================
+    import mplang.v2.libs.mpc.psi.okvs_gct as okvs_gct
+
+    expansion = okvs_gct.get_okvs_expansion(n)
+    M = int(n * expansion)
+    if M % 128 != 0:
+        M = ((M // 128) + 1) * 128
+
+    # =========================================================================
+    # Phase 2. Correlated Randomness (VOLE)
+    # =========================================================================
+    # In the original paper logic (Fig 4), the PSI Sender holds Delta.
+    # Therefore, we swap the roles in the OT call.
+    #
+    # silent_vole_random_u(A, B) gives:
+    #   A (OT Sender):   U, V
+    #   B (OT Receiver): W, Delta
+    #
+    # We want PSI Sender to be OT Receiver.
+    res_tuple = silent_ot.silent_vole_random_u(receiver, sender, M, base_k=1024)
+
+    # PSI Receiver gets U, V
+    v_recv, w_sender, u_recv, delta_sender = res_tuple[:4]
+
+    # =========================================================================
+    # Phase 3. Receiver Encoding & Masking
+    # =========================================================================
+    # Receiver computes P such that P(y) = H(y).
+    # Receiver masks P with U (Paper's A vector).
+    # Q = P ^ U
+
+    from mplang.v2.dialects import crypto
+    from mplang.v2.edsl import typing as elt
+
+    def _gen_seed() -> Any:
+        return crypto.random_tensor((2,), elt.u64)
+
+    okvs_seed = simp.pcall_static((receiver,), _gen_seed)
+    okvs_seed_sender = simp.shuffle_static(okvs_seed, {sender: receiver})
+
+    okvs = okvs_gct.SparseOKVS(M)
+
+    def _recv_ops(y: Any, u: Any, seed: Any) -> Any:
+        # 3.1 Compute H(y)
+        def _reshape_seeds(items: Any) -> Any:
+            lo = items
+            hi = jnp.zeros_like(items)
+            return jnp.stack([lo, hi], axis=1)  # (N, 2)
+
+        seeds = tensor.run_jax(_reshape_seeds, y)
+        res_exp = field.aes_expand(seeds, 1)  # (N, 1, 2)
+
+        def _davies_meyer(enc: Any, s: Any) -> Any:
+            enc_flat = enc.reshape(enc.shape[0], 2)
+            return jnp.bitwise_xor(enc_flat, s)
+
+        h_y = tensor.run_jax(_davies_meyer, res_exp, seeds)
+
+        # 3.2 Encode P
+        p_storage = okvs.encode(y, h_y, seed)
+
+        # 3. Mask with U (instead of W)
+        # Q = P ^ U
+        q_storage = field.add(p_storage, u)
+
+        return q_storage
+
+    # Receiver uses U to mask
+    q_shared = simp.pcall_static(
+        (receiver,), _recv_ops, receiver_items, u_recv, okvs_seed
+    )
+
+    q_sender_view = simp.shuffle_static(q_shared, {sender: receiver})
+
+    # =========================================================================
+    # Phase 4. Sender "One Decode" & Tag Generation
+    # =========================================================================
+    # Sender holds W, Delta. Receives Q.
+    # W = V + U * Delta
+    #
+    # Derivation:
+    # K = Q * Delta + W
+    #   = (P + U) * Delta + (V + U * Delta)
+    #   = P * Delta + U * Delta + V + U * Delta
+    #   = P * Delta + V
+    #
+    # Sender computes Tag = Decode(K, x) - H(x) * Delta
+    # If x in Intersection: P(x) = H(x)
+    # Tag = (P(x) * Delta + V(x)) - P(x) * Delta
+    # Tag = V(x)
+
+    def _sender_ops(x: Any, q: Any, w: Any, delta: Any, seed: Any) -> Any:
+        # q, w: (M, 2)
+        # delta: (2,)
+
+        # Safe tiling assuming M is aligned
+        def _tile_m_simple(d: Any) -> Any:
+            return jnp.tile(d, (M, 1))
+
+        delta_expanded_m = tensor.run_jax(_tile_m_simple, delta)
+
+        # 4.2. Compute Global K = Q * Delta + W
+        # This is the O(M) multiplication mentioned in the paper
+        q_times_delta = field.mul(q, delta_expanded_m)
+        k_storage = field.add(q_times_delta, w)
+
+        # 4.3 One Decode
+        # decoded_val = P(x)*Delta + V(x)
+        decoded_k = okvs.decode(x, k_storage, seed)
+
+        # 4.4 Remove H(x)*Delta
+        def _reshape_seeds(items: Any) -> Any:
+            lo = items
+            hi = jnp.zeros_like(items)
+            return jnp.stack([lo, hi], axis=1)
+
+        seeds_x = tensor.run_jax(_reshape_seeds, x)
+        res_exp_x = field.aes_expand(seeds_x, 1)
+
+        def _davies_meyer(enc: Any, s: Any) -> Any:
+            enc_flat = enc.reshape(enc.shape[0], 2)
+            return jnp.bitwise_xor(enc_flat, s)
+
+        h_x = tensor.run_jax(_davies_meyer, res_exp_x, seeds_x)
+
+        # Expand delta for batch N
+        def _tile_n(d: Any) -> Any:
+            return jnp.tile(d, (n, 1))
+
+        delta_expanded_n = tensor.run_jax(_tile_n, delta)
+
+        h_x_times_delta = field.mul(h_x, delta_expanded_n)
+
+        # Final Tag = (P*Delta + V) + H*Delta = V(x)
+        tag = field.add(decoded_k, h_x_times_delta)
+
+        return tag
+
+    # Execute on Sender
+    sender_tags = simp.pcall_static(
+        (sender,),
+        _sender_ops,
+        sender_items,
+        q_sender_view,
+        w_sender,
+        delta_sender,
+        okvs_seed_sender,
+    )
+    # =========================================================================
+    # Phase 5. Verification (Receiver Side)
+    # =========================================================================
+    # Sender sends Tags (which should be V(x)) to Receiver. To reduce
+    # communication we hash and truncate on the sender side and only send
+    # the truncated hash (first 16 bytes).
+
+    # 5.1 Compute hashed & truncated tags on Sender
+    from mplang.v2.libs.mpc.ot import extension as ot_extension
+
+    def _hash_and_trunc(tags: Any) -> Any:
+        # Compute batched hash on sender and truncate to 16 bytes
+        full_h = ot_extension.vec_hash(tags, domain_sep=0x1111, num_rows=n)
+        # Use tensor.slice_tensor to slice TraceObjects (start=(0,0), end=(n,16))
+        return tensor.slice_tensor(full_h, (0, 0), (n, 16))
+
+    h_sender_trunc = simp.pcall_static((sender,), _hash_and_trunc, sender_tags)
+
+    # 5.2 Send truncated hashes to Receiver (much smaller payload)
+    tags_at_recv = simp.shuffle_static(h_sender_trunc, {receiver: sender})
+
+    # 5.3 Receiver computes local V(y) and compares
+    def _recv_verify(y: Any, v: Any, seed: Any, remote_tags: Any) -> Any:
+        # 1. Decode V locally: target = V(y)
+        local_v_y = okvs.decode(y, v, seed)
+
+        # 2. Hash local V(y) and compare with received truncated sender hashes
+        # Note: `remote_tags` here is already the truncated hash (16 bytes)
+        # sent from the Sender.
+        h_local = ot_extension.vec_hash(local_v_y, domain_sep=0x1111, num_rows=n)
+
+        def _core(h_r16: Any, h_l_full: Any) -> Any:
+            # h_r16: (n, 16) truncated bytes from sender
+            # h_l_full: (n, k) full hash bytes; truncate to 16
+            h_l16 = h_l_full[:, :16]
+
+            eq_matrix = jnp.all(h_r16[:, None, :] == h_l16[None, :, :], axis=2)
+            membership = jnp.any(eq_matrix, axis=0)
+            return membership.astype(jnp.uint8)
+
+        return tensor.run_jax(_core, remote_tags, h_local)
+
+    intersection_mask = simp.pcall_static(
+        (receiver,), _recv_verify, receiver_items, v_recv, okvs_seed, tags_at_recv
+    )
+
+    return cast(el.Object, intersection_mask)

{mplang_nightly-0.1.dev267 → mplang_nightly-0.1.dev269}/tests/v2/libs/mpc/psi/test_rr22.py

@@ -201,7 +201,7 @@ def _run_psi_simulation(
 
         traced = mp.compile(job)
         mask_obj = mp.evaluate(traced)
-        mask_val = mp.fetch(mask_obj)[SENDER]
+        mask_val = mp.fetch(mask_obj)[RECEIVER]
     return mask_val
 
 
@@ -225,8 +225,9 @@ def test_rr22_partial_intersection() -> None:
 
     mask_val = _run_psi_simulation(sender_items, receiver_items)
 
-    # Compute expected mask: 1 if sender_item is in common_items, else 0
-    expected_mask = np.isin(sender_items, common_items).astype(np.uint8)
+    # Compute expected mask: 1 if receiver_item is in common_items, else 0
+    # (PSI returns mask located at Receiver)
+    expected_mask = np.isin(receiver_items, common_items).astype(np.uint8)
 
     np.testing.assert_array_equal(
         mask_val,
@@ -275,8 +276,8 @@ def test_rr22_single_element_intersection() -> None:
 
     mask_val = _run_psi_simulation(sender_items, receiver_items)
 
-    # Exactly one element should match
-    expected_mask = (sender_items == common_item).astype(np.uint8)
+    # Exactly one element should match (mask is for receiver's items)
+    expected_mask = (receiver_items == common_item).astype(np.uint8)
 
     np.testing.assert_array_equal(
         mask_val,