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

mplang/v2/libs/mpc/__init__.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""MPC (Multi-Party Computation) library for MPLang2.
+
+Subpackages:
+- ot: Oblivious Transfer protocols
+- vole: Vector OLE protocols
+- psi: Private Set Intersection
+- analytics: Privacy-preserving analytics
+
+Example usage:
+    from mplang.v2.libs.mpc import ot_transfer, apply_permutation
+    from mplang.v2.libs.mpc.vole import silver_vole
+    from mplang.v2.libs.mpc.psi import psi_intersect
+"""
+
+from .analytics.aggregation import rotate_and_sum
+from .analytics.groupby import oblivious_groupby_sum_bfv, oblivious_groupby_sum_shuffle
+from .analytics.permutation import apply_permutation, secure_switch
+from .ot.base import transfer as ot_transfer
+
+__all__ = [
+    "apply_permutation",
+    "oblivious_groupby_sum_bfv",
+    "oblivious_groupby_sum_shuffle",
+    "ot_transfer",
+    "rotate_and_sum",
+    "secure_switch",
+]

mplang/v2/libs/mpc/_utils.py

@@ -0,0 +1,99 @@
+# 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.
+
+"""Utilities for MPC protocols."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import jax.numpy as jnp
+
+import mplang.v2.edsl as el
+from mplang.v2.dialects import tensor
+
+
+def bytes_to_bits(data: el.Object) -> el.Object:
+    """Convert bytes (uint8 tensor) to bits (uint8 tensor of 0s and 1s).
+
+    Output shape logic: (..., N) -> (..., N * 8)
+    """
+
+    def _to_bits(arr: Any) -> Any:
+        # View as u8
+        y_u8 = arr.view(jnp.uint8)
+        # Unpack produces Big Endian bits [b7, b6, ..., b0] per byte
+        bits = jnp.unpackbits(y_u8)
+        # Reshape to (N, 8) and flip to get [b0, ..., b7]
+        bits = bits.reshape(-1, 8)
+        bits = jnp.fliplr(bits)
+        return bits.reshape(-1)
+
+    return cast(el.Object, tensor.run_jax(_to_bits, data))
+
+
+def bits_to_bytes(bits: el.Object) -> el.Object:
+    """Convert bits to bytes.
+
+    Output shape logic: (..., N * 8) -> (..., N)
+    """
+
+    def _to_bytes(arr: Any) -> Any:
+        return jnp.packbits(arr, axis=-1)
+
+    return cast(el.Object, tensor.run_jax(_to_bytes, bits))
+
+
+def transpose_128(matrix_bits: el.Object) -> el.Object:
+    """Transpose a bit matrix.
+
+    Just a wrapper for tensor.transpose currently.
+    """
+    return tensor.transpose(matrix_bits, perm=(1, 0))
+
+
+class CuckooHash:
+    """Simple Cuckoo Hashing simulation."""
+
+    def __init__(self, num_bins: int, num_hash_functions: int = 3, stash_size: int = 0):
+        self.num_bins = num_bins
+        self.num_functions = num_hash_functions
+        self.stash_size = stash_size
+
+    def hash(self, items: el.Object, seed: int) -> el.Object:
+        """Hash items to bin indices."""
+
+        # We perform hashing.
+        # Note: We return hashes for each function?
+        # Usually simplest cuckoo uses 3 hash functions.
+        # We can return (num_funcs, N) or (N, num_funcs)
+
+        def _hash_fn(xs: Any, s: int) -> Any:
+            # xs: array of items
+
+            # Simple hash: (x * s + s) % bins
+            # We want multiple hashes?
+            # For now, let's just return one hash per seed provided (assuming call per seed)
+            # Or if seed is a single int, we might mix it.
+
+            # Let's assume this function handles one hash instance.
+            res = (xs * s + s) % self.num_bins
+            return res.astype(jnp.int32)
+
+        # Passing self.num_bins as constant implementation detail inside _hash_fn closure is fine
+        # if using run_jax (as it's compiled).
+        # Actually run_jax recompiles if closure changes?
+        # run_jax supports closures.
+
+        return cast(el.Object, tensor.run_jax(_hash_fn, items, seed))

mplang/v2/libs/mpc/analytics/__init__.py

