mplang-nightly 0.1.dev284__py3-none-any.whl → 0.1.dev285__py3-none-any.whl

mplang/edsl/README.md

@@ -1,279 +1,252 @@
-# MPLang EDSL - Experimental Architecture
+# MPLang EDSL - Embedded Domain-Specific Language
 
-**⚠️ Status**: Experimental / Work in Progress
+The EDSL (Embedded Domain-Specific Language) module is the core infrastructure for MPLang's graph-based IR system. It provides the foundational components for tracing Python functions into SSA-based graphs, type checking, and program compilation.
 
-This directory contains the next-generation EDSL (Embedded Domain-Specific Language) architecture for MPLang.
+## Overview
 
-## Why a New Architecture?
+MPLang EDSL uses a modern **Operation List + SSA** approach for better analyzability and optimization, similar to PyTorch FX and JAX. The EDSL captures Python function execution into an explicit graph IR that can be analyzed, optimized, and executed across distributed parties and devices.
 
-The current `mplang.core` architecture (Expr Tree + @primitive) has served us well, but we're hitting limitations:
+## Architecture
 
-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
+### Key Components
 
-Modern EDSLs (torch.fx, JAX) use **Operation List + SSA** for better analyzability and optimization.
+```
+mplang/edsl/
+├── typing.py         # Unified type system (MPType hierarchy)
+├── graph.py          # IR: Operation List + SSA (Graph, Operation, Value)
+├── primitive.py      # Primitive abstraction and registration
+├── object.py         # Object hierarchy (TraceObject, runtime values)
+├── context.py        # Context management (tracing vs execution)
+├── tracer.py         # Explicit tracer for graph construction
+├── jit.py            # @jit decorator for function compilation
+├── program.py        # Compiled program representation
+├── printer.py        # Graph visualization and debugging
+├── registry.py       # Op registry and dialect management
+└── serde.py          # Serialization/deserialization
+```
 
-## Goals
+### Type System
 
-### 1. Modern IR (Operation List)
+**Single source of truth**: `mplang.edsl.typing.MPType`
+
+The EDSL provides a unified, extensible type system that supports:
+
+- **ScalarType**: `i32`, `i64`, `f32`, `f64`, `bool`, etc.
+- **TensorType**: Multi-dimensional arrays with shape and dtype
+- **TableType**: Structured tabular data
+- **VectorType**: Encrypted vectors (for FHE/MPC)
+- **SSType**: Secret-shared values (for MPC)
+- **CustomType**: Extensible type system for domain-specific types
 
-**From** (Expr Tree):
 ```python
-CallExpr(
-    func=add,
-    args=[VariableExpr("x"), VariableExpr("y")]
-)
+from mplang.edsl.typing import TensorType, VectorType, f32, i64
+
+# Define types
+plaintext: TensorType = TensorType(f32, (4096,))
+ciphertext: VectorType = VectorType(f32, 4096)
+counter: TensorType = TensorType(i64, ())  # scalar
 ```
 
-**To** (Operation List):
+### Graph IR
+
+The Graph IR uses SSA (Static Single Assignment) form with explicit operations:
+
 ```python
-%0 = input "x"
-%1 = input "y"
-%2 = add %0, %1
-return %2
+# Example graph structure:
+# %0 = input "x"
+# %1 = input "y"
+# %2 = add %0, %1
+# return %2
 ```
 
-### 2. Unified Type System
+Each operation has:
+- **opcode**: Operation type (e.g., "add", "mul", "simp.pcall")
+- **inputs**: List of input Values (SSA variables)
+- **outputs**: List of output Values
+- **attributes**: Operation-specific metadata
 
-**Single source of truth**: `mplang.edsl.typing.MPType`
+### Tracing and Compilation
 
 ```python
-from mplang2.edsl.typing import Tensor, Vector, MPType, f32
+from mplang.edsl import jit, trace, Tracer
 
-# All types use BaseType
-plaintext: MPType = Tensor[f32, (4096,)]
-ciphertext: MPType = Vector[f32, 4096]
-```
+# Method 1: @jit decorator (automatic tracing and execution)
+@jit
+def my_program(x, y):
+    return x + y
 
-### 3. Explicit Tracing
+result = my_program(data_x, data_y)
 
