mplang-nightly 0.1.dev169__py3-none-any.whl → 0.1.dev170__py3-none-any.whl

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/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/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/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/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/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.dist-info → mplang_nightly-0.1.dev170.dist-info}/METADATA

@@ -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.dist-info → mplang_nightly-0.1.dev170.dist-info}/RECORD

@@ -10,19 +10,19 @@ mplang/core/context_mgr.py,sha256=R0QJAod-1nYduVoOknLfAsxZiy-RtmuQcp-07HABYZU,15
 mplang/core/dtype.py,sha256=0rZqFaFikFu9RxtdO36JLEgFL-E-lo3hH10whwkTVVY,10213
 mplang/core/interp.py,sha256=JKjKJGWURU5rlHQ2yG5XNKWzN6hLZsmo--hZuveQgxI,5915
 mplang/core/mask.py,sha256=14DFxaA446lGjN4dzTuQgm9Shcn34rYI87YJHg0YGNQ,10693
-mplang/core/mpir.py,sha256=V6S9RqegaI0yojhLkHla5nGBi27ASoxlrEs1k4tGubM,37980
+mplang/core/mpir.py,sha256=3NyHa1cDnUaw3XWIUgyOMXfZ9JS-30COb29AoXYcRtM,38251
 mplang/core/mpobject.py,sha256=0pHSd7SrAFTScCFcB9ziDztElYQn-oIZOKBx47B3QX0,3732
 mplang/core/mptype.py,sha256=7Cp2e58uUX-uqTp6QxuioOMJ8BzLBPXlWG5rRakv2uo,13773
 mplang/core/pfunc.py,sha256=PAr8qRhVveWO5HOI0TgdsWjpi4PFi2iEyuTlr9UVKSY,5106
-mplang/core/primitive.py,sha256=-IkGqdbwtbMkLEOOTghXfuFtFvxu5jFQBupm5nPV-RI,40569
+mplang/core/primitive.py,sha256=MxnHr12BorscFiz_sv-iVAbFWd4mdy-WYSg4ORtspOM,43871
 mplang/core/table.py,sha256=BqTBZn7Tfwce4vzl3XYhaX5hVmKagVq9-YoERDta6d8,5892
 mplang/core/tensor.py,sha256=86u6DogSZMoL0w5XjtTmQm2PhA_VjwybN1b6U4Zzphg,2361
 mplang/core/tracer.py,sha256=dVMfUeCMmPz4o6tLXewGMW1Kpy5gpZORvr9w4MhwDtM,14288
 mplang/core/expr/__init__.py,sha256=qwiSTUOcanFJLyK8HZ13_L1ZDrybqpPXIlTHAyeumE8,1988
-mplang/core/expr/ast.py,sha256=KE46KTtlH9RA2V_EzWVKCKolsycgTmt7SotUrOc8Qxs,20923
-mplang/core/expr/evaluator.py,sha256=EFy71vYUL2xLHCtMkWlYJpyGyujDdVSAx8ByET-62qQ,23297
-mplang/core/expr/printer.py,sha256=VblKGnO0OUfzH7EBkszwRNxQUB8QyyC7BlJWJEUv9so,9546
-mplang/core/expr/transformer.py,sha256=TyL-8FjrVvDq_C9X7kAuKkiqt2XdZM-okjzVQj0A33s,4893
+mplang/core/expr/ast.py,sha256=K-rNqlpgkdjVzwSrLgunYnL4zWl1USJGLOgfz0qJNO4,20959
+mplang/core/expr/evaluator.py,sha256=rpzZQPPVtxBvUuCx-9_bFmzr_7tfAQjPlP_rqpWjgIo,23313
+mplang/core/expr/printer.py,sha256=RIHoG0gTDr6jPke391KekmPLu4AdeIfrhlROWa0XseQ,9883
+mplang/core/expr/transformer.py,sha256=gez9eedVsWoLasSgWvPmGR8WfQnGXPlldWeVFEjqyYo,4904
 mplang/core/expr/utils.py,sha256=VDTJ_-CsdHtVy9wDaGa7XdFxQ7o5lYYaeqcgsAhkbNI,2625
 mplang/core/expr/visitor.py,sha256=2Ge-I5N-wH8VVXy8d2WyNaEv8x6seiRx9peyH9S2BYU,2044
 mplang/core/expr/walk.py,sha256=lXkGJEEuvKGDqQihbxXPxfz2RfR1Q1zYUlt11iooQW0,11889
@@ -73,8 +73,8 @@ mplang/utils/crypto.py,sha256=rvPomBFtznRHc3RPi6Aip9lsU8zW2oxBqGv1K3vn7Rs,1052
 mplang/utils/func_utils.py,sha256=vCJcZmu0bEbqhOQKdpttV2_MBllIcPSN0b8U4WjNGGo,5164
 mplang/utils/spu_utils.py,sha256=S3L9RBkBe2AvSuMSQQ12cBY5Y1NPthubvErSX_7nj1A,4158
 mplang/utils/table_utils.py,sha256=aC-IZOKkSmFkpr3NZchLM0Wt0GOn-rg_xHBHREWBwAU,2202
-mplang_nightly-0.1.dev169.dist-info/METADATA,sha256=3Ml9Mvi3n9iBnvcVp7dp7lFJABk1hUoIP5AB5BbeQFE,16547
-mplang_nightly-0.1.dev169.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mplang_nightly-0.1.dev169.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
-mplang_nightly-0.1.dev169.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mplang_nightly-0.1.dev169.dist-info/RECORD,,
+mplang_nightly-0.1.dev170.dist-info/METADATA,sha256=xh004xC6U_s5O350oQsAdZ3yFFJ5BvZz1dw2Kw3Tzuw,16547
+mplang_nightly-0.1.dev170.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev170.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev170.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev170.dist-info/RECORD,,