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

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

@@ -15,16 +15,38 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import spu.api as spu_api
 import spu.libspu as libspu
 
-from mplang.core.mptype import TensorLike
-from mplang.core.pfunc import PFunction
-from mplang.kernels.base import cur_kctx, kernel_def
-from mplang.runtime.link_comm import LinkCommunicator
+from mplang.v1.core import (
+    BOOL,
+    FLOAT32,
+    FLOAT64,
+    INT8,
+    INT16,
+    INT32,
+    INT64,
+    UINT8,
+    UINT16,
+    UINT32,
+    UINT64,
+    DType,
+    PFunction,
+)
+from mplang.v1.kernels.base import cur_kctx, 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
+from mplang.v1.runtime.link_comm import LinkCommunicator
 
 
 def shape_spu_to_np(spu_shape: Any) -> tuple[int, ...]:
@@ -32,36 +54,106 @@ def shape_spu_to_np(spu_shape: Any) -> tuple[int, ...]:
     return tuple(spu_shape.dims)
 
 
-def dtype_spu_to_np(spu_dtype: Any) -> np.dtype:
-    """Convert SPU dtype to numpy dtype."""
+def dtype_spu_to_mpl(spu_dtype: libspu.DataType) -> DType:
+    """Convert libspu.DataType to MPLang DType."""
     MAP = {
-        libspu.DataType.DT_F32: np.float32,
-        libspu.DataType.DT_F64: np.float64,
-        libspu.DataType.DT_I1: np.bool_,
-        libspu.DataType.DT_I8: np.int8,
-        libspu.DataType.DT_U8: np.uint8,
-        libspu.DataType.DT_I16: np.int16,
-        libspu.DataType.DT_U16: np.uint16,
-        libspu.DataType.DT_I32: np.int32,
-        libspu.DataType.DT_U32: np.uint32,
-        libspu.DataType.DT_I64: np.int64,
-        libspu.DataType.DT_U64: np.uint64,
+        libspu.DataType.DT_F32: FLOAT32,
+        libspu.DataType.DT_F64: FLOAT64,
+        libspu.DataType.DT_I1: BOOL,
+        libspu.DataType.DT_I8: INT8,
+        libspu.DataType.DT_U8: UINT8,
+        libspu.DataType.DT_I16: INT16,
+        libspu.DataType.DT_U16: UINT16,
+        libspu.DataType.DT_I32: INT32,
+        libspu.DataType.DT_U32: UINT32,
+        libspu.DataType.DT_I64: INT64,
+        libspu.DataType.DT_U64: UINT64,
     }
-    return MAP[spu_dtype]  # type: ignore[return-value]
+    return MAP[spu_dtype]
 
 
+@register_value
 @dataclass
-class SpuValue:
-    """SPU value container for secure computation."""
+class SpuValue(Value):
+    """SPU value container for secure computation (Value type)."""
+
+    KIND: ClassVar[str] = "mplang.spu.SpuValue"
+    WIRE_VERSION: ClassVar[int] = 1
 
     shape: tuple[int, ...]
-    dtype: Any
+    dtype: DType  # Now uses MPLang's unified DType
     vtype: libspu.Visibility
     share: libspu.Share
 
     def __repr__(self) -> str:
         return f"SpuValue({self.shape},{self.dtype},{self.vtype})"
 
+    def to_proto(self) -> _value_pb2.ValueProto:
+        """Serialize SpuValue to wire format.
+
+        libspu.Share has two attributes:
+        - meta: bytes (protobuf serialized metadata)
+        - share_chunks: list[bytes] (the actual secret share data)
+
+        Strategy: Store shape/dtype/vtype in runtime_attrs, concatenate share.meta + all chunks in payload.
+        """
+        # Store metadata in runtime_attrs; keep chunk lengths for payload splitting
+        chunk_lengths = [len(chunk) for chunk in self.share.share_chunks]
+
+        # Payload contains only share chunks (meta stored in attrs)
+        payload = b""
+        for chunk in self.share.share_chunks:
+            payload += chunk
+
+        return (
+            ValueProtoBuilder(self.KIND, self.WIRE_VERSION)
+            .set_attr("shape", list(self.shape))
+            .set_attr("dtype", self.dtype.name)  # Serialize DType name
+            .set_attr("vtype", int(self.vtype))
+            .set_attr("share_meta", self.share.meta)
+            .set_attr("chunk_lengths", chunk_lengths)
+            .set_payload(payload)
+            .build()
+        )
+
+    @classmethod
+    def from_proto(cls, proto: _value_pb2.ValueProto) -> SpuValue:
+        """Deserialize SpuValue from wire format."""
+        reader = ValueProtoReader(proto)
+        if reader.version != cls.WIRE_VERSION:
+            raise ValueDecodeError(f"Unsupported SpuValue version {reader.version}")
+
+        # Read metadata from runtime_attrs
+        shape = tuple(reader.get_attr("shape"))
+        dtype_name = reader.get_attr("dtype")
+        # Reconstruct DType from serialized name (numpy dtype string)
+        dtype = DType.from_numpy(dtype_name)
+        vtype = libspu.Visibility(reader.get_attr("vtype"))
+        share_meta = reader.get_attr("share_meta")
+        chunk_lengths = reader.get_attr("chunk_lengths")
+
+        # Parse payload: [chunk_0][chunk_1]...
+        payload = reader.payload
+        offset = 0
+
+        share_chunks: list[bytes] = []
+        for chunk_len in chunk_lengths:
+            chunk = payload[offset : offset + chunk_len]
+            offset += chunk_len
+            share_chunks.append(chunk)
+
+        # Reconstruct libspu.Share
+        share = libspu.Share()
+        share.meta = share_meta
+        share.share_chunks = share_chunks
+
+        return cls(
+            shape=shape,
+            dtype=dtype,
+            vtype=vtype,
+            share=share,
+        )
+
 
 def _get_spu_config_and_world() -> tuple[libspu.RuntimeConfig, int]:
     kctx = cur_kctx()
