mplang-nightly 0.1.dev253__tar.gz → 0.1.dev255__tar.gz

{mplang_nightly-0.1.dev253 → mplang_nightly-0.1.dev255}/PKG-INFO

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

{mplang_nightly-0.1.dev253 → mplang_nightly-0.1.dev255}/mplang/v2/backends/simp_driver/mem.py

@@ -154,6 +154,10 @@ class SimpMemDriver(SimpDriver):
         self._workers = workers
         self._mesh = mesh
 
+    def shutdown(self) -> None:
+        """Shutdown the local memory driver and its mesh."""
+        self._mesh.shutdown()
+
     @property
     def world_size(self) -> int:
         return self._world_size

{mplang_nightly-0.1.dev253 → mplang_nightly-0.1.dev255}/mplang/v2/backends/spu_impl.py

@@ -27,6 +27,7 @@ import spu.api as spu_api
 import spu.libspu as libspu
 
 from mplang.v2.backends.simp_worker import SimpWorker
+from mplang.v2.backends.spu_state import SPUState
 from mplang.v2.backends.tensor_impl import TensorValue
 from mplang.v2.dialects import spu
 from mplang.v2.edsl import serde
@@ -107,84 +108,6 @@ def to_runtime_config(config: spu.SPUConfig) -> libspu.RuntimeConfig:
     return runtime_config
 
 
-# Global cache for SPU runtimes per (local_rank, world_size) pair
-# Key: (local_rank, spu_world_size, protocol, field, link_mode), Value: (Runtime, Io)
-_SPU_RUNTIMES: dict[
-    tuple[int, int, str, str, str], tuple[spu_api.Runtime, spu_api.Io]
-] = {}
-
-
-def _create_mem_link(local_rank: int, spu_world_size: int) -> libspu.link.Context:
-    """Create in-memory link for simulation."""
-    desc = libspu.link.Desc()  # type: ignore
-    desc.recv_timeout_ms = 30 * 1000
-    for i in range(spu_world_size):
-        desc.add_party(f"P{i}", f"mem:{i}")
-    return libspu.link.create_mem(desc, local_rank)
-
-
-def _create_brpc_link(local_rank: int, spu_endpoints: list[str]) -> libspu.link.Context:
-    """Create BRPC link for distributed execution.
-
-    Args:
-        local_rank: The local rank within the SPU device (0-indexed).
-        spu_endpoints: List of BRPC endpoints for all SPU parties.
-
-    Returns:
-        A libspu.link.Context for BRPC communication.
-    """
-    desc = libspu.link.Desc()  # type: ignore
-    desc.recv_timeout_ms = 100 * 1000  # 100 seconds
-    desc.http_max_payload_size = 32 * 1024 * 1024  # 32MB
-
-    for i, endpoint in enumerate(spu_endpoints):
-        desc.add_party(f"P{i}", endpoint)
-
-    return libspu.link.create_brpc(desc, local_rank)
-
-
-def _get_spu_ctx(
-    local_rank: int,
-    spu_world_size: int,
-    config: spu.SPUConfig,
-    spu_endpoints: list[str] | None = None,
-) -> tuple[spu_api.Runtime, spu_api.Io]:
-    """Get or create SPU runtime and IO for the given local rank within SPU.
-
-    Args:
-        local_rank: The local rank within the SPU device (0-indexed).
-        spu_world_size: The number of parties in the SPU device.
-        config: SPU configuration including protocol settings.
-        spu_endpoints: Optional list of BRPC endpoints. If None, use mem link.
-
-    Returns:
-        A tuple of (Runtime, Io) for this party.
-    """
-    # Determine link mode
-    link_mode = "brpc" if spu_endpoints else "mem"
-
-    # Include protocol, field, and link_mode in cache key
-    cache_key = (local_rank, spu_world_size, config.protocol, config.field, link_mode)
-    if cache_key in _SPU_RUNTIMES:
-        return _SPU_RUNTIMES[cache_key]
-
-    # Create Link
-    if spu_endpoints:
-        link = _create_brpc_link(local_rank, spu_endpoints)
-    else:
-        link = _create_mem_link(local_rank, spu_world_size)
-
-    # Use config from SPUConfig
-    runtime_config = to_runtime_config(config)
-
-    # Create Runtime and Io
-    runtime = spu_api.Runtime(link, runtime_config)
-    io = spu_api.Io(spu_world_size, runtime_config)
-
-    _SPU_RUNTIMES[cache_key] = (runtime, io)
-    return runtime, io
-
-
 @spu.makeshares_p.def_impl
 def makeshares_impl(
     interpreter: Interpreter, op: Operation, data: TensorValue
@@ -293,7 +216,15 @@ def exec_impl(interpreter: Interpreter, op: Operation, *args: Any) -> Any:
                 )
             spu_endpoints.append(spu_endpoints_map[party_rank])
 
