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

mplang/{kernels → v1/kernels}/mock_tee.py

@@ -20,8 +20,9 @@ import warnings
 import numpy as np
 from numpy.typing import NDArray
 
-from mplang.core.pfunc import PFunction
-from mplang.kernels.base import cur_kctx, kernel_def
+from mplang.v1.core import PFunction
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import TensorValue
 
 __all__: list[str] = []
 
@@ -46,28 +47,26 @@ def _quote_from_pk(pk: np.ndarray) -> NDArray[np.uint8]:
 
 
 @kernel_def("mock_tee.quote_gen")
-def _tee_quote_gen(pfunc: PFunction, pk: object) -> NDArray[np.uint8]:
+def _tee_quote_gen(pfunc: PFunction, pk: TensorValue) -> TensorValue:
     warnings.warn(
         "Insecure mock TEE kernel 'mock_tee.quote_gen' in use. NOT secure; for local testing only.",
         stacklevel=3,
     )
-    pk = np.asarray(pk, dtype=np.uint8)
+    pk_arr = pk.to_numpy().astype(np.uint8, copy=False)
     # rng access ensures deterministic seeding per rank even if unused now
     _rng()
-    return _quote_from_pk(pk)
+    quote = _quote_from_pk(pk_arr)
+    return TensorValue(np.array(quote, copy=True))
 
 
 @kernel_def("mock_tee.attest")
-def _tee_attest(pfunc: PFunction, quote: object) -> NDArray[np.uint8]:
+def _tee_attest(pfunc: PFunction, quote: TensorValue) -> TensorValue:
     warnings.warn(
         "Insecure mock TEE kernel 'mock_tee.attest' in use. NOT secure; for local testing only.",
         stacklevel=3,
     )
-    quote = np.asarray(quote, dtype=np.uint8)
-    platform = pfunc.attrs.get("platform")
-    if platform is None:
-        raise ValueError("missing required 'platform' attribute in PFunction")
-
-    if quote.size != 33:
+    quote_arr = quote.to_numpy().astype(np.uint8, copy=False)
+    if quote_arr.size != 33:
         raise ValueError("mock quote must be 33 bytes (1 header + 32 pk)")
-    return quote[1:33].astype(np.uint8)
+    attest = quote_arr[1:33].astype(np.uint8, copy=True)
+    return TensorValue(attest)

mplang/{kernels → v1/kernels}/phe.py

@@ -14,15 +14,26 @@
 
 """PHE (Partially Homomorphic Encryption) backend implementation using lightPHE."""
 
-from typing import Any
+from __future__ import annotations
+
+import json
+from typing import Any, ClassVar
 
 import numpy as np
 from lightphe import LightPHE
-
-from mplang.core.dtype import DType
-from mplang.core.mptype import TensorLike
-from mplang.core.pfunc import PFunction
-from mplang.kernels.base import kernel_def
+from lightphe.models.Ciphertext import Ciphertext
+
+from mplang.v1.core import DType, PFunction
+from mplang.v1.kernels.base import kernel_def
+from mplang.v1.kernels.value import (
+    TensorValue,
+    Value,
+    ValueDecodeError,
+    ValueProtoBuilder,
+    ValueProtoReader,
+    register_value,
+)
+from mplang.v1.protos.v1alpha1 import value_pb2 as _value_pb2
 
 # This controls the decimal precision used in lightPHE for float operations
 # we force it to 0 to only support integer operations
@@ -30,8 +41,12 @@ from mplang.kernels.base import kernel_def
 PRECISION = 0
 
 
