mplang-nightly 0.1.dev143__tar.gz → 0.1.dev145__tar.gz

{mplang_nightly-0.1.dev143 → mplang_nightly-0.1.dev145}/PKG-INFO

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

mplang_nightly-0.1.dev145/mplang/backend/base.py

@@ -0,0 +1,175 @@
+# 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.
+
+"""Backend kernel registry & per-participant runtime (explicit op->kernel binding).
+
+This version decouples *kernel implementation registration* from *operation binding*.
+
+Concepts:
+    * kernel_id: unique identifier of a concrete backend implementation.
+    * op_type: semantic operation name carried by ``PFunction.fn_type``.
+    * bind_op(op_type, kernel_id): performed by higher layer (see ``backend.context``)
+        to select which implementation handles an op. Runtime dispatch is now a 2-step:
+        pfunc.fn_type -> active kernel_id -> KernelSpec.fn
+
+The previous implicit "import == register+bind" coupling is removed. Kernel
+modules only call ``@kernel_def(kernel_id)``. Default bindings are established
+centrally (lazy) the first time a runtime executes a kernel.
+"""
+
+from __future__ import annotations
+
+import contextvars
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+__all__ = [
+    "KernelContext",
+    "KernelSpec",
+    "bind_op",
+    "cur_kctx",
+    "get_kernel_for_op",
+    "list_kernels",
+    "list_ops",
+    "unbind_op",
+]
+
+
+@dataclass
+class KernelContext:
+    """Ephemeral call context set via contextvar while a kernel runs."""
+
+    rank: int
+    world_size: int
+    state: dict[str, dict[str, Any]]  # backend namespace -> pocket
+    cache: dict[str, Any]  # runtime-level shared cache (per BackendRuntime)
+
+
+_CTX_VAR: contextvars.ContextVar[KernelContext | None] = contextvars.ContextVar(
+    "_flat_backend_ctx", default=None
+)
+
+
+def cur_kctx() -> KernelContext:
+    """Return the current kernel execution context (only valid inside a kernel).
+
+    Two storages:
+      - state: namespaced pockets (dict[str, dict]) for backend-local mutable helpers
+      - cache: global (per runtime) shared dict; prefer state unless truly cross-backend
+
+    Examples:
+      1) Compile cache::
+            @kernel_def("mlir.stablehlo")
+            def _exec(pfunc, args):
+                ctx = cur_kctx()
+                pocket = ctx.state.setdefault("stablehlo", {})
+                cache = pocket.setdefault("compile_cache", {})
+                text = pfunc.fn_text
+                mod = cache.get(text)
+                if mod is None:
+                    mod = compile_mlir(text)
+                    cache[text] = mod
+                return run(mod, args)
+
+      2) Deterministic RNG::
+            @kernel_def("crypto.keygen")
+            def _keygen(pfunc, args):
+                ctx = cur_kctx()
+                pocket = ctx.state.setdefault("crypto", {})
+                rng = pocket.get("rng")
+                if rng is None:
+                    rng = np.random.default_rng(1234 + ctx.rank * 7919)
+                    pocket["rng"] = rng
+                return (rng.integers(0, 256, size=(32,), dtype=np.uint8),)
+    """
+    ctx = _CTX_VAR.get()
+    if ctx is None:
+        raise RuntimeError("cur_kctx() called outside backend kernel execution")
+    return ctx
+
+
+# ---------------- Registry ----------------
+
+# Kernel callable signature: (pfunc, *args) -> Any | sequence (no **kwargs)
+KernelFn = Callable[..., Any]
+
+
+@dataclass
+class KernelSpec:
+    kernel_id: str
+    fn: KernelFn
+    meta: dict[str, Any]
+
+
+# All registered kernel implementations: kernel_id -> spec
+_KERNELS: dict[str, KernelSpec] = {}
+
+# Active op bindings: op_type -> kernel_id
+_BINDINGS: dict[str, str] = {}
+
+
+def kernel_def(kernel_id: str, /, **meta: Any) -> Callable[[KernelFn], KernelFn]:
+    """Decorator to register a concrete kernel implementation.
+
+    This ONLY registers the implementation (kernel_id -> fn). It does NOT bind
+    any op. Higher layer must call ``bind_op(op_type, kernel_id)`` explicitly.
+    """
+
+    def _decorator(fn: KernelFn) -> KernelFn:
+        if kernel_id in _KERNELS:
+            raise ValueError(f"duplicate kernel_id={kernel_id}")
+        _KERNELS[kernel_id] = KernelSpec(kernel_id=kernel_id, fn=fn, meta=dict(meta))
+        return fn
+
+    return _decorator
+
+
+def bind_op(op_type: str, kernel_id: str, *, force: bool = True) -> None:
+    """Bind an op_type to a registered kernel implementation.
+
+    Args:
+        op_type: Semantic operation name.
+        kernel_id: Previously registered kernel identifier.
+        force: If False and op_type already bound, keep existing binding.
+               If True (default), overwrite.
+    """
+    if kernel_id not in _KERNELS:
+        raise KeyError(f"kernel_id {kernel_id} not registered")
+    if not force and op_type in _BINDINGS:
+        return
+    _BINDINGS[op_type] = kernel_id
+
+
+def unbind_op(op_type: str) -> None:
+    _BINDINGS.pop(op_type, None)
+
+
+def get_kernel_for_op(op_type: str) -> KernelSpec:
+    kid = _BINDINGS.get(op_type)
+    if kid is None:
+        # Tests expect NotImplementedError for unsupported operations
+        raise NotImplementedError(f"no backend kernel registered for op {op_type}")
+    spec = _KERNELS.get(kid)
+    if spec is None:  # inconsistent state
+        raise RuntimeError(f"active kernel_id {kid} missing spec")
+    return spec
+
+
+def list_kernels() -> list[str]:
+    return sorted(_KERNELS.keys())
+
+
+def list_ops() -> list[str]:
+    return sorted(_BINDINGS.keys())