-    runtime, io = _get_spu_ctx(local_rank, spu_world_size, config, spu_endpoints)
+    # Get or create SPUState for caching Runtime/Io
+    spu_state = interpreter.get_dialect_state(SPUState.dialect_name)
+    if not isinstance(spu_state, SPUState):
+        spu_state = SPUState()
+        interpreter.set_dialect_state(SPUState.dialect_name, spu_state)
+
+    runtime, io = spu_state.get_or_create(
+        local_rank, spu_world_size, config, spu_endpoints
+    )
 
     executable_code = op.attrs["executable"]
     input_names = op.attrs["input_names"]

mplang_nightly-0.1.dev255/mplang/v2/backends/spu_state.py

@@ -0,0 +1,124 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SPU Dialect State.
+
+Manages SPU Runtime lifecycle as a dialect state, enabling reuse across
+multiple executions while binding to the Interpreter's lifecycle.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import spu.api as spu_api
+import spu.libspu as libspu
+
+from mplang.v2.runtime.dialect_state import DialectState
+
+if TYPE_CHECKING:
+    from mplang.v2.dialects import spu
+
+
+class SPUState(DialectState):
+    """SPU Runtime cache as dialect state.
+
+    Caches SPU Runtime and Io objects per (local_rank, world_size, config, link_mode)
+    to enable reuse across multiple SPU kernel executions.
+
+    This replaces the previous global `_SPU_RUNTIMES` cache with a properly
+    lifecycle-managed dialect state.
+    """
+
+    dialect_name: str = "spu"
+
+    def __init__(self) -> None:
+        # Key: (local_rank, world_size, protocol, field, link_mode)
+        # Value: (Runtime, Io)
+        self._runtimes: dict[
+            tuple[int, int, str, str, str], tuple[spu_api.Runtime, spu_api.Io]
+        ] = {}
+
+    def get_or_create(
+        self,
+        local_rank: int,
+        spu_world_size: int,
+        config: spu.SPUConfig,
+        spu_endpoints: list[str] | None = None,
+    ) -> tuple[spu_api.Runtime, spu_api.Io]:
+        """Get or create SPU Runtime and Io for the given configuration.
+
+        Args:
+            local_rank: The local rank within the SPU device (0-indexed).
+            spu_world_size: The number of parties in the SPU device.
+            config: SPU configuration including protocol settings.
+            spu_endpoints: Optional list of BRPC endpoints. If None, use mem link.
+
+        Returns:
+            A tuple of (Runtime, Io) for this party.
+        """
+        from mplang.v2.backends.spu_impl import to_runtime_config
+
+        link_mode = "brpc" if spu_endpoints else "mem"
+        cache_key = (
+            local_rank,
+            spu_world_size,
+            config.protocol,
+            config.field,
+            link_mode,
+        )
+
+        if cache_key in self._runtimes:
+            return self._runtimes[cache_key]
+
+        # Create Link
+        if spu_endpoints:
+            link = self._create_brpc_link(local_rank, spu_endpoints)
+        else:
+            link = self._create_mem_link(local_rank, spu_world_size)
+
+        # Create Runtime and Io
+        runtime_config = to_runtime_config(config)
+        runtime = spu_api.Runtime(link, runtime_config)
+        io = spu_api.Io(spu_world_size, runtime_config)
+
+        self._runtimes[cache_key] = (runtime, io)
+        return runtime, io
+
+    def _create_mem_link(
+        self, local_rank: int, spu_world_size: int
+    ) -> libspu.link.Context:
+        """Create in-memory link for simulation."""
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 30 * 1000
+        for i in range(spu_world_size):
+            desc.add_party(f"P{i}", f"mem:{i}")
+        return libspu.link.create_mem(desc, local_rank)
+
+    def _create_brpc_link(
+        self, local_rank: int, spu_endpoints: list[str]
+    ) -> libspu.link.Context:
+        """Create BRPC link for distributed execution."""
+        desc = libspu.link.Desc()  # type: ignore
+        desc.recv_timeout_ms = 100 * 1000  # 100 seconds
+        desc.http_max_payload_size = 32 * 1024 * 1024  # 32MB
+
+        for i, endpoint in enumerate(spu_endpoints):
+            desc.add_party(f"P{i}", endpoint)
+
+        return libspu.link.create_brpc(desc, local_rank)
+
+    def shutdown(self) -> None:
+        """Clear all cached runtimes."""
+        self._runtimes.clear()

