mplang-nightly 0.1.dev269__py3-none-any.whl → 0.1.dev270__py3-none-any.whl

mplang/v1/ops/base.py

@@ -1,424 +0,0 @@
-# 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 abc import ABC, abstractmethod
-from collections.abc import Callable
-from typing import Any
-
-from jax.tree_util import PyTreeDef, tree_flatten
-
-from mplang.v1.core import MPContext, MPObject, PFunction, TableType, TensorType
-
-# -----------------------------------------------------------------------------
-# Triad ABI
-# The standard return contract for frontend operations (FeOperation.trace).
-#
-# Triad := (PFunction, list[MPObject], PyTreeDef)
-# - PFunction: Captures fn_type (routing key, e.g., "mlir.stablehlo", "sql.run"),
-#              input/output MPTypes and optional attributes.
-# - list[MPObject]: The flat positional MPObjects captured under the current
-#                   context (Trace/Interp). Order matches pfunc.ins_info.
-# - PyTreeDef: The output pytree structure to unflatten results after execution.
-#
-# Error modes:
-# - Type errors if non-MPObject positional args provided to simple ops.
-# - Kernel/type builder must produce TensorType/TableType leaves for outs.
-# - Context errors propagate from cur_ctx() usage if called outside capture.
-# -----------------------------------------------------------------------------
-Triad = tuple[PFunction, list[MPObject], PyTreeDef]
-
-
-# -----------------------------------------------------------------------------
-# Lightweight fe module/feop system (new FeOperation based)
-# -----------------------------------------------------------------------------
-
-
-# Global registry for frontend modules and operations
-class FeRegistry:
-    """Registry for FeModules and FeOperations.
-
-    Maintains:
-    - modules: name -> FeModule
-    - ops: (module, name) -> FeOperation (callable) returning Triad
-    """
-
-    __slots__ = ("_modules", "_ops")
-
-    def __init__(self) -> None:
-        # Typed registries
-        self._modules: dict[str, FeModule] = {}
-        self._ops: dict[tuple[str, str], FeOperation] = {}
-
-    # ----------------------------- Modules -----------------------------
-    def register_module(self, mod: FeModule, *, replace: bool = False) -> None:
-        if not replace and mod.name in self._modules:
-            raise ValueError(f"Module already registered: {mod.name}")
-        self._modules[mod.name] = mod
-
-    def get_module(self, name: str) -> FeModule:
-        if name not in self._modules:
-            raise KeyError(f"Unknown module: {name}")
-        return self._modules[name]
-
-    def has_module(self, name: str) -> bool:
-        return name in self._modules
-
-    def list_modules(self) -> dict[str, FeModule]:
-        return dict(self._modules)
-
-    # ------------------------------ Ops -------------------------------
-    def register_op(
-        self, module: str, name: str, op: FeOperation, *, replace: bool = False
-    ) -> None:
-        key = (module, name)
-        if not replace and key in self._ops:
-            raise ValueError(f"Op already registered: {module}.{name}")
-        self._ops[key] = op
-
-    def get_op(self, module: str, name: str) -> FeOperation:
-        key = (module, name)
-        if key not in self._ops:
-            raise KeyError(f"Unknown op: {module}.{name}")
-        return self._ops[key]
-
-    def list_ops(self, module: str | None = None) -> dict[tuple[str, str], FeOperation]:
-        if module is None:
-            return dict(self._ops)
-        return {k: v for k, v in self._ops.items() if k[0] == module}
-
-
-_REGISTRY = FeRegistry()
-
-
-def get_registry() -> FeRegistry:
-    return _REGISTRY
-
-
-def is_feop(x: Any) -> bool:
-    """Return True if x is a frontend operation instance."""
-    return isinstance(x, FeOperation)
-
-
-class FeModule(ABC):
-    """Frontend module with feop/typed_op decorators.
-
-    When to use which:
-    - Use typed_op (SimpleFeOperation) when:
-        - You know the backend routing key up front via pfunc_name, and the kernel is pure type logic.
-        - Inputs are MPObjects (positional/kwargs). Attributes are simple Python values (int/float/str/bytes/tuples/lists of primitives) passed as keywords.
-        - The kernel returns TensorType/TableType (or a PyTree thereof); no IR construction inside.
-    - Use feop (InlineFeOperation) when:
-        - You already build and return the Triad explicitly, or need custom packing/attrs/multi-output composition.
-    - Subclass FeOperation when:
-        - You need compilation/stateful behavior/dynamic routing, multiple PFunctions, or complex capture flows.
-
-    Tips:
-    - Keep routing information in PFunction.fn_type (e.g., "basic.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)).
-    """
-
-    def __init__(self, name: str):
-        self.name = name
-        get_registry().register_module(self)
-
-    @abstractmethod
-    def initialize(self, ctx: MPContext) -> None: ...
-
-    def op_def(self) -> Callable[[Callable[..., Triad]], FeOperation]:
-        """Decorator for inline/complex ops which already return a Triad.
-
-        Usage:
-            @mymod.feop()
-            def scale(x: MPObject, factor: int) -> Triad:
-                # build PFunction and return triad directly
-                ...
-                return pfunc, [x], out_tree
-        """
-
-        def _decorator(trace_fn: Callable[..., Triad]) -> FeOperation:
-            name = trace_fn.__name__
-            op = InlineFeOperation(self, name, trace_fn)
-            get_registry().register_op(self.name, name, op)
-            return op
-
-        return _decorator
-
-    def simple_op(
-        self, pfunc_name: str | None = None
-    ) -> Callable[[Callable[..., Any]], FeOperation]:
-        """Decorator for type-driven ops that return only types/schemas.
-
-        The decorated kernel should compute and return a TensorType/TableType (or PyTree thereof).
-        Positional inputs may be MPObjects (captured as inputs) or data-like values (TableLike/TensorLike)
-        used for type inference/validation. Keyword arguments are PFunction attributes and must be plain
-        Python values (int/float/str/bytes/tuples/lists of primitives). Passing MPObjects via kwargs is not allowed.
-
-        SSOT naming: The operation name is derived from the kernel function name (kernel.__name__),
-        ensuring there's a single source of truth and improving readability. Use clear, concise
-        function names to define the public op names.
-
-        Example:
-            @mymod.typed_op(pfunc_name="builtin.add")
-            def add_kernel(x: MPObject, y: MPObject) -> TensorType:
-                return x.mptype._type  # same shape/type as x
-
-        Bad vs Good (signatures and calls):
-        - Bad:  def op(x: MPObject, **kwargs): ...               # disallowed: **kwargs
-          Good: def op(x: MPObject, *, attr: int): ...
-
-        - Bad:  def op(*args, **kwargs): ...                     # disallowed: *args/**kwargs
-          Good: def op(x: MPObject, y: MPObject, *, k: str): ...
-
-        - Bad:  enc(plaintext=pt, key=mp_key)                    # MPObject via kwargs (disallowed)
-          Good: enc(pt, mp_key)                                  # pass MPObjects positionally
-
-        - Good: hkdf(secret, "info")                             # data-like positional mapped to kw-only attr
-          Also good: hkdf(secret, info="info")
-
-        - Good: phe.mul(jnp.array(...), jnp.array(...))          # data-like positionals allowed for type inference
-        """
-
-        def _decorator(kernel: Callable[..., Any]) -> FeOperation:
-            # Default PFunction routing when not provided: "<module>.<kernel_name>"
-            final_pfunc_name = pfunc_name or f"{self.name}.{kernel.__name__}"
-            op = SimpleFeOperation(self, final_pfunc_name, kernel)
-            # Use kernel function name as SSOT for op name
-            get_registry().register_op(self.name, op.name, op)
-            return op
-
-        return _decorator
-
-
-class StatelessFeModule(FeModule):
-    """Stateless frontend module with no ctx-level state."""
-
-    def initialize(self, ctx: MPContext) -> None:
-        pass
-
-
-# -----------------------------------------------------------------------------
-# Class-based contracts and adapters
-# -----------------------------------------------------------------------------
-
-
-class FeOperation(ABC):
-    """Class-based frontend operation contract.
-
-    Subclasses implement trace() to produce a standard triad. __call__ delegates to trace().
-    """
-
-    module: FeModule
-    name: str
-
-    def __init__(self, module: FeModule, name: str):
-        self.module = module
-        self.name = name
-
-    @abstractmethod
-    def trace(self, *args: Any, **kwargs: Any) -> Triad:
-        """Produce a standard triad for this operation."""
-
-    # Convenience: allow calling an FeOperation like a function.
-    def __call__(self, *args: Any, **kwargs: Any) -> Triad:
-        return self.trace(*args, **kwargs)
-
-
-class InlineFeOperation(FeOperation):
-    """FeOperation that delegates tracing to a provided triad-returning function."""
-
-    def __init__(self, module: FeModule, name: str, trace_fn: Callable[..., Triad]):
-        super().__init__(module, name)
-        self._trace_fn = trace_fn
-
-    # override
-    def trace(self, *args: Any, **kwargs: Any) -> Triad:
-        return self._trace_fn(*args, **kwargs)
-
-
-class SimpleFeOperation(FeOperation):
-    """FeOperation that builds Triad from a type-only kernel.
-
-    Contract (keep it simple):
-    - Kernel computes and returns TensorType/TableType or a PyTree thereof.
-    - Positional inputs may be MPObjects (captured as inputs) or data-like values (TableLike/TensorLike)
-        used for type inference/validation. Keyword arguments are attributes and must be plain Python
-        values (TensorType/TableType are also excluded from attrs). MPObject kwargs are disallowed.
-    - Prefer keyword-only attributes in the kernel signature for explicitness. For convenience, non-MPObject
-        positional values that are not data-like will be mapped to keyword-only parameters by order when possible.
-    - No IR building inside the kernel; PFunction is assembled here with fn_type=pfunc_name.
-    """
-
-    def __init__(
-        self,
-        module: FeModule,
-        pfunc_name: str,
-        kernel: Callable[..., Any],
-    ):
-        # Derive operation name from kernel function name for SSOT
-        super().__init__(module, kernel.__name__)
-        self.pfunc_name = pfunc_name
-        self._kernel = kernel
-
-        # Validate kernel signature: typed_op kernels must not use *args/**kwargs.
-        import inspect
-
-        sig = inspect.signature(kernel)
-        for p in sig.parameters.values():
-            if p.kind == inspect.Parameter.VAR_KEYWORD:
-                raise TypeError(
-                    f"typed_op kernel '{module.name}.{kernel.__name__}' must not use **kwargs; define explicit keywords instead"
-                )
-            if p.kind == inspect.Parameter.VAR_POSITIONAL:
-                raise TypeError(
-                    f"typed_op kernel '{module.name}.{kernel.__name__}' must not use *args; define explicit parameters instead"
-                )
-
-        # Cache signature and kw-only parameter names for fast path in trace
-        self._kernel_sig = sig
-        self._kwonly_names = [
-            p.name
-            for p in sig.parameters.values()
-            if p.kind == inspect.Parameter.KEYWORD_ONLY
-        ]
-
-    # override
-    def trace(self, *args: MPObject, **kwargs: Any) -> Triad:
-        # Actual params may not match kernel signature exactly, so we do flexible binding.
-        sig = self._kernel_sig
-
-        # Inputs at PFunction layer are MPObjects captured from positional args only.
-        pos_mp_inputs: list[MPObject] = [a for a in args if isinstance(a, MPObject)]
-
-        # Enforce: no MPObject kwargs per simplified contract
-        for k, v in kwargs.items():
-            if isinstance(v, MPObject):
-                raise TypeError(
-                    f"typed_op does not accept MPObject kwargs: {k}; pass MPObjects positionally"
-                )
-
-        # Try original call; if it binds, keep it as-is to support data-like positionals
-        try:
-            sig.bind_partial(*args, **kwargs)
-            call_pos = args
-            call_kwargs = kwargs
-        except TypeError as _bind_err:
-            # Fallback: For convenience, map non-MPObject positional arguments to
-            # keyword-only parameters by order. This allows ergonomic calls like
-            # `crypto.keygen(32)` where the kernel is `def keygen(*, length: int)`.
-            # The direct binding `sig.bind_partial(32)` would fail, so we manually
-            # map the positional `32` to the `length` keyword.
-            non_mp_positional = [a for a in args if not isinstance(a, MPObject)]
-            call_kwargs = dict(kwargs)
-            filled = 0
-            for _i, name in enumerate(self._kwonly_names):
-                if filled < len(non_mp_positional) and name not in call_kwargs:
-                    call_kwargs[name] = non_mp_positional[filled]
-                    filled += 1
-            if filled < len(non_mp_positional):
-                leftover = non_mp_positional[filled:]
-                raise TypeError(
-                    f"too many non-MPObject positional values for typed_op '{self.module.name}.{self.name}': {leftover}. "
-                    "Pass attributes explicitly by keyword (e.g., foo(x, *, attr=...))."
-                ) from None
-            call_pos = tuple(pos_mp_inputs)
-
-        # Compute PFunction attrs from the call kwargs (exclude MPObject and type objects)
-        attr_kwargs: dict[str, Any] = {
-            k: v
-            for k, v in call_kwargs.items()
-            if not isinstance(v, MPObject)
-            and not isinstance(v, (TensorType, TableType))
-        }
-
-        # Prepare kernel positional arguments: replace MPObject with its underlying type so
-        # the kernel always sees TensorType/TableType (never TraceVar/InterpVar).
-        call_pos_types = tuple(a.mptype._type for a in call_pos)
-
-        # Sanity: no MPObject should appear in kwargs (enforced earlier), but be safe.
-        if any(isinstance(v, MPObject) for v in call_kwargs.values()):
-            raise TypeError("kernel kwargs should not be MPObject")
-
-        # Execute kernel to compute return types
-        result = self._kernel(*call_pos_types, **call_kwargs)
-
-        outs_info, out_tree = tree_flatten(result)
-
-        # ensure all out_vars are TensorType or TableType.
-        # TODO(jint), theoretically we can also python constants here.
-        for o in outs_info:
-            if not isinstance(o, (TensorType, TableType)):
-                raise TypeError(
-                    f"simple op kernel must return TensorType or TableType, got {type(o).__name__}"
-                )
-
-        # Build input types from positional MPObjects only
-        ins_info = [a.mptype._type for a in pos_mp_inputs]
-
-        # Compose PFunction and return triad
-        pfunc = PFunction(
-            fn_type=self.pfunc_name,
-            ins_info=tuple(ins_info),
-            outs_info=tuple(outs_info),
-            **attr_kwargs,
-        )
-        return pfunc, pos_mp_inputs, out_tree
-
-
-def stateless_mod(mod_name: str) -> FeModule:
-    return StatelessFeModule(mod_name)
-
-
-def list_ops(module: str | None = None) -> dict[tuple[str, str], FeOperation]:
-    """Return a view of registered feops, optionally filtered by module name."""
-    return get_registry().list_ops(module)
-
-
-# -----------------------------------------------------------------------------
-# Guidance: complex ops via subclassing
-# -----------------------------------------------------------------------------
-
-# Example pattern (non-executable) showing how a complex op (e.g., jax_cc) could
-# capture a Python callable and compile it to a Triad by subclassing FeOperation.
-#
-# class JaxCompileOp(FeOperation):
-#     def __init__(self, module: FeModule, name: str, func: Callable[..., Any], *,
-#                  fn_type: str = "mlir.stablehlo", **options: Any) -> None:
-#         super().__init__(module, name)
-#         self.func = func
-#         self.fn_type = fn_type
-#         self.options = dict(options)
-#
-#     def trace(self, *args: MPObject, **kwargs: Any) -> Triad:
-#         # 1) Infer output types from func and args, respecting current ctx/masks.
-#         # 2) Build PFunction with fn_type=self.fn_type and any attributes.
-#         # 3) Return (pfunc, list(args), out_tree)
-#         raise NotImplementedError
-
-
-# -----------------------------------------------------------------------------
-# Migration notes (checklist)
-# -----------------------------------------------------------------------------
-
-# - Replace any isinstance(FEOp)/metadata checks with isinstance(x, FeOperation).
-# - Define a FeModule via femod("module_name") and register it in FeRegistry automatically.
-# - For inline ops that already produce a triad, use @module.feop()(trace_fn). The op name is derived from the function name.
-# - For type-only kernels, use @module.typed_op(pfunc_name)(kernel). The op name is derived from the kernel function name.
-# - For complex ops (with Python callables/closures), subclass FeOperation and register
-#   using get_registry().register_op(module, name, op_instance) or use @module.feop with InlineFeOperation.
-# - Ensure PFunction.fn_type is set as the routing key (e.g., "mlir.stablehlo", "sql.run").
-# - Keep device selection/routing out of frontend code; only set fn_type and attributes.
-# - Avoid moving MPObjects across contexts directly; capture within current ctx in trace().