mplang-nightly 0.1.dev147__py3-none-any.whl → 0.1.dev149__py3-none-any.whl

mplang/backend/base.py

@@ -12,20 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Backend kernel registry & per-participant runtime (explicit op->kernel binding).
+"""Backend kernel registry: mapping kernel_id -> implementation.
 
-This version decouples *kernel implementation registration* from *operation binding*.
+This module provides a lightweight registry for backend kernel implementations.
+It does not track or decide which kernel handles a given semantic operation;
+that policy (op -> kernel_id) is managed externally by each ``RuntimeContext``.
 
-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
+Exposed primitives:
+* ``@kernel_def(kernel_id)``: decorator to register a kernel implementation.
+* ``get_kernel_spec(kernel_id)``: look up a previously registered kernel.
+* ``cur_kctx()`` / ``KernelContext``: execution context available only
+    inside a kernel body (rank, world_size, per-backend state pockets, and a
+    runtime-wide cache shared by kernels of the same runtime instance).
 
-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.
+No global op binding table exists here; callers resolve an op to a kernel_id
+before invoking the kernel function.
 """
 
 from __future__ import annotations
@@ -38,12 +39,10 @@ from typing import Any
 __all__ = [
     "KernelContext",
     "KernelSpec",
-    "bind_op",
     "cur_kctx",
-    "get_kernel_for_op",
+    "get_kernel_spec",
+    "kernel_exists",
     "list_kernels",
-    "list_ops",
-    "unbind_op",
 ]
 
 
@@ -116,9 +115,6 @@ class KernelSpec:
 # 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.
@@ -136,34 +132,11 @@ def kernel_def(kernel_id: str, /, **meta: Any) -> Callable[[KernelFn], KernelFn]
     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:
+def get_kernel_spec(kernel_id: str) -> KernelSpec:
+    """Return KernelSpec for a registered kernel_id (no op binding lookup)."""
+    spec = _KERNELS.get(kernel_id)
+    if spec is None:
         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
 
 
@@ -171,5 +144,6 @@ def list_kernels() -> list[str]:
     return sorted(_KERNELS.keys())
 
 
-def list_ops() -> list[str]:
-    return sorted(_BINDINGS.keys())
+def kernel_exists(kernel_id: str) -> bool:
+    """Return True if a kernel_id has been registered."""
+    return kernel_id in _KERNELS

mplang/backend/context.py

@@ -15,11 +15,10 @@
 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.backend.base import KernelContext, get_kernel_spec, kernel_exists
 from mplang.core.dtype import UINT8, DType
 from mplang.core.pfunc import PFunction
 from mplang.core.table import TableLike, TableType
@@ -100,30 +99,57 @@ _DEFAULT_BINDINGS: dict[str, str] = {
 # --- 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:
+    """Per-runtime execution context with isolated op->kernel bindings.
+
+    Parameters
+    ----------
+    rank : int
+        Local rank of this participant.
+    world_size : int
+        Total number of participants.
+    initial_bindings : Mapping[str, str] | None, optional
+        Optional partial overrides applied on top of the default binding table
+        during construction (override semantics, not replace). After
+        initialization, all (re)binding must go through ``bind_op`` /
+        ``rebind_op``.
+    state / cache / stats : dict, optional
+        Mutable pockets reused across kernel invocations. If omitted, new
+        dictionaries are created.
+    """
+
+    __slots__ = ("_ibindings", "cache", "rank", "state", "stats", "world_size")
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        initial_bindings: Mapping[str, str] | None = None,
+        *,
+        state: dict[str, dict[str, Any]] | None = None,
+        cache: dict[str, Any] | None = None,
+        stats: dict[str, Any] | None = None,
+    ) -> 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.rank = rank
+        self.world_size = world_size
+        # Merge defaults with user overrides (override semantics)
+        self._ibindings: dict[str, str] = {
+            **_DEFAULT_BINDINGS,
+            **(initial_bindings or {}),
+        }
+        self.state = state if state is not None else {}
+        self.cache = cache if cache is not None else {}
+        self.stats = stats if stats is not None else {}
         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
+        kid = self._ibindings.get(fn_type)
+        if kid is None:
+            raise NotImplementedError(f"no backend kernel registered for op {fn_type}")
+        spec = get_kernel_spec(kid)
+        fn = spec.fn  # kernel implementation
         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)}"
@@ -186,17 +212,32 @@ class RuntimeContext:
 
     # ---- 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.
+        """Bind an operation to a kernel for THIS context only.
 
-        force=False (default) preserves any existing binding to avoid accidental
-        silent overrides. Use ``rebind_op`` or ``force=True`` to intentionally
-        change a binding.
+        force=False (default) keeps existing binding (no silent override).
         """
-        base.bind_op(op_type, kernel_id, force=force)
+        if not kernel_exists(kernel_id):
+            raise KeyError(f"kernel_id {kernel_id} not registered")
+        if not force and op_type in self._ibindings:
+            return
+        self._ibindings[op_type] = kernel_id
 
     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)
