mplang-nightly 0.1.dev169__tar.gz → 0.1.dev170__tar.gz

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/PKG-INFO

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

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/mplang/core/expr/ast.py

@@ -528,8 +528,9 @@ class FuncDefExpr(Expr):
 class CallExpr(Expr):
     """Expression for function call."""
 
-    def __init__(self, fn: FuncDefExpr, args: list[Expr]):
+    def __init__(self, name: str, fn: FuncDefExpr, args: list[Expr]):
         super().__init__()
+        self.name = name
         self.fn = fn
         self.args = args
 

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/mplang/core/expr/evaluator.py

@@ -341,10 +341,10 @@ class RecursiveEvaluator(EvalSemantic, ExprVisitor):
 
         # Only evaluate selected branch locally
         if bool(pred):
-            then_call = CallExpr(expr.then_fn, expr.args)
+            then_call = CallExpr("then", expr.then_fn, expr.args)
             return self._values(then_call)
         else:
-            else_call = CallExpr(expr.else_fn, expr.args)
+            else_call = CallExpr("else", expr.else_fn, expr.args)
             return self._values(else_call)
 
     def visit_call(self, expr: CallExpr) -> Any:

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/mplang/core/expr/printer.py

@@ -50,11 +50,13 @@ class Printer(ExprVisitor):
         compact_format: bool = True,
         *,
         verbose_peval: bool = False,
+        inline_pcall: bool = True,
     ):
         super().__init__()  # Initialize MemorizedVisitor
         self.indent_size = indent_size
         self.compact_format = compact_format
         self.verbose_peval = verbose_peval
+        self.inline_pcall = inline_pcall
         self._cur_indent = 0
         self._output: list[str] = []
         self._visited: dict[Expr, str] = {}
@@ -92,6 +94,7 @@ class Printer(ExprVisitor):
                 body_printer = Printer(
                     indent_size=self.indent_size,
                     compact_format=self.compact_format,
+                    inline_pcall=self.inline_pcall,
                 )
                 func_def_expr.accept(body_printer)
                 regions_str += f"{indent}{r_name}: "
@@ -209,12 +212,19 @@ class Printer(ExprVisitor):
 
     def visit_call(self, expr: CallExpr) -> str:
         arg_names = [self._var_name(arg) for arg in expr.args]
-        return self._do_print(
-            "pcall",
-            arg_names,
-            regions={"fn": expr.fn},
-            mptypes=expr.mptypes,
-        )
+        if self.inline_pcall:
+            return self._do_print(
+                expr.name,
+                arg_names,
+                mptypes=expr.mptypes,
+            )
+        else:
+            return self._do_print(
+                "pcall",
+                arg_names,
+                regions={"fn": expr.fn},
+                mptypes=expr.mptypes,
+            )
 
     def visit_while(self, expr: WhileExpr) -> str:
         arg_names = [self._var_name(arg) for arg in expr.args]

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/mplang/core/expr/transformer.py

@@ -79,7 +79,7 @@ class ExprTransformer(ExprVisitor):
     def visit_call(self, expr: CallExpr) -> Expr:
         # Transform child expressions first
         transformed_args = [arg.accept(self) for arg in expr.args]
-        new_expr = CallExpr(expr.fn, transformed_args)
+        new_expr = CallExpr(expr.name, expr.fn, transformed_args)
 
         if "call" in self.trans_rules:
             return self.trans_rules["call"](new_expr)

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/mplang/core/mpir.py

@@ -491,6 +491,7 @@ class Writer:
             op = self._create_node_proto(expr, "call")
             self._add_single_expr_inputs(op, expr.fn)
             self._add_expr_inputs(op, *expr.args)
+            self._add_attrs(op, name=expr.name)
             self._finalize_node(op, expr)
         elif isinstance(expr, WhileExpr):
             op = self._create_node_proto(expr, "while")
@@ -822,8 +823,12 @@ class Reader:
                 arg_exprs.append(self._value_cache[dep_name])
             else:
                 raise ValueError(f"Input {input_name} not found for call node")
+        # Optional call-site name attribute
+        call_name = None
+        if "name" in node_proto.attrs:
+            call_name = self._proto_to_attr(node_proto.attrs["name"])  # type: ignore[assignment]
 
-        return CallExpr(fn_expr, arg_exprs)
+        return CallExpr(call_name or "", fn_expr, arg_exprs)
 
     def _proto_to_mptype(self, type_proto: mpir_pb2.MPTypeProto) -> MPType:
         """Convert MPTypeProto to MPType."""

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/mplang/core/primitive.py