@@ -128,33 +220,25 @@ def _spu_seed_env(pfunc: PFunction, *args: Any) -> Any:
 
 
 @kernel_def("spu.makeshares")
-def _spu_makeshares(pfunc: PFunction, *args: Any) -> Any:
-    """Create SPU shares from input data.
-
-    Args:
-        pfunc: PFunction containing makeshares metadata
-        args: Input data to be shared (single tensor)
-
-    Returns:
-        Tuple of SPU shares (SpuValue), one for each party.
-    """
-    assert len(args) == 1
-
+def _spu_makeshares(pfunc: PFunction, tensor: TensorValue) -> tuple[SpuValue, ...]:
+    """Create SPU shares from input TensorValue data."""
     visibility_value = pfunc.attrs.get("visibility", libspu.Visibility.VIS_SECRET.value)
     if isinstance(visibility_value, int):
         visibility = libspu.Visibility(visibility_value)
     else:
         visibility = visibility_value
 
-    arg = np.array(args[0], copy=False)
+    arg = tensor.to_numpy()
     cfg, world = _get_spu_config_and_world()
     spu_io = spu_api.Io(world, cfg)
     shares = spu_io.make_shares(arg, visibility)
     assert len(shares) == world, f"Expected {world} shares, got {len(shares)}"
