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

mplang/v1/kernels/basic.py

@@ -0,0 +1,240 @@
+# 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.
+
+from __future__ import annotations
+
+import numpy as np
+
+from mplang.v1.core import PFunction, TableType, TensorType
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import TableValue, TensorValue, Value
+from mplang.v1.runtime.data_providers import get_provider, resolve_uri
+from mplang.v1.utils import table_utils
+
+
+@kernel_def("basic.identity")
+def _identity(pfunc: PFunction, value: Value) -> Value:
+    # Runtime guarantees exactly one argument; no extra arity checks here.
+    return value
+
+
+@kernel_def("basic.read")
+def _read(pfunc: PFunction) -> Value:
+    path = pfunc.attrs.get("path")
+    if path is None:
+        raise ValueError("missing path attr for basic.read")
+    out_t = pfunc.outs_info[0]
+    uri = resolve_uri(str(path))
+    prov = get_provider(uri.scheme)
+    if prov is None:
+        raise NotImplementedError(f"no resource provider for scheme: {uri.scheme}")
+    ctx = cur_kctx()
+    try:
+        data = prov.read(uri, out_t, ctx=ctx)
+    except Exception as e:  # pragma: no cover - provider errors
+        raise RuntimeError(f"basic.read failed: {e}") from e
+
+    if isinstance(data, Value):
+        return data
+
+    if isinstance(out_t, TableType):
+        return TableValue(data)
+    elif isinstance(out_t, TensorType):
+        return TensorValue(np.asarray(data))
+    else:
+        raise TypeError(
+            f"basic.read only supports TableType/TensorType outputs, got {type(out_t).__name__}"
+        )
+
+
+@kernel_def("basic.write")
+def _write(pfunc: PFunction, obj: Value) -> Value:
+    path = pfunc.attrs.get("path")
+    if path is None:
+        raise ValueError("missing path attr for basic.write")
+    uri = resolve_uri(str(path))
+    prov = get_provider(uri.scheme)
+    if prov is None:
+        raise NotImplementedError(f"no resource provider for scheme: {uri.scheme}")
+    # Pass Value object directly to provider - let provider decide how to handle it
+    ctx = cur_kctx()
+    try:
+        prov.write(uri, obj, ctx=ctx)
+    except Exception as e:  # pragma: no cover
+        raise RuntimeError(f"basic.write failed: {e}") from e
+    return obj
+
+
+@kernel_def("basic.constant")
+def _constant(pfunc: PFunction) -> Value:
+    """Return constants as Value types (TensorValue or TableValue)."""
+    data_bytes = pfunc.attrs.get("data_bytes")
+    if data_bytes is None:
+        raise ValueError("missing data_bytes attr for basic.constant")
+    out_t = pfunc.outs_info[0]
+    fmt = pfunc.attrs.get("data_format")
+    if isinstance(out_t, TableType):
+        if fmt != "bytes[parquet]":
+            raise ValueError(f"unsupported table constant format {fmt}")
+        df = table_utils.decode_table(data_bytes, format="parquet")
+        return TableValue(df)
+    # tensor path
+    shape = out_t.shape  # type: ignore[attr-defined,union-attr]
+    dtype = out_t.dtype.numpy_dtype()  # type: ignore[attr-defined,union-attr]
+    arr = np.frombuffer(data_bytes, dtype=dtype).reshape(shape)
+    return TensorValue(arr)
+
+
+@kernel_def("basic.rank")
+def _rank(pfunc: PFunction) -> TensorValue:
+    """Return rank as TensorValue."""
+    ctx = cur_kctx()
+    arr = np.array(ctx.rank, dtype=np.uint64)
+    return TensorValue(arr)
+
+
+@kernel_def("basic.prand")
+def _prand(pfunc: PFunction) -> TensorValue:
+    """Return random data as TensorValue."""
+    shape = pfunc.attrs.get("shape", ())
+    rng = np.random.default_rng()
+    info = np.iinfo(np.uint64)
+    data = rng.integers(
+        low=info.min, high=info.max, size=shape, dtype=np.uint64, endpoint=True
+    )
+    return TensorValue(data)
+
+
+@kernel_def("basic.table_to_tensor")
+def _table_to_tensor(pfunc: PFunction, table: TableValue) -> TensorValue:
+    """Convert table to tensor, return as TensorValue."""
+    arrow_table = table.to_arrow()
+    if arrow_table.num_columns == 0:
+        raise ValueError("cannot pack empty table")
+    # Convert Arrow columns to numpy arrays and stack
+    mat = np.column_stack([
+        arrow_table.column(i).to_numpy() for i in range(arrow_table.num_columns)
+    ])
+    return TensorValue(mat)
+
+
+@kernel_def("basic.tensor_to_table")
+def _tensor_to_table(pfunc: PFunction, tensor: TensorValue) -> TableValue:
+    """Convert tensor to table, return as TableValue."""
+    import pyarrow as pa  # type: ignore
+
+    arr = tensor.to_numpy()
+    if arr.ndim != 2:
+        raise ValueError("tensor_to_table expects rank-2 array")
+    col_names = pfunc.attrs.get("column_names")
+    if col_names is None:
+        raise ValueError("missing column_names attr")
+    # Create Arrow table directly from numpy array columns
+    arrays = [pa.array(arr[:, i]) for i in range(arr.shape[1])]
+    arrow_table = pa.table(dict(zip(col_names, arrays, strict=True)))
+    return TableValue(arrow_table)
+
+
+def _summ(v: Value) -> str:
+    try:
+        if isinstance(v, TableValue):
+            # Use Arrow's native string representation (more efficient)
+            arrow_table = v.to_arrow()
+            # Show first 8 rows
+            preview = arrow_table.slice(0, min(8, arrow_table.num_rows))
+            return str(preview)
+        if isinstance(v, TensorValue):
+            arr = v.to_numpy()
+            return str(
+                np.array2string(
+                    arr, threshold=64, edgeitems=3, precision=6, suppress_small=True
+                )
+            )
+        return repr(v)
+    except Exception as e:  # pragma: no cover
+        return f"<unprintable {type(v).__name__}: {e}>"
+
+
+@kernel_def("basic.debug_print")
+def _debug_print(pfunc: PFunction, val: Value) -> Value:
+    prefix = pfunc.attrs.get("prefix", "")
+    ctx = cur_kctx()
+    print(f"[debug_print][rank={ctx.rank}] {prefix}{_summ(val)}")
+    return val
+
+
+@kernel_def("basic.pack")
+def _pack(pfunc: PFunction, value: Value) -> TensorValue:
+    outs_info = pfunc.outs_info
+    if len(outs_info) != 1:
+        raise ValueError("basic.pack expects single output type")
+    out_ty = outs_info[0]
+    if not isinstance(out_ty, TensorType):
+        raise TypeError("basic.pack must return TensorType")
+    if out_ty.dtype.numpy_dtype() != np.uint8:
+        raise TypeError("basic.pack output dtype must be uint8")
+
+    if isinstance(value, TableValue):
+        # Serialize Arrow table using IPC stream for consistency with Value serde
+        import pyarrow as pa  # type: ignore
+        import pyarrow.ipc as pa_ipc  # type: ignore
+
+        arrow_table = value.to_arrow()
+        sink = pa.BufferOutputStream()
+        with pa_ipc.new_stream(sink, arrow_table.schema) as writer:  # type: ignore[arg-type]
+            writer.write_table(arrow_table)  # type: ignore[arg-type]
+        ipc_bytes = sink.getvalue().to_pybytes()
+        return TensorValue(np.frombuffer(ipc_bytes, dtype=np.uint8))
+
+    if isinstance(value, TensorValue):
+        arr = value.to_numpy()
+        return TensorValue(np.frombuffer(arr.tobytes(order="C"), dtype=np.uint8))
+
+    raise TypeError(f"basic.pack does not support Value type {type(value).__name__}")
+
+
+@kernel_def("basic.unpack")
+def _unpack(pfunc: PFunction, packed: TensorValue) -> Value:
+    outs_info = pfunc.outs_info
+    if len(outs_info) != 1:
+        raise ValueError("basic.unpack expects single output type")
+    out_ty = outs_info[0]
+
+    b = packed.to_numpy().astype(np.uint8, copy=False).reshape(-1)
+
+    if isinstance(out_ty, TensorType):
+        np_dtype = out_ty.dtype.numpy_dtype()
+        shape = tuple(out_ty.shape)
+        if any(dim < 0 for dim in shape):
+            raise ValueError("basic.unpack does not support dynamic tensor shapes")
+        elem_count = int(np.prod(shape))
+        expected = elem_count * np.dtype(np_dtype).itemsize
+        if b.size != expected:
+            raise ValueError(
+                f"unpack size mismatch: got {b.size} bytes, expect {expected} for {np_dtype} {shape}"
+            )
+        arr = np.frombuffer(b.tobytes(), dtype=np_dtype)
+        return TensorValue(arr.reshape(shape))
+
+    if isinstance(out_ty, TableType):
+        # Deserialize Arrow IPC stream back to TableValue
+        import pyarrow as pa  # type: ignore
+        import pyarrow.ipc as pa_ipc  # type: ignore
+
+        buf = pa.py_buffer(b.tobytes())
+        reader = pa_ipc.open_stream(buf)
+        table = reader.read_all()
+        return TableValue(table)
+
+    raise TypeError("basic.unpack output type must be TensorType or TableType")

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