@@ -32,6 +32,7 @@ from mplang.core.context_mgr import cur_ctx
 from mplang.core.dtype import BOOL
 from mplang.core.expr.ast import (
     AccessExpr,
+    CallExpr,
     CondExpr,
     ConvExpr,
     EvalExpr,
@@ -87,30 +88,106 @@ P = ParamSpec("P")
 R = TypeVar("R")
 
 
-def primitive(fn: Callable[P, R]) -> Callable[P, R]:
+def trace_before_apply(fn: Callable[P, R], make_call: bool) -> Callable[P, R]:
     """A decorator to make all primitive call in trace context."""
 
     @wraps(fn)
     def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
         current_ctx = cur_ctx()
         if isinstance(current_ctx, TraceContext):
-            # If we are in a tracer context, just call the function.
-            # Note: switch_ctx will do the capture if needed.
-            args, kwargs = tree_map(partial(_switch_ctx, current_ctx), (args, kwargs))
-            return fn(*args, **kwargs)
+            # If we are already in a tracer context
+            if make_call:
+                # make a primitive call
+                tracer = current_ctx
+                tfn = trace(tracer.fork(), fn, *args, **kwargs)
+                is_mpobj = lambda x: isinstance(x, MPObject)
+                in_vars, in_imms, in_struct = var_morph((args, kwargs), is_mpobj)
+                assert in_struct == tfn.in_struct and in_imms == tfn.in_imms
+                arg_exprs = [arg.expr for arg in in_vars]
+                # re-capture all captured variables into current context if needed.
+                cap_exprs = [tracer.capture(var).expr for var in tfn.capture_map.keys()]
+                caller_expr = CallExpr(
+                    name=fn.__name__, fn=tfn.make_expr(), args=arg_exprs + cap_exprs
+                )
+                out_vars = [
+                    TraceVar(tracer, AccessExpr(caller_expr, idx))
+                    for idx in range(caller_expr.num_outputs)
+                ]
+                return cast(R, var_demorph(out_vars, tfn.out_imms, tfn.out_struct))
+            else:
+                # embed the function call in the current tracer context
+                # Note: switch_ctx will do the capture if needed.
+                args, kwargs = tree_map(
+                    partial(_switch_ctx, current_ctx), (args, kwargs)
+                )
+                return fn(*args, **kwargs)
         elif isinstance(current_ctx, InterpContext):
             trace_ctx = TraceContext(current_ctx.cluster_spec, parent=current_ctx)
             # TODO(jint): should we add trace_and_apply to improve the performance?
-            traced_fn = trace(trace_ctx, fn, *args, **kwargs)
+            tfn = trace(trace_ctx, fn, *args, **kwargs)
             # Return back to the original context.
-            return cast(R, apply(current_ctx, traced_fn, *args, **kwargs))
+            return cast(R, apply(current_ctx, tfn, *args, **kwargs))
         else:
             raise ValueError(f"Unsupported context type: {type(current_ctx)}")
 
     return wrapped
 
 
-function = primitive
+def primitive(fn: Callable[P, R]) -> Callable[P, R]:
+    """Decorator to trace a Python function as an opaque primitive call (`CallExpr`).
+
+    When a function decorated with `@primitive` is called within a `TraceContext`, it is
+    not inlined. Instead, it is traced separately in a forked context, and a `CallExpr`
+    node is inserted into the main graph. This is useful for encapsulating complex
+    operations or third-party library calls as single, opaque nodes.
+
+    **Implementation Note:**
+    A `CallExpr` represents a call to a single inline lambda (non-recursive, as we don't
+    have Y-combinator support). This single lambda call can be treated as a "primitive call"
+    by the printer/visualizer - hence the name "primitive". The function body is captured
+    once during tracing and represented as an opaque callable unit in the expression graph,
+    maintaining a clear boundary between the caller and callee contexts.
+
+    Args:
+        fn: The function to be traced as a primitive operation.
+
+    Returns:
+        A wrapped function that creates a `CallExpr` node when called in a trace context.
+
+    Example:
+        ```python
+        @primitive
+        def my_op(x: MPObject) -> MPObject:
+            # Complex logic traced as a single CallExpr node
+            return x + 1
+        ```
+    """
+    return trace_before_apply(fn, make_call=True)
+
+
+def function(fn: Callable[P, R]) -> Callable[P, R]:
+    """Decorator to trace a Python function by inlining its body.
+
+    When a function decorated with `@function` is called within a `TraceContext`, its
+    underlying primitive operations are expanded and inserted directly into the caller's
+    graph. This is the default tracing behavior and is suitable for most pure-Python
+    multi-party functions.
+
+    Args:
+        fn: The function to be traced and inlined.
+
+    Returns:
+        A wrapped function that inlines its operations into the caller's trace context.
+
+    Example:
+        ```python
+        @function
+        def my_func(x: MPObject, y: MPObject) -> MPObject:
+            # Operations are inlined into the caller's trace
+            return x + y * constant(2)
+        ```
+    """
+    return trace_before_apply(fn, make_call=False)
 
 
 # ============================================================================
@@ -126,18 +203,15 @@ def _tracer() -> TraceContext:
     return ctx
 
 
-@primitive
 def psize() -> int:
     """Get the size of the current party world.
 
     Returns:
         int: The total number of parties in the current multi-party computation context.
     """
-    ctx = _tracer()
-    return ctx.world_size()
+    return cur_ctx().world_size()
 
 
-@primitive
 def pmask() -> Mask:
     """Get the current party mask in this computation context.
 
@@ -145,8 +219,7 @@ def pmask() -> Mask:
         Mask: The current party mask indicating which parties are active
               in the current computation context.
     """
-    ctx = _tracer()
-    return ctx.mask
+    return _tracer().mask
 
 
 @primitive
@@ -203,7 +276,6 @@ def prand(shape: Shape = ()) -> MPObject:
     return out_tree.unflatten(results)  # type: ignore[no-any-return]
 
 
-@primitive
 def constant(data: TensorLike | ScalarType | TableLike) -> MPObject:
     """Create a constant tensor or table from data.
 
@@ -250,7 +322,7 @@ def debug_print(obj: MPObject, prefix: str = "") -> MPObject:
     return out_tree.unflatten(results)  # type: ignore[no-any-return]
 
 
-@primitive
+@function
 def peval(
     pfunc: PFunction,
     args: list[MPObject],
@@ -378,7 +450,7 @@ def set_mask(arg: MPObject, mask: Mask) -> MPObject:
     return out_tree.unflatten(results)  # type: ignore[no-any-return]
 
 
-@primitive
+@function
 def uniform_cond(
     pred: MPObject,
     then_fn: Callable[..., Any],
@@ -588,7 +660,7 @@ def uniform_cond(
     return var_demorph(out_vars, then_tfn.out_imms, then_tfn.out_struct)  # type: ignore[no-any-return]
 
 
-@primitive
+@function
 def while_loop(
     cond_fn: Callable[[Any], MPObject],
     body_fn: Callable[[Any], Any],
@@ -781,7 +853,7 @@ def while_loop(
     return var_demorph(out_vars, body_tfn.out_imms, body_tfn.out_struct)
 
 
-@primitive
+@function
 def pshfl(src: MPObject, index: MPObject) -> MPObject:
     """Shuffle the input tensor to the specified index (dynamic version).
 
@@ -851,7 +923,7 @@ def pshfl(src: MPObject, index: MPObject) -> MPObject:
     return TraceVar(_tracer(), shfl_expr)
 
 
-@primitive
+@function
 def pshfl_s(src_val: MPObject, pmask: Mask, src_ranks: list[Rank]) -> MPObject:
     """Shuffle the input tensor to the specified rank, static version.
 
@@ -910,7 +982,7 @@ def pshfl_s(src_val: MPObject, pmask: Mask, src_ranks: list[Rank]) -> MPObject:
     return TraceVar(_tracer(), shfl_s_expr)
 
 
-@primitive
+@function
 def pconv(vars: list[MPObject]) -> MPObject:
     """Combine multiple variables that share the same dtype and shape into one.
 

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/tests/core/expr/test_ast.py

@@ -233,7 +233,7 @@ class TestCallExpr:
         arg_mptype = MPType.tensor(FLOAT32, (2, 3), pmask_2p)
         arg = VariableExpr("arg", arg_mptype)
 
-        expr = CallExpr(func, [arg])
+        expr = CallExpr("test", func, [arg])
 
         # The call should have the correct structure
         assert isinstance(expr.fn, FuncDefExpr)

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/tests/core/expr/test_printer.py

@@ -321,7 +321,7 @@ class TestPrinterComplexExpressions:
 
     def test_all_expr_types_printing(self, pmask_2p, pfunc_1i1o):
         """Test printing of a complex expression involving all Expr types with meaningful parameter usage."""
-        printer = Printer(compact_format=False)
+        printer = Printer(compact_format=False, inline_pcall=False)
 
         # 1. Variable expressions
         var1_mptype = MPType.tensor(FLOAT32, (2,), pmask_2p)
@@ -357,7 +357,7 @@ class TestPrinterComplexExpressions:
         var_expr = VariableExpr("input_data", param_type)
         func_body = TupleExpr([var_expr, access_expr])
         func_def = FuncDefExpr(["input_data"], func_body)
-        call_expr = CallExpr(func_def, [var1])
+        call_expr = CallExpr("test", func_def, [var1])
 
         # Access the first output of the call expression to get a single-output expr
         call_expr_first = AccessExpr(call_expr, 0)
@@ -405,7 +405,7 @@ class TestPrinterComplexExpressions:
         assert result == expected.strip()
 
         # Test with optimize_variables=True for comparison
-        printer_optimized = Printer(compact_format=True)
+        printer_optimized = Printer(compact_format=True, inline_pcall=False)
         result_optimized = printer_optimized.print_expr(final_expr)
 
         # Expected output with variable optimization (compact mode uses variable names directly)
@@ -469,7 +469,7 @@ class TestPrinterComplexExpressions:
 
     def test_call_expr_printing(self, pmask_2p):
         """Test printing of CallExpr with nested function that uses its parameters."""
-        printer = Printer(compact_format=False)
+        printer = Printer(compact_format=False, inline_pcall=False)
 
         # Create function body that actually uses the parameters
         x_mptype = MPType.tensor(FLOAT32, (1,), pmask_2p)
@@ -489,7 +489,7 @@ class TestPrinterComplexExpressions:
         arg2 = VariableExpr("arg2", arg2_mptype)
 
         # Create call expression
-        expr = CallExpr(fn, [arg1, arg2])
+        expr = CallExpr("test", fn, [arg1, arg2])
 
         result = printer.print_expr(expr)
 
@@ -617,7 +617,7 @@ class TestPrinterEdgeCases:
 
     def test_deep_nesting_indentation(self, pmask_2p):
         """Test printer with deeply nested expressions that use parameters meaningfully."""
-        printer = Printer(compact_format=False)
+        printer = Printer(compact_format=False, inline_pcall=False)
 
         # Create nested function definitions with meaningful parameter usage
         inner_param_type = MPType.tensor(FLOAT32, (1,), pmask_2p)
@@ -630,13 +630,13 @@ class TestPrinterEdgeCases:
 
         # Middle function: takes parameter and calls inner function with it
         middle_param = VariableExpr("middle_param", middle_param_type)
-        middle_body = CallExpr(inner_fn, [middle_param])
+        middle_body = CallExpr("middle", inner_fn, [middle_param])
         middle_fn = FuncDefExpr(["middle_param"], middle_body)
 
         # Outer expression: call middle function with a variable expression
         arg_mptype = MPType.tensor(UINT64, (), pmask_2p)
         arg = VariableExpr("arg", arg_mptype)
-        outer_expr = CallExpr(middle_fn, [arg])
+        outer_expr = CallExpr("outer", middle_fn, [arg])
 
         result = printer.print_expr(outer_expr)
 

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/tests/core/test_mpir.py

@@ -374,7 +374,7 @@ class TestComplexExpressions:
         # Create arguments using VariableExpr
         arg = VariableExpr("input", MPType.tensor(FLOAT32, (2,), pmask=Mask(7)))
 
-        original = CallExpr(fn, [arg])
+        original = CallExpr("original", fn, [arg])
 
         writer = Writer()
         proto = writer.dumps(original)
@@ -887,7 +887,7 @@ class TestComplexExpressionRoundtrip:
             MPType.tensor(FLOAT32, (2,), pmask=Mask(7)),
         )
 
