mplang-nightly 0.1.dev169__py3-none-any.whl → 0.1.dev171__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}: "
@@ -161,13 +164,9 @@ class Printer(ExprVisitor):
         arg_names = [self._var_name(arg) for arg in expr.args]
         fn_type = expr.pfunc.fn_type
 
-        # for well known builtin functions
-        if fn_type == "builtin.constant":
+        # for well known basic functions
+        if fn_type == "basic.constant":
             return self._print_const(expr.pfunc, expr.mptypes)
-        elif fn_type == "builtin.rank":
-            return self._do_print("prank", [], mptypes=expr.mptypes)
-        elif fn_type == "builtin.prand":
-            return self._do_print("prand", [], mptypes=expr.mptypes)
 
         attrs = {"fn_type": fn_type}
         if expr.pfunc.fn_name:
@@ -209,12 +208,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/pfunc.py

@@ -33,7 +33,7 @@ class PFunction:
 
     PFunction serves as a unified interface for describing single-party computations
     in multi-party computing scenarios. It can represent both:
-    1. Built-in operations (e.g., "spu.makeshares", "builtin.read")
+    1. Built-in operations (e.g., "spu.makeshares", "basic.read")
     2. User-defined programmable functions with custom code
 
     The PFunction accepts a list of typed inputs (TensorType/TableType). For
@@ -47,7 +47,7 @@ class PFunction:
 
     Args:
         fn_type: The type/category identifier of this PFunction, indicating which
