JSTprove 1.0.0__py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

python/core/model_processing/onnx_quantizer/onnx_op_quantizer.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+from typing import Callable
+
+import onnx
+
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    HandlerImplementationError,
+    MissingHandlerError,
+    UnsupportedOpError,
+)
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    PassthroughQuantizer,
+    ScaleConfig,
+)
+from python.core.model_processing.onnx_quantizer.layers.constant import (
+    ConstantQuantizer,
+)
+from python.core.model_processing.onnx_quantizer.layers.conv import ConvQuantizer
+from python.core.model_processing.onnx_quantizer.layers.gemm import GemmQuantizer
+from python.core.model_processing.onnx_quantizer.layers.maxpool import MaxpoolQuantizer
+from python.core.model_processing.onnx_quantizer.layers.relu import ReluQuantizer
+
+
+class ONNXOpQuantizer:
+    """
+    Registry for ONNX operator quantizers.
+    This should be used to obtain the quantized
+    layer based on any provided operation of that layer type
+
+    Attributes
+    ----------
+    handlers : Dict[str, Callable]
+        Maps ONNX op_type strings to quantizer handler instances.
+    new_initializers : List[onnx.TensorProto]
+        A list of newly created ONNX initializers
+        (weights or biases typically) during quantization.
+        This is shared with handlers that may add new constants.
+
+    Methods
+    -------
+    register(op_type, handler)
+        Registers a handler for an ONNX op_type.
+    quantize(node, rescale, graph, scale_exponent, scale_base, initializer_map)
+        Apply quantization to a specific ONNX node using its registered handler.
+    check_model(model)
+        Ensure all operations in the model are supported and validate
+        each layer's parameters are valid and supported.
+    check_layer(node, initializer_map)
+        Validate a single ONNX node using its handler's check_supported method,
+        to check that the given layers parameters and structure is supported.
+    get_initializer_map(model)
+        Build a {name: TensorProto} mapping for the model's initializers.
+    """
+
+    def __init__(self: ONNXOpQuantizer) -> None:
+        self.handlers: dict[
+            str,
+            Callable[
+                [onnx.NodeProto, bool],
+                onnx.NodeProto | list[onnx.NodeProto],
+            ],
+        ] = {}
+        self.new_initializers = []
+
+        # Register handlers
+        self.register("Conv", ConvQuantizer(self.new_initializers))
+        self.register("Relu", ReluQuantizer())
+        self.register("Reshape", PassthroughQuantizer())
+        self.register("Gemm", GemmQuantizer(self.new_initializers))
+        self.register("Constant", ConstantQuantizer())
+        self.register("MaxPool", MaxpoolQuantizer())
+        self.register("Flatten", PassthroughQuantizer())
+
+    def register(
+        self: ONNXOpQuantizer,
+        op_type: str,
+        handler: Callable[
+            [onnx.NodeProto, bool],
+            onnx.NodeProto | list[onnx.NodeProto],
+        ],
+    ) -> None:
+        """Register a quantizer handler for a given ONNX op_type.
+
+        Args:
+            op_type (str): Name of the ONNX operator type (e.g., "Conv", "Relu").
+            handler (Callable[[onnx.NodeProto, bool],
+                    Union[onnx.NodeProto, list[onnx.NodeProto]]]):
+                - Handler instance implementing `quantize()`
+                    (and optionally `check_supported()`).
+
+        Raises:
+            HandlerImplementationError: If handler has not properly implemented
+                `quantize` method
+        """
+        if not hasattr(handler, "quantize") or not callable(handler.quantize):
+            raise HandlerImplementationError(op_type, "Missing 'quantize' method.")
+
+        self.handlers[op_type] = handler
+
+    def quantize(  # noqa: PLR0913
+        self: ONNXOpQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_exponent: int,
+        scale_base: int,
+        initializer_map: dict[str, onnx.TensorProto],
+        *,
+        rescale: bool = True,
+    ) -> onnx.NodeProto | list[onnx.NodeProto]:
+        """Quantize an ONNX node using its registered handler.
+
+        Args:
+            node (onnx.NodeProto): The ONNX node to quantize.
+            rescale (bool): Whether to apply rescaling.
+            graph (onnx.GraphProto): The ONNX graph containing the node.
+            scale_exponent (int): Quantization scale value.
+                The scaling becomes scale_base**scale_exponent.
+            scale_base (int): Base for the quantization scale.
+                The scaling becomes scale_base**scale.
+            initializer_map (dict[str, onnx.TensorProto]):
+                Mapping of initializer names (typically weights and biases) to tensors.
+
+        Returns:
+            Union[onnx.NodeProto, List[onnx.NodeProto]]: The quantized node +
+                any additional nodes created in the process.
+        """
+        handler = self.handlers.get(node.op_type)
+        if handler:
+            result = handler.quantize(
+                node=node,
+                graph=graph,
+                scale_config=ScaleConfig(scale_exponent, scale_base, rescale),
+                initializer_map=initializer_map,
+            )
+            if isinstance(result, onnx.NodeProto):
+                return [result]
+            return result
+
+        raise UnsupportedOpError(node.op_type)
+
+    def check_model(self: ONNXOpQuantizer, model: onnx.ModelProto) -> None:
+        """Verify that all nodes in the model are supported and valid.
+
+        Args:
+            model (onnx.ModelProto): The ONNX model to check.
+
+        Raises:
+            UnsupportedOpError: If the model contains unsupported operators.
+        """
+        initializer_map = self.get_initializer_map(model)
+
+        model_ops = {node.op_type for node in model.graph.node}
+        unsupported = model_ops - self.handlers.keys()
+
+        if unsupported:
+            raise UnsupportedOpError(unsupported)
+
+        # Call check_layer on each node (e.g., for param validation)
+        for node in model.graph.node:
+            self.check_layer(node, initializer_map)
+
+    def check_layer(
+        self: ONNXOpQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> None:
+        """
+        Check an individual node using its handler.
+        Parameters for the node will be checked that they
+        meet the supported parameter requirements.
+
+        Args:
+            node (onnx.NodeProto): The node to check.
+            initializer_map (dict[str, onnx.TensorProto]): Mapping of initializer names
+                to tensor typically used in weights and biases.
+
+        Raises:
+            MissingHandlerError: If no handler is registered for the given node.
+        """
+        handler = self.handlers.get(node.op_type)
+        if not handler:
+            raise MissingHandlerError(node.op_type)
+
+        if hasattr(handler, "check_supported") and callable(handler.check_supported):
+            handler.check_supported(node, initializer_map)
+
+    def get_initializer_map(
+        self: ONNXOpQuantizer,
+        model: onnx.ModelProto,
+    ) -> dict[str, onnx.TensorProto]:
+        """Build a dictionary mapping initializer names to tensors in graph.
+
+        Args:
+            model (onnx.ModelProto): The ONNX model.
+
+        Returns:
+            dict[str, onnx.TensorProto]: Map from initializer name to tensors in graph.
+        """
+        return {init.name: init for init in model.graph.initializer}

python/core/model_templates/circuit_template.py

@@ -0,0 +1,57 @@
+import torch.nn as nn
+from python.core.circuits.base import Circuit
+from random import randint
+
+class SimpleCircuit(Circuit):
+    '''
+    Note: This template is irrelevant if using the onnx circuit builder. 
+    The template only helps developers if they choose to incorporate other circuit builders into the framework.
+
+    To begin, we need to specify some basic attributes surrounding the circuit we will be using.
+    required_keys - specify the variables in the input dictionary (and input file).
+    name - name of the rust bin to be run by the circuit.
+
+    scale_base - specify the base of the scaling applied to each value
+    scale_exponent - the exponent applied to the base to get the scaling factor. Scaling factor will be multiplied by each input
+
+    Other default inputs can be defined below
+    '''
+    def __init__(self, file_name):
+        # Initialize the base class
+        super().__init__()
+        
+        # Circuit-specific parameters
+        self.required_keys = ["input_a", "input_b", "nonce"]
+        self.name = "simple_circuit"  # Use exact name that matches the binary
+
+        self.scale_exponent = 1
+        self.scale_base = 1
+    
+        self.input_a = 100
+        self.input_b = 200
+        self.nonce = randint(0,10000)
+    
+    '''
+    The following are some important functions used by the model. get inputs should be defined to specify the inputs to the circuit
+    '''
+    def get_inputs(self):
+        '''
+        Specify the inputs to the circuit, based on what was specified in the __init__. Can also have inputs to this function for the inputs.
+        '''
+        return {'input_a': self.input_a, 'input_b': self.input_b, 'nonce': self.nonce}
+    
+    def get_outputs(self, inputs = None):
+        """
+        Compute the output of the circuit.
+        This is overwritten from the base class to ensure computation happens only once.
+        """
+        if inputs == None:
+            inputs = {'input_a': self.input_a, 'input_b': self.input_b, 'nonce': self.nonce}
+        print(f"Performing addition operation: {inputs['input_a']} + {inputs['input_b']}")
+        return inputs['input_a'] + inputs['input_b']
+    
+    # def format_inputs(self, inputs):
+    #     return {"input": inputs.long().tolist()}
+    
+    # def format_outputs(self, outputs):
+    #     return {"output": outputs.long().tolist()}

python/core/utils/benchmarking_helpers.py

@@ -0,0 +1,163 @@
+# python/core/utils/benchmarking_helpers.py
+from __future__ import annotations
+
+# --- Standard library --------------------------------------------------------
+import threading
+import time
+
+# --- Third-party -------------------------------------------------------------
+import psutil
+
+"""
+Lightweight helpers to measure peak memory usage of child processes during
+benchmarks. Uses a background thread that periodically sums the RSS of all
+descendants of the current process (optionally filtered by a name keyword).
+"""
+
+
+def _safe_rss_kb(pid: int) -> int:
+    """
+    Return RSS for a PID in KB. On errors/missing process, return 0.
+    """
+    try:
+        proc = psutil.Process(pid)
+        rss_bytes = proc.memory_info().rss  # type: ignore[attr-defined]
+        return int(rss_bytes // 1024)
+    except (psutil.NoSuchProcess, psutil.AccessDenied):
+        return 0
+
+
+def _list_children(parent_pid: int) -> list[psutil.Process]:
+    """
+    Return all descendant processes of a parent PID.
+    Empty list if the parent is gone or access is denied.
+    """
+    try:
+        parent = psutil.Process(parent_pid)
+        return parent.children(recursive=True)
+    except (psutil.NoSuchProcess, psutil.AccessDenied):
+        return []
+
+
+def _safe_name_lower(proc: psutil.Process) -> str | None:
+    """
+    Lowercased process name, or None if unavailable.
+    """
+    try:
+        return proc.name().lower()
+    except (psutil.NoSuchProcess, psutil.AccessDenied):
+        return None
+
+
+def monitor_subprocess_memory(
+    parent_pid: int,
+    process_name_keyword: str,
+    results: dict[str, int],
+    stop_event: threading.Event,
+    *,
+    poll_interval_s: float = 0.1,
+) -> None:
+    """
+    Track the peak sum of RSS across child processes of `parent_pid`.
+
+    If `process_name_keyword` is non-empty, only include children whose
+    name contains that lowercase keyword.
+
+    Writes peaks (KB) into `results` under:
+      - 'peak_subprocess_mem'   (RSS)
+      - 'peak_subprocess_swap'  (0; swap not collected here)
+      - 'peak_subprocess_total' (mem + swap)
+    """
+    keyword = process_name_keyword.strip().lower()
+    peak_rss_kb = 0
+
+    # Initialize keys so callers can inspect mid-run safely
+    results["peak_subprocess_mem"] = 0
+    results["peak_subprocess_swap"] = 0
+    results["peak_subprocess_total"] = 0
+
+    while not stop_event.is_set():
+        children = _list_children(parent_pid)
+        if not children and not psutil.pid_exists(parent_pid):
+            break
+
+        if keyword:
+            filtered: list[psutil.Process] = []
+            for c in children:
+                nm = _safe_name_lower(c)
+                if nm and keyword in nm:
+                    filtered.append(c)
+        else:
+            filtered = children
+
+        rss_sum_kb = 0
+        for c in filtered:
+            rss_sum_kb += _safe_rss_kb(c.pid)
+
+        if rss_sum_kb > peak_rss_kb:
+            peak_rss_kb = rss_sum_kb
+            results["peak_subprocess_mem"] = peak_rss_kb
+            results["peak_subprocess_swap"] = 0
+            results["peak_subprocess_total"] = peak_rss_kb
+
+        time.sleep(poll_interval_s)
+
+    # Final write (covers the case where peak never changed inside the loop)
+    results["peak_subprocess_mem"] = max(
+        results.get("peak_subprocess_mem", 0),
+        peak_rss_kb,
+    )
+    results["peak_subprocess_swap"] = 0
+    results["peak_subprocess_total"] = results["peak_subprocess_mem"]
+
+
+def start_memory_collection(
+    process_name: str,
+) -> tuple[threading.Event, threading.Thread, dict[str, int]]:
+    """
+    Spawn and start a monitoring thread for the current process' children.
+
+    Args:
+        process_name:
+            Optional substring to filter child process names (case-insensitive).
+            Pass "" to include all children.
+
+    Returns:
+        (stop_event, monitor_thread, monitor_results_dict)
+    """
+    parent_pid = psutil.Process().pid
+    monitor_results: dict[str, int] = {}
+    stop_event = threading.Event()
+    monitor_thread = threading.Thread(
+        target=monitor_subprocess_memory,
+        args=(parent_pid, process_name, monitor_results, stop_event),
+        kwargs={"poll_interval_s": 0.02},
+        daemon=True,
+    )
+    monitor_thread.start()
+    time.sleep(0.05)  # allow thread to start and populate initial keys
+    return stop_event, monitor_thread, monitor_results
+
+
+def end_memory_collection(
+    stop_event: threading.Event,
+    monitor_thread: threading.Thread,
+    monitor_results: dict[str, int],
+) -> dict[str, float]:
+    """
+    Stop the monitor thread and return a summary dict in MB:
+      {'ram': <MB>, 'swap': <MB>, 'total': <MB>}
+    """
+    stop_event.set()
+    monitor_thread.join(timeout=5.0)
+
+    rss_kb = int(monitor_results.get("peak_subprocess_mem", 0))
+    swap_kb = int(monitor_results.get("peak_subprocess_swap", 0))
+    total_kb = int(monitor_results.get("peak_subprocess_total", rss_kb + swap_kb))
+
+    kb_to_mb = 1.0 / 1024.0
+    return {
+        "ram": rss_kb * kb_to_mb,
+        "swap": swap_kb * kb_to_mb,
+        "total": total_kb * kb_to_mb,
+    }

python/core/utils/constants.py

@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+MODEL_SOURCE_ONNX: str = "onnx"
+MODEL_SOURCE_CLASS: str = "class"

python/core/utils/errors.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+
+class CircuitExecutionError(Exception):
+    """Base exception for all circuit execution-related errors."""
+
+    def __init__(self: CircuitExecutionError, message: str) -> None:
+        super().__init__(message)
+        self.message = message
+
+
+class MissingFileError(CircuitExecutionError):
+    """Raised when cant find file"""
+
+    def __init__(self: MissingFileError, message: str, path: str | None = None) -> None:
+        full_message = message if path is None else f"{message} [Path: {path}]"
+        super().__init__(full_message)
+        self.path = path
+
+
+class FileCacheError(CircuitExecutionError):
+    """Raised when reading or writing cached output fails."""
+
+    def __init__(self: FileCacheError, message: str, path: str | None = None) -> None:
+        full_message = message if path is None else f"{message} [Path: {path}]"
+        super().__init__(full_message)
+        self.path = path
+
+
+class ProofBackendError(CircuitExecutionError):
+    """Raised when a Cargo command fails."""
+
+    def __init__(  # noqa: PLR0913
+        self: ProofBackendError,
+        message: str,
+        command: list[str] | None = None,
+        returncode: int | None = None,
+        stdout: str | None = None,
+        stderr: str | None = None,
+    ) -> None:
+        parts = [message]
+        if command is not None:
+            command2 = [str(c) for c in command]
+            parts.append(f"Command: {' '.join(command2)}")
+            command = command2
+        if returncode is not None:
+            parts.append(f"Exit code: {returncode}")
+        if stdout:
+            parts.append(f"STDOUT:\n{stdout}")
+        if stderr:
+            parts.append(f"STDERR:\n{stderr}")
+        full_message = "\n".join(parts)
+        super().__init__(full_message)
+        self.command = command
+        self.returncode = returncode
+        self.stdout = stdout
+        self.stderr = stderr
+
+
+class ProofSystemNotImplementedError(CircuitExecutionError):
+    """Raised when a proof system is not implemented."""
+
+    def __init__(self: ProofSystemNotImplementedError, proof_system: object) -> None:
+        message = f"Proof system '{proof_system}' is not implemented."
+        super().__init__(message)
+        self.proof_system = proof_system
+
+
+class CircuitUtilsError(Exception):
+    """Base exception for layer utility errors."""
+
+
+class InputFileError(CircuitUtilsError):
+    """Raised when reading an input file fails."""
+
+    def __init__(
+        self: InputFileError,
+        file_path: str,
+        message: str,
+        *,
+        cause: Exception | None = None,
+    ) -> None:
+        full_msg = f"Failed to read input file '{file_path}': {message}"
+        super().__init__(full_msg)
+        self.file_path = file_path
+        self.__cause__ = cause
+
+
+class MissingCircuitAttributeError(CircuitUtilsError):
+    """Raised when a required attribute is missing or not set."""
+
+    def __init__(
+        self: MissingCircuitAttributeError,
+        attribute_name: str,
+        context: str | None = None,
+    ) -> None:
+        msg = f"Required attribute '{attribute_name}' is missing"
+        if context:
+            msg += f" ({context})"
+        super().__init__(msg)
+        self.attribute_name = attribute_name
+
+
+class ShapeMismatchError(CircuitUtilsError):
+    """Raised when reshaping tensors fails due to incompatible shapes."""
+
+    def __init__(
+        self: ShapeMismatchError,
+        expected_shape: list[int],
+        actual_shape: list[int],
+    ) -> None:
+        super().__init__(
+            f"Cannot reshape tensor of shape {actual_shape}"
+            f" to expected shape {expected_shape}",
+        )
+        self.expected_shape = expected_shape
+        self.actual_shape = actual_shape