-        original = CallExpr(fn, [arg])
+        original = CallExpr("original", fn, [arg])
 
         # Test roundtrip
         writer = Writer()

{mplang_nightly-0.1.dev169 → mplang_nightly-0.1.dev170}/tests/core/test_primitive.py

@@ -37,11 +37,11 @@ from mplang.core.mptype import Rank
 from mplang.core.primitive import (
     _switch_ctx,
     constant,
+    function,
     pconv,
     peval,
     prand,
     prank,
-    primitive,
     pshfl,
     pshfl_s,
     set_mask,
@@ -76,7 +76,7 @@ class TestPrimitiveDecorator:
     def test_primitive_decorator_basic(self, trace_context):
         """Test basic primitive decorator functionality."""
 
-        @primitive
+        @function
         def simple_func():
             return constant(42)
 
@@ -1156,7 +1156,7 @@ class TestCompleteExample:
     def test_complete_primitive_example(self, trace_context):
         """Complete example showing the complete testing pattern."""
 
-        @primitive
+        @function
         def example_computation():
             """Example multi-party computation function."""
             my_rank = prank()
@@ -1193,7 +1193,7 @@ class TestComplexExpressions:
     def test_complex_function(self, trace_context):
         """Test function combining multiple primitives."""
 
-        @primitive
+        @function
         def complex_func():
             prank()
             prand((2, 2))
@@ -1226,9 +1226,9 @@ class TestComplexExpressions:
     def test_nested_primitives(self, trace_context):
         """Test nested primitive calls."""
 
-        @primitive
+        @function
         def outer_func():
-            @primitive
+            @function
             def inner_func():
                 return constant(1)