@@ -17,12 +17,12 @@ from __future__ import annotations
 from collections.abc import Mapping
 from typing import Any
 
-from mplang.core.dtype import UINT8, DType
-from mplang.core.pfunc import PFunction
-from mplang.core.table import TableLike, TableType
-from mplang.core.tensor import TensorLike, TensorType
-from mplang.kernels import base
-from mplang.kernels.base import KernelContext, get_kernel_spec, kernel_exists
+from mplang.v1.core.dtypes import UINT8, DType
+from mplang.v1.core.pfunc import PFunction
+from mplang.v1.core.table import PandasTableLike, TableLike, TableType
+from mplang.v1.core.tensor import TensorLike, TensorType
+from mplang.v1.kernels import base
+from mplang.v1.kernels.base import KernelContext, get_kernel_spec, kernel_exists
 
 # Default bindings
 # Import kernel implementation modules explicitly so their @kernel_def entries
@@ -35,13 +35,14 @@ def _ensure_impl_imported() -> None:
     global _IMPL_IMPORTED
     if _IMPL_IMPORTED:
         return
-    from mplang.kernels import builtin as _impl_builtin  # noqa: F401
-    from mplang.kernels import crypto as _impl_crypto  # noqa: F401
-    from mplang.kernels import mock_tee as _impl_tee  # noqa: F401
-    from mplang.kernels import phe as _impl_phe  # noqa: F401
-    from mplang.kernels import spu as _impl_spu  # noqa: F401
-    from mplang.kernels import sql_duckdb as _impl_sql_duckdb  # noqa: F401
-    from mplang.kernels import stablehlo as _impl_stablehlo  # noqa: F401
+    from mplang.v1.kernels import basic as _impl_basic  # noqa: F401
+    from mplang.v1.kernels import crypto as _impl_crypto  # noqa: F401
+    from mplang.v1.kernels import fhe as _impl_fhe  # noqa: F401
+    from mplang.v1.kernels import mock_tee as _impl_tee  # noqa: F401
+    from mplang.v1.kernels import phe as _impl_phe  # noqa: F401
+    from mplang.v1.kernels import spu as _impl_spu  # noqa: F401
+    from mplang.v1.kernels import sql_duckdb as _impl_sql_duckdb  # noqa: F401
+    from mplang.v1.kernels import stablehlo as _impl_stablehlo  # noqa: F401
 
     _IMPL_IMPORTED = True
 