mplang_nightly-0.1.dev145/mplang/backend/context.py

@@ -0,0 +1,255 @@
+# 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
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any
+
+from mplang.backend import base
+from mplang.backend.base import KernelContext, bind_op, get_kernel_for_op
+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
+
+# Default bindings
+# Import kernel implementation modules explicitly so their @kernel_def entries
+# register at import time. Keep imports grouped; alias with leading underscore
+# to silence unused variable warnings without F401 pragmas.
+_IMPL_IMPORTED = False
+
+
+def _ensure_impl_imported() -> None:
+    global _IMPL_IMPORTED
+    if _IMPL_IMPORTED:
+        return
+    from mplang.backend import builtin as _impl_builtin  # noqa: F401
+    from mplang.backend import crypto as _impl_crypto  # noqa: F401
+    from mplang.backend import phe as _impl_phe  # noqa: F401
+    from mplang.backend import spu as _impl_spu  # noqa: F401
+    from mplang.backend import sql_duckdb as _impl_sql_duckdb  # noqa: F401
+    from mplang.backend import stablehlo as _impl_stablehlo  # noqa: F401
+    from mplang.backend import tee as _impl_tee  # noqa: F401
+
+    _IMPL_IMPORTED = True
+
+
+# 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",
+    # crypto
+    "crypto.keygen": "crypto.keygen",
+    "crypto.enc": "crypto.enc",
+    "crypto.dec": "crypto.dec",
+    "crypto.kem_keygen": "crypto.kem_keygen",
+    "crypto.kem_derive": "crypto.kem_derive",
+    "crypto.hkdf": "crypto.hkdf",
+    # phe
+    "phe.keygen": "phe.keygen",
+    "phe.encrypt": "phe.encrypt",
+    "phe.mul": "phe.mul",
+    "phe.add": "phe.add",
+    "phe.decrypt": "phe.decrypt",
+    "phe.dot": "phe.dot",
+    "phe.gather": "phe.gather",
+    "phe.scatter": "phe.scatter",
+    "phe.concat": "phe.concat",
+    "phe.reshape": "phe.reshape",
+    "phe.transpose": "phe.transpose",
+    # spu
+    "spu.seed_env": "spu.seed_env",
+    "spu.makeshares": "spu.makeshares",
+    "spu.reconstruct": "spu.reconstruct",
+    "spu.run_pphlo": "spu.run_pphlo",
+    # stablehlo
+    "mlir.stablehlo": "mlir.stablehlo",
+    # sql
+    # generic SQL op; backend-specific kernel id for duckdb
+    "sql.run": "duckdb.run_sql",
+    # tee
+    "tee.quote": "tee.quote",
+    "tee.attest": "tee.attest",
+}
+
+
+# --- RuntimeContext ---
+
+
+@dataclass
+class RuntimeContext:
+    rank: int
+    world_size: int
+    bindings: Mapping[str, str] | None = None  # optional overrides
+    state: dict[str, dict[str, Any]] = field(default_factory=dict)
+    cache: dict[str, Any] = field(default_factory=dict)
+    stats: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        _ensure_impl_imported()
+        if self.bindings is not None:
+            for op, kid in self.bindings.items():
+                bind_op(op, kid)
+        else:
+            for op, kid in _DEFAULT_BINDINGS.items():
+                bind_op(op, kid)
+        # Initialize stats pocket
+        self.stats.setdefault("op_calls", {})
+
+    def run_kernel(self, pfunc: PFunction, arg_list: list[Any]) -> list[Any]:
+        fn_type = pfunc.fn_type
+        spec = get_kernel_for_op(fn_type)
+        fn = spec.fn
+        if len(arg_list) != len(pfunc.ins_info):
+            raise ValueError(
+                f"kernel {fn_type} arg count mismatch: got {len(arg_list)}, expect {len(pfunc.ins_info)}"
+            )
+        for idx, (ins_spec, val) in enumerate(
+            zip(pfunc.ins_info, arg_list, strict=True)
+        ):
+            if isinstance(ins_spec, TableType):
+                _validate_table_arg(fn_type, idx, ins_spec, val)
+                continue
+            if isinstance(ins_spec, TensorType):
+                _validate_tensor_arg(fn_type, idx, ins_spec, val)
+                continue
+        # install kernel context
+        kctx = KernelContext(
+            rank=self.rank,
+            world_size=self.world_size,
+            state=self.state,
+            cache=self.cache,
+        )
+        token = base._CTX_VAR.set(kctx)  # type: ignore[attr-defined]
+        try:
+            raw = fn(pfunc, *arg_list)
+        finally:
+            base._CTX_VAR.reset(token)  # type: ignore[attr-defined]
+        # Stats (best effort)
+        try:
+            op_calls = self.stats.setdefault("op_calls", {})
+            op_calls[fn_type] = op_calls.get(fn_type, 0) + 1
+        except Exception:  # pragma: no cover - never raise due to stats
+            pass
+        expected = len(pfunc.outs_info)
+        if expected == 0:
+            if raw in (None, (), []):
+                return []
+            raise ValueError(
+                f"kernel {fn_type} should return no values; got {type(raw).__name__}"
+            )
+        if expected == 1:
+            if isinstance(raw, (tuple, list)):
+                if len(raw) != 1:
+                    raise ValueError(
+                        f"kernel {fn_type} produced {len(raw)} outputs, expected 1"
+                    )
+                return [raw[0]]
+            return [raw]
+        if not isinstance(raw, (tuple, list)):
+            raise TypeError(
+                f"kernel {fn_type} must return sequence (len={expected}), got {type(raw).__name__}"
+            )
+        if len(raw) != expected:
+            raise ValueError(
+                f"kernel {fn_type} produced {len(raw)} outputs, expected {expected}"
+            )
+        return list(raw)
+
+    def reset(self) -> None:
+        self.state.clear()
+        self.cache.clear()
+
+    # ---- explicit (re)binding API ----
+    def bind_op(self, op_type: str, kernel_id: str, *, force: bool = False) -> None:
+        """Bind an operation to a kernel at runtime.
+
+        force=False (default) preserves any existing binding to avoid accidental
+        silent overrides. Use ``rebind_op`` or ``force=True`` to intentionally
+        change a binding.
+        """
+        base.bind_op(op_type, kernel_id, force=force)
+
+    def rebind_op(self, op_type: str, kernel_id: str) -> None:
+        """Force rebind an operation to a different kernel (shorthand)."""
+        base.bind_op(op_type, kernel_id, force=True)
+
+
+def _validate_table_arg(
+    fn_type: str, arg_index: int, spec: TableType, value: Any
+) -> None:
+    if not isinstance(value, TableLike):
+        raise TypeError(
+            f"kernel {fn_type} input[{arg_index}] expects TableLike, got {type(value).__name__}"
+        )
+    if len(value.columns) != len(spec.columns):
+        raise ValueError(
+            f"kernel {fn_type} input[{arg_index}] column count mismatch: got {len(value.columns)}, expected {len(spec.columns)}"
+        )
+
+
+def _validate_tensor_arg(
+    fn_type: str, arg_index: int, spec: TensorType, value: Any
+) -> None:
+    # Backend-only handle sentinel (e.g., PHE keys) bypasses all structural checks
+    if tuple(spec.shape) == (-1, 0) and spec.dtype == UINT8:
+        return
+
+    if isinstance(value, (int, float, bool, complex)):
+        val_shape: tuple[Any, ...] = ()
+        duck_dtype: Any = type(value)
+    else:
+        if not isinstance(value, TensorLike):
+            raise TypeError(
+                f"kernel {fn_type} input[{arg_index}] expects TensorLike, got {type(value).__name__}"
+            )
+        val_shape = getattr(value, "shape", ())
+        duck_dtype = getattr(value, "dtype", None)
+
+    if len(spec.shape) != len(val_shape):
+        raise ValueError(
+            f"kernel {fn_type} input[{arg_index}] rank mismatch: got {val_shape}, expected {spec.shape}"
+        )
+
+    for dim_idx, (spec_dim, val_dim) in enumerate(
+        zip(spec.shape, val_shape, strict=True)
+    ):
+        if spec_dim >= 0 and spec_dim != val_dim:
+            raise ValueError(
+                f"kernel {fn_type} input[{arg_index}] shape mismatch at dim {dim_idx}: got {val_dim}, expected {spec_dim}"
+            )
+
+    try:
+        val_dtype = DType.from_any(duck_dtype)
+    except (ValueError, TypeError):  # pragma: no cover
+        raise TypeError(
+            f"kernel {fn_type} input[{arg_index}] has unsupported dtype object {duck_dtype!r}"
+        ) from None
+    if val_dtype != spec.dtype:
+        raise ValueError(
+            f"kernel {fn_type} input[{arg_index}] dtype mismatch: got {val_dtype}, expected {spec.dtype}"
+        )

