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

mplang/v2/edsl/README.md

@@ -0,0 +1,279 @@
+# MPLang EDSL - Experimental Architecture
+
+**⚠️ Status**: Experimental / Work in Progress
+
+This directory contains the next-generation EDSL (Embedded Domain-Specific Language) architecture for MPLang.
+
+## Why a New Architecture?
+
+The current `mplang.core` architecture (Expr Tree + @primitive) has served us well, but we're hitting limitations:
+
+1. **Expr Tree** is hard to optimize (visitor pattern, nested structure)
+2. **@primitive decorators** hide complexity and limit flexibility
+3. **Type system** is split between `mptype.MPType` and `typing.BaseType`
+4. **No clear separation** between IR, frontend, and backend
+
+Modern EDSLs (torch.fx, JAX) use **Operation List + SSA** for better analyzability and optimization.
+
+## Goals
+
+### 1. Modern IR (Operation List)
+
+**From** (Expr Tree):
+```python
+CallExpr(
+    func=add,
+    args=[VariableExpr("x"), VariableExpr("y")]
+)
+```
+
+**To** (Operation List):
+```python
+%0 = input "x"
+%1 = input "y"
+%2 = add %0, %1
+return %2
+```
+
+### 2. Unified Type System
+
+**Single source of truth**: `mplang.edsl.typing.MPType`
+
+```python
+from mplang2.edsl.typing import Tensor, Vector, MPType, f32
+
+# All types use BaseType
+plaintext: MPType = Tensor[f32, (4096,)]
+ciphertext: MPType = Vector[f32, 4096]
+```
+
+### 3. Explicit Tracing
+
+**Clean context management**:
+```python
+from mplang2.edsl import Tracer
+
+tracer = Tracer()
+with tracer:  # Context manager protocol
+    result = my_function(x, y)
+graph = tracer.finalize(result)
+```
+
+### 4. Extensibility
+
+Easy to add new backends:
+- FHE (Fully Homomorphic Encryption)
+- TEE (Trusted Execution Environment)
+- Custom accelerators
+
+### 5. Layered API Architecture
+
+The EDSL provides two distinct API layers:
+
+1.  **Low-Level API (Graph Manipulation)**:
+    - Direct manipulation of the `Graph` IR.
+    - Generic `add_op` method (pure graph API, no op semantics).
+    - Analogous to MLIR's generic operation construction.
+    - Used by compiler passes and backend implementations.
+
+2.  **High-Level API (Tracing)**:
+    - Uses `Tracer` + `Primitive` (with `abstract_eval`).
+    - Pythonic interface (functions, operators).
+    - Automatic type inference and graph construction.
+    - The primary interface for users.
+
+## Directory Structure
+
+```
+mplang/edsl/
+├── __init__.py          # Public API
+├── README.md            # This file
+│
+├── design/              # Design documents
+│   ├── architecture.md  # Complete architecture overview
+│   ├── type_system.md   # Type system design
+│   └── migration.md     # Migration from mplang.core
+│
+├── typing.py            # ✅ Unified type system
+├── graph.py             # ✅ IR: Operation List + SSA
+├── primitive.py         # ✅ Primitive abstraction
+├── object.py            # ✅ TraceObject/InterpObject
+├── context.py           # ✅ Context management
+├── tracer.py            # ✅ Explicit tracer
+├── interpreter.py       # ✅ Interpreter + GraphInterpreter
+└── jit.py               # ✅ @jit decorator
+```
+
+## Implementation Status
+
+### ✅ Completed (Phase 1-4)
+
+- [x] Type system (`typing.py`) - 649 lines
+- [x] Graph IR (`graph.py`) - 388 lines
+- [x] Primitive abstraction (`primitive.py`) - 338 lines
+- [x] Object hierarchy (`object.py`) - 153 lines
+- [x] Context system (`context.py`) - 117 lines
+- [x] Tracer (`tracer.py`) - 201 lines
+- [x] Interpreter (`interpreter.py`) - 66 lines
+- [x] JIT decorator (`jit.py`) - 42 lines
+- [x] Design documents
+- [x] **153 tests passing** (140 edsl + 13 core2)
+
+### 🚧 In Progress
+- [ ] Integration with existing ops/kernels
+- [ ] Migration utilities
+- [ ] Performance benchmarks
+
+### ❌ Dropped / Deprecated
+- [x] Builder API (`builder.py`) - Integrated into `Tracer`
+
+### 📋 Planned
+- [ ] Advanced optimizations
+- [ ] More backends (TEE, MPC)
+
+## Quick Start
+
+### Using the New Type System
+
+```python
+from mplang2.edsl.typing import Tensor, Vector, CustomType, f32
+
+# Define types
+PlaintextVec = Tensor[f32, (4096,)]
+CiphertextVec = Vector[f32, 4096]
+EncryptionKey = CustomType("EncryptionKey")
+
+# Type annotations
+def encrypt(data: PlaintextVec, key: EncryptionKey) -> CiphertextVec:
+    ...
+```
+
+### Using the Tracer (Graph Construction)
+
+```python
+from mplang2.edsl import Tracer
+from mplang2.dialects.simp import pcall_static
+
+def my_program(x, y):
+    # This function is traced into a Graph
+    return pcall_static((0, 1), lambda a, b: a + b, x, y)
+
+tracer = Tracer()
+with tracer:
+    # Inputs are automatically lifted to TraceObjects
+    result = my_program(x, y)
+
+# Finalize graph
+graph = tracer.finalize(result)
+```
+
+## Design Documents
+
+Detailed design documents are in the `design/` subdirectory:
+
+### 1. [architecture.md](design/architecture.md)
+
+Complete EDSL architecture overview covering:
+- Core components (Tracer, Graph)
+- Design principles (Closed-World, TracedFunction vs First-Class Functions)
+- Control flow handling (Dialect-specific, e.g., `simp.uniform_cond`)
+- Comparison with JAX, PyTorch, TensorFlow
+
+### 2. [type_system.md](design/type_system.md)
+
+New type system design:
+- Three orthogonal dimensions (Layout, Encryption, Distribution)
+- Type composition examples
+- Ops writing guide
+- Migration strategy
+
+### 3. [migration.md](design/migration.md)
+
+Migration path from `mplang.core` to `mplang.edsl`:
+- 6-phase migration plan
+- Backward compatibility strategy
+- Type conversion utilities
+
+## Relationship with mplang.core
+
+```
+mplang/
+├── core/           # Stable API (current production)
+│   ├── primitive.py
+│   ├── tracer.py
+│   └── expr/
+│
+├── edsl/           # Experimental (this directory)
+│   ├── typing.py   # Can be used independently
+│   ├── graph.py    # Future replacement for core.expr
+│   └── tracer.py   # Future replacement for core.tracer
+│
+├── ops/            # Shared between core and edsl
+├── kernels/        # Shared between core and edsl
+└── runtime/        # Shared between core and edsl
+```
+
+**Migration Strategy**:
+1. Develop `edsl` in parallel (no breaking changes to `core`)
+2. Gradually move internal code to use `edsl.typing`
+3. Add adapters between `core` and `edsl`
+4. Deprecate `core` in future major version
+
+## Contributing
+
+We welcome contributions! Since this is experimental:
+
+1. **Read the design docs first**: Understand the architecture
+2. **Start small**: Pick a specific component (e.g., Graph IR)
+3. **Discuss early**: Open an issue before implementing
+4. **Test thoroughly**: Add unit tests for new code
+
+### Development Workflow
+
+```bash
+# Install dev dependencies
+uv sync --group dev
+
+# Run tests (future)
+uv run pytest mplang/edsl/
+
+# Lint
+uv run ruff check mplang/edsl/
+uv run ruff format mplang/edsl/
+
+# Type check
+uv run mypy mplang/edsl/
+```
+
+## FAQ
+
+### Q: Should I use `mplang.edsl` in production?
+
+**A**: No, use `mplang.core`. `mplang.edsl` is experimental.
+
+### Q: Can I use `mplang.edsl.typing` independently?
+
+**A**: Yes! The type system is stable and can be used for type annotations.
+
+### Q: When will `edsl` replace `core`?
+
+**A**: No timeline yet. We need to:
+1. Complete the implementation
+2. Validate performance
+3. Migrate all tests
+4. Get community feedback
+
+### Q: How can I help?
+
+**A**: Check the implementation status above and pick an unimplemented component. Open an issue to discuss!
+
+## References
+
+- **torch.fx**: https://pytorch.org/docs/stable/fx.html
+- **JAX jaxpr**: https://jax.readthedocs.io/en/latest/jaxpr.html
+- **MLIR**: https://mlir.llvm.org/
+
+---
+
+**Last Updated**: 2025-01-11
+**Maintainers**: MPLang Team