@@ -0,0 +1,35 @@
+# 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.
+
+"""Privacy-preserving analytics operations.
+
+Submodules:
+- aggregation: BFV homomorphic aggregation
+- groupby: Oblivious Group-By operations
+- permutation: Secure permutation (Bitonic Sort)
+"""
+
+from .aggregation import aggregate_sparse, batch_bucket_aggregate, rotate_and_sum
+from .groupby import oblivious_groupby_sum_bfv, oblivious_groupby_sum_shuffle
+from .permutation import apply_permutation, secure_switch
+
+__all__ = [
+    "aggregate_sparse",
+    "apply_permutation",
+    "batch_bucket_aggregate",
+    "oblivious_groupby_sum_bfv",
+    "oblivious_groupby_sum_shuffle",
+    "rotate_and_sum",
+    "secure_switch",
+]

mplang/v2/libs/mpc/analytics/aggregation.py

@@ -0,0 +1,372 @@
+# 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.
+
+"""Homomorphic Aggregation library.
+
+This module implements efficient aggregation algorithms using BFV rotation.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import numpy as np
+
+from mplang.v2.dialects import bfv, tensor
+
+
+def _safe_rotate(
+    ciphertext: Any, step: int, galois_keys: Any, max_step: int = 1024
+) -> Any:
+    """Rotate ciphertext by step, decomposing large steps if needed.
+
+    SEAL's rotate_rows requires step to be in range (-slot_count/2, slot_count/2).
+    For poly_modulus_degree=4096, slot_count=4096, max valid step is 2047.
+    We use a conservative max_step=1024 by default for safety.
+
+    For large steps, we decompose into multiple rotations:
+    - rotate(x, 3000) = rotate(rotate(rotate(x, 1024), 1024), 952)
+    """
+    if step == 0:
+        return ciphertext
+    if abs(step) <= max_step:
+        return bfv.rotate(ciphertext, step, galois_keys)
+
+    # Decompose large step into multiple rotations
+    current = ciphertext
+    remaining = abs(step)
+    sign = 1 if step > 0 else -1
+
+    while remaining > 0:
+        rot = min(remaining, max_step)
+        current = bfv.rotate(current, sign * rot, galois_keys)
+        remaining -= rot
+
+    return current
+
+
+def _rotate_and_sum_row(
+    ciphertext: Any, k: int, galois_keys: Any, max_step: int = 1024
+) -> Any:
+    """Sum first k elements within a single row (k <= row_size).
+
+    Uses the recursive doubling algorithm with O(log k) rotations.
+    """
+    if k <= 1:
+        return ciphertext
+
+    num_steps = math.ceil(math.log2(k))
+    current = ciphertext
+
+    for i in range(num_steps):
+        step = 1 << i
+        if step >= k:
+            break
+        rotated = _safe_rotate(current, step, galois_keys, max_step)
+        current = bfv.add(current, rotated)
+
+    return current
+
+
+def rotate_and_sum(
+    ciphertext: Any, k: int, galois_keys: Any, slot_count: int = 4096
+) -> Any:
+    """Aggregate the first k elements of a ciphertext using O(log k) rotations.
+
+    The result is placed in the 0-th slot.
+    This assumes the input ciphertext has relevant data in slots 0..k-1
+    and zeros (or irrelevant data) elsewhere, OR that the caller will mask the result.
+
+    Args:
+        ciphertext: The BFV ciphertext.
+        k: The number of elements to sum.
+        galois_keys: Keys required for rotation.
+        slot_count: Total number of slots (default 4096 for poly_degree=4096).
+
+    Returns:
+        Ciphertext where slot 0 contains sum(ciphertext[0..k-1]).
+
+    Note:
+        SEAL batching arranges slots as 2 rows of slot_count/2 each.
+        - rotate_rows rotates within each row (circular)
+        - rotate_columns swaps the two rows
+
+        For k <= row_size (2048), only row rotations are needed.
+        For k > row_size, we use rotate_columns to aggregate across rows.
+    """
+    row_size = slot_count // 2
+
+    if k <= row_size:
+        # Simple case: all elements in row 0
+        return _rotate_and_sum_row(ciphertext, k, galois_keys)
+
+    # k > row_size: data spans both rows
+    # Strategy:
+    # 1. Sum row 0 completely (row_size elements)
+    # 2. rotate_columns to bring row 1 to row 0 position
+    # 3. Sum the first (k - row_size) elements of what was row 1
+    # 4. Add the two partial sums
+
+    # Sum row 0 completely
+    row0_sum = _rotate_and_sum_row(ciphertext, row_size, galois_keys)
+
+    # Rotate columns: swap row 0 <-> row 1
+    # Now row 1's data is in row 0 position
+    col_rotated = bfv.rotate_columns(ciphertext, galois_keys)
+
+    # Sum the first (k - row_size) elements (originally in row 1)
+    row1_count = k - row_size
+    row1_sum = _rotate_and_sum_row(col_rotated, row1_count, galois_keys)
+
+    # Both row0_sum and row1_sum have their results in slot 0
+    # Add them together
+    return bfv.add(row0_sum, row1_sum)
+
+
+def aggregate_sparse(
+    ciphertext: Any,
+    aggregations: list[tuple[int, list[int]]],
+    galois_keys: Any,
+    encoder: Any,
+    vector_size: int,
+) -> Any:
+    """Perform sparse aggregation.
+
+    Args:
+        ciphertext: Input ciphertext.
+        aggregations: List of (target_slot, [source_slots]).
+                      e.g. [(0, [0, 3, 8]), (1, [1, 5])]
+        galois_keys: Rotation keys.
+        encoder: BFV encoder for encoding masks.
+        vector_size: Total size of the vector (slots).
+
+    Returns:
+        Ciphertext with aggregated results in target slots.
+    """
+    # Naive approach: For each target, sum sources.
+    # Optimized approach:
+    # 1. Decompose into rotations.
+    #    For target t, source s: need rotation by (t-s).
+    #    Group by rotation amount.
+    # 2. Apply rotations and accumulate.
+
+    # Map: rotation_amount -> mask
+    # We want to compute: Result = Sum( Rotate(Input, r) * Mask_r )
+    # where Mask_r has 1 at slot t if (t - r) is a source for t.
+
+    # Example: t=0, s={0, 3, 8}.
+    #   s=0: rot=0. Mask[0]=1.
+    #   s=3: rot=-3. Mask[0]=1.
+    #   s=8: rot=-8. Mask[0]=1.
+    # Example: t=1, s={1, 5}.
+    #   s=1: rot=0. Mask[1]=1.
+    #   s=5: rot=-4. Mask[1]=1.
+
+    # Combined:
+    # Rot 0: Mask[0]=1, Mask[1]=1. -> Mask = [1, 1, 0...]
+    # Rot -3: Mask[0]=1. -> Mask = [1, 0...]
+    # Rot -8: Mask[0]=1. -> Mask = [1, 0...]
+    # Rot -4: Mask[1]=1. -> Mask = [0, 1, 0...]
+
+    rotations = {}  # shift -> mask_list
+
+    for target, sources in aggregations:
+        for src in sources:
+            shift = src - target
+            if shift not in rotations:
+                rotations[shift] = [0] * vector_size
+            rotations[shift][target] = 1
+
+    final_result = None
+
+    for shift, mask_list in rotations.items():
+        # Optimization: Skip if mask is all zeros (no contribution from this rotation)
+        if not any(mask_list):
+            continue
+
+        # Create mask plaintext
+        # In a real implementation, we encode this list to a Plaintext
+        mask_tensor = tensor.constant(np.array(mask_list, dtype=np.int64))
+        mask_pt = bfv.encode(mask_tensor, encoder)
+
+        # Rotate
+        if shift == 0:
+            rotated_c = ciphertext
+        else:
+            rotated_c = bfv.rotate(ciphertext, shift, galois_keys)
+
+        # Mask
+        masked_c = bfv.mul(rotated_c, mask_pt)
+
+        # Accumulate
+        if final_result is None:
+            final_result = masked_c
+        else:
+            final_result = bfv.add(final_result, masked_c)
+
+    return final_result
+
+
+def masked_aggregate(ciphertexts: list[Any], masks: list[Any]) -> Any:
+    """Aggregate multiple partial results using masks.
+
+    Args:
+        ciphertexts: List of ciphertexts.
+        masks: List of plaintexts (masks).
+
+    Returns:
+        Sum(ct * mask)
+    """
+    if not ciphertexts or not masks:
+        raise ValueError("Empty input lists")
+    if len(ciphertexts) != len(masks):
+        raise ValueError("Mismatch in ciphertexts and masks length")
+
+    total = None
+
+    for ct, mask in zip(ciphertexts, masks, strict=True):
+        # ct * mask
+        masked = bfv.mul(ct, mask)
+
+        if total is None:
+            total = masked
+        else:
+            total = bfv.add(total, masked)
+
+    return total
+
+
+# ==============================================================================
+# SIMD Bucket Packing for Histogram Computation
+# ==============================================================================
+
+
+def strided_rotate_and_sum(
+    ciphertext: Any,
+    stride: int,
+    n_elements: int,
+    galois_keys: Any,
+    max_step: int = 1024,
+) -> Any:
+    """Aggregate elements at positions [0, stride, 2*stride, ...] into slot 0.
+
+    This is used for SIMD bucket packing where each bucket's values are
+    placed at strided positions.
+
+    Args:
+        ciphertext: The BFV ciphertext with values at strided positions.
+        stride: Distance between consecutive elements to sum.
+        n_elements: Number of elements to aggregate (at positions 0, stride, ..., (n-1)*stride).
+        galois_keys: Rotation keys.
+        max_step: Maximum rotation step for safe_rotate.
+
+    Returns:
+        Ciphertext where slot 0 contains sum of strided elements.
+
+    Example:
+        stride=64, n_elements=47 (bucket has 47 samples)
+        Values at slots: 0, 64, 128, 192, ...
+        Result: slot[0] = sum of all these values
+    """
+    if n_elements <= 1:
+        return ciphertext
+
+    # Use recursive doubling with strided rotations
+    # Step 1: rotate by stride, add -> pairs summed at even positions
+    # Step 2: rotate by 2*stride, add -> quads summed at positions 0, 4*stride, ...
+    # ...
+    num_steps = math.ceil(math.log2(n_elements))
+    current = ciphertext
+
+    for i in range(num_steps):
+        step = stride * (1 << i)
+        if step >= n_elements * stride:
+            break
+        rotated = _safe_rotate(current, step, galois_keys, max_step)
+        current = bfv.add(current, rotated)
+
+    return current
+
+
+def batch_bucket_aggregate(
+    ciphertext: Any,
+    n_buckets: int,
+    samples_per_bucket: int,
+    galois_keys: Any,
+    slot_count: int = 4096,
+) -> Any:
+    """Aggregate samples within each bucket region in a packed ciphertext.
+
+    Assumes the ciphertext has the following layout:
+    - slot_count is divided into n_buckets regions of size `stride = slot_count // n_buckets`
+    - Each bucket b occupies slots [b*stride, b*stride + samples_per_bucket)
+    - Samples are placed at consecutive positions within their bucket region
+
+    After aggregation, slot[b * stride] contains sum of bucket b.
+
+    Args:
+        ciphertext: Packed ciphertext with samples in bucket regions.
+        n_buckets: Number of buckets.
+        samples_per_bucket: Max samples per bucket (for rotation count).
+        galois_keys: Rotation keys.
+        slot_count: Total BFV slots.
+
+    Returns:
+        Ciphertext where slot[b * stride] = sum of bucket b's values.
+    """
+    if samples_per_bucket <= 1:
+        return ciphertext
+
+    # Use recursive doubling within each bucket region
+    # Since all buckets use the same relative positions, one set of rotations
+    # aggregates ALL buckets simultaneously!
+    num_steps = math.ceil(math.log2(samples_per_bucket))
+    current = ciphertext
+
+    for i in range(num_steps):
+        step = 1 << i
+        if step >= samples_per_bucket:
+            break
+        # Rotating by `step` shifts values within each bucket region
+        # Add original + rotated to sum pairs/quads/etc.
+        rotated = _safe_rotate(current, step, galois_keys)
+        current = bfv.add(current, rotated)
+
+    return current
+
+
+def extract_bucket_results(
+    vector: Any,
+    n_buckets: int,
+    slot_count: int = 4096,
+) -> Any:
+    """Extract bucket sums from a packed result vector.
+
+    After batch_bucket_aggregate, each bucket's sum is at slot[b * stride].
+    This function extracts those values.
+
+    Args:
+        vector: Decoded vector from packed ciphertext.
+        n_buckets: Number of buckets.
+        slot_count: Total slots.
+
+    Returns:
+        (n_buckets,) array of bucket sums.
+    """
+    import jax.numpy as jnp
+
+    stride = slot_count // n_buckets
+    indices = jnp.arange(n_buckets) * stride
+    return vector[indices]

mplang/v2/libs/mpc/analytics/groupby.md

@@ -0,0 +1,99 @@
+# Oblivious Group-by Sum Design
+
+This document outlines the design for Oblivious Group-by Sum algorithms in MPLang. The goal is to compute the sum of values in `data` (held by P0) grouped by `bins` (held by P1), such that:
+- P0 learns nothing about the `bins` (permutation/grouping).
+- P1 learns nothing about the `data` values (except the final aggregated sums).
+- The result is revealed to P1 (or shared).
+
+We propose two approaches based on the trade-off between communication and computation, and the cardinality of groups ($K$).
+
+## Interface
+
+```python
+def oblivious_groupby_sum(
+    data: Plaintext[P0],
+    bins: Plaintext[P1],
+    K: int,
+    method: str = "auto"
+) -> Plaintext[P1]:
+    """
+    Args:
+        data: Input data vector held by P0.
+        bins: Bin assignments for each data element held by P1.
+              Values must be in [0, K).
+        K: The number of bins (groups).
+        method: "bfv" (HE-based) or "shuffle" (OT-based).
+
+    Returns:
+        A vector of length K held by P1 containing the sum of data for each bin.
+    """
+```
+
+## Approach 1: HE-based (BFV SIMD)
+
+Best for: **Small K** (e.g., $K < 1000$), Low Bandwidth.
+
+### Algorithm
+
+1.  **Encryption (P0)**:
+    - P0 encrypts `data` using a BFV scheme with SIMD packing.
+    - Sends ciphertext(s) `Enc(data)` to P1.
+
+2.  **Aggregation (P1)**:
+    - P1 holds `bins`. For each bin $k \in [0, K)$:
+        - Construct a plaintext mask vector $M_k$ where $M_k[i] = 1$ if $bins[i] == k$, else $0$.
+        - Compute homomorphic multiplication: $Enc(Sum_k) = Enc(data) \otimes M_k$.
+        - Sum the slots in $Enc(Sum_k)$ to get the total sum for bin $k$.
+          - *Optimization*: Instead of full slot summation for every bin (which is expensive), P1 can just compute the element-wise product and accumulate. The final reduction can be done by sending back to P0 or using rotations if $K$ is small enough to pack into result ciphertexts.
+          - *Simplified Flow*: P1 computes $Enc(Partial_k) = Enc(data) \cdot M_k$. P1 sends these $K$ ciphertexts (or batched versions) back to P0.
+
+3.  **Decryption & Finalize (P0 -> P1)**:
+    - P0 decrypts the partial sums.
+    - P0 computes the sum of the vector for each bin.
+    - P0 sends the final $K$ sums to P1.
+    - *Privacy Note*: To prevent P0 from learning the partial sums (which reveals data distribution), P1 should add a random mask to the result before sending to P0, or use a proper threshold decryption if available. For the "Simplified Flow" above, P0 sees the masked data values. This might leak info.
+    - *Refined Privacy Flow*:
+        - P1 computes $Enc(V_k) = Enc(data) \cdot M_k$.
+        - P1 computes $Enc(S_k) = \text{TotalSum}(Enc(V_k))$ using rotations and additions.
+        - P1 masks $Enc(S_k)$ with a random value $r_k$: $Enc(O_k) = Enc(S_k) + Enc(r_k)$.
+        - P1 sends $Enc(O_k)$ to P0.
+        - P0 decrypts to get $O_k = S_k + r_k$ and sends back to P1.
+        - P1 subtracts $r_k$ to get $S_k$.
+
+### Complexity
+- **Comm**: $O(N/B)$ ciphertexts (P0->P1) + $O(K)$ ciphertexts (P1->P0). ($B$ is batch size).
+- **Comp**: $O(K \cdot N/B)$ homomorphic multiplications and additions.
+
+## Approach 2: OT-based (Shuffle + Prefix Sum)
+
+Best for: **Large K**, High Bandwidth.
+
+### Algorithm
+
+1.  **Sort Permutation (P1)**:
+    - P1 calculates a permutation $\pi$ that sorts `data` according to `bins`.
+    - P1 calculates boundary indices for each bin.
+
+2.  **Oblivious Shuffle (P0, P1)**:
+    - Use a Benes network or similar switching network.
+    - P0 inputs `data`. P1 inputs control bits derived from $\pi$.
+    - Output: Secret shares of permuted data $\langle D' \rangle_0, \langle D' \rangle_1$.
+
+3.  **Secret Shared Prefix Sum (P0, P1)**:
+    - Locally compute prefix sums of shares: $\langle S \rangle_0 = \text{cumsum}(\langle D' \rangle_0)$, $\langle S \rangle_1 = \text{cumsum}(\langle D' \rangle_1)$.
+
+4.  **Oblivious Gather (P0, P1)**:
+    - P1 knows the boundary indices $idx_k$.
+    - P1 needs $S[idx_k] = \langle S \rangle_0[idx_k] + \langle S \rangle_1[idx_k]$.
+    - P1 has $\langle S \rangle_1[idx_k]$ locally.
+    - To get $\langle S \rangle_0[idx_k]$ obliviously:
+        - Use another permutation network or ORAM to fetch these values without revealing $idx_k$ to P0.
+        - Or, since P1 is the result receiver, we can use a simpler selection protocol if we don't hide the access pattern from P0 (but we must hide it to protect bin sizes).
+        - A second shuffle network mapping $idx_k \to k$ is secure.
+
+5.  **Difference (P1)**:
+    - P1 computes $Result[k] = S[idx_k] - S[idx_{k-1}]$.
+
+### Complexity
+- **Comm**: $O(N \log N)$ bits for shuffle.
+- **Comp**: Low (symmetric crypto).