{mplang_nightly-0.1.dev253 → mplang_nightly-0.1.dev255}/mplang/v2/cli.py

@@ -62,7 +62,7 @@ def run_worker(
     signal.signal(signal.SIGINT, signal.SIG_DFL)
     signal.signal(signal.SIGTERM, signal.SIG_DFL)
 
-    from mplang.v2.backends.simp_http_worker import create_worker_app
+    from mplang.v2.backends.simp_worker.http import create_worker_app
 
     app = create_worker_app(rank, world_size, endpoints, spu_endpoints)
 

{mplang_nightly-0.1.dev253 → mplang_nightly-0.1.dev255}/mplang/v2/edsl/context.py

@@ -69,7 +69,7 @@ class Context(ABC):
     """
 
     def __init__(self) -> None:
-        self._state: dict[str, Any] = {}
+        self._states: dict[str, Any] = {}
 
     # =========================================================================
     # State Management
@@ -82,7 +82,7 @@ class Context(ABC):
             key: State key (e.g., "dialect.simp", "device.cluster")
             value: State value
         """
-        self._state[key] = value
+        self._states[key] = value
 
     def get_state(self, key: str, default: Any = None) -> Any:
         """Get attached state by key.
@@ -94,7 +94,7 @@ class Context(ABC):
         Returns:
             State value or default
         """
-        return self._state.get(key, default)
+        return self._states.get(key, default)
 
     def has_state(self, key: str) -> bool:
         """Check if state exists.
@@ -105,7 +105,7 @@ class Context(ABC):
         Returns:
             True if state exists
         """
-        return key in self._state
+        return key in self._states
 
     # =========================================================================
     # Abstract Methods

mplang_nightly-0.1.dev255/mplang/v2/libs/ml/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2025 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Machine Learning algorithms for secure multi-party computation."""
+
+from mplang.v2.libs.ml.sgb import SecureBoost, Tree, TreeEnsemble
+
+__all__ = [
+    "SecureBoost",
+    "Tree",
+    "TreeEnsemble",
+]

{mplang_nightly-0.1.dev253/examples/v2 → mplang_nightly-0.1.dev255/mplang/v2/libs/ml}/sgb.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# mypy: disable-error-code="no-untyped-def,no-any-return,var-annotated"
+
 """SecureBoost v2: Optimized implementation using mplang.v2 low-level BFV APIs.
 
 This implementation improves upon v1 by leveraging BFV SIMD slots and the
@@ -879,7 +881,7 @@ def _find_splits_pps(
         subgroup_map_fn = make_get_subgroup_map(n_level)
         subgroup_map = simp.pcall_static(
             (pp_rank,),
-            lambda fn=subgroup_map_fn: tensor.run_jax(fn, bt_level_pp),
+            lambda fn=subgroup_map_fn, bt_lv=bt_level_pp: tensor.run_jax(fn, bt_lv),
         )
 
         n_pp_features = n_features_per_party[pp_idx + 1]
@@ -978,8 +980,8 @@ def _find_splits_pps(
 
             pp_hists = simp.pcall_static(
                 (ap_rank,),
-                lambda: tensor.run_jax(
-                    derive_right_and_combine, left_hists, parent_hists
+                lambda lh=left_hists, ph=parent_hists: tensor.run_jax(
+                    derive_right_and_combine, lh, ph
                 ),
             )
 
@@ -996,7 +998,7 @@ def _find_splits_pps(
 
         pp_gains, pp_feats, pp_threshs = simp.pcall_static(
             (ap_rank,),
-            lambda: tensor.run_jax(find_splits, pp_hists),
+            lambda h=pp_hists: tensor.run_jax(find_splits, h),
         )
 
         pp_gains_list.append(pp_gains)
@@ -1292,7 +1294,7 @@ def build_tree(
         # Local bt for this level
         bt_level = simp.pcall_static(
             (ap_rank,),
-            lambda off=level_offset: tensor.run_jax(lambda b: b - off, bt),
+            lambda off=level_offset, b=bt: tensor.run_jax(lambda x: x - off, b),
         )
 
         # === AP: Local histogram computation ===
@@ -1352,7 +1354,7 @@ def build_tree(
 
         best_gains, best_party = simp.pcall_static(
             (ap_rank,),
-            lambda: tensor.run_jax(find_global_best, *all_gains),
+            lambda gains=all_gains: tensor.run_jax(find_global_best, *gains),
         )
 
         # === Update Tree State ===
@@ -1498,8 +1500,9 @@ def predict_tree(
             (rank,),
             lambda d=all_datas[i],
             f=tree.feature[i],
-            t=tree.threshold[i]: predict_tree_single_party(
-                d, f, t, tree.is_leaf, tree.owned_party_id, i, n_nodes
+            t=tree.threshold[i],
+            idx=i: predict_tree_single_party(
+                d, f, t, tree.is_leaf, tree.owned_party_id, idx, n_nodes
             ),
         )
         # Transfer to AP
@@ -1564,7 +1567,7 @@ def predict_ensemble(
 
         y_pred_logits = simp.pcall_static(
             (ap_rank,),
-            lambda: tensor.run_jax(update_pred, y_pred_logits, tree_pred),
+            lambda yp=y_pred_logits, tp=tree_pred: tensor.run_jax(update_pred, yp, tp),
         )
 
     # Convert logits to probabilities
@@ -1634,7 +1637,9 @@ def fit_tree_ensemble(
 
         gh, qg, qh = simp.pcall_static(
             (ap_rank,),
-            lambda: tensor.run_jax(compute_gh_quantized, y_data, y_pred, fxp_scale),
+            lambda yp=y_pred: tensor.run_jax(
+                compute_gh_quantized, y_data, yp, fxp_scale
+            ),
         )
 
         # FHE encrypt only if we have passive parties
@@ -1678,7 +1683,7 @@ def fit_tree_ensemble(
 
         y_pred = simp.pcall_static(
             (ap_rank,),
-            lambda: tensor.run_jax(update_pred_fn, y_pred, tree_pred),
+            lambda yp=y_pred, tp=tree_pred: tensor.run_jax(update_pred_fn, yp, tp),
         )
 
     return TreeEnsemble(
@@ -1866,275 +1871,3 @@ class SecureBoost:
             lambda: tensor.run_jax(compute_metrics, y_prob, y_data),
         )
         return accuracy
-
-    # ==============================================================================
-    # Main
-    # ==============================================================================
-
-    print("=" * 60)
-
-
-def run_sgb_demo(sim):
-    """Test SecureBoost with 2 parties, enabling BFV FHE.
-
-    AP (party 0) holds labels + some features.
-    PP (party 1) holds remaining features.
-
-    Flow:
-    1. PP encrypts its histogram buckets with BFV
-    2. AP receives encrypted histogram, multiplies with gradient mask
-    3. PP decrypts and aggregates
-    """
-    import mplang.v2 as mp
-
-    print("=" * 60)
-    print("Multi-Party SecureBoost Test (with BFV FHE)")
-    print("=" * 60)
-
-    # Load BFV backend (registers implementations)
-    from mplang.v2.backends import load_backend
-
-    try:
-        load_backend("mplang.v2.backends.bfv_impl")
-        print("✓ BFV backend loaded successfully")
-    except ImportError as e:
-        print(f"✗ Failed to load BFV backend: {e}")
-        print("  Make sure 'tenseal' is installed: pip install tenseal")
-        return
-
-    # Generate synthetic data for 2 parties
-    np.random.seed(42)
-    n_samples = 100
-    n_features_ap = 3  # AP's features
-    n_features_pp = 2  # PP's features
-
-    # Generate linearly separable data
-    X_all = np.random.randn(n_samples, n_features_ap + n_features_pp).astype(np.float32)
-    y = (X_all[:, 0] + X_all[:, 1] + X_all[:, 3] > 0).astype(np.float32)
-
-    # Split features by party
-    X_ap = X_all[:, :n_features_ap]  # Party 0's features
-    X_pp = X_all[:, n_features_ap:]  # Party 1's features
-
-    print("\nData Distribution:")
-    print(f"  Party 0 (AP): {X_ap.shape} features + labels")
-    print(f"  Party 1 (PP): {X_pp.shape} features (encrypted)")
-
-    # Model parameters
-    n_estimators = 2  # Small for test
-    max_depth = 2
-
-    def job():
-        # Put data on respective parties
-        data_ap = mp.put("P0", X_ap)
-        data_pp = mp.put("P1", X_pp)
-        y_data = mp.put("P0", y)
-
-        # Create model with passive party
-        model = SecureBoost(
-            n_estimators=n_estimators,
-            max_depth=max_depth,
-            learning_rate=0.1,
-            max_bin=8,
-            ap_rank=0,
-            pp_ranks=[1],  # Party 1 is passive party (uses FHE)
-        )
-
-        print("\nTraining with BFV FHE...")
-        model.fit(
-            [data_ap, data_pp],  # Both parties' data
-            y_data,
-            n_samples=n_samples,
-            n_features_per_party=[n_features_ap, n_features_pp],
-        )
-        print("Training complete!")
-
-        # Predict
-        print("\nPredicting...")
-        y_prob = model.predict([data_ap, data_pp], n_samples=n_samples)
-        print(f"Predictions type: {y_prob.type}")
-        return y_prob
-
-    # Execute with 2 parties
-    print("\nExecuting graph with 2 parties...")
-    y_prob_obj = mp.evaluate(job, context=sim)
-
-    # Calculate accuracy
-    y_pred_probs = mp.fetch(y_prob_obj, context=sim)
-    if isinstance(y_pred_probs, list):
-        y_pred_probs = y_pred_probs[0]
-    y_pred_class = (y_pred_probs > 0.5).astype(np.float32)
-    accuracy = np.mean(y_pred_class == y)
-    print(f"\nPredictions (first 10): {y_pred_probs[:10]}")
-    print(f"True labels (first 10): {y[:10]}")
-    print(f"Training Accuracy: {accuracy:.4f}")
-
-    print("\n" + "=" * 60)
-    print("Multi-party test completed!")
-    print("=" * 60)
-
-
-def run_sgb_bench(sim):
-    """Benchmark SecureBoost with BFV FHE for performance analysis."""
-    import time
-
-    print("=" * 70)
-    print("SecureBoost v2 - Multi-Party FHE Performance Benchmark")
-    print("=" * 70)
-    # Benchmark configurations
-    configs = [
-        # Test m > 4096 case (multi-CT support)
-        {
-            "n_samples": 1000,
-            "n_features_ap": 50,
-            "n_features_pp": 50,
-            "n_trees": 2,
-            "max_depth": 3,
-        },
-    ]
-
-    results = []
-
-    for cfg in configs:
-        n_samples = cfg["n_samples"]
-        n_features_ap = cfg["n_features_ap"]
-        n_features_pp = cfg["n_features_pp"]
-        n_trees = cfg["n_trees"]
-        max_depth = cfg["max_depth"]
-
-        print(f"\n{'─' * 70}")
-        print(
-            f"Config: samples={n_samples}, features=({n_features_ap}+{n_features_pp}), "
-            f"trees={n_trees}, depth={max_depth}"
-        )
-        print(f"{'─' * 70}")
-
-        # Generate data
-        np.random.seed(42)
-        X_all = np.random.randn(n_samples, n_features_ap + n_features_pp).astype(
-            np.float32
-        )
-        y = (X_all[:, 0] + X_all[:, 1] > 0).astype(np.float32)
-        X_ap = X_all[:, :n_features_ap]
-        X_pp = X_all[:, n_features_ap:]
-
-        def job():
-            data_ap = mp.put("P0", X_ap)
-            data_pp = mp.put("P1", X_pp)
-            y_data = mp.put("P0", y)
-
-            model = SecureBoost(
-                n_estimators=n_trees,
-                max_depth=max_depth,
-                learning_rate=0.1,
-                max_bin=8,
-                ap_rank=0,
-                pp_ranks=[1],
-            )
-
-            model.fit(
-                [data_ap, data_pp],
-                y_data,
-                n_samples=n_samples,
-                n_features_per_party=[n_features_ap, n_features_pp],
-            )
-
-            y_prob = model.predict([data_ap, data_pp], n_samples=n_samples)
-            return y_prob
-
-        # Measure tracing time
-        t0 = time.perf_counter()
-        traced = mp.compile(job, context=sim)
-        trace_time = time.perf_counter() - t0
-
-        graph = traced.graph
-        n_ops = len(graph.operations)
-
-        print(f"  Tracing:    {trace_time:.3f}s ({n_ops} ops)")
-
-        # Measure execution time
-        t0 = time.perf_counter()
-
-        y_prob_obj = mp.evaluate(traced, context=sim)
-        exec_time = time.perf_counter() - t0
-
-        # Calculate accuracy
-        y_pred_probs = mp.fetch(y_prob_obj, context=sim)
-        if isinstance(y_pred_probs, list):
-            y_pred_probs = y_pred_probs[0]
-        y_pred_class = (y_pred_probs > 0.5).astype(np.float32)
-        accuracy = np.mean(y_pred_class == y)
-
-        print(f"  Execution:  {exec_time:.3f}s")
-        print(f"  Total:      {trace_time + exec_time:.3f}s")
-        print(f"  Accuracy:   {accuracy:.2%}")
-
-        results.append({
-            "config": cfg,
-            "trace_time": trace_time,
-            "exec_time": exec_time,
-            "n_ops": n_ops,
-            "accuracy": accuracy,
-        })
-
-    # Summary
-    print("\n" + "=" * 70)
-    print("BENCHMARK SUMMARY")
-    print("=" * 70)
-    print(f"{'Config':<40} {'Trace':<10} {'Exec':<10} {'Ops':<8} {'Acc':<8}")
-    print("-" * 70)
-    for r in results:
-        cfg = r["config"]
-        cfg_str = f"n={cfg['n_samples']}, f={cfg['n_features_ap']}+{cfg['n_features_pp']}, t={cfg['n_trees']}, d={cfg['max_depth']}"
-        print(
-            f"{cfg_str:<40} {r['trace_time']:.3f}s    {r['exec_time']:.3f}s    {r['n_ops']:<8} {r['accuracy']:.2%}"
-        )
-
-    # Analyze operation distribution for last config
-    print("\n" + "=" * 70)
-    print("OPERATION DISTRIBUTION (last config)")
-    print("=" * 70)
-    from collections import Counter
-
-    op_counts = Counter(op.opcode for op in graph.operations)
-    for op_name, count in sorted(op_counts.items(), key=lambda x: -x[1])[:15]:
-        pct = count / n_ops * 100
-        bar = "█" * int(pct / 2)
-        print(f"  {op_name:<30} {count:>5}  ({pct:>5.1f}%) {bar}")
-
-    # BFV-specific operations
-    bfv_ops = {k: v for k, v in op_counts.items() if k.startswith("bfv.")}
-    if bfv_ops:
-        bfv_total = sum(bfv_ops.values())
-        print(f"\n  BFV operations total: {bfv_total} ({bfv_total / n_ops * 100:.1f}%)")
-        for op_name, count in sorted(bfv_ops.items(), key=lambda x: -x[1]):
-            print(f"    {op_name:<28} {count:>5}")
-
-    # Print profiling results
-    if hasattr(sim, "backend") and getattr(sim.backend, "profiler", None) is not None:
-        sim.backend.profiler.stop(filename_prefix="sgb_trace")
-
-    profiler = mp.get_profiler()
-    profiler.print_summary()  # All ops (includes container overhead)
-
-
-__mp_main__ = run_sgb_demo
-
-
-if __name__ == "__main__":
-    import argparse
-
-    import mplang.v2 as mp
-
-    parser = argparse.ArgumentParser(description="SecureBoost v2 Example")
-    parser.add_argument("--benchmark", action="store_true", help="Run benchmark")
-    parser.add_argument("--profile", action="store_true", help="Enable profiling")
-    args = parser.parse_args()
-
-    # Use high-level Simulator API with profiling if requested
-    sim = mp.make_simulator(2, enable_profiling=args.profile)
-
-    if args.benchmark:
-        run_sgb_bench(sim)
-    else:
-        run_sgb_demo(sim)