-class PublicKey:
-    """PHE Public Key that implements TensorLike protocol."""
+@register_value
+class PublicKey(Value):
+    """PHE Public Key Value type."""
+
+    KIND: ClassVar[str] = "mplang.phe.PublicKey"
+    WIRE_VERSION: ClassVar[int] = 1
 
     def __init__(
         self,
@@ -62,12 +77,56 @@ class PublicKey:
         """Maximum float value that can be encoded."""
         return float(self.max_value / (2**self.fxp_bits))
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize PublicKey to wire format."""
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("scheme", self.scheme)
+            .set_attr("key_size", self.key_size)
+            .set_attr("max_value", self.max_value)
+            .set_attr("fxp_bits", self.fxp_bits)
+            .set_attr("modulus", str(self.modulus) if self.modulus is not None else "")
+            .set_payload(json.dumps(self.key_data).encode("utf-8"))
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> PublicKey:
+        """Deserialize PublicKey from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported PublicKey version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        scheme = reader.get_attr("scheme")
+        key_size = reader.get_attr("key_size")
+        max_value = reader.get_attr("max_value")
+        fxp_bits = reader.get_attr("fxp_bits")
+        modulus_str = reader.get_attr("modulus")
+        modulus = None if modulus_str == "" else int(modulus_str)
+
+        # JSON deserialize the public key dict
+        key_data = json.loads(reader.payload.decode("utf-8"))
+
+        return cls(
+            key_data=key_data,
+            scheme=scheme,
+            key_size=key_size,
+            max_value=max_value,
+            fxp_bits=fxp_bits,
+            modulus=modulus,
+        )
+
     def __repr__(self) -> str:
         return f"PublicKey(scheme={self.scheme}, key_size={self.key_size}, max_value={self.max_value}, fxp_bits={self.fxp_bits})"
 
 
-class PrivateKey:
-    """PHE Private Key that implements TensorLike protocol."""
+@register_value
+class PrivateKey(Value):
+    """PHE Private Key Value type."""
+
+    KIND: ClassVar[str] = "mplang.phe.PrivateKey"
+    WIRE_VERSION: ClassVar[int] = 1
 
     def __init__(
         self,
@@ -100,12 +159,63 @@ class PrivateKey:
         """Maximum float value that can be encoded."""
         return float(self.max_value / (2**self.fxp_bits))
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize PrivateKey to wire format."""
+        # JSON serialize both key dicts (contain int values)
+        # Store both keys in a single dict to avoid needing length metadata
+        keys_dict = {"sk": self.sk_data, "pk": self.pk_data}
+
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("scheme", self.scheme)
+            .set_attr("key_size", self.key_size)
+            .set_attr("max_value", self.max_value)
+            .set_attr("fxp_bits", self.fxp_bits)
+            .set_attr("modulus", str(self.modulus) if self.modulus is not None else "")
+            .set_payload(json.dumps(keys_dict).encode("utf-8"))
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> PrivateKey:
+        """Deserialize PrivateKey from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported PrivateKey version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        scheme = reader.get_attr("scheme")
+        key_size = reader.get_attr("key_size")
+        max_value = reader.get_attr("max_value")
+        fxp_bits = reader.get_attr("fxp_bits")
+        modulus_str = reader.get_attr("modulus")
+        modulus = None if modulus_str == "" else int(modulus_str)
+
+        # JSON deserialize both key dicts
+        keys_dict = json.loads(reader.payload.decode("utf-8"))
+        sk_data = keys_dict["sk"]
+        pk_data = keys_dict["pk"]
+
+        return cls(
+            sk_data=sk_data,
+            pk_data=pk_data,
+            scheme=scheme,
+            key_size=key_size,
+            max_value=max_value,
+            fxp_bits=fxp_bits,
+            modulus=modulus,
+        )
+
     def __repr__(self) -> str:
         return f"PrivateKey(scheme={self.scheme}, key_size={self.key_size}, max_value={self.max_value}, fxp_bits={self.fxp_bits})"
 
 
-class CipherText:
-    """PHE CipherText that implements TensorLike protocol."""
+@register_value
+class CipherText(Value):
+    """PHE CipherText Value type."""
+
+    KIND: ClassVar[str] = "mplang.phe.CipherText"
+    WIRE_VERSION: ClassVar[int] = 1
 
     def __init__(
         self,
@@ -142,6 +252,106 @@ class CipherText:
         """Maximum float value that can be encoded."""
         return float(self.max_value / (2**self.fxp_bits))
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize CipherText to wire format.
+
+        WARNING: This serialization is tightly coupled to lightphe.Ciphertext
+        internal attributes (value, algorithm_name, keys). Any changes to these
+        attributes in future lightphe versions will break serialization.
+
+        TODO: Check if lightphe provides official serialization methods and
+        migrate to them if available. Consider adding version compatibility checks.
+        """
+        # JSON serialize ciphertext components
+        # ct_data is a list of lightPHE Ciphertext objects
+        # Each Ciphertext has: value, algorithm_name, keys
+        # We need to serialize the list of ciphertexts
+        if not isinstance(self.ct_data, list):
+            raise ValueError(f"ct_data should be a list, got {type(self.ct_data)}")
+
+        ct_list = []
+        for ct in self.ct_data:
+            if not isinstance(ct, Ciphertext):
+                raise TypeError(
+                    f"ct_data must contain lightphe.Ciphertext objects, got {type(ct).__name__}"
+                )
+            ct_list.append({
+                "value": ct.value,
+                "algorithm_name": ct.algorithm_name,
+                "keys": ct.keys,
+            })
+
+        # Combine ct_data and pk_data into single dict
+        payload_dict = {
+            "ct_list": ct_list,
+            "pk": self.pk_data if self.pk_data is not None else None,
+        }
+
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("semantic_dtype", str(self.semantic_dtype))
+            .set_attr("semantic_shape", list(self.semantic_shape))
+            .set_attr("scheme", self.scheme)
+            .set_attr("key_size", self.key_size)
+            .set_attr("max_value", self.max_value)
+            .set_attr("fxp_bits", self.fxp_bits)
+            .set_attr("modulus", str(self.modulus) if self.modulus is not None else "")
+            .set_payload(json.dumps(payload_dict).encode("utf-8"))
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> CipherText:
+        """Deserialize CipherText from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported CipherText version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        semantic_dtype_str = reader.get_attr("semantic_dtype")
+        semantic_shape = reader.get_attr("semantic_shape")
+        scheme = reader.get_attr("scheme")
+        key_size = reader.get_attr("key_size")
+        max_value = reader.get_attr("max_value")
+        fxp_bits = reader.get_attr("fxp_bits")
+        modulus_str = reader.get_attr("modulus")
+        modulus = None if modulus_str == "" else int(modulus_str)
+
+        # JSON deserialize ciphertext and public key
+        payload_dict = json.loads(reader.payload.decode("utf-8"))
+        ct_list = payload_dict["ct_list"]
+        pk_data = payload_dict["pk"]
+
+        # Reconstruct ct_data: list of Ciphertext objects
+        ct_data = []
+        for ct_dict in ct_list:
+            if ct_dict["keys"] is None or ct_dict["algorithm_name"] is None:
+                raise ValueDecodeError(
+                    "Invalid CipherText: missing keys or algorithm_name in serialized data"
+                )
+            ct_data.append(
+                Ciphertext(
+                    algorithm_name=ct_dict["algorithm_name"],
+                    keys=ct_dict["keys"],
+                    value=ct_dict["value"],
+                )
+            )
+
+        # Parse dtype string back to DType
+        dtype = DType.from_any(semantic_dtype_str)
+
+        return cls(
+            ct_data=ct_data,
+            semantic_dtype=dtype,
+            semantic_shape=tuple(semantic_shape),
+            scheme=scheme,
+            key_size=key_size,
+            pk_data=pk_data,
+            max_value=max_value,
+            fxp_bits=fxp_bits,
+            modulus=modulus,
+        )
+
     def __repr__(self) -> str:
         return f"CipherText(dtype={self.semantic_dtype}, shape={self.semantic_shape}, scheme={self.scheme})"
 
@@ -257,33 +467,15 @@ def _range_decode_mixed(
         return _range_decode_integer(encoded_value, max_value, modulus)
 
 
-def _convert_to_numpy(obj: TensorLike) -> np.ndarray:
-    """Convert a TensorLike object to numpy array."""
-    if isinstance(obj, np.ndarray):
-        return obj
-
-    # Try to use .numpy() method if available
-    if hasattr(obj, "numpy"):
-        numpy_method = getattr(obj, "numpy", None)
-        if callable(numpy_method):
-            try:
-                return np.asarray(numpy_method())
-            except Exception:
-                pass
-
-    return np.asarray(obj)
-
-
 @kernel_def("phe.keygen")
 def _phe_keygen(pfunc: PFunction) -> Any:
     scheme = pfunc.attrs.get("scheme", "paillier")
     # use small key_size to speed up tests
     # in production use at least 2048 bits or 3072 bits for better security
     key_size = pfunc.attrs.get("key_size", 2048)
-    max_value = pfunc.attrs.get(
-        "max_value", 2**32
-    )  # Use larger range to avoid overflow
-    fxp_bits = pfunc.attrs.get("fxp_bits", 12)
+    # Accept very large max_value; allow decimal string input, kept simple like other attrs
+    max_value = int(pfunc.attrs.get("max_value", 2**32))
+    fxp_bits = int(pfunc.attrs.get("fxp_bits", 12))
 
     # Validate scheme
     if scheme.lower() not in ["paillier"]:
@@ -334,14 +526,16 @@ def _phe_keygen(pfunc: PFunction) -> Any:
 
 
 @kernel_def("phe.encrypt")
-def _phe_encrypt(pfunc: PFunction, plaintext: Any, public_key: PublicKey) -> Any:
+def _phe_encrypt(
+    pfunc: PFunction, plaintext: TensorValue, public_key: PublicKey
+) -> Any:
     # Validate public_key type
     if not isinstance(public_key, PublicKey):
         raise ValueError("Second argument must be a PublicKey instance")
 
     try:
         # Convert plaintext to numpy to get semantic type info
-        plaintext_np = _convert_to_numpy(plaintext)
+        plaintext_np = plaintext.to_numpy()
         semantic_dtype = DType.from_numpy(plaintext_np.dtype)
         semantic_shape = plaintext_np.shape
 
@@ -403,14 +597,14 @@ def _phe_encrypt(pfunc: PFunction, plaintext: Any, public_key: PublicKey) -> Any
 
 
 @kernel_def("phe.mul")
-def _phe_mul(pfunc: PFunction, ciphertext: CipherText, plaintext: Any) -> Any:
+def _phe_mul(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorValue) -> Any:
     # Validate that first argument is a CipherText
     if not isinstance(ciphertext, CipherText):
         raise ValueError("First argument must be a CipherText instance")
 
     try:
         # Convert plaintext to numpy
-        plaintext_np = _convert_to_numpy(plaintext)
+        plaintext_np = plaintext.to_numpy()
 
         # Check if plaintext is floating point type - multiplication not supported
         if np.issubdtype(plaintext_np.dtype, np.floating):
@@ -443,7 +637,8 @@ def _phe_mul(pfunc: PFunction, ciphertext: CipherText, plaintext: Any) -> Any:
             # Use numpy to create a properly broadcasted index mapping
             # Create a dummy array with same shape as ciphertext, fill with indices
             dummy_ct = (
-                np.arange(np.prod(ciphertext.semantic_shape))
+                np
+                .arange(np.prod(ciphertext.semantic_shape))
                 .reshape(ciphertext.semantic_shape)
                 .astype(np.int64)
             )
@@ -511,7 +706,7 @@ def _phe_add(pfunc: PFunction, lhs: Any, rhs: Any) -> Any:
         elif isinstance(rhs, CipherText):
             return _phe_add_ct2pt(rhs, lhs)
         else:
-            return _convert_to_numpy(lhs) + _convert_to_numpy(rhs)
+            return TensorValue(lhs.to_numpy() + rhs.to_numpy())
     except ValueError:
         raise
     except Exception as e:  # pragma: no cover
@@ -550,7 +745,8 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     # Broadcast ct1 if needed
     if ct1.semantic_shape != result_shape:
         dummy_ct1 = (
-            np.arange(np.prod(ct1.semantic_shape))
+            np
+            .arange(np.prod(ct1.semantic_shape))
             .reshape(ct1.semantic_shape)
             .astype(np.int64)
         )
@@ -563,7 +759,8 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     # Broadcast ct2 if needed
     if ct2.semantic_shape != result_shape:
         dummy_ct2 = (
-            np.arange(np.prod(ct2.semantic_shape))
+            np
+            .arange(np.prod(ct2.semantic_shape))
             .reshape(ct2.semantic_shape)
             .astype(np.int64)
         )
@@ -593,9 +790,9 @@ def _phe_add_ct2ct(ct1: CipherText, ct2: CipherText) -> CipherText:
     )
 
 
-def _phe_add_ct2pt(ciphertext: CipherText, plaintext: TensorLike) -> CipherText:
+def _phe_add_ct2pt(ciphertext: CipherText, plaintext: TensorValue) -> CipherText:
     # Convert plaintext to numpy
-    plaintext_np = _convert_to_numpy(plaintext)
+    plaintext_np = plaintext.to_numpy()
     plaintext_dtype = DType.from_numpy(plaintext_np.dtype)
 
     # Check for mixed precision issue: floating point ciphertext + integer plaintext
@@ -636,7 +833,8 @@ def _phe_add_ct2pt(ciphertext: CipherText, plaintext: TensorLike) -> CipherText:
     # Broadcast ciphertext if needed
     if ciphertext.semantic_shape != result_shape:
         dummy_ct = (
-            np.arange(np.prod(ciphertext.semantic_shape))
+            np
+            .arange(np.prod(ciphertext.semantic_shape))
             .reshape(ciphertext.semantic_shape)
             .astype(np.int64)
         )
@@ -802,12 +1000,17 @@ def _phe_decrypt(
         # Convert to target dtype
         if target_dtype.kind in "iu":  # integer types
             # Convert floats back to integers for integer semantic types
-            processed_data = [round(val) for val in decoded_data]
-            # Handle overflow for smaller integer types
-            info = np.iinfo(target_dtype)
-            processed_data = [
-                max(info.min, min(info.max, val)) for val in processed_data
-            ]
+            # decoded_data are numeric (ints or floats); normalize to Python int
+            ints = [round(v) if isinstance(v, float) else v for v in decoded_data]
+            if np.issubdtype(target_dtype, np.unsignedinteger):
+                # Reduce modulo 2^k for unsigned to preserve ring semantics
+                width = np.iinfo(target_dtype).bits
+                mod = 1 << width
+                processed_data = [v % mod for v in ints]
+            else:
+                # Signed integers: clamp to dtype range
+                info = np.iinfo(target_dtype)
+                processed_data = [max(info.min, min(info.max, v)) for v in ints]
         else:  # float types
             processed_data = decoded_data
 
@@ -816,14 +1019,14 @@ def _phe_decrypt(
             ciphertext.semantic_shape
         )
 
-        return [plaintext_np]
+        return [TensorValue(plaintext_np)]
 
     except Exception as e:
         raise RuntimeError(f"Failed to decrypt data: {e}") from e
 
 
 @kernel_def("phe.dot")
-def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorLike) -> Any:
+def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorValue) -> Any:
     """Execute homomorphic dot product with zero-value optimization.
 
     Supports various dot product operations:
@@ -844,7 +1047,7 @@ def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorLike) ->
 
     try:
         # Convert plaintext to numpy
-        plaintext_np = _convert_to_numpy(plaintext)
+        plaintext_np = plaintext.to_numpy()
 
         # Check if plaintext is floating point type - dot product not supported
         if np.issubdtype(plaintext_np.dtype, np.floating):
@@ -1109,7 +1312,7 @@ def _phe_dot(pfunc: PFunction, ciphertext: CipherText, plaintext: TensorLike) ->
 
 
 @kernel_def("phe.gather")
-def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: Any) -> Any:
+def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: TensorValue) -> Any:
     """Execute gather operation on CipherText.
 
     Supports gathering from multidimensional CipherText using multidimensional indices.
@@ -1126,7 +1329,7 @@ def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: Any) -> Any:
 
     try:
         # Convert indices to numpy
-        indices_np = _convert_to_numpy(indices)
+        indices_np = indices.to_numpy()
 
         if not np.issubdtype(indices_np.dtype, np.integer):
             raise ValueError("Indices must be of integer type")
@@ -1224,7 +1427,10 @@ def _phe_gather(pfunc: PFunction, ciphertext: CipherText, indices: Any) -> Any:
 
 @kernel_def("phe.scatter")
 def _phe_scatter(
-    pfunc: PFunction, ciphertext: CipherText, indices: TensorLike, updated: CipherText
+    pfunc: PFunction,
+    ciphertext: CipherText,
+    indices: TensorValue,
+    updated: CipherText,
 ) -> Any:
     """Execute scatter operation on CipherText.
 
@@ -1252,7 +1458,7 @@ def _phe_scatter(
 
     try:
         # Convert indices to numpy
-        indices_np = _convert_to_numpy(indices)
+        indices_np = indices.to_numpy()
 
         if not np.issubdtype(indices_np.dtype, np.integer):
             raise ValueError("Indices must be of integer type")