+    # Store MPLang DType instead of libspu.DataType
+    dtype = DType.from_numpy(arg.dtype)
     return tuple(
         SpuValue(
             shape=arg.shape,
-            dtype=arg.dtype,
+            dtype=dtype,
             vtype=visibility,
             share=share,
         )
@@ -163,24 +247,29 @@ def _spu_makeshares(pfunc: PFunction, *args: Any) -> Any:
 
 
 @kernel_def("spu.reconstruct")
-def _spu_reconstruct(pfunc: PFunction, *args: Any) -> Any:
+def _spu_reconstruct(pfunc: PFunction, *shares: SpuValue) -> TensorValue:
     """Reconstruct plaintext data from SPU shares."""
     cfg, world = _get_spu_config_and_world()
-    assert len(args) == world, f"Expected {world} shares, got {len(args)}"
-    for i, arg in enumerate(args):
-        if not isinstance(arg, SpuValue):
+    assert len(shares) == world, f"Expected {world} shares, got {len(shares)}"
+    for i, share in enumerate(shares):
+        if not isinstance(share, SpuValue):
             raise ValueError(
-                f"Input {i} must be SpuValue, got {type(arg)}. Reconstruction requires SPU shares as input."
+                f"Input {i} must be SpuValue, got {type(share)}. Reconstruction requires SPU shares as input."
             )
-    spu_args: list[SpuValue] = list(args)  # type: ignore
-    shares = [spu_arg.share for spu_arg in spu_args]
+    spu_args: list[SpuValue] = list(shares)  # type: ignore
+    share_payloads = [spu_arg.share for spu_arg in spu_args]
     spu_io = spu_api.Io(world, cfg)
-    reconstructed = spu_io.reconstruct(shares)
-    return reconstructed
+    reconstructed = spu_io.reconstruct(share_payloads)
+    base = np.array(reconstructed, copy=False)
+    # Respect semantic dtype/shape recorded on shares (all shares share same meta).
+    semantic_dtype = shares[0].dtype.to_numpy()  # DType now has to_numpy() method
+    semantic_shape = shares[0].shape
+    restored = np.asarray(base, dtype=semantic_dtype).reshape(semantic_shape)
+    return TensorValue(np.array(restored, copy=False))
 
 
 @kernel_def("spu.run_pphlo")
-def _spu_run_mlir(pfunc: PFunction, *args: Any) -> Any:
+def _spu_run_mlir(pfunc: PFunction, *args: SpuValue) -> tuple[SpuValue, ...]:
     """Execute compiled SPU function (spu.run_pphlo) and return SpuValue outputs.
 
     Participation rule: a rank participates iff its entry in the stored
@@ -240,10 +329,10 @@ def _spu_run_mlir(pfunc: PFunction, *args: Any) -> Any:
     spu_rt.run(executable)
     shares = [spu_rt.get_var(out_name) for out_name in output_names]
     metas = [spu_rt.get_var_meta(out_name) for out_name in output_names]
-    results: list[TensorLike] = [
+    results: list[SpuValue] = [
         SpuValue(
             shape=shape_spu_to_np(meta.shape),
-            dtype=dtype_spu_to_np(meta.data_type),
+            dtype=dtype_spu_to_mpl(meta.data_type),
             vtype=meta.visibility,
             share=shares[idx],
         )

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

@@ -14,16 +14,14 @@
 
 from __future__ import annotations
 
-from typing import Any
-
-from mplang.core.pfunc import PFunction
-from mplang.kernels.base import kernel_def
+from mplang.v1.core import PFunction
+from mplang.v1.kernels.base import kernel_def
+from mplang.v1.kernels.value import TableValue
 
 
 @kernel_def("duckdb.run_sql")
-def _duckdb_sql(pfunc: PFunction, *args: Any) -> Any:
+def _duckdb_sql(pfunc: PFunction, *args: TableValue) -> TableValue:
     import duckdb
-    import pandas as pd
 
     # TODO: maybe we could translate the sql to duckdb dialect
     # instead of raising an exception
@@ -36,12 +34,11 @@ def _duckdb_sql(pfunc: PFunction, *args: Any) -> Any:
         if in_names is None:
             raise ValueError("duckdb sql missing in_names attr")
         for arg, name in zip(args, in_names, strict=True):
-            if isinstance(arg, pd.DataFrame):
-                df = arg
-            elif isinstance(arg, list):  # const list-of-dict for tests
-                df = pd.DataFrame.from_records(arg)
-            else:
-                raise ValueError(f"unsupported duckdb input type {type(arg)}")
-            conn.register(name, df)
-    res_df = conn.execute(pfunc.fn_text).fetchdf()
-    return res_df
+            # Use Arrow directly for zero-copy data transfer
+            arrow_table = arg.to_arrow()
+            conn.register(name, arrow_table)
+    # Fetch result as Arrow table for consistency
+    if pfunc.fn_text is None:
+        raise ValueError("SQL function text is None")
+    res_arrow = conn.execute(pfunc.fn_text).fetch_arrow_table()
+    return TableValue(res_arrow)

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

@@ -17,12 +17,14 @@ from __future__ import annotations
 from typing import Any
 
 import jax
+import jax.extend as jxt
 import jax.numpy as jnp
-from jax._src import xla_bridge
-from jax.lib import xla_client as xc
+import numpy as np
+from jax._src import compiler
 
-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
 
 
 @kernel_def("mlir.stablehlo")
@@ -45,11 +47,13 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
     key = f"stablehlo.compile_cache.{h}"
     compiled = rt.get_state(key)
     if compiled is None:
-        backend = jax.default_backend()
-        client = xla_bridge.get_backend(backend)
-        compile_options = xc.CompileOptions()
+        client = jxt.backend.get_backend()
+        compile_options = compiler.get_compile_options(num_replicas=1, num_partitions=1)
+
         try:
-            compiled = client.compile(mlir_text, compile_options)
+            compiled = client.compile_and_load(
+                mlir_text, client.devices(), compile_options
+            )
         except Exception as e:  # pragma: no cover
             raise RuntimeError(f"StableHLO compile failed: {e}") from e
         rt.set_state(key, compiled)
@@ -61,23 +65,26 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
         # Filter out arguments that were eliminated by JAX during compilation
         runtime_args = tuple(args[i] for i in keep_indices)
 
-    jax_args = []
-    for arg in runtime_args:
-        if hasattr(arg, "numpy"):
-            jax_arg = jnp.array(arg.numpy())  # type: ignore
-        else:
-            jax_arg = jnp.array(arg)
-        jax_args.append(jax.device_put(jax_arg))
+    tensor_args: list[TensorValue] = []
+    for idx, arg in enumerate(runtime_args):
+        if not isinstance(arg, TensorValue):
+            raise TypeError(
+                f"StableHLO kernel expects TensorValue inputs, got {type(arg).__name__} at position {idx}"
+            )
+        tensor_args.append(arg)
+
+    jax_args = [
+        jax.device_put(jnp.asarray(tensor.to_numpy())) for tensor in tensor_args
+    ]
 
     try:
-        result = compiled.execute_sharded(jax_args)
-        arrays = result.disassemble_into_single_device_arrays()
-        flat: list[Any] = []
-        for lst in arrays:
-            if isinstance(lst, list) and len(lst) == 1:
-                flat.append(jnp.array(lst[0]))
-            else:
-                flat.extend([jnp.array(a) for a in lst])
+        # Execute with the new LoadedExecutable interface
+        result = compiled.execute(jax_args)
+
+        # Use jax.tree_util.tree_flatten to robustly handle any PyTree structure
+        flat_results, _ = jax.tree_util.tree_flatten(result)
+        flat = [TensorValue(np.asarray(item)) for item in flat_results]
+
         return tuple(flat)
     except Exception as e:  # pragma: no cover
         raise RuntimeError(f"StableHLO execute failed: {e}") from e