mplang-nightly 0.1.dev142__py3-none-any.whl → 0.1.dev144__py3-none-any.whl

mplang/backend/__init__.py

@@ -11,10 +11,3 @@
 # 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 module for mplang.
-
-This module contains handlers that execute serialized functions on individual
-parties in a multi-party computation system.
-"""

mplang/backend/base.py

@@ -12,12 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Flat backend kernel registry & per-participant runtime.
+"""Backend kernel registry & per-participant runtime (explicit op->kernel binding).
 
-Design revision:
-- Global, stateless kernel function catalog (fn_type -> callable).
-- BackendRuntime: per-rank state & cache; executes kernels.
-- Legacy global helpers removed after full migration to explicit runtimes.
+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
@@ -27,22 +35,17 @@ from collections.abc import Callable
 from dataclasses import dataclass
 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
-
 __all__ = [
-    "BackendRuntime",
     "KernelContext",
-    "create_runtime",
+    "KernelSpec",
+    "bind_op",
     "cur_kctx",
-    "kernel_def",
-    "list_registered_kernels",
+    "get_kernel_for_op",
+    "list_kernels",
+    "list_ops",
+    "unbind_op",
 ]
 
-# ---------------- Context ----------------
-
 
 @dataclass
 class KernelContext:
@@ -99,189 +102,74 @@ def cur_kctx() -> KernelContext:
 
 # ---------------- Registry ----------------
 
-# Canonical kernel callable signature (new style): (pfunc, *args) -> Any | sequence
-# - No **kwargs (explicitly disallowed)
-# - Return normalization handled by BackendRuntime.run_kernel
+# Kernel callable signature: (pfunc, *args) -> Any | sequence (no **kwargs)
 KernelFn = Callable[..., Any]
 
-_KERNELS: dict[str, KernelFn] = {}
 
+@dataclass
+class KernelSpec:
+    kernel_id: str
+    fn: KernelFn
+    meta: dict[str, Any]
 
-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)}"
-        )
 
+# All registered kernel implementations: kernel_id -> spec
+_KERNELS: dict[str, KernelSpec] = {}
+
+# Active op bindings: op_type -> kernel_id
+_BINDINGS: dict[str, str] = {}
 
-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}"
-        )
-
-
-def kernel_def(fn_type: str) -> Callable[[KernelFn], KernelFn]:
-    """Decorator to register a backend kernel (new signature).
-
-    Expected Python signature form:
-
-        @kernel_def("namespace.op")
-        def _op(pfunc: PFunction, *args): ...
-
-    Rules:
-      * First parameter MUST be the PFunction object.
-      * Positional arguments correspond 1:1 to pfunc.ins_info order.
-      * **kwargs are NOT supported (will raise at call site if used).
-      * Return value forms accepted (n = len(pfunc.outs_info)):
-          - n == 0: return None / () / []
-          - n == 1: return scalar/object OR (value,) / [value]
-          - n > 1 : return tuple/list of length n
-        Anything else raises a ValueError.
+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 fn_type in _KERNELS:
-            raise ValueError(f"duplicate backend kernel fn_type={fn_type}")
-        _KERNELS[fn_type] = fn
+        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 list_registered_kernels() -> list[str]:  # public API unchanged
-    return sorted(_KERNELS.keys())
+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
 
-class BackendRuntime:
-    """Per-rank backend execution environment.
 
-    Holds mutable backend state (namespaced pockets) and a cache. Stateless
-    kernel implementations look up their state through cur_kctx().
-    """
+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 __init__(self, rank: int, world_size: int):
-        self.rank = rank
-        self.world_size = world_size
-        self.state: dict[str, dict[str, Any]] = {}
-        self.cache: dict[str, Any] = {}
-
-    # Main entry
-    def run_kernel(self, pfunc: PFunction, arg_list: list[Any]) -> list[Any]:
-        fn_type = pfunc.fn_type
-        fn = _KERNELS.get(fn_type)
-        if fn is None:
-            raise NotImplementedError(f"no backend kernel registered for {fn_type}")
-
-        # Strict positional arg count validation (no kernel-managed arity bypass)
-        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, (spec, val) in enumerate(zip(pfunc.ins_info, arg_list, strict=True)):
-            if isinstance(spec, TableType):
-                _validate_table_arg(fn_type, idx, spec, val)
-                continue
-
-            if isinstance(spec, TensorType):
-                _validate_tensor_arg(fn_type, idx, spec, val)
-                continue
-
-            # Unknown spec type: silently skip validation (legacy behavior)
-            continue
-
-        kctx = KernelContext(
-            rank=self.rank,
-            world_size=self.world_size,
-            state=self.state,
-            cache=self.cache,
-        )
-        token = _CTX_VAR.set(kctx)
-        try:
-            raw = fn(pfunc, *arg_list)
-        finally:
-            _CTX_VAR.reset(token)
-
-        # Normalize return values
-        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 multi-output expected, raw must be sequence of right length
-        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]]
-            # Single object
-            return [raw]
-
-        # expected > 1
-        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)
-
-    # Optional helper
-    def reset(self) -> None:  # pragma: no cover - simple
-        self.state.clear()
-        self.cache.clear()
-
-
-def create_runtime(rank: int, world_size: int) -> BackendRuntime:
-    """Factory for BackendRuntime (allows future policy injection)."""
-    return BackendRuntime(rank, world_size)
+def list_ops() -> list[str]:
+    return sorted(_BINDINGS.keys())

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}"
+        )