+        self.bind_op(op_type, kernel_id, force=True)
+
+    # Introspection helpers
+    def list_bound_ops(self) -> list[str]:  # pragma: no cover - convenience
+        return sorted(self._ibindings.keys())
+
+    def get_binding(self, op_type: str) -> str | None:  # pragma: no cover
+        return self._ibindings.get(op_type)
+
+    def __repr__(self) -> str:  # pragma: no cover - debug aid
+        return (
+            f"RuntimeContext(rank={self.rank}, world_size={self.world_size}, "
+            f"bound_ops={len(self._ibindings)})"
+        )
 
 
 def _validate_table_arg(

mplang/backend/stablehlo.py

@@ -51,8 +51,15 @@ def _stablehlo_exec(pfunc: PFunction, *args: Any) -> Any:
             raise RuntimeError(f"StableHLO compile failed: {e}") from e
         cache[mlir_text] = compiled
 
+    # Handle JAX's unused parameter elimination via arg_keep_map
+    runtime_args = args
+    if "arg_keep_map" in pfunc.attrs:
+        keep_indices = pfunc.attrs["arg_keep_map"]
+        # 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 args:
+    for arg in runtime_args:
         if hasattr(arg, "numpy"):
             jax_arg = jnp.array(arg.numpy())  # type: ignore
         else:

mplang/frontend/jax_cc.py

@@ -106,14 +106,46 @@ def jax2stablehlo(
     out_info_flat, out_tree = tree_flatten(lowered.out_info)
     out_info_flat = [TensorType.from_obj(info) for info in out_info_flat]
 
+    # Extract argument keep mapping to handle JAX's unused parameter elimination
+    # JAX can eliminate unused parameters during compilation, but the runtime still
+    # receives all original arguments. We need the mapping to filter them correctly.
+    arg_keep_map = None
+    original_arg_count = len(in_vars)
+
+    try:
+        # Access JAX internal kept_var_idx - the authoritative source
+        # This tells us exactly which original parameters survived compilation
+        compile_args = lowered._lowering.compile_args
+        kept_var_idx = compile_args["kept_var_idx"]
+
+        kept_indices = sorted(kept_var_idx)
+        if len(kept_indices) < original_arg_count:
+            arg_keep_map = kept_indices
+
+    except (AttributeError, KeyError, TypeError) as e:
+        # JAX internal API is not available or changed
+        # This is a hard error - we cannot reliably handle unused parameters
+        # without knowing exactly which ones were kept
+        raise RuntimeError(
+            f"Cannot access JAX's kept_var_idx to handle unused parameter elimination. "
+            f"This function may have unused parameters that JAX optimized away, "
+            f"but we cannot determine which ones without the internal API. "
+            f"Original error: {e}"
+        ) from e
+
     # This format tells JaxRT how to handle the compiled result
-    pfn = PFunction(
-        fn_type="mlir.stablehlo",  # Key: specify StableHLO MLIR format
-        ins_info=tuple(TensorType.from_obj(x) for x in in_vars),
-        outs_info=tuple(out_info_flat),
-        fn_name=get_fn_name(flat_fn),
-        fn_text=mlir_text,  # MLIR text, serializable for transmission
-    )
+    pfn_kwargs: dict[str, Any] = {
+        "fn_type": "mlir.stablehlo",  # Key: specify StableHLO MLIR format
+        "ins_info": tuple(TensorType.from_obj(x) for x in in_vars),
+        "outs_info": tuple(out_info_flat),
+        "fn_name": get_fn_name(flat_fn),
+        "fn_text": mlir_text,  # MLIR text, serializable for transmission
+    }
+
+    if arg_keep_map is not None:
+        pfn_kwargs["arg_keep_map"] = arg_keep_map
+
+    pfn = PFunction(**pfn_kwargs)
     return pfn, in_vars, out_tree
 
 

{mplang_nightly-0.1.dev147.dist-info → mplang_nightly-0.1.dev149.dist-info}/METADATA

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

{mplang_nightly-0.1.dev147.dist-info → mplang_nightly-0.1.dev149.dist-info}/RECORD

@@ -4,14 +4,14 @@ mplang/device.py,sha256=Iz_YFKkrbTFKtTQdGqkQZfc0NQH9dIxXP7-fUkIQOa4,12568
 mplang/analysis/__init__.py,sha256=CTHFvRsi-nFngojqjn08UaR3RY9i7CJ7T2UdR95kCrk,1056
 mplang/analysis/diagram.py,sha256=ffwgD12gL1_KH1uJ_EYkjmIlDrfxYJJkWj-wHl09_Xk,19520
 mplang/backend/__init__.py,sha256=2WE4cmW96Xkzyq2yRRYNww4kZ5o6u6NbPV0BxqZG698,581
-mplang/backend/base.py,sha256=Rp4Ze1aF5dypdMOskXnQWE-rM5nc4vkdDogEoxQp7FU,5712
+mplang/backend/base.py,sha256=eizxj16sWkUvBvXWS0Zl-S9uuqalJmNUzB1xLhBg8S8,4920
 mplang/backend/builtin.py,sha256=Mk1uUO2Vpw3meqZ0B7B0hG-wndea6cmFv2Uk1vM_uTg,7052
-mplang/backend/context.py,sha256=e22JzsHydi1DLkjzOEYwGBaAoVHyf8aEFYH_SuN7oR4,9351
+mplang/backend/context.py,sha256=fVJ0w0cw15JEqJO048dncWg7DGNWqbHSUjq42Jsyvos,10952
 mplang/backend/crypto.py,sha256=H_s5HI7lUP7g0xz-a9qMbSn6dhJStUilKbn3-7SIh0I,3812
 mplang/backend/phe.py,sha256=uNqmrbDAbd97TWS_O6D5sopastHy6J20R7knFE4M4uc,65247
 mplang/backend/spu.py,sha256=QT1q5uv-5P_nBGtTvtA_yI2h3h3zIqNSnvzGT7Shua4,9307
 mplang/backend/sql_duckdb.py,sha256=U_KzEUinxrBRDoUz2Vh597-N4I3hPOBT0RT3tX-ZqKE,1502
-mplang/backend/stablehlo.py,sha256=GOxy-qgOxyEdtBkt6LASKzaZnPewZhvHYSPOgFFXgIM,2612
+mplang/backend/stablehlo.py,sha256=RhKf6TUvjLrRvgtdVY2HxcRDGtjpKBobuBFMfsvZQOI,2937
 mplang/backend/tee.py,sha256=6kc7qTe8nWc3pr6iYtozEGLO8Umg-UBQLDiz6p3pdVg,1918
 mplang/core/__init__.py,sha256=lWxlEKfRwX7FNDzgyKZ1fiDMaCiqkyg0j5mKlZD_v7g,2244
 mplang/core/cluster.py,sha256=gqMJenvXUfHhE181Dd5JiUkD4nT07RLoicBnvsGmRkE,8598
@@ -41,7 +41,7 @@ mplang/frontend/base.py,sha256=rGtfBejcDh9mTRxOdJK5VUlG5vYiVJSir8X72X0Huvc,18264
 mplang/frontend/builtin.py,sha256=8qrlbe_SSy6QTXTnMG6_ADB8jSklVZGFBrkoR-p02FE,9368
 mplang/frontend/crypto.py,sha256=Nf8zT4Eko7MIs4R2tgZecKVd7d6Hvd_CGGmANhs3Ghs,3651
 mplang/frontend/ibis_cc.py,sha256=CTTbPPZ9hFnHuFDDIfgJHie1EdNnHmi5Ha1KsX0iYh8,4235
-mplang/frontend/jax_cc.py,sha256=ssP6rCvyWQ5VAr80-7z9QZUE2mWXyozJCGpq1dYQYY8,6374
+mplang/frontend/jax_cc.py,sha256=lMqaYD1tyM5DsStTNYifAXzhzsNM5nDiG3a61ygbWyc,7807
 mplang/frontend/phe.py,sha256=tDsCvStjVJ1Fs07yF3idkFnugUCA1zdFApPx7Uuulik,6795
 mplang/frontend/spu.py,sha256=7G6DaEfC5APSDhfeWSISTG_8tEcVbWth3XmjL8QUrVA,4994
 mplang/frontend/sql.py,sha256=DFdvjEPQX28VCRgUMeHYR0rwwOaoCH15bpvvlclLtHA,1999
@@ -70,8 +70,8 @@ mplang/utils/crypto.py,sha256=rvPomBFtznRHc3RPi6Aip9lsU8zW2oxBqGv1K3vn7Rs,1052
 mplang/utils/func_utils.py,sha256=vCJcZmu0bEbqhOQKdpttV2_MBllIcPSN0b8U4WjNGGo,5164
 mplang/utils/spu_utils.py,sha256=S3L9RBkBe2AvSuMSQQ12cBY5Y1NPthubvErSX_7nj1A,4158
 mplang/utils/table_utils.py,sha256=aC-IZOKkSmFkpr3NZchLM0Wt0GOn-rg_xHBHREWBwAU,2202
-mplang_nightly-0.1.dev147.dist-info/METADATA,sha256=WxU1KBbckH3JcYQF1jnDD8ph-EE9wKhbompdIg0KTP8,16547
-mplang_nightly-0.1.dev147.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev147.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev147.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev147.dist-info/RECORD,,
+mplang_nightly-0.1.dev149.dist-info/METADATA,sha256=eZ_qGx1500gbGY5Ms8smHjQsyZ2Wxfa11z2mSeF0CSk,16547
+mplang_nightly-0.1.dev149.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev149.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev149.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev149.dist-info/RECORD,,