-**Clean context management**:
-```python
-from mplang2.edsl import Tracer
+# Method 2: Explicit tracing (returns TracedFunction with .graph)
+traced_fn = trace(my_program, x_obj, y_obj)
+graph = traced_fn.graph
 
+# Method 3: Manual tracing with Tracer context
 tracer = Tracer()
-with tracer:  # Context manager protocol
-    result = my_function(x, y)
+with tracer:
+    result = my_program(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
+## API Layers
 
 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.
+### 1. High-Level API (User-Facing)
 
-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.
+- **Tracing**: `@jit`, `trace()`, `Tracer` context manager
+- **Primitives**: `@primitive` decorator for defining new operations
+- **Types**: Type annotations and type inference
+- **Objects**: Automatic wrapping of Python values
 
-## Directory Structure
+This is the **primary interface for users**.
 
-```
-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
-```
+### 2. Low-Level API (Compiler/Backend)
 
-## Implementation Status
+- **Graph manipulation**: Direct `Graph.add_op()` calls
+- **Op registry**: `register_impl()`, `get_impl()`
+- **Serialization**: Graph to/from protobuf
+- **Execution**: Direct graph interpretation via runtime
 
-### ✅ Completed (Phase 1-4)
+Used by compiler passes and backend implementations.
 
-- [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)
+## Integration with Dialects
 
-### 🚧 In Progress
-- [ ] Integration with existing ops/kernels
-- [ ] Migration utilities
-- [ ] Performance benchmarks
+The EDSL is dialect-agnostic. Dialects provide domain-specific operations:
 
-### ❌ Dropped / Deprecated
-- [x] Builder API (`builder.py`) - Integrated into `Tracer`
+- **mplang.dialects.simp**: SPMD/MPI-style operations (pcall_static, pcall_dynamic)
+- **mplang.dialects.tensor**: Tensor operations (run_jax for JAX-backed computation, structural ops)
+- **mplang.dialects.table**: Table operations (run_sql, read/write, conversions)
+- **mplang.dialects.spu**: Secure multi-party computation
+- **mplang.dialects.tee**: Trusted execution environment
+- **mplang.dialects.bfv**: Homomorphic encryption (BFV scheme)
+- **mplang.dialects.phe**: Paillier homomorphic encryption
 
-### 📋 Planned
-- [ ] Advanced optimizations
-- [ ] More backends (TEE, MPC)
+Each dialect registers its operations and type implementations with the EDSL.
 
-## Quick Start
+## Examples
 
-### Using the New Type System
+### Basic Tracing
 
 ```python
-from mplang2.edsl.typing import Tensor, Vector, CustomType, f32
+import mplang.edsl as el
+from mplang.dialects.simp import pcall_static
+
+@el.jit
+def distribute_computation(x, y):
+    # Execute computation on parties 0 and 1
+    result = pcall_static((0, 1), lambda a, b: a + b, x, y)
+    return result
+```
 
-# Define types
-PlaintextVec = Tensor[f32, (4096,)]
-CiphertextVec = Vector[f32, 4096]
-EncryptionKey = CustomType("EncryptionKey")
+### Custom Primitives
 
-# Type annotations
-def encrypt(data: PlaintextVec, key: EncryptionKey) -> CiphertextVec:
-    ...
+```python
+from mplang.edsl import primitive
+from mplang.edsl.typing import TensorType, f32
+
+@primitive("custom_op")
+def custom_op_abstract(x: TensorType, y: TensorType) -> TensorType:
+    # Type inference (abstract evaluation) for the "custom_op" primitive.
+    # In this simple example, we say the result has the same abstract type as `x`.
+    return x
 ```
 
-### Using the Tracer (Graph Construction)
+### Graph Inspection
 
 ```python
-from mplang2.edsl import Tracer
-from mplang2.dialects.simp import pcall_static
+from mplang.edsl import trace, format_graph
 
-def my_program(x, y):
-    # This function is traced into a Graph
-    return pcall_static((0, 1), lambda a, b: a + b, x, y)
+def my_fn(x):
+    return x * 2 + 1
 
-tracer = Tracer()
-with tracer:
-    # Inputs are automatically lifted to TraceObjects
-    result = my_program(x, y)
+# Trace the function
+traced_fn = trace(my_fn, x_obj)
 