{mplang_nightly-0.1.dev143 → mplang_nightly-0.1.dev145}/mplang/backend/spu.py

@@ -186,16 +186,18 @@ def _spu_reconstruct(pfunc: PFunction, *args: Any) -> Any:
     return reconstructed
 
 
-@kernel_def("mlir.pphlo")
+@kernel_def("spu.run_pphlo")
 def _spu_run_mlir(pfunc: PFunction, *args: Any) -> Any:
-    """Execute compiled SPU function (mlir.pphlo) and return SpuValue outputs.
+    """Execute compiled SPU function (spu.run_pphlo) and return SpuValue outputs.
 
     Participation rule: a rank participates iff its entry in the stored
     link_ctx list is non-None. This allows us to allocate a world-sized list
     (indexed by global rank) and simply assign None for non-SPU parties.
     """
-    if pfunc.fn_type != "mlir.pphlo":
-        raise ValueError(f"Unsupported format: {pfunc.fn_type}. Expected 'mlir.pphlo'")
+    if pfunc.fn_type != "spu.run_pphlo":
+        raise ValueError(
+            f"Unsupported format: {pfunc.fn_type}. Expected 'spu.run_pphlo'"
+        )
 
     cfg, _ = _get_spu_config_and_world()
     pocket = _get_spu_pocket()

{mplang_nightly-0.1.dev143 → mplang_nightly-0.1.dev145}/mplang/backend/sql_duckdb.py

@@ -20,7 +20,7 @@ from mplang.backend.base import kernel_def
 from mplang.core.pfunc import PFunction
 
 
-@kernel_def("sql[duckdb]")
+@kernel_def("duckdb.run_sql")
 def _duckdb_sql(pfunc: PFunction, *args: Any) -> Any:
     import duckdb
     import pandas as pd

{mplang_nightly-0.1.dev143 → mplang_nightly-0.1.dev145}/mplang/core/expr/evaluator.py

@@ -27,7 +27,7 @@ from __future__ import annotations
 from dataclasses import dataclass
 from typing import Any, Protocol
 
-from mplang.backend.base import BackendRuntime
+from mplang.backend.context import RuntimeContext
 from mplang.core.comm import ICommunicator
 from mplang.core.expr.ast import (
     AccessExpr,
@@ -56,7 +56,7 @@ class IEvaluator(Protocol):
     backend state via evaluator.runtime.run_kernel(...).
     """
 
