mplang-nightly 0.1.dev284__tar.gz → 0.1.dev286__tar.gz

{mplang_nightly-0.1.dev284 → mplang_nightly-0.1.dev286}/PKG-INFO

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

mplang_nightly-0.1.dev286/mplang/edsl/README.md

@@ -0,0 +1,252 @@
+# MPLang EDSL - Embedded Domain-Specific Language
+
+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.
+
+## Overview
+
+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.
+
+## Architecture
+
+### Key Components
+
+```
+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
+```
+
+### Type System
+
+**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
+
+```python
+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
+```
+
+### Graph IR
+
+The Graph IR uses SSA (Static Single Assignment) form with explicit operations:
+
+```python
+# Example graph structure:
+# %0 = input "x"
+# %1 = input "y"
+# %2 = add %0, %1
+# return %2
+```
+
+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
+
+### Tracing and Compilation
+
+```python
+from mplang.edsl import jit, trace, Tracer
+
+# Method 1: @jit decorator (automatic tracing and execution)
+@jit
+def my_program(x, y):
+    return x + y
+
+result = my_program(data_x, data_y)
+
+# 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:
+    result = my_program(x, y)
+graph = tracer.finalize(result)
+```
+
+## API Layers
+
+The EDSL provides two distinct API layers:
+
+### 1. High-Level API (User-Facing)
+
+- **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
+
+This is the **primary interface for users**.
+
+### 2. Low-Level API (Compiler/Backend)
+
+- **Graph manipulation**: Direct `Graph.add_op()` calls
+- **Op registry**: `register_impl()`, `get_impl()`
+- **Serialization**: Graph to/from protobuf
+- **Execution**: Direct graph interpretation via runtime
+
+Used by compiler passes and backend implementations.
+
+## Integration with Dialects
+
+The EDSL is dialect-agnostic. Dialects provide domain-specific operations:
+
+- **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
+
+Each dialect registers its operations and type implementations with the EDSL.
+
+## Examples
+
+### Basic Tracing
+
+```python
+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
+```
+
+### Custom Primitives
+
+```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
+```
+
+### Graph Inspection
+
+```python
+from mplang.edsl import trace, format_graph
+
+def my_fn(x):
+    return x * 2 + 1
+
+# Trace the function
+traced_fn = trace(my_fn, x_obj)
+
+# Print graph IR
+print(format_graph(traced_fn.graph))
+```
+
+## Testing
+
+The EDSL has comprehensive test coverage:
+
+```bash
+# Run all EDSL tests
+uv run pytest tests/edsl/
+
+# 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
+```
+
+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
+
+## Development
+
+### Adding New Operations
+
+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/`
+
+### Type System Extension
+
+To add a new type:
+
+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`
+
+### Code Style
+
+```bash
+# Format code
+uv run ruff format mplang/edsl/
+
+# Lint
+uv run ruff check mplang/edsl/ --fix
+
+# Type check
+uv run mypy mplang/edsl/
+```
+
+## Relationship with Other Components
+
+```
+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
+```
+
+**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
+
+- **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**: 2026-02-02

{mplang_nightly-0.1.dev284 → mplang_nightly-0.1.dev286}/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`).
 
@@ -82,11 +82,18 @@ underlying HE libraries.
     - **Core Type**: `Tensor[Scalar, ...]`
     - **API Standard**: Follows NumPy/JAX conventions. All layout and arithmetic operations are valid.
 
+-   **World 2: The Element-wise HE World**
     - **Core Type**: `Tensor[EncryptedScalar, ...]` (e.g., `Tensor[phe.CiphertextType, ...]`)
     - **API Standard**: Follows TenSEAL-like (Tensor-level) conventions. Layout operations
       (`transpose`, `reshape`) are valid as they merely shuffle independent ciphertext objects.
       Arithmetic operations are overloaded for element-wise HE computation.
 
+-   **World 3: The SIMD HE World**
+    - **Core Type**: `Tensor[Vector[...], ...]` where the inner `Vector` holds SIMD-encrypted slots
+    - **API Standard**: SIMD-HE specific (e.g., BFV/CKKS). Only specific operations are supported
+      due to the batched nature of SIMD encryption. Layout operations must account for slot packing,
+      and arithmetic operates on batched ciphertexts with slot-wise semantics.
+
 ===========================
 Principle 3: Contracts via Protocols
 ===========================
@@ -137,7 +144,15 @@ class BaseType:
     """Base class for all MPLang types."""
 
     def __repr__(self) -> str:
-        return str(self)
+        # Prevent infinite recursion: only call __str__ if it's overridden
+        if type(self).__str__ is not BaseType.__str__:
+            return str(self)
+        # Fallback to default object repr if __str__ not implemented
+        return object.__repr__(self)
+
+    def __str__(self) -> str:
+        # Default implementation for subclasses that don't override
+        return f"{self.__class__.__name__}()"
 
 
 # ==============================================================================
@@ -175,11 +190,11 @@ class ScalarType(BaseType):
 
 @serde.register_class
 class IntegerType(ScalarType):
-    """Represents a variable-length integer type.
+    """Represents a fixed-width integer type with configurable bitwidth.
 
-    This is a standard integer type with configurable bit width, used for
-    arbitrary-precision arithmetic. It can represent integers that exceed
-    the range of fixed-width types like i64.
+    This is a standard integer type with parameterized bit width, used for
+    arbitrary-precision arithmetic. By configuring larger bitwidths (e.g., 128, 256),
+    this type can represent integers that exceed the range of standard types like i64.
 
     Examples:
         >>> i128 = IntegerType(bitwidth=128, signed=True)  # i128
@@ -229,15 +244,16 @@ class IntegerType(ScalarType):
 
 @serde.register_class
 class FloatType(ScalarType):
-    """Represents a floating-point type.
+    """Represents a fixed-width floating-point type with configurable bitwidth.
 
-    This supports standard IEEE 754 floating-point types with configurable
-    precision (bitwidth).
+    This supports standard IEEE 754 floating-point types with parameterized
+    bit width for different precision requirements.
 
     Examples:
         >>> f16 = FloatType(bitwidth=16)  # half precision
         >>> f32 = FloatType(bitwidth=32)  # single precision
         >>> f64 = FloatType(bitwidth=64)  # double precision
+        >>> f128 = FloatType(bitwidth=128)  # quadruple precision
     """
 
     def __init__(self, *, bitwidth: int = 32):
@@ -245,7 +261,7 @@ class FloatType(ScalarType):
 
         Args:
             bitwidth: Number of bits for the float representation.
-                     Standard values: 16 (half), 32 (single), 64 (double).
+                     Standard values: 16 (half), 32 (single), 64 (double), 128 (quadruple).
         """
         if bitwidth not in (16, 32, 64, 128):
             raise ValueError(f"bitwidth must be 16, 32, 64, or 128, got {bitwidth}")

mplang_nightly-0.1.dev284/mplang/edsl/README.md

@@ -1,279 +0,0 @@
-# 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