-# Finalize graph
-graph = tracer.finalize(result)
+# Print graph IR
+print(format_graph(traced_fn.graph))
 ```
 
-## Design Documents
-
-Detailed design documents are in the `design/` subdirectory:
-
-### 1. [architecture.md](design/architecture.md)
+## Testing
 
-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
+The EDSL has comprehensive test coverage:
 
-### 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
+```bash
+# Run all EDSL tests
+uv run pytest tests/edsl/
 
-### 3. [migration.md](design/migration.md)
+# Run specific test files
+uv run pytest tests/edsl/test_tracer.py
+uv run pytest tests/edsl/test_typing.py
+uv run pytest tests/edsl/test_graph.py
+```
 
-Migration path from `mplang.core` to `mplang.edsl`:
-- 6-phase migration plan
-- Backward compatibility strategy
-- Type conversion utilities
+Test files:
+- `test_typing.py`: Type system tests
+- `test_graph.py`: Graph IR tests
+- `test_tracer.py`: Tracing functionality
+- `test_primitive.py`: Primitive operations
+- `test_context.py`: Context management
+- `test_printer.py`: Graph visualization
+- `test_serde.py`: Serialization/deserialization
+- `test_compiled_program_artifact.py`: Program compilation
 
-## Relationship with mplang.core
+## Development
 
-```
-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
-```
+### Adding New Operations
 
-**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
+1. Define the operation in the appropriate dialect (e.g., `mplang/dialects/tensor.py`)
+2. Register the operation with `@primitive` or explicit registry
+3. Implement backend execution in `mplang/backends/`
+4. Add tests in `tests/dialects/` and `tests/backends/`
 
-## Contributing
+### Type System Extension
 
-We welcome contributions! Since this is experimental:
+To add a new type:
 
-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
+1. Subclass `MPType` in `mplang/edsl/typing.py`
+2. Implement required methods (`__repr__`, `__eq__`, `__hash__`)
+3. Add serialization support in `typing.py` and `serde.py`
+4. Add tests in `tests/edsl/test_typing.py`
 
-### Development Workflow
+### Code Style
 
 ```bash
-# Install dev dependencies
-uv sync --group dev
-
-# Run tests (future)
-uv run pytest mplang/edsl/
+# Format code
+uv run ruff format mplang/edsl/
 
 # Lint
-uv run ruff check mplang/edsl/
-uv run ruff format mplang/edsl/
+uv run ruff check mplang/edsl/ --fix
 
 # 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?
+## Relationship with Other Components
 
-**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?
+```
+mplang/
+├── edsl/              # Core IR and tracing (this module)
+├── dialects/          # Domain-specific operations
+├── backends/          # Execution implementations
+├── runtime/           # Runtime execution engine
+├── libs/              # High-level libraries (device, ml, mpc)
+└── kernels/           # Low-level kernel implementations
+```
 
-**A**: Check the implementation status above and pick an unimplemented component. Open an issue to discuss!
+**Design documents** are in the project root `design/` directory:
+- [`design/architecture.md`](../../design/architecture.md): Overall architecture
+- [`design/control_flow.md`](../../design/control_flow.md): Control flow handling
+- [`design/compile_execute_decoupling.md`](../../design/compile_execute_decoupling.md): Compilation model
 
 ## References
 
-- **torch.fx**: https://pytorch.org/docs/stable/fx.html
+- **MPLang Documentation**: See repository root README.md and AGENTS.md
+- **Design Documents**: `/design/` directory
+- **PyTorch 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
+**Last Updated**: 2026-02-02

mplang/edsl/typing.py

@@ -59,7 +59,7 @@ rather than by creating a large, monolithic set of specific types.
     - `Tensor`: A multi-dimensional array of a `ScalarType` element type.
     - `Table`: A dictionary-like structure with named columns of any type.
 
-    2.  **Encryption Types**: Wrap other types to confer privacy properties by making them opaque.
+2.  **Encryption Types**: Wrap other types to confer privacy properties by making them opaque.
     - `SS`: A single share of a secret-shared value.
     - Note: Element-wise HE types (like `phe.CiphertextType`) are defined in their respective dialects (e.g., `phe`).
 

{mplang_nightly-0.1.dev284.dist-info → mplang_nightly-0.1.dev285.dist-info}/METADATA

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

{mplang_nightly-0.1.dev284.dist-info → mplang_nightly-0.1.dev285.dist-info}/RECORD