-    runtime: BackendRuntime
+    runtime: RuntimeContext
 
     def evaluate(self, root: Expr, env: dict[str, Any] | None = None) -> list[Any]: ...
 
@@ -72,7 +72,7 @@ class EvalSemantic:
     rank: int
     env: dict[str, Any]
     comm: ICommunicator
-    runtime: BackendRuntime
+    runtime: RuntimeContext
 
     # ------------------------------ Shared helpers (semantics) ------------------------------
     def _should_run(self, rmask: Mask | None, args: list[Any]) -> bool:
@@ -205,7 +205,7 @@ class RecursiveEvaluator(EvalSemantic, ExprVisitor):
         rank: int,
         env: dict[str, Any],
         comm: ICommunicator,
-        runtime: BackendRuntime,
+        runtime: RuntimeContext,
     ) -> None:
         super().__init__(rank, env, comm, runtime)
         self._cache: dict[int, Any] = {}  # Cache based on expr id
@@ -380,7 +380,7 @@ class IterativeEvaluator(EvalSemantic):
         rank: int,
         env: dict[str, Any],
         comm: ICommunicator,
-        runtime: BackendRuntime,
+        runtime: RuntimeContext,
     ) -> None:
         super().__init__(rank, env, comm, runtime)
 
@@ -501,7 +501,7 @@ def create_evaluator(
     rank: int,
     env: dict[str, Any],
     comm: ICommunicator,
-    runtime: BackendRuntime,
+    runtime: RuntimeContext,
     kind: str | None = "iterative",
 ) -> IEvaluator:
     """Factory to create an evaluator engine.

{mplang_nightly-0.1.dev143 → mplang_nightly-0.1.dev145}/mplang/frontend/base.py

@@ -129,7 +129,7 @@ class FeModule(ABC):
         - You need compilation/stateful behavior/dynamic routing, multiple PFunctions, or complex capture flows.
 
     Tips:
-    - Keep routing information in PFunction.fn_type (e.g., "builtin.read", "sql[duckdb]", "mlir.stablehlo").
+    - Keep routing information in PFunction.fn_type (e.g., "builtin.read", "sql.run", "mlir.stablehlo").
     - Avoid backend-specific logic in kernels; only validate and shape types.
     - Prefer keyword-only attributes in typed_op kernels for clarity (def op(x: MPObject, *, attr: int)).
     """