@@ -49,18 +50,18 @@ def _ensure_impl_imported() -> None:
 # imports consolidated above
 
 _DEFAULT_BINDINGS: dict[str, str] = {
-    # builtin
-    "builtin.identity": "builtin.identity",
-    "builtin.read": "builtin.read",
-    "builtin.write": "builtin.write",
-    "builtin.constant": "builtin.constant",
-    "builtin.rank": "builtin.rank",
-    "builtin.prand": "builtin.prand",
-    "builtin.table_to_tensor": "builtin.table_to_tensor",
-    "builtin.tensor_to_table": "builtin.tensor_to_table",
-    "builtin.debug_print": "builtin.debug_print",
-    "builtin.pack": "builtin.pack",
-    "builtin.unpack": "builtin.unpack",
+    # basic
+    "basic.identity": "basic.identity",
+    "basic.read": "basic.read",
+    "basic.write": "basic.write",
+    "basic.constant": "basic.constant",
+    "basic.rank": "basic.rank",
+    "basic.prand": "basic.prand",
+    "basic.table_to_tensor": "basic.table_to_tensor",
+    "basic.tensor_to_table": "basic.tensor_to_table",
+    "basic.debug_print": "basic.debug_print",
+    "basic.pack": "basic.pack",
+    "basic.unpack": "basic.unpack",
     # crypto
     "crypto.keygen": "crypto.keygen",
     "crypto.enc": "crypto.enc",
@@ -80,6 +81,17 @@ _DEFAULT_BINDINGS: dict[str, str] = {
     "phe.concat": "phe.concat",
     "phe.reshape": "phe.reshape",
     "phe.transpose": "phe.transpose",
+    # fhe
+    "fhe.keygen": "fhe.keygen",
+    "fhe.encrypt": "fhe.encrypt",
+    "fhe.decrypt": "fhe.decrypt",
+    "fhe.add": "fhe.add",
+    "fhe.mul": "fhe.mul",
+    "fhe.dot": "fhe.dot",
+    "fhe.polyval": "fhe.polyval",
+    "fhe.sub": "fhe.sub",
+    "fhe.negate": "fhe.negate",
+    "fhe.square": "fhe.square",
     # spu
     "spu.seed_env": "spu.seed_env",
     "spu.makeshares": "spu.makeshares",
@@ -305,9 +317,12 @@ def _validate_table_arg(
         raise TypeError(
             f"kernel {fn_type} input[{arg_index}] expects TableLike, got {type(value).__name__}"
         )
-    if len(value.columns) != len(spec.columns):
+    columns = (
+        value.columns if isinstance(value, PandasTableLike) else value.column_names
+    )
+    if len(columns) != len(spec.columns):
         raise ValueError(
-            f"kernel {fn_type} input[{arg_index}] column count mismatch: got {len(value.columns)}, expected {len(spec.columns)}"
+            f"kernel {fn_type} input[{arg_index}] column count mismatch: got {len(columns)}, expected {len(spec.columns)}"
         )
 
 

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

@@ -15,15 +15,15 @@
 from __future__ import annotations
 
 import os
-from typing import Any
 
 import numpy as np
 
-from mplang.core.pfunc import PFunction
-from mplang.kernels.base import cur_kctx, kernel_def
-from mplang.utils.crypto import blake2b
+from mplang.v1.core import PFunction
+from mplang.v1.kernels.base import cur_kctx, kernel_def
+from mplang.v1.kernels.value import TensorValue
+from mplang.v1.utils.crypto import blake2b
 
-__all__: list[str] = []  # flat kernels only
+__all__: list[str] = []  # No public exports currently
 
 
 def _get_rng() -> np.random.Generator:
@@ -45,71 +45,78 @@ def _get_rng() -> np.random.Generator:
 def _keystream(key: bytes, nonce: bytes, length: int) -> bytes:
     # WARNING (INSECURE): hash-based keystream (key||nonce||counter)
     out = bytearray()
-    counter = 0
     while len(out) < length:
-        chunk = blake2b(key + nonce + counter.to_bytes(4, "little"))
+        chunk = blake2b(key + nonce)
         out.extend(chunk)
-        counter += 1
     return bytes(out[:length])
 
 
 @kernel_def("crypto.keygen")
-def _crypto_keygen(pfunc: PFunction) -> Any:
+def _crypto_keygen(pfunc: PFunction) -> TensorValue:
     length = int(pfunc.attrs.get("length", 32))
     rng = _get_rng()
     key = rng.integers(0, 256, size=(length,), dtype=np.uint8)
-    return key
+    return TensorValue(key)
 
 
 @kernel_def("crypto.enc")
-def _crypto_encrypt(pfunc: PFunction, pt_bytes: Any, key: Any) -> Any:
-    pt_bytes = np.asarray(pt_bytes, dtype=np.uint8)
-    key = np.asarray(key, dtype=np.uint8)
+def _crypto_encrypt(
+    pfunc: PFunction, pt_bytes: TensorValue, key: TensorValue
+) -> TensorValue:
+    pt_bytes_np = pt_bytes.to_numpy().astype(np.uint8, copy=False)
+    key_np = key.to_numpy().astype(np.uint8, copy=False)
     rng = _get_rng()
-    nonce = rng.integers(0, 256, size=(12,), dtype=np.uint8)
+    nonce = rng.integers(0, 256, size=(16,), dtype=np.uint8)
     stream = np.frombuffer(
-        _keystream(key.tobytes(), nonce.tobytes(), pt_bytes.size), dtype=np.uint8
+        _keystream(key_np.tobytes(), nonce.tobytes(), pt_bytes_np.size), dtype=np.uint8
     )
-    ct = (pt_bytes ^ stream).astype(np.uint8)
+    ct = (pt_bytes_np ^ stream).astype(np.uint8)
     out = np.concatenate([nonce, ct]).astype(np.uint8)
-    return out
+    return TensorValue(out)
 
 
 @kernel_def("crypto.dec")
-def _crypto_decrypt(pfunc: PFunction, ct_with_nonce: Any, key: Any) -> Any:
-    ct_with_nonce = np.asarray(ct_with_nonce, dtype=np.uint8)
-    key = np.asarray(key, dtype=np.uint8)
-    nonce = ct_with_nonce[:12]
-    ct = ct_with_nonce[12:]
+def _crypto_decrypt(
+    pfunc: PFunction, ct_with_nonce: TensorValue, key: TensorValue
+) -> TensorValue:
+    ct_np = ct_with_nonce.to_numpy().astype(np.uint8, copy=False)
+    key_np = key.to_numpy().astype(np.uint8, copy=False)
+    nonce = ct_np[:16]
+    ct = ct_np[16:]
     stream = np.frombuffer(
-        _keystream(key.tobytes(), nonce.tobytes(), len(ct)), dtype=np.uint8
+        _keystream(key_np.tobytes(), nonce.tobytes(), len(ct)), dtype=np.uint8
     )
     pt_bytes = (ct ^ stream).astype(np.uint8)
-    return pt_bytes
+    return TensorValue(pt_bytes)
 
 
 @kernel_def("crypto.kem_keygen")
-def _crypto_kem_keygen(pfunc: PFunction) -> Any:
+def _crypto_kem_keygen(pfunc: PFunction) -> tuple[TensorValue, TensorValue]:
     rng = _get_rng()
     sk = rng.integers(0, 256, size=(32,), dtype=np.uint8)
-    pk = np.frombuffer(blake2b(sk.tobytes())[:32], dtype=np.uint8)
-    return (sk, pk)
+    pk_bytes = blake2b(sk.tobytes())[:32]
+    pk = np.frombuffer(pk_bytes, dtype=np.uint8)
+    return (TensorValue(sk), TensorValue(pk))
 
 
 @kernel_def("crypto.kem_derive")
-def _crypto_kem_derive(pfunc: PFunction, sk: Any, peer_pk: Any) -> Any:
-    sk = np.asarray(sk, dtype=np.uint8)
-    peer_pk = np.asarray(peer_pk, dtype=np.uint8)
-    self_pk = np.frombuffer(blake2b(sk.tobytes())[:32], dtype=np.uint8)
-    xored = (self_pk ^ peer_pk).astype(np.uint8)
+def _crypto_kem_derive(
+    pfunc: PFunction, sk: TensorValue, peer_pk: TensorValue
+) -> TensorValue:
+    sk_np = sk.to_numpy().astype(np.uint8, copy=False)
+    peer_pk_np = peer_pk.to_numpy().astype(np.uint8, copy=False)
+
+    self_pk_bytes = blake2b(sk_np.tobytes())[:32]
+    self_pk_arr = np.frombuffer(self_pk_bytes, dtype=np.uint8)
+    xored = (self_pk_arr ^ peer_pk_np).astype(np.uint8)
     secret = np.frombuffer(blake2b(xored.tobytes())[:32], dtype=np.uint8)
-    return secret
+    return TensorValue(secret)
 
 
 @kernel_def("crypto.hkdf")
-def _crypto_hkdf(pfunc: PFunction, secret: Any) -> Any:
-    secret = np.asarray(secret, dtype=np.uint8)
+def _crypto_hkdf(pfunc: PFunction, secret: TensorValue) -> TensorValue:
+    secret_np = secret.to_numpy().astype(np.uint8, copy=False)
     info_str = str(pfunc.attrs.get("info", ""))
     info = info_str.encode("utf-8")
-    out = np.frombuffer(blake2b(secret.tobytes() + info)[:32], dtype=np.uint8)
-    return out
+    out = np.frombuffer(blake2b(secret_np.tobytes() + info)[:32], dtype=np.uint8)
+    return TensorValue(out)