@@ -40,7 +40,7 @@ mplang/dialects/store.py,sha256=FkvTOdcNC_oCgJPMZfdpGLlKZAlelEqRIZ2WgMgYH-8,2108
 mplang/dialects/table.py,sha256=i9ruyh91_tSWu9rsLomrBUfqRdbHiZMMMJzNKfMrAUc,13534
 mplang/dialects/tee.py,sha256=BMFSbeK-Ck2jQP4qY9bZeNYTxEa7uEtUWLZLC4BPQxk,10111
 mplang/dialects/tensor.py,sha256=7aAYKaMaFjJ8N25yPFnmVhUuUdKJYy-M-a4NsZGE7kY,39893
-mplang/edsl/README.md,sha256=viflvdRojOa6Xk_UMRPqpuPGXcPGmdlv2-XR6LO7B58,7592
+mplang/edsl/README.md,sha256=T_gmLkcWvN5b8jU6DX59LMZdgdgGupD9YlpPzH9FwlQ,7750
 mplang/edsl/__init__.py,sha256=WL4efo6uY1br781_8IaCkSi7yCUldcfJfbtFsn6Fdj4,2698
 mplang/edsl/context.py,sha256=Ln8n3bDe8_ISe42TAGzUuz8fw57-tu1APuihMfAtW1Y,10075
 mplang/edsl/graph.py,sha256=nCeCN7-bxfzyv40fmxcEXOaVUx14cOCaHfFb7A9OBnE,14968
@@ -52,7 +52,7 @@ mplang/edsl/program.py,sha256=_JdEU2-nb79VlFLcgMJf4JS30TARBeUIzno0y0SFVsg,4467
 mplang/edsl/registry.py,sha256=hudXZPUrUUueEwgksDKN0cnE3iiXucuTaDdDK8uSPmk,6822
 mplang/edsl/serde.py,sha256=8K94laE8ObeGuBoF6m7g3A-xEe98EvqQ_6ZPPspddAY,11641
 mplang/edsl/tracer.py,sha256=WQFNL2ZgXSLjxD4JA7cXIDUKIQXe3aZ94qer57IKPXc,23128
-mplang/edsl/typing.py,sha256=Vp0r_oTyFrOhwwpVD_6XAX0sSsKvR04X-KdihkmD4QA,29287
+mplang/edsl/typing.py,sha256=Hmn_w04XWepmvurN1ZEb6rVNW1kW9bOBYe3AkD8CAJc,29283
 mplang/kernels/Makefile,sha256=5PoPpajcb_8ByPGNHzVytmovXUwkjJs_K8MbXX9qDYs,1033
 mplang/kernels/__init__.py,sha256=J_rDl9lAXd7QL3Nt_P3YX6j9yge7ssguSaHuafPZNKE,876
 mplang/kernels/gf128.cpp,sha256=WIvCr3MijzwJxMi1Wnfhm8aWT8oL0fia6FeyTmFJtPQ,5975
@@ -99,8 +99,8 @@ mplang/tool/program.py,sha256=W3H8bpPirnoJ4ZrmyPYuMCPadJis20o__n_1MKqCsWU,11058
 mplang/utils/__init__.py,sha256=Hwrwti2nfPxWUXV8DN6T1QaqXH_Jsd27k8UMSdBGUns,1073
 mplang/utils/func_utils.py,sha256=aZ-X43w8JKJgiF-IUMS0G7QqrNeoTM5ZPzRNd-tKxpw,5180
 mplang/utils/logging.py,sha256=9dMhwprVbx1WMGJrgoQbWmV50vyYuLU4NSPnetcl1Go,7237
-mplang_nightly-0.1.dev284.dist-info/METADATA,sha256=ufhEEGe9Ieui9mISxEFm0yXQfzhSS0T0x3TurFsHwRM,16783
-mplang_nightly-0.1.dev284.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mplang_nightly-0.1.dev284.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev284.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev284.dist-info/RECORD,,
+mplang_nightly-0.1.dev285.dist-info/METADATA,sha256=cXD5_vMkz7YYbixonHu4GN49cnfgzuGgaofeV2VJ9z4,16783
+mplang_nightly-0.1.dev285.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mplang_nightly-0.1.dev285.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev285.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev285.dist-info/RECORD,,