{mplang_nightly-0.1.dev143 → mplang_nightly-0.1.dev145}/mplang/frontend/ibis_cc.py

@@ -57,8 +57,9 @@ def ibis2sql(
     outs_info = [_convert(expr.schema())]
 
     sql = ibis.to_sql(expr, dialect="duckdb")
+    # Emit generic sql.run op; runtime maps to backend-specific kernel.
     pfn = PFunction(
-        fn_type="sql[duckdb]",
+        fn_type="sql.run",
         fn_name=fn_name,
         fn_text=sql,
         ins_info=tuple(ins_info),

{mplang_nightly-0.1.dev143 → mplang_nightly-0.1.dev145}/mplang/frontend/spu.py

@@ -94,9 +94,10 @@ def _compile_jax(
     *args: Any,
     **kwargs: Any,
 ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
-    """Compile a JAX function into SPU pphlo MLIR representation."""
+    """Compile a JAX function into SPU pphlo MLIR and wrap as PFunction.
 
-    """Compile a JAX function into SPU pphlo MLIR representation."""
+    Resulting PFunction uses fn_type 'spu.run_pphlo'.
+    """
 
     def is_variable(arg: Any) -> bool:
         return isinstance(arg, MPObject)
@@ -132,7 +133,7 @@ def _compile_jax(
     executable_code = executable_code.decode("utf-8")
 
     pfunc = PFunction(
-        fn_type="mlir.pphlo",
+        fn_type="spu.run_pphlo",
         ins_info=tuple(TensorType.from_obj(x) for x in in_vars),
         outs_info=tuple(output_tensor_infos),
         fn_name=get_fn_name(fn),