-            backend or handler should process it (e.g., "spu.makeshares", "builtin.read",
+            backend or handler should process it (e.g., "spu.makeshares", "basic.read",
             "mlir.stablehlo"). This serves as a routing mechanism for execution.
         ins_info: Type information for input parameters (TensorType or TableType)
         outs_info: Type information for output values (TensorType or TableType)

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,
@@ -47,7 +48,7 @@ from mplang.core.pfunc import PFunction
 from mplang.core.table import TableLike
 from mplang.core.tensor import ScalarType, Shape, TensorLike
 from mplang.core.tracer import TraceContext, TraceVar, trace
-from mplang.ops import builtin
+from mplang.ops import basic
 from mplang.utils.func_utils import var_demorph, var_morph
 
 
@@ -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 bltin_function(fn: Callable[P, R]) -> Callable[P, R]:
+    """Decorator to trace a Python function as an opaque primitive call (`CallExpr`).
+
+    When a function decorated with `@bltin_function` 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
+        @bltin_function
+        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,11 +219,10 @@ 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
+@bltin_function
 def prank() -> MPObject:
     """Multi-party get the rank (party identifier) of each party.
 
@@ -170,12 +243,12 @@ def prank() -> MPObject:
     Note:
         Each party in the current party mask independently produces its own rank value.
     """
-    pfunc, eval_args, out_tree = builtin.rank()
+    pfunc, eval_args, out_tree = basic.rank()
     results = peval(pfunc, eval_args)
     return out_tree.unflatten(results)  # type: ignore[no-any-return]
 
 
-@primitive
+@bltin_function
 def prand(shape: Shape = ()) -> MPObject:
     """Multi-party generate a private random (uint64) tensor with the given shape.
 
@@ -198,12 +271,11 @@ def prand(shape: Shape = ()) -> MPObject:
         private random values. The randomness is local to each party and is
         not shared or revealed to other parties.
     """
-    pfunc, eval_args, out_tree = builtin.prand(shape)
+    pfunc, eval_args, out_tree = basic.prand(shape)
     results = peval(pfunc, eval_args)
     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.
 
@@ -233,24 +305,24 @@ def constant(data: TensorLike | ScalarType | TableLike) -> MPObject:
         Note that the constant primitive is not designed to carry large tables efficiently -
         consider using dedicated table loading mechanisms for substantial datasets.
     """
-    pfunc, eval_args, out_tree = builtin.constant(data)
+    pfunc, eval_args, out_tree = basic.constant(data)
     results = peval(pfunc, eval_args)
     return out_tree.unflatten(results)  # type: ignore[no-any-return]
 
 
-@primitive
+@bltin_function
 def debug_print(obj: MPObject, prefix: str = "") -> MPObject:
     """Print local value of obj on owning parties and pass it through.
 
     Returns the same MPObject value to keep it alive against DCE and to
     support usage like: x = debug_print(x, prefix="x=").
     """
-    pfunc, eval_args, out_tree = builtin.debug_print(obj, prefix=prefix)
+    pfunc, eval_args, out_tree = basic.debug_print(obj, prefix=prefix)
     results = peval(pfunc, eval_args)
     return out_tree.unflatten(results)  # type: ignore[no-any-return]
 
 
-@primitive
+@function
 def peval(
     pfunc: PFunction,
     args: list[MPObject],
@@ -373,12 +445,12 @@ def set_mask(arg: MPObject, mask: Mask) -> MPObject:
         The underlying implementation uses JAX identity function with the
         specified execution mask.
     """
-    pfunc, eval_args, out_tree = builtin.identity(arg)
+    pfunc, eval_args, out_tree = basic.identity(arg)
     results = peval(pfunc, eval_args, mask)
     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/device.py

@@ -35,7 +35,7 @@ from mplang.core import InterpContext, MPObject, primitive
 from mplang.core.cluster import ClusterSpec, Device
 from mplang.core.context_mgr import cur_ctx
 from mplang.core.tensor import TensorType
-from mplang.ops import builtin, crypto, ibis_cc, jax_cc, tee
+from mplang.ops import basic, crypto, ibis_cc, jax_cc, tee
 from mplang.ops.base import FeOperation
 from mplang.ops.ibis_cc import IbisRunner
 from mplang.ops.jax_cc import JaxRunner
@@ -209,11 +209,11 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         sess_p, sess_t = _ensure_tee_session(frm_dev_id, to_dev_id, frm_rank, tee_rank)
         # Bytes-only path: pack -> enc -> p2p -> dec -> unpack (with static out type)
         obj_ty = TensorType.from_obj(obj)
-        b = simp.runAt(frm_rank, builtin.pack)(obj)
+        b = simp.runAt(frm_rank, basic.pack)(obj)
         ct = simp.runAt(frm_rank, crypto.enc)(b, sess_p)
         ct_at_tee = mpi.p2p(frm_rank, tee_rank, ct)
         b_at_tee = simp.runAt(tee_rank, crypto.dec)(ct_at_tee, sess_t)
-        pt_at_tee = simp.runAt(tee_rank, builtin.unpack)(b_at_tee, out_ty=obj_ty)
+        pt_at_tee = simp.runAt(tee_rank, basic.unpack)(b_at_tee, out_ty=obj_ty)
         return tree_map(partial(_set_devid, dev_id=to_dev_id), pt_at_tee)  # type: ignore[no-any-return]
     elif frm_to_pair == ("TEE", "PPU"):
         # Transparent encryption from TEE to a specific PPU using the reverse-direction session key
@@ -223,11 +223,11 @@ def _d2d(to_dev_id: str, obj: MPObject) -> MPObject:
         # Ensure bidirectional session established for this pair
         sess_p, sess_t = _ensure_tee_session(to_dev_id, frm_dev_id, ppu_rank, tee_rank)
         obj_ty = TensorType.from_obj(obj)
-        b = simp.runAt(tee_rank, builtin.pack)(obj)
+        b = simp.runAt(tee_rank, basic.pack)(obj)
         ct = simp.runAt(tee_rank, crypto.enc)(b, sess_t)
         ct_at_ppu = mpi.p2p(tee_rank, ppu_rank, ct)
         b_at_ppu = simp.runAt(ppu_rank, crypto.dec)(ct_at_ppu, sess_p)
-        pt_at_ppu = simp.runAt(ppu_rank, builtin.unpack)(b_at_ppu, out_ty=obj_ty)
+        pt_at_ppu = simp.runAt(ppu_rank, basic.unpack)(b_at_ppu, out_ty=obj_ty)
         return tree_map(partial(_set_devid, dev_id=to_dev_id), pt_at_ppu)  # type: ignore[no-any-return]
     else:
         supported = [

mplang/kernels/{builtin.py → basic.py}

@@ -25,17 +25,17 @@ from mplang.runtime.data_providers import get_provider, resolve_uri
 from mplang.utils import table_utils
 
 
-@kernel_def("builtin.identity")
+@kernel_def("basic.identity")
 def _identity(pfunc: PFunction, value: Value) -> Value:
     # Runtime guarantees exactly one argument; no extra arity checks here.
     return value
 
 
-@kernel_def("builtin.read")
+@kernel_def("basic.read")
 def _read(pfunc: PFunction) -> Value:
     path = pfunc.attrs.get("path")
     if path is None:
-        raise ValueError("missing path attr for builtin.read")
+        raise ValueError("missing path attr for basic.read")
     out_t = pfunc.outs_info[0]
     uri = resolve_uri(str(path))
     prov = get_provider(uri.scheme)
@@ -45,7 +45,7 @@ def _read(pfunc: PFunction) -> Value:
     try:
         data = prov.read(uri, out_t, ctx=ctx)
     except Exception as e:  # pragma: no cover - provider errors
-        raise RuntimeError(f"builtin.read failed: {e}") from e
+        raise RuntimeError(f"basic.read failed: {e}") from e
 
     if isinstance(out_t, TableType):
         if isinstance(data, TableValue):
@@ -56,15 +56,15 @@ def _read(pfunc: PFunction) -> Value:
             return data
         return TensorValue(np.asarray(data))
     raise TypeError(
-        f"builtin.read only supports TableType/TensorType outputs, got {type(out_t).__name__}"
+        f"basic.read only supports TableType/TensorType outputs, got {type(out_t).__name__}"
     )
 
 
-@kernel_def("builtin.write")
+@kernel_def("basic.write")
 def _write(pfunc: PFunction, obj: Value) -> Value:
     path = pfunc.attrs.get("path")
     if path is None:
-        raise ValueError("missing path attr for builtin.write")
+        raise ValueError("missing path attr for basic.write")
     uri = resolve_uri(str(path))
     prov = get_provider(uri.scheme)
     if prov is None:
@@ -74,16 +74,16 @@ def _write(pfunc: PFunction, obj: Value) -> Value:
     try:
         prov.write(uri, obj, ctx=ctx)
     except Exception as e:  # pragma: no cover
-        raise RuntimeError(f"builtin.write failed: {e}") from e
+        raise RuntimeError(f"basic.write failed: {e}") from e
     return obj
 
 
-@kernel_def("builtin.constant")
+@kernel_def("basic.constant")
 def _constant(pfunc: PFunction) -> Value:
     """Return constants as Value types (TensorValue or TableValue)."""
     data_bytes = pfunc.attrs.get("data_bytes")
     if data_bytes is None:
-        raise ValueError("missing data_bytes attr for builtin.constant")
+        raise ValueError("missing data_bytes attr for basic.constant")
     out_t = pfunc.outs_info[0]
     fmt = pfunc.attrs.get("data_format")
     if isinstance(out_t, TableType):
@@ -98,7 +98,7 @@ def _constant(pfunc: PFunction) -> Value:
     return TensorValue(arr)
 
 
-@kernel_def("builtin.rank")
+@kernel_def("basic.rank")
 def _rank(pfunc: PFunction) -> TensorValue:
     """Return rank as TensorValue."""
     ctx = cur_kctx()
@@ -106,7 +106,7 @@ def _rank(pfunc: PFunction) -> TensorValue:
     return TensorValue(arr)
 
 
-@kernel_def("builtin.prand")
+@kernel_def("basic.prand")
 def _prand(pfunc: PFunction) -> TensorValue:
     """Return random data as TensorValue."""
     shape = pfunc.attrs.get("shape", ())
@@ -118,7 +118,7 @@ def _prand(pfunc: PFunction) -> TensorValue:
     return TensorValue(data)
 
 
-@kernel_def("builtin.table_to_tensor")
+@kernel_def("basic.table_to_tensor")
 def _table_to_tensor(pfunc: PFunction, table: TableValue) -> TensorValue:
     """Convert table to tensor, return as TensorValue."""
     arrow_table = table.to_arrow()
@@ -131,7 +131,7 @@ def _table_to_tensor(pfunc: PFunction, table: TableValue) -> TensorValue:
     return TensorValue(mat)
 
 
-@kernel_def("builtin.tensor_to_table")
+@kernel_def("basic.tensor_to_table")
 def _tensor_to_table(pfunc: PFunction, tensor: TensorValue) -> TableValue:
     """Convert tensor to table, return as TableValue."""
     import pyarrow as pa  # type: ignore
@@ -168,7 +168,7 @@ def _summ(v: Value) -> str:
         return f"<unprintable {type(v).__name__}: {e}>"
 
 
-@kernel_def("builtin.debug_print")
+@kernel_def("basic.debug_print")
 def _debug_print(pfunc: PFunction, val: Value) -> Value:
     prefix = pfunc.attrs.get("prefix", "")
     ctx = cur_kctx()
@@ -176,16 +176,16 @@ def _debug_print(pfunc: PFunction, val: Value) -> Value:
     return val
 
 
-@kernel_def("builtin.pack")
+@kernel_def("basic.pack")
 def _pack(pfunc: PFunction, value: Value) -> TensorValue:
     outs_info = pfunc.outs_info
     if len(outs_info) != 1:
-        raise ValueError("builtin.pack expects single output type")
+        raise ValueError("basic.pack expects single output type")
     out_ty = outs_info[0]
     if not isinstance(out_ty, TensorType):
-        raise TypeError("builtin.pack must return TensorType")
+        raise TypeError("basic.pack must return TensorType")
     if out_ty.dtype.numpy_dtype() != np.uint8:
-        raise TypeError("builtin.pack output dtype must be uint8")
+        raise TypeError("basic.pack output dtype must be uint8")
 
     if isinstance(value, TableValue):
         # Serialize Arrow table using IPC stream for consistency with Value serde
@@ -203,14 +203,14 @@ def _pack(pfunc: PFunction, value: Value) -> TensorValue:
         arr = value.to_numpy()
         return TensorValue(np.frombuffer(arr.tobytes(order="C"), dtype=np.uint8))
 
-    raise TypeError(f"builtin.pack does not support Value type {type(value).__name__}")
+    raise TypeError(f"basic.pack does not support Value type {type(value).__name__}")
 
 
-@kernel_def("builtin.unpack")
+@kernel_def("basic.unpack")
 def _unpack(pfunc: PFunction, packed: TensorValue) -> Value:
     outs_info = pfunc.outs_info
     if len(outs_info) != 1:
-        raise ValueError("builtin.unpack expects single output type")
+        raise ValueError("basic.unpack expects single output type")
     out_ty = outs_info[0]
 
     b = packed.to_numpy().astype(np.uint8, copy=False).reshape(-1)
@@ -219,7 +219,7 @@ def _unpack(pfunc: PFunction, packed: TensorValue) -> Value:
         np_dtype = out_ty.dtype.numpy_dtype()
         shape = tuple(out_ty.shape)
         if any(dim < 0 for dim in shape):
-            raise ValueError("builtin.unpack does not support dynamic tensor shapes")
+            raise ValueError("basic.unpack does not support dynamic tensor shapes")
         elem_count = int(np.prod(shape))
         expected = elem_count * np.dtype(np_dtype).itemsize
         if b.size != expected:
@@ -239,4 +239,4 @@ def _unpack(pfunc: PFunction, packed: TensorValue) -> Value:
         table = reader.read_all()
         return TableValue(table)
 
-    raise TypeError("builtin.unpack output type must be TensorType or TableType")
+    raise TypeError("basic.unpack output type must be TensorType or TableType")

mplang/kernels/context.py

@@ -35,7 +35,7 @@ def _ensure_impl_imported() -> None:
     global _IMPL_IMPORTED
     if _IMPL_IMPORTED:
         return
-    from mplang.kernels import builtin as _impl_builtin  # noqa: F401
+    from mplang.kernels import basic as _impl_basic  # noqa: F401
     from mplang.kernels import crypto as _impl_crypto  # noqa: F401
     from mplang.kernels import mock_tee as _impl_tee  # noqa: F401
     from mplang.kernels import phe as _impl_phe  # noqa: F401
@@ -49,18 +49,18 @@ def _ensure_impl_imported() -> None:
 # imports consolidated above
 
 _DEFAULT_BINDINGS: dict[str, str] = {
-    # builtin
-    "builtin.identity": "builtin.identity",
-    "builtin.read": "builtin.read",
-    "builtin.write": "builtin.write",
-    "builtin.constant": "builtin.constant",
-    "builtin.rank": "builtin.rank",
-    "builtin.prand": "builtin.prand",
-    "builtin.table_to_tensor": "builtin.table_to_tensor",
-    "builtin.tensor_to_table": "builtin.tensor_to_table",
-    "builtin.debug_print": "builtin.debug_print",
-    "builtin.pack": "builtin.pack",
-    "builtin.unpack": "builtin.unpack",
+    # basic
+    "basic.identity": "basic.identity",
+    "basic.read": "basic.read",
+    "basic.write": "basic.write",
+    "basic.constant": "basic.constant",
+    "basic.rank": "basic.rank",
+    "basic.prand": "basic.prand",
+    "basic.table_to_tensor": "basic.table_to_tensor",
+    "basic.tensor_to_table": "basic.tensor_to_table",
+    "basic.debug_print": "basic.debug_print",
+    "basic.pack": "basic.pack",
+    "basic.unpack": "basic.unpack",
     # crypto
     "crypto.keygen": "crypto.keygen",
     "crypto.enc": "crypto.enc",

mplang/ops/__init__.py

@@ -19,12 +19,12 @@ This module contains compilers that transform high-level functions into
 portable, serializable intermediate representations.
 """
 
-from mplang.ops import builtin, crypto, ibis_cc, jax_cc, phe, spu, sql_cc, tee
+from mplang.ops import basic, crypto, ibis_cc, jax_cc, phe, spu, sql_cc, tee
 from mplang.ops.base import FeOperation as FeOperation
 
 __all__ = [
     "FeOperation",
-    "builtin",
+    "basic",
     "crypto",
     "ibis_cc",
     "jax_cc",

mplang/ops/base.py

@@ -129,7 +129,7 @@ class FeModule(ABC):
         - You need compilation/stateful behavior/dynamic routing, multiple PFunctions, or complex capture flows.
 
     Tips:
-    - Keep routing information in PFunction.fn_type (e.g., "builtin.read", "sql.run", "mlir.stablehlo").
+    - Keep routing information in PFunction.fn_type (e.g., "basic.read", "sql.run", "mlir.stablehlo").
     - Avoid backend-specific logic in kernels; only validate and shape types.
     - Prefer keyword-only attributes in typed_op kernels for clarity (def op(x: MPObject, *, attr: int)).
     """

mplang/ops/{builtin.py → basic.py}

@@ -23,10 +23,10 @@ from mplang.core.tensor import ScalarType, Shape, TensorLike, TensorType
 from mplang.ops.base import stateless_mod
 from mplang.utils import table_utils
 
-_BUILTIN_MOD = stateless_mod("builtin")
+_BASIC_MOD = stateless_mod("basic")
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def identity(x: TensorType) -> TensorType:
     """Return the input type unchanged.
 
@@ -40,7 +40,7 @@ def identity(x: TensorType) -> TensorType:
     return x
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def read(*, path: str, ty: TensorType) -> TensorType:
     """Declare reading a value of type ``ty`` from ``path`` (type-only).
 
@@ -63,7 +63,7 @@ def read(*, path: str, ty: TensorType) -> TensorType:
     return ty
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def write(x: TensorType, *, path: str) -> TensorType:
     """Declare writing the input value to ``path`` and return the same type.
 
@@ -77,7 +77,7 @@ def write(x: TensorType, *, path: str) -> TensorType:
     return x
 
 
-@_BUILTIN_MOD.op_def()
+@_BASIC_MOD.op_def()
 def constant(
     data: TensorLike | ScalarType | TableLike,
 ) -> tuple[PFunction, list[MPObject], PyTreeDef]:
@@ -89,7 +89,7 @@ def constant(
 
     Returns:
         Tuple[PFunction, list[MPObject], PyTreeDef]:
-        - PFunction: ``fn_type='builtin.constant'`` with one output whose type
+        - PFunction: ``fn_type='basic.constant'`` with one output whose type
             matches ``data``; payload serialized via ``data_bytes`` with
             ``data_format`` ('bytes[numpy]' or 'bytes[csv]').
         - list[MPObject]: Empty (no inputs captured).
@@ -120,7 +120,7 @@ def constant(
         data_format = "bytes[numpy]"
 
     pfunc = PFunction(
-        fn_type="builtin.constant",
+        fn_type="basic.constant",
         ins_info=(),
         outs_info=(out_type,),
         data_bytes=data_bytes,
@@ -130,7 +130,7 @@ def constant(
     return pfunc, [], treedef
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def rank() -> TensorType:
     """Return the scalar UINT64 tensor type for the current party rank.
 
@@ -140,7 +140,7 @@ def rank() -> TensorType:
     return TensorType(UINT64, ())
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def prand(*, shape: Shape = ()) -> TensorType:
     """Declare a private random UINT64 tensor with the given shape.
 
@@ -153,7 +153,7 @@ def prand(*, shape: Shape = ()) -> TensorType:
     return TensorType(UINT64, shape)
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def debug_print(
     x: TensorType | TableType, *, prefix: str = ""
 ) -> TableType | TensorType:
@@ -169,7 +169,7 @@ def debug_print(
     return x
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def pack(x: TensorType | TableType) -> TensorType:
     """Serialize a tensor/table into a byte vector (type-only).
 
@@ -189,7 +189,7 @@ def pack(x: TensorType | TableType) -> TensorType:
     return TensorType(UINT8, (-1,))
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def unpack(b: TensorType, *, out_ty: TensorType | TableType) -> TensorType | TableType:
     """Deserialize a byte vector into the explicit output type.
 
@@ -215,7 +215,7 @@ def unpack(b: TensorType, *, out_ty: TensorType | TableType) -> TensorType | Tab
     return out_ty
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def table_to_tensor(table: TableType, *, number_rows: int) -> TensorType:
     """Convert a homogeneous-typed table to a dense 2D tensor.
 
@@ -248,7 +248,7 @@ def table_to_tensor(table: TableType, *, number_rows: int) -> TensorType:
     return TensorType(first, shape)  # type: ignore[arg-type]
 
 
-@_BUILTIN_MOD.simple_op()
+@_BASIC_MOD.simple_op()
 def tensor_to_table(tensor: TensorType, *, column_names: list[str]) -> TableType:
     """Convert a rank-2 tensor into a table with named columns.
 

mplang/simp/__init__.py

@@ -105,7 +105,7 @@ def run_impl(
         The result of evaluating the function through the appropriate handler
 
     Raises:
-        ValueError: If builtin.write is called without required arguments
+        ValueError: If basic.write is called without required arguments
         TypeError: If the function compilation or evaluation fails
         RuntimeError: If the underlying peval execution encounters errors
 
@@ -114,11 +114,11 @@ def run_impl(
 
         >>> tensor_info = TensorType(shape=(10, 10), dtype=np.float32)
         >>> attrs = {"format": "binary"}
-        >>> result = run_impl(builtin.read, "data/input.bin", tensor_info, attrs)
+        >>> result = run_impl(basic.read, "data/input.bin", tensor_info, attrs)
 
         Writing data to a file:
 
-        >>> run_impl(builtin.write, data, "data/output.bin")
+        >>> run_impl(basic.write, data, "data/output.bin")
 
         Running a JAX function:
 

{mplang_nightly-0.1.dev169.dist-info → mplang_nightly-0.1.dev171.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mplang-nightly
-Version: 0.1.dev169
+Version: 0.1.dev171
 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.dev171.dist-info}/RECORD

@@ -1,6 +1,6 @@
 mplang/__init__.py,sha256=ofO-F-CNoVIxpMpTJtTJoQtKegJcHwcOJLzoVispiyc,1852
 mplang/api.py,sha256=ssmv0_CyZPFORhOUJ84Jo6NwRJSK7_Ono3n7ZjEg4sA,3058
-mplang/device.py,sha256=7X2rPp3hnHFVAYrvBLqdB5XEbURd3VGzXWbw-6bJ1sU,12488
+mplang/device.py,sha256=SgR3LF-vno-TeogSz9zplUpDspogbBm9ZhLnkAO-mI4,12478
 mplang/analysis/__init__.py,sha256=CTHFvRsi-nFngojqjn08UaR3RY9i7CJ7T2UdR95kCrk,1056
 mplang/analysis/diagram.py,sha256=ffwgD12gL1_KH1uJ_EYkjmIlDrfxYJJkWj-wHl09_Xk,19520
 mplang/core/__init__.py,sha256=lWxlEKfRwX7FNDzgyKZ1fiDMaCiqkyg0j5mKlZD_v7g,2244
@@ -10,26 +10,26 @@ 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/pfunc.py,sha256=WOGmMr4HCUELED5QxYbhhyQJRDXrA5Bk3tPbZWpwmw8,5102
+mplang/core/primitive.py,sha256=vu60-k0fSAWWidcWDC0_FGvrRZww12oGXjB8CR9F6Yo,43889
 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=Ec6tCLtOUYqu0i1ZmtRvuSLjGpMqB30SM5CYZ_l2CqA,9660
+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
 mplang/kernels/__init__.py,sha256=eooIUklLSg-cvyGk6uDSwZ3bUAjM6AXtHw_YdbUamYo,1052
 mplang/kernels/base.py,sha256=-YV4Aj5fs6GT4ehS6Tyi8WQ-amxn5edHTFJRQzyjHXY,3826
-mplang/kernels/builtin.py,sha256=4XhIu7yLGuVMJeED207ZGtecwV6zvQ2Zg61Nm5DK2co,9015
-mplang/kernels/context.py,sha256=n-Z7fz7HjHb3UY380iZcasmn2sK-OQUGEIWJk2-fT18,13602
+mplang/kernels/basic.py,sha256=thE4jAoozsk_3_t7ahyQj8J9QQlC79CyNRPCqHnyQhk,8967
+mplang/kernels/context.py,sha256=OTXuqZ8ziu1fVXem6lIn4DMygOAeSgySm_-dfzYakEA,13552
 mplang/kernels/crypto.py,sha256=y5epCht71QenQnSbn5xRB0DCnb55Wm83iBJ1KRUedUU,4323
 mplang/kernels/mock_tee.py,sha256=IMTIy5-tEMqB8bD1FG0Ki5UcH8dLAQUFIlFnktXUDX0,2492
 mplang/kernels/phe.py,sha256=xglWuxPiclKVL1_YHgeN4KSUudqKwP5c4IskB-QjC1Y,72711
@@ -37,9 +37,9 @@ mplang/kernels/spu.py,sha256=vh-WG-uDVPvK11CDzc-f58sUalGts_eWJPjeuxLANfY,12508
 mplang/kernels/sql_duckdb.py,sha256=vq4UCth_PCsH8dxcpx7mMnsq54KtjD8kxwISH7tj3qg,1631
 mplang/kernels/stablehlo.py,sha256=gQyg-2ANyA1TjRg90MZ79mn4cHoXhU7g5GFUCuYXyKs,3231
 mplang/kernels/value.py,sha256=2NZ0UvaLyRYb3bTCqL_fQXMzbHk1qbQ3j9xiv4v5E0A,20726
-mplang/ops/__init__.py,sha256=ryPlNziAfc4rwVazdJ6VG7kuUD0bqlRZ2smY5HNIRGI,1018
-mplang/ops/base.py,sha256=h67_SHWNZGUuTCuMll-9kDgGvlPhlFov7WAQCHTmUvw,18258
-mplang/ops/builtin.py,sha256=D7T8rRF9g05VIw9T72lsncF5cDQqaT37eapBieRKvRI,9363
+mplang/ops/__init__.py,sha256=QIRs49KTG-vUavhUiRHuYFZWrlKh97kwb2nVGHmxII0,1014
+mplang/ops/base.py,sha256=HohQ0I39dda576mMhlzQp53oOWdRdUTsvn42t8QP_68,18256
+mplang/ops/basic.py,sha256=Mj_4x9SUs0m0F3lzZYU-LIzDdBPyrsDvPJSrUUkIzPE,9333
 mplang/ops/crypto.py,sha256=9CeFJrYmvjmgx-3WQl6jHXh8VafRpT4QBunbzsPF8Uc,3646
 mplang/ops/ibis_cc.py,sha256=a5OqZVRZ1NzugQPYigdlJcGKbMZHqKh1xkiJen-LtCU,4242
 mplang/ops/jax_cc.py,sha256=kVhJM8i8oPd-yPqyeaZ1hfVxcZPzNhTwjhltDh50hyY,7809
@@ -64,7 +64,7 @@ mplang/runtime/link_comm.py,sha256=ZHNcis8QDu2rcyyF3rhpxMiJDkczoMA_c0iZ2GDW_bA,2
 mplang/runtime/server.py,sha256=CdmBmpbylEl7XeZj26i0rUmTrPTvl2CVdRgbtR02gcg,16543
 mplang/runtime/session.py,sha256=I2711V-pPRCYibNgBhjboUURdubnL6ltCoh5RvFVabs,10641
 mplang/runtime/simulation.py,sha256=1I_8dIqxivxtYnQK0ofz0oXk3lXXh3-zN0lmNFnWucA,11615
-mplang/simp/__init__.py,sha256=U-JUvaO3EEg0cvEtU1Zihh674vJKQBCDy19gI9t1f-0,11581
+mplang/simp/__init__.py,sha256=_X1kpq9qhoPUL2gQRVNkSjS5jNSxYRu7-_1GSwJ9PK8,11575
 mplang/simp/mpi.py,sha256=Wv_Q16TQ3rdLam6OzqXiefIGSMmagGkso09ycyOkHEs,4774
 mplang/simp/random.py,sha256=7PVgWNL1j7Sf3MqT5PRiWplUu-0dyhF3Ub566iqX86M,3898
 mplang/simp/smpc.py,sha256=tdH54aU4T-GIDPhpmf9NCeJC0G67PdOYc04cyUkOnwE,7119
@@ -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.dev171.dist-info/METADATA,sha256=Lgh9kvCVpPpkkDPlGylgBfo82Kweei-CFce9B2w8sMI,16547
+mplang_nightly-0.1.dev171.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mplang_nightly-0.1.dev171.dist-info/entry_points.txt,sha256=mG1oJT-GAjQR834a62_QIWb7litzWPPyVnwFqm-rWuY,55
+mplang_nightly-0.1.dev171.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mplang_nightly-0.1.dev171.dist-info/RECORD,,