mplang/v2/edsl/__init__.py

@@ -0,0 +1,99 @@
+# 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.
+
+"""Public entrypoint for the MPLang EDSL.
+
+This module keeps the surface area intentionally small so downstream code can
+simply write::
+
+    import mplang.v2.edsl as el
+    import mplang.v2.edsl.typing as elt
+
+The `el` namespace re-exports the commonly used building blocks (context,
+graph, tracer, primitives, etc.), while the full type system lives under
+``mplang.edsl.typing``.
+"""
+
+from __future__ import annotations
+
+# Re-export the typing module so callers can `import mplang.v2.edsl.typing as elt`
+from . import typing as typing
+
+# Context management
+from .context import (
+    Context,
+    find_context,
+    find_context_with_state,
+    find_interpreter,
+    get_current_context,
+    get_default_context,
+    is_tracing,
+    pop_context,
+    push_context,
+    register_default_context_factory,
+    set_root_context,
+)
+
+# Graph IR
+from .graph import Graph, Operation, Value
+
+# High-level helpers
+from .jit import jit
+from .object import Object
+from .primitive import Primitive, primitive
+from .printer import GraphPrinter, format_graph
+from .tracer import TracedFunction, TraceObject, Tracer, trace
+from .typing import MPType, ScalarType, SSType, TableType, TensorType, VectorType
+
+# Type Aliases for strong typing
+MPObject = Object[MPType]
+ScalarObject = Object[ScalarType]
+SSObject = Object[SSType]
+TableObject = Object[TableType]
+TensorObject = Object[TensorType]
+VectorObject = Object[VectorType]
+
+__all__ = [
+    "Context",
+    "Graph",
+    "GraphPrinter",
+    "MPObject",
+    "Object",
+    "Operation",
+    "Primitive",
+    "SSObject",
+    "ScalarObject",
+    "TableObject",
+    "TensorObject",
+    "TraceObject",
+    "TracedFunction",
+    "Tracer",
+    "Value",
+    "VectorObject",
+    "find_context",
+    "find_context_with_state",
+    "find_interpreter",
+    "format_graph",
+    "get_current_context",
+    "get_default_context",
+    "is_tracing",
+    "jit",
+    "pop_context",
+    "primitive",
+    "push_context",
+    "register_default_context_factory",
+    "set_root_context",
+    "trace",
+    "typing",
+]

mplang/v2/edsl/context.py

@@ -0,0 +1,311 @@
+# 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.
+
+"""Context: EDSL Execution Context Abstraction.
+
+This module defines the Context hierarchy:
+- Context: Base class for EDSL execution contexts (with bind_primitive method)
+- Tracer: Tracing context (records operations to Graph IR)
+- Interpreter: Execution context (executes operations immediately)
+
+Contexts can be used directly with Python's 'with' statement:
+
+    from mplang.v2.edsl import Tracer
+
+    tracer = Tracer()
+    with tracer:
+        # Operations run under tracer context
+        result = primitive.bind(x, y)
+
+State Management:
+    Contexts can carry arbitrary named state via set_state/get_state.
+    This allows different layers (device, ml, analytics) to attach their
+    own state without the EDSL layer knowing about specific state types.
+
+    State key conventions:
+    - "dialect.{name}": Dialect runtime state (e.g., "dialect.simp")
+    - "device.cluster": Device/cluster configuration
+    - "ml.{component}": ML pipeline components
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Self
+
+if TYPE_CHECKING:
+    from mplang.v2.edsl.graph import Graph
+    from mplang.v2.edsl.object import Object
+    from mplang.v2.edsl.primitive import Primitive
+
+
+class Context(ABC):
+    """Base class for EDSL execution contexts with extensible state slots.
+
+    A Context represents an environment where primitives are executed.
+    There are two types of contexts:
+    - Tracer: Records operations to Graph IR (compile-time)
+    - Interpreter: Execution context (executes operations immediately)
+
+    State Management:
+        Contexts can carry arbitrary named state. Different layers can attach
+        their own state without the EDSL layer knowing specifics:
+
+        >>> ctx.set_state("device.cluster", cluster_spec)
+        >>> ctx.set_state("dialect.simp", simp_driver)
+        >>> cluster = ctx.get_state("device.cluster")
+    """
+
+    def __init__(self) -> None:
+        self._states: dict[str, Any] = {}
+
+    # =========================================================================
+    # State Management
+    # =========================================================================
+
+    def set_state(self, key: str, value: Any) -> None:
+        """Attach state to this context.
+
+        Args:
+            key: State key (e.g., "dialect.simp", "device.cluster")
+            value: State value
+        """
+        self._states[key] = value
+
+    def get_state(self, key: str, default: Any = None) -> Any:
+        """Get attached state by key.
+
+        Args:
+            key: State key
+            default: Default value if key not found
+
+        Returns:
+            State value or default
+        """
+        return self._states.get(key, default)
+
+    def has_state(self, key: str) -> bool:
+        """Check if state exists.
+
+        Args:
+            key: State key
+
+        Returns:
+            True if state exists
+        """
+        return key in self._states
+
+    # =========================================================================
+    # Abstract Methods
+    # =========================================================================
+
+    @abstractmethod
+    def bind_primitive(
+        self, primitive: Primitive, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> Any:
+        """Execute a primitive in this context.
+
+        Args:
+            primitive: The primitive to execute
+            args: Positional arguments (Objects)
+            kwargs: Keyword arguments (plain values)
+
+        Returns:
+            Result Object (TraceObject in Tracer, InterpObject in Interpreter)
+        """
+
+    @abstractmethod
+    def lift(self, obj: Any) -> Object:
+        """Lift an object to this context's native Object type.
+
+        Converts objects to the appropriate type for this context:
+        - Tracer: InterpObject → TraceObject (via promote), constants → TraceObject
+        - Interpreter: keeps InterpObject as-is, may convert constants
+
+        Args:
+            obj: Object to lift (Object, constant, etc.)
+
+        Returns:
+            Object in the context's native type (TraceObject or InterpObject)
+        """
+
+    # =========================================================================
+    # Context Manager
+    # =========================================================================
+
+    def __enter__(self) -> Self:
+        """Enter context manager (push context onto stack)."""
+        push_context(self)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:  # type: ignore[no-untyped-def]
+        """Exit context manager (pop context from stack)."""
+        pop_context()
+
+
+# =============================================================================
+# Abstract Interpreter Interface
+# =============================================================================
+
+
+class AbstractInterpreter(Context):
+    """Abstract interface for Interpreters.
+
+    This allows EDSL components (like JIT) to depend on the Interpreter interface
+    without depending on the concrete Runtime implementation (which may depend on
+    ObjectStore, Backends, etc.).
+    """
+
+    @abstractmethod
+    def evaluate_graph(self, graph: Graph, inputs: list[Any]) -> Any:
+        """Execute a Graph IR with given inputs."""
+
+    @abstractmethod
+    def lift(self, obj: Any) -> Any:
+        """Lift a python object to an interpreter object."""
+
+
+# =============================================================================
+# Global Context Stack Management
+# =============================================================================
+
+_context_stack: list[Context] = []
+_default_context: Context | None = None
+_default_context_factory: Callable[[], Context] | None = None
+
+
+def get_current_context() -> Context | None:
+    """Get the current active context (top of stack).
+
+    Returns None if no context is active.
+    """
+    return _context_stack[-1] if _context_stack else None
+
+
+def push_context(context: Context) -> None:
+    """Push a context onto the stack (enter context)."""
+    _context_stack.append(context)
+
+
+def pop_context() -> Context | None:
+    """Pop a context from the stack (exit context).
+
+    Returns:
+        The popped context, or None if stack was empty.
+    """
+    return _context_stack.pop() if _context_stack else None
+
+
+def find_context(predicate: Callable[[Context], bool]) -> Context | None:
+    """Find a context in the stack that satisfies the predicate.
+
+    Traverses from top (most recent) to bottom of the context stack,
+    returning the first context for which predicate(ctx) returns True.
+
+    Args:
+        predicate: A callable that takes a Context and returns True if it matches.
+
+    Returns:
+        The first matching Context, or None if no match found.
+
+    Example:
+        >>> # Find context with simp dialect state
+        >>> ctx = find_context(lambda c: c.has_state("dialect.simp"))
+    """
+    for ctx in reversed(_context_stack):
+        if predicate(ctx):
+            return ctx
+    return None
+
+
+def find_context_with_state(key: str) -> Context | None:
+    """Find first context that has the specified state.
+
+    Args:
+        key: State key to look for
+
+    Returns:
+        First context with the state, or None
+    """
+    return find_context(lambda c: c.has_state(key))
+
+
+def find_interpreter() -> Context | None:
+    """Find first Interpreter in the context stack.
+
+    Returns:
+        First Interpreter context, or None if not found.
+    """
+    return find_context(lambda c: isinstance(c, AbstractInterpreter))
+
+
+def is_tracing() -> bool:
+    """Check if current context is a Tracer.
+
+    Returns:
+        True if the top of the context stack is a Tracer.
+    """
+    from mplang.v2.edsl.tracer import Tracer
+
+    return isinstance(get_current_context(), Tracer)
+
+
+# =============================================================================
+# Default Context Management
+# =============================================================================
+
+
+def register_default_context_factory(factory: Callable[[], Context]) -> None:
+    """Register a factory function to create the default context."""
+    global _default_context_factory
+    _default_context_factory = factory
+
+
+def get_default_context() -> Context:
+    """Get the default context for eager execution."""
+    global _default_context
+    if _default_context is None:
+        if _default_context_factory is None:
+            raise RuntimeError(
+                "No default context factory registered. "
+                "Ensure mplang.v2.edsl is imported or register a factory manually."
+            )
+        _default_context = _default_context_factory()
+    return _default_context
+
+
+def set_root_context(context: Context, force: bool = False) -> None:
+    """Set the root/default execution context.
+
+    This sets the provided context as the base of the context stack.
+    All subsequent operations will use this context as the default environment.
+
+    Args:
+        context: Context to set as root.
+        force: If True, clears the existing context stack before setting.
+               If False (default), raises error if stack is not empty.
+    """
+    if force:
+        _context_stack.clear()
+        _context_stack.append(context)
+        return
+
+    if get_current_context() is not None:
+        raise RuntimeError(
+            "Cannot set root context: Context stack is not empty. "
+            "Use force=True to overwrite the existing context."
+        )
+
+    push_context(context)