torch-prisma 0.1.0__py3-none-any.whl

torch_prisma/__init__.py

@@ -0,0 +1,20 @@
+"""Inference latency and model complexity measurement for PyTorch."""
+
+from .complexity import (
+    ModelComplexity,
+    OperationCounts,
+    calculate_model_complexity,
+    count_operations,
+    count_parameters,
+)
+from .latency import LatencyStats, measure_module_latency
+
+__all__ = [
+    "LatencyStats",
+    "ModelComplexity",
+    "OperationCounts",
+    "calculate_model_complexity",
+    "count_operations",
+    "count_parameters",
+    "measure_module_latency",
+]

torch_prisma/complexity.py

@@ -0,0 +1,303 @@
+"""Parameter counts, nominal FLOPs, dense MACs, and coverage diagnostics."""
+
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+from torch import nn
+from torch.autograd.profiler_util import FunctionEvent
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.profiler import ProfilerActivity, profile, record_function
+from torch.utils.flop_counter import FlopCounterMode, flop_registry
+
+from .execution import call_arguments, tensor_leaves
+from .formulas import (
+    FLOP_COMPOSITES,
+    FLOP_FREE_OPS,
+    MAC_COMPOSITES,
+    MAC_FORMULAS,
+    MAC_FREE_OPS,
+    arithmetic_formulas,
+    biased_matrix_vector_flops,
+    cudnn_convolution_flops,
+    dot_flops,
+    efficient_attention_flops,
+    is_native_pointwise,
+    matrix_extra_flops,
+    matrix_vector_flops,
+    slow_conv2d_flops,
+)
+
+_FORMULA_PREFIX = "torch_prisma.formula/"
+
+
+def count_parameters(model: nn.Module) -> int:
+    """Count elements in the model's unique registered parameters.
+
+    Counts use tensor dimensions, including for frozen or masked parameters.
+
+    Args:
+        model: Module whose registered parameters are counted.
+
+    Returns:
+        Total number of parameter elements.
+    """
+    return sum(parameter.numel() for parameter in model.parameters())
+
+
+@dataclass(frozen=True)
+class OperationCounts:
+    """Dense MACs, nominal FLOPs, and coverage gaps for a forward."""
+
+    macs: int
+    flops: int
+    uncovered_ops: dict[str, tuple[str, ...]]
+
+
+@dataclass(frozen=True)
+class ModelComplexity:
+    """Model parameter count and operation counts for one forward.
+
+    Attributes:
+        macs: Known dense matrix and convolution multiply-accumulates.
+        params: Number of elements in unique registered parameters.
+        flops: Known nominal floating-point operations.
+        uncovered_ops: Sorted names of operations with incomplete counts,
+            grouped under `macs` and `flops`.
+    """
+
+    macs: int
+    params: int
+    flops: int
+    uncovered_ops: dict[str, tuple[str, ...]] = field(
+        default_factory=lambda: {"macs": (), "flops": ()}
+    )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return counts and coverage as JSON-ready data."""
+        return {
+            "macs": self.macs,
+            "params": self.params,
+            "flops": self.flops,
+            "uncovered_ops": {key: list(ops) for key, ops in self.uncovered_ops.items()},
+        }
+
+
+@contextmanager
+def _counting_attention() -> Iterator[None]:
+    """Expose attention matrix work and restore public backend switches on exit."""
+    fastpath = torch.backends.mha.get_fastpath_enabled()
+    try:
+        torch.backends.mha.set_fastpath_enabled(False)
+        with sdpa_kernel(SDPBackend.MATH):
+            yield
+    finally:
+        torch.backends.mha.set_fastpath_enabled(fastpath)
+
+
+def _find_coverage_gaps(
+    events: list[FunctionEvent], mac_free: frozenset[str]
+) -> tuple[set[str], set[str]]:
+    """Match each formula marker to one invocation and its backend children.
+
+    A C++ composite can invoke an operation below Python dispatch. Follow
+    profiler parents to distinguish hidden calls from separately counted calls.
+    """
+    # Native pointwise semantics prove the absence of dense matrix/convolution
+    # MACs without implying a FLOP count. Inspect each observed packet once.
+    mac_free = mac_free | {
+        name
+        for name in {event.name for event in events}
+        if name not in mac_free and name not in MAC_FORMULAS and is_native_pointwise(name)
+    }
+    counted: dict[int, str] = {}
+    recorded: set[str] = set()
+    for event in events:
+        if not event.name.startswith(_FORMULA_PREFIX):
+            continue
+        name = event.name.removeprefix(_FORMULA_PREFIX)
+        recorded.add(name)
+        parent = event.cpu_parent
+        while parent is not None and parent.name != name:
+            parent = parent.cpu_parent
+        if (
+            parent is None
+            and not name.startswith("aten::")
+            and event.cpu_parent is not None
+            and event.cpu_parent.name == "PythonDispatchMode"
+        ):
+            # Python custom operators can lack their own profiler event. The
+            # immediate dispatch scope still contains this call and its formula.
+            parent = event.cpu_parent
+        if parent is not None:
+            counted[parent.id] = name
+    mac_gaps = recorded - MAC_FORMULAS - mac_free
+    flop_gaps: set[str] = set()
+    mac_ignored = (mac_free | MAC_COMPOSITES) - mac_gaps
+    flop_ignored = FLOP_FREE_OPS | MAC_COMPOSITES | FLOP_COMPOSITES
+    for event in events:
+        if "::" not in event.name or event.name.startswith(_FORMULA_PREFIX):
+            continue
+        mac_covered = flop_covered = False
+        parent = event
+        while parent is not None:
+            if parent.id in counted:
+                flop_covered = True
+                mac_covered |= counted[parent.id] in MAC_FORMULAS
+            parent = parent.cpu_parent
+        if not mac_covered and event.name not in mac_ignored:
+            mac_gaps.add(event.name)
+        if not flop_covered and event.name not in flop_ignored:
+            flop_gaps.add(event.name)
+    return mac_gaps, flop_gaps
+
+
+def calculate_model_complexity(
+    target: Callable[..., Any],
+    input_args: Any,
+    input_kwargs: dict[str, Any] | None = None,
+) -> ModelComplexity:
+    """Count eager inference FLOPs, MACs and parameters for one forward.
+
+    Prepare the model mode and device placement before calling this function.
+    Changes made by the forward persist. Pass a separate model instance to
+    preserve the original model state.
+
+    Args:
+        target: An eager module or a prepared callable.
+        input_args: One positional input or a tuple/list of positional inputs.
+        input_kwargs: Optional keyword arguments.
+
+    Returns:
+        Counts for the supplied batch and executed path, with coverage gaps
+        listed in `uncovered_ops`. A plain callable has zero registered parameters.
+    """
+    args, kwargs = call_arguments(input_args, input_kwargs)
+    params = count_parameters(target) if isinstance(target, nn.Module) else 0
+    operations = count_operations(lambda: target(*args, **kwargs))
+    return ModelComplexity(operations.macs, params, operations.flops, operations.uncovered_ops)
+
+
+def count_operations(forward: Callable[[], Any]) -> OperationCounts:
+    """Count the MACs and FLOPs of a prepared eager forward.
+
+    Args:
+        forward: A zero-argument callable that runs the desired forward pass.
+            The caller prepares the model and inputs and manages changes
+            made by the forward.
+
+    Returns:
+        MACs, FLOPs, and coverage gaps. Counting uses math scaled dot-product
+        attention and temporarily disables the multihead attention fast path.
+    """
+    with torch.inference_mode(False), torch.no_grad():
+        # Preserve external non-MAC formulas. The mappings and validation
+        # wrappers belong to this invocation, never the global registry.
+        arithmetic = arithmetic_formulas()
+        additions = {op: formula for op, formula in arithmetic.items() if op not in flop_registry}
+        counter = FlopCounterMode(
+            display=False,
+            custom_mapping={
+                **additions,
+                torch.ops.aten.dot: dot_flops,
+                torch.ops.aten.vdot: dot_flops,
+                torch.ops.aten.mv: matrix_vector_flops,
+                torch.ops.aten.addmv: biased_matrix_vector_flops,
+                torch.ops.aten.addmv_: biased_matrix_vector_flops,
+                torch.ops.aten.cudnn_convolution: cudnn_convolution_flops,
+                torch.ops.aten._slow_conv2d_forward: slow_conv2d_flops,
+            },
+        )
+        # Reuse the existing dense formulas for identical operand layouts.
+        # Assign functions already wrapped by the registry directly.
+        for alias, source in (
+            (torch.ops.aten.addmm_, torch.ops.aten.addmm),
+            (torch.ops.aten.baddbmm_, torch.ops.aten.baddbmm),
+            (torch.ops.aten.addbmm, torch.ops.aten.baddbmm),
+            (torch.ops.aten.addbmm_, torch.ops.aten.baddbmm),
+            (
+                torch.ops.aten._scaled_dot_product_flash_attention_for_cpu,
+                torch.ops.aten._scaled_dot_product_flash_attention,
+            ),
+        ):
+            counter.flop_registry[alias] = counter.flop_registry[source]
+        counter.flop_registry[torch.ops.aten._efficient_attention_forward] = (
+            efficient_attention_flops
+        )
+        macs = 0
+        uncovered_flops: set[str] = set()
+
+        def checked_formula(op: object, formula: Callable, added: bool) -> Callable:
+            """Validate each invocation before aggregation and retain its dense MAC count."""
+            name = str(op).replace(".", "::", 1)
+
+            def record(*values: Any, out_val: Any = None, **options: Any) -> int:
+                """Check the formula result and add separately defined non-MAC arithmetic."""
+                nonlocal macs
+                output_tensors = tensor_leaves(out_val)
+                if added and any(t.is_complex() for t in tensor_leaves((values, options))):
+                    uncovered_flops.add(name)
+                    return 0
+                if added and output_tensors and not output_tensors[0].is_floating_point():
+                    if output_tensors[0].is_complex():
+                        uncovered_flops.add(name)
+                    with record_function(f"{_FORMULA_PREFIX}{name}"):
+                        pass
+                    return 0
+                value = formula(*values, out_val=out_val, **options)
+                if value is None and added:
+                    uncovered_flops.add(name)
+                    return 0
+                is_mac = name in MAC_FORMULAS
+                if type(value) is not int or value < 0 or (is_mac and value % 2):
+                    expected = "nonnegative even" if is_mac else "nonnegative integer"
+                    raise ValueError(f"{name}: expected a {expected} FLOP count")
+                if (
+                    name in ("aten::_scaled_mm", "aten::convolution_backward")
+                    or (name.startswith("aten::") and "attention" in name)
+                    or name in ("flex_attention", "flex_attention_backward")
+                ):
+                    # Native fused formulas omit arithmetic outside their
+                    # matrix products, such as softmax or bias gradients.
+                    uncovered_flops.add(name)
+                if is_mac:
+                    macs += value // 2
+                    value += matrix_extra_flops(name, values, options, out_val)
+                    if output_tensors and not output_tensors[0].is_floating_point():
+                        if output_tensors[0].is_complex():
+                            uncovered_flops.add(name)
+                        # Count integer and complex MACs as logical products.
+                        value = 0
+                with record_function(f"{_FORMULA_PREFIX}{name}"):
+                    pass
+                return value
+
+            return record
+
+        counter.flop_registry = {
+            op: checked_formula(op, formula, op in additions)
+            for op, formula in counter.flop_registry.items()
+        }
+        with (
+            profile(activities=[ProfilerActivity.CPU]) as trace,
+            _counting_attention(),
+            counter,
+        ):
+            forward()
+        mac_free = (
+            MAC_FREE_OPS
+            | FLOP_FREE_OPS
+            | {str(op).replace("aten.", "aten::", 1) for op in arithmetic}
+        )
+        uncovered_macs, flop_gaps = _find_coverage_gaps(trace.events(), mac_free)
+        uncovered_flops.update(flop_gaps)
+        return OperationCounts(
+            macs=macs,
+            flops=counter.get_total_flops(),
+            uncovered_ops={
+                "macs": tuple(sorted(uncovered_macs)),
+                "flops": tuple(sorted(uncovered_flops)),
+            },
+        )

torch_prisma/execution.py

@@ -0,0 +1,19 @@
+"""Argument and tensor helpers shared by the measurement APIs."""
+
+from typing import Any
+
+import torch
+from torch.utils._pytree import tree_leaves
+
+
+def call_arguments(
+    input_args: Any, input_kwargs: dict[str, Any] | None
+) -> tuple[tuple[Any, ...], dict[str, Any]]:
+    """Normalize positional and keyword arguments for a forward call."""
+    args = tuple(input_args) if isinstance(input_args, (tuple, list)) else (input_args,)
+    return args, {} if input_kwargs is None else input_kwargs
+
+
+def tensor_leaves(value: Any) -> list[torch.Tensor]:
+    """Find the tensors in a nested value."""
+    return [leaf for leaf in tree_leaves(value) if isinstance(leaf, torch.Tensor)]

torch_prisma/formulas.py

@@ -0,0 +1,677 @@
+"""Shape formulas for nominal real arithmetic, independent of pruning."""
+
+import math
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch.utils.flop_counter import flop_registry
+
+
+def is_native_pointwise(name: str) -> bool:
+    """Identify ATen packets whose overloads all declare pointwise semantics.
+
+    Pointwise work is outside the dense matrix/convolution MAC convention.
+    These tags do not determine FLOPs. Require every overload to carry the tag
+    because profiler names identify packets rather than individual overloads.
+    """
+    if not name.startswith("aten::"):
+        return False
+    packet = getattr(torch.ops.aten, name.removeprefix("aten::"), None)
+    if packet is None or not callable(getattr(packet, "overloads", None)):
+        return False
+    overloads = packet.overloads()
+    return bool(overloads) and all(
+        torch.Tag.pointwise in getattr(packet, overload).tags for overload in overloads
+    )
+
+
+def elementwise_flops(*args: Any, out_shape: tuple[int, ...], **kwargs: Any) -> int:
+    """Count one arithmetic operation per output element, including broadcasting."""
+    return math.prod(out_shape)
+
+
+def add_flops(*args: Any, out_shape: tuple[int, ...], alpha: float = 1, **kwargs: Any) -> int:
+    """Count addition and an explicit nonunit scale of the second operand."""
+    return math.prod(out_shape) * (1 + (alpha != 1))
+
+
+def add_product_flops(
+    *args: Any, out_shape: tuple[int, ...], value: float = 1, **kwargs: Any
+) -> int:
+    """Count addcmul/addcdiv arithmetic and an optional nonunit scalar scale."""
+    return math.prod(out_shape) * (2 + (value != 1))
+
+
+def sum_flops(input: tuple[int, ...], *args: Any, out_shape: tuple[int, ...], **kwargs: Any) -> int:
+    """Count one fewer addition than input elements in each reduction group."""
+    return max(0, math.prod(input) - math.prod(out_shape))
+
+
+def mean_flops(
+    input: tuple[int, ...], *args: Any, out_shape: tuple[int, ...], **kwargs: Any
+) -> int:
+    """Count reduction additions and one division per output."""
+    return sum_flops(input, out_shape=out_shape) + math.prod(out_shape)
+
+
+def softmax_flops(input: tuple[int, ...], dim: int, *args: Any, **kwargs: Any) -> int:
+    """Count stable softmax arithmetic, excluding maximum comparisons."""
+    size = input[dim] if input else 1
+    total = math.prod(input)
+    return 4 * total - total // size if size else 0
+
+
+def log_softmax_flops(input: tuple[int, ...], *args: Any, **kwargs: Any) -> int:
+    """Count stable log-softmax arithmetic with one unit per exp/log."""
+    return 4 * math.prod(input)
+
+
+def normalization_flops(total: int, groups: int, weight: object, bias: object) -> int:
+    """Count centered variance with reused centered values and optional affine work."""
+    return 5 * total + 2 * groups + total * ((weight is not None) + (bias is not None))
+
+
+def layer_norm_flops(
+    input: tuple[int, ...],
+    normalized_shape: tuple[int, ...],
+    weight: object,
+    bias: object,
+    *args: Any,
+    **kwargs: Any,
+) -> int | None:
+    """Count mean, centered variance, epsilon, rsqrt and affine arithmetic."""
+    width = math.prod(normalized_shape)
+    return (
+        normalization_flops(math.prod(input), math.prod(input) // width, weight, bias)
+        if width
+        else None
+    )
+
+
+def group_norm_flops(
+    input: tuple[int, ...],
+    weight: object,
+    bias: object,
+    N: int,
+    C: int,
+    HxW: int,
+    group: int,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
+    """Count one normalization domain per sample and group."""
+    return normalization_flops(math.prod(input), N * group, weight, bias)
+
+
+def batch_norm_eval_flops(
+    input: tuple[int, ...],
+    weight: object,
+    bias: object,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
+    """Count inference normalization from stored statistics, without algebraic folding."""
+    total = math.prod(input)
+    return 2 * input[1] + total * (2 + (weight is not None) + (bias is not None))
+
+
+def batch_norm_flops(
+    input: tuple[int, ...],
+    weight: object,
+    bias: object,
+    running_mean: object = None,
+    running_var: object = None,
+    training: bool = True,
+    *args: Any,
+    **kwargs: Any,
+) -> int | None:
+    """Count inference with stored statistics and flag training updates as uncovered."""
+    # The no_stats overload places training/momentum/eps after bias. Those
+    # arguments are not running-statistic shapes and are outside this formula.
+    if training or not isinstance(running_mean, tuple) or not isinstance(running_var, tuple):
+        return None
+    return batch_norm_eval_flops(input, weight, bias)
+
+
+def dot_flops(a: tuple[int, ...], b: tuple[int, ...], **kwargs: Any) -> int:
+    """Count a vector contraction under the dense two-FLOPs-per-MAC convention."""
+    return 2 * a[0]
+
+
+def matrix_vector_flops(matrix: tuple[int, ...], vector: tuple[int, ...], **kwargs: Any) -> int:
+    """Count matrix-vector work without bias additions or scalar scaling."""
+    return 2 * math.prod(matrix)
+
+
+def biased_matrix_vector_flops(
+    bias: tuple[int, ...], matrix: tuple[int, ...], vector: tuple[int, ...], **kwargs: Any
+) -> int:
+    """Use the same dense matrix-vector formula for the biased form."""
+    return matrix_vector_flops(matrix, vector)
+
+
+def conv_flops(input: tuple[int, ...], weight: tuple[int, ...], output: tuple[int, ...]) -> int:
+    """Count ordinary convolution products, including implicit padding positions."""
+    return 2 * input[0] * math.prod(weight) * math.prod(output[2:])
+
+
+def cudnn_convolution_flops(
+    input: tuple[int, ...],
+    weight: tuple[int, ...],
+    padding: object,
+    stride: object,
+    dilation: object,
+    groups: int,
+    benchmark: bool,
+    deterministic: bool,
+    allow_tf32: bool,
+    *,
+    out_shape: tuple[int, ...],
+    **kwargs: Any,
+) -> int:
+    """Bind cuDNN's distinct signature without interpreting benchmark as transposed."""
+    return conv_flops(input, weight, out_shape)
+
+
+def slow_conv2d_flops(
+    input: tuple[int, ...],
+    weight: tuple[int, ...],
+    kernel_size: object,
+    bias: object,
+    stride: object,
+    padding: object,
+    *,
+    out_shape: tuple[int, ...],
+    **kwargs: Any,
+) -> int:
+    """Bind the six-argument CPU convolution signature."""
+    return conv_flops(input, weight, out_shape)
+
+
+def efficient_attention_flops(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    bias: torch.Tensor | None,
+    cu_seqlens_q: torch.Tensor | None,
+    cu_seqlens_k: torch.Tensor | None,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
+    """Count dense BMHD attention and reuse the native variable-length formula.
+
+    This raw kernel uses batch, sequence, heads, and width order. The public
+    SDPA kernels use batch, heads, sequence, and width order instead.
+    """
+    if cu_seqlens_q is not None or cu_seqlens_k is not None:
+        return flop_registry[torch.ops.aten._efficient_attention_forward](
+            query, key, value, bias, cu_seqlens_q, cu_seqlens_k, *args, **kwargs
+        )
+    batch, queries, heads, width = query.shape
+    return 2 * batch * heads * queries * key.shape[1] * (width + value.shape[-1])
+
+
+def arithmetic_formulas() -> dict[object, Callable[..., int | None]]:
+    """Return local additions to the counter, without changing its global registry.
+
+    None means the formula does not cover this invocation. Comparisons and
+    data movement do not contribute floating-point arithmetic.
+    """
+    result = {
+        getattr(torch.ops.aten, name): elementwise_flops
+        for name in (
+            "mul",
+            "mul_",
+            "div",
+            "div_",
+            "divide",
+            "divide_",
+            "true_divide",
+            "true_divide_",
+            "multiply",
+            "multiply_",
+            "neg",
+            "neg_",
+            "negative",
+            "negative_",
+            "abs",
+            "abs_",
+            "absolute",
+            "absolute_",
+            "sqrt",
+            "sqrt_",
+            "rsqrt",
+            "rsqrt_",
+            "exp",
+            "exp_",
+            "log",
+            "log_",
+            "reciprocal",
+            "reciprocal_",
+            "square",
+            "square_",
+            "pow",
+            "pow_",
+        )
+    }
+    result.update(
+        {
+            torch.ops.aten.add: add_flops,
+            torch.ops.aten.add_: add_flops,
+            torch.ops.aten.sub: add_flops,
+            torch.ops.aten.sub_: add_flops,
+            torch.ops.aten.subtract: add_flops,
+            torch.ops.aten.subtract_: add_flops,
+            torch.ops.aten.rsub: add_flops,
+            torch.ops.aten.addcmul: add_product_flops,
+            torch.ops.aten.addcmul_: add_product_flops,
+            torch.ops.aten.addcdiv: add_product_flops,
+            torch.ops.aten.addcdiv_: add_product_flops,
+            torch.ops.aten.sum: sum_flops,
+            torch.ops.aten.mean: mean_flops,
+            torch.ops.aten._softmax: softmax_flops,
+            torch.ops.aten._safe_softmax: softmax_flops,
+            torch.ops.aten._log_softmax: log_softmax_flops,
+            torch.ops.aten.native_layer_norm: layer_norm_flops,
+            torch.ops.aten.native_group_norm: group_norm_flops,
+            torch.ops.aten._native_batch_norm_legit_no_training: batch_norm_eval_flops,
+            torch.ops.aten.native_batch_norm: batch_norm_flops,
+            torch.ops.aten._native_batch_norm_legit: batch_norm_flops,
+            torch.ops.aten.cudnn_batch_norm: batch_norm_flops,
+        }
+    )
+    return result
+
+
+def matrix_extra_flops(op: str, args: tuple[Any, ...], kwargs: dict[str, Any], output: Any) -> int:
+    """Count fused bias and scalar scaling separately from the dense MAC formula."""
+    if op in (
+        "aten::addmm",
+        "aten::addmm_",
+        "aten::baddbmm",
+        "aten::baddbmm_",
+        "aten::addmv",
+        "aten::addmv_",
+        "aten::addbmm",
+        "aten::addbmm_",
+    ):
+        beta, alpha = kwargs.get("beta", 1), kwargs.get("alpha", 1)
+        return output.numel() * ((beta != 0) + (beta not in (0, 1)) + (alpha != 1))
+    positions = {
+        "aten::convolution": 2,
+        "aten::_convolution": 2,
+        "aten::convolution_overrideable": 2,
+        "aten::_slow_conv2d_forward": 3,
+    }
+    if op in positions:
+        bias = args[positions[op]] if len(args) > positions[op] else kwargs.get("bias")
+        return output.numel() if bias is not None else 0
+    return 0
+
+
+# Data movement, indexing, comparisons and selection are free under this
+# arithmetic convention.
+FLOP_FREE_OPS = frozenset(
+    f"aten::{name}"
+    for name in (
+        "t",
+        "transpose",
+        "permute",
+        "view",
+        "_unsafe_view",
+        "reshape",
+        "flatten",
+        "unflatten",
+        "unsqueeze",
+        "squeeze",
+        "squeeze_",
+        "unsqueeze_",
+        "expand",
+        "expand_as",
+        "repeat",
+        "repeat_interleave",
+        "slice",
+        "select",
+        "narrow",
+        "split",
+        "split_with_sizes",
+        "chunk",
+        "cat",
+        "stack",
+        "unbind",
+        "as_strided",
+        "as_strided_",
+        "contiguous",
+        "clone",
+        "copy_",
+        "detach",
+        "alias",
+        "to",
+        "_to_copy",
+        "type_as",
+        "result_type",
+        "resolve_conj",
+        "resolve_neg",
+        "empty",
+        "empty_like",
+        "empty_strided",
+        "zeros",
+        "zeros_like",
+        "ones",
+        "ones_like",
+        "full",
+        "full_like",
+        "fill_",
+        "zero_",
+        "resize_",
+        "arange",
+        "scalar_tensor",
+        "lift_fresh",
+        "embedding",
+        "index",
+        "index_select",
+        "gather",
+        "masked_fill",
+        "masked_fill_",
+        "triu",
+        "tril",
+        "_nnpack_available",
+        "_local_scalar_dense",
+        "item",
+        "is_nonzero",
+        "_reshape_alias",
+        "_autocast_to_reduced_precision",
+        "_autocast_to_full_precision",
+        "eq",
+        "ne",
+        "lt",
+        "le",
+        "gt",
+        "ge",
+        "all",
+        "any",
+        "logical_not",
+        "logical_and",
+        "where",
+        "isneginf",
+        "isinf",
+        "isnan",
+        "relu",
+        "relu_",
+        "clamp",
+        "clamp_min",
+        "clamp_min_",
+        "amax",
+        "amin",
+        "max",
+        "min",
+        # Forward max pooling selects values and indices. Comparisons are free
+        # under this convention. Backward scatter/add kernels require separate rules.
+        "max_pool1d",
+        "max_pool2d",
+        "max_pool3d",
+        "max_pool1d_with_indices",
+        "max_pool2d_with_indices",
+        "max_pool3d_with_indices",
+        "adaptive_max_pool1d",
+        "adaptive_max_pool2d",
+        "adaptive_max_pool3d",
+        # Coordinate transforms and padding only copy, select, or fill values.
+        "pixel_shuffle",
+        "pixel_unshuffle",
+        "channel_shuffle",
+        "native_channel_shuffle",
+        "unfold",
+        "im2col",
+        "max_unpool2d",
+        "max_unpool3d",
+        "pad",
+        "constant_pad_nd",
+        "_pad_circular",
+        "reflection_pad1d",
+        "reflection_pad2d",
+        "reflection_pad3d",
+        "replication_pad1d",
+        "replication_pad2d",
+        "replication_pad3d",
+    )
+)
+
+FLOP_COMPOSITES = frozenset(
+    f"aten::{name}"
+    for name in (
+        "softmax",
+        "_safe_softmax",
+        "square",
+        "square_",
+        "absolute",
+        "absolute_",
+        "negative",
+        "negative_",
+        "multiply",
+        "multiply_",
+        "subtract",
+        "subtract_",
+        "divide",
+        "divide_",
+        "true_divide",
+        "true_divide_",
+        "log_softmax",
+        "batch_norm",
+        "_batch_norm_impl_index",
+        "native_batch_norm",
+        "layer_norm",
+        "group_norm",
+    )
+)
+
+
+# Only these audited forward formulas use two FLOPs per dense MAC. Do not
+# convert the counter's total: its registry also accepts arbitrary FLOP formulas.
+MAC_FORMULAS = frozenset(
+    f"aten::{name}"
+    for name in (
+        "mm",
+        "addmm",
+        "addmm_",
+        "bmm",
+        "baddbmm",
+        "baddbmm_",
+        "addbmm",
+        "addbmm_",
+        "dot",
+        "vdot",
+        "mv",
+        "addmv",
+        "addmv_",
+        "_scaled_mm",
+        "convolution",
+        "_convolution",
+        "cudnn_convolution",
+        "_slow_conv2d_forward",
+        "convolution_overrideable",
+        "_scaled_dot_product_efficient_attention",
+        "_scaled_dot_product_flash_attention",
+        "_scaled_dot_product_cudnn_attention",
+        "_scaled_dot_product_flash_attention_for_cpu",
+        "_flash_attention_forward",
+        "_efficient_attention_forward",
+    )
+)
+
+# These wrappers contribute through their constituent matrix/conv calls, not
+# through formulas registered directly on the wrappers.
+MAC_COMPOSITES = frozenset(
+    f"aten::{name}"
+    for name in (
+        "linear",
+        "matmul",
+        "einsum",
+        "conv1d",
+        "conv2d",
+        "conv3d",
+        "_convolution_mode",
+        "conv_transpose1d",
+        "conv_transpose2d",
+        "conv_transpose3d",
+        "scaled_dot_product_attention",
+        "_scaled_dot_product_attention_math",
+    )
+)
+
+# Explicit exclusions under the matrix/conv MAC convention. Unfamiliar kernels
+# must appear in uncovered_ops["macs"], not silently become zero-cost operations.
+MAC_FREE_OPS = frozenset(
+    f"aten::{name}"
+    for name in (
+        "add",
+        "add_",
+        "sub",
+        "mul",
+        "div",
+        "div_",
+        "neg",
+        "pow",
+        "sqrt",
+        "rsqrt",
+        "abs",
+        "exp",
+        "log",
+        "tanh",
+        "sigmoid",
+        "relu",
+        "relu_",
+        "gelu",
+        "silu",
+        "clamp_min",
+        "clamp_min_",
+        "clamp",
+        "where",
+        "eq",
+        "ne",
+        "lt",
+        "le",
+        "gt",
+        "ge",
+        "all",
+        "any",
+        "logical_not",
+        "logical_and",
+        "softmax",
+        "_softmax",
+        "_safe_softmax",
+        "log_softmax",
+        "_log_softmax",
+        "batch_norm",
+        "_batch_norm_impl_index",
+        "_native_batch_norm_legit",
+        "_native_batch_norm_legit_no_training",
+        "native_batch_norm",
+        "cudnn_batch_norm",
+        "miopen_batch_norm",
+        "layer_norm",
+        "native_layer_norm",
+        "group_norm",
+        "native_group_norm",
+        "avg_pool1d",
+        "avg_pool2d",
+        "avg_pool3d",
+        "adaptive_avg_pool1d",
+        "adaptive_avg_pool2d",
+        "adaptive_avg_pool3d",
+        "_adaptive_avg_pool2d",
+        "_adaptive_avg_pool3d",
+        "max_pool1d",
+        "max_pool2d",
+        "max_pool3d",
+        "max_pool1d_with_indices",
+        "max_pool2d_with_indices",
+        "max_pool3d_with_indices",
+        "adaptive_max_pool1d",
+        "adaptive_max_pool2d",
+        "adaptive_max_pool3d",
+        "mean",
+        "sum",
+        "amax",
+        "amin",
+        "max",
+        "min",
+        "t",
+        "transpose",
+        "permute",
+        "view",
+        "_unsafe_view",
+        "reshape",
+        "flatten",
+        "unflatten",
+        "unsqueeze",
+        "squeeze",
+        "squeeze_",
+        "unsqueeze_",
+        "lgamma",
+        "special_gammaln",
+        "expand",
+        "expand_as",
+        "repeat",
+        "repeat_interleave",
+        "slice",
+        "select",
+        "narrow",
+        "split",
+        "split_with_sizes",
+        "chunk",
+        "cat",
+        "stack",
+        "unbind",
+        "as_strided",
+        "as_strided_",
+        "contiguous",
+        "clone",
+        "copy_",
+        "detach",
+        "alias",
+        "to",
+        "_to_copy",
+        "type_as",
+        "resolve_conj",
+        "resolve_neg",
+        "empty",
+        "empty_like",
+        "empty_strided",
+        "zeros",
+        "zeros_like",
+        "ones",
+        "ones_like",
+        "full",
+        "full_like",
+        "fill_",
+        "zero_",
+        "resize_",
+        "arange",
+        "scalar_tensor",
+        "lift_fresh",
+        "dropout",
+        "native_dropout",
+        "embedding",
+        "index",
+        "index_select",
+        "gather",
+        "masked_fill",
+        "isneginf",
+        "isinf",
+        "isnan",
+        "masked_fill_",
+        "triu",
+        "tril",
+        "_nnpack_available",
+        "_local_scalar_dense",
+        "item",
+        "is_nonzero",
+        "_reshape_alias",
+        "_autocast_to_reduced_precision",
+        "_autocast_to_full_precision",
+        # Fold adds overlapping patches but performs no dense matrix product.
+        "col2im",
+    )
+)

torch_prisma/latency.py

@@ -0,0 +1,246 @@
+"""Eager and compiled inference latency on CPU, CUDA and Ascend."""
+
+import functools
+import gc
+import math
+import statistics
+import time
+from collections.abc import Callable, Iterable, Iterator
+from contextlib import AbstractContextManager, ExitStack, contextmanager
+from dataclasses import asdict, dataclass, field
+from importlib import import_module
+from types import ModuleType
+from typing import Any
+
+import torch
+
+from .execution import call_arguments
+
+_COMPILE_ITERS = 3
+_UNLIMITED_RECOMPILES = 2**31 - 1
+
+
+@dataclass(frozen=True)
+class LatencyStats:
+    """Summarize forward latency in milliseconds.
+
+    Attributes:
+        median_ms: Median sample duration.
+        mean_ms: Arithmetic mean of the sample durations.
+        std_ms: Population standard deviation of the sample durations.
+        p90_ms: 90th percentile selected from the measured samples.
+        p99_ms: 99th percentile selected from the measured samples.
+        min_ms: Shortest sample duration.
+        max_ms: Longest sample duration.
+        samples_ms: Sample durations in measurement order.
+    """
+
+    median_ms: float
+    mean_ms: float
+    std_ms: float
+    p90_ms: float
+    p99_ms: float
+    min_ms: float
+    max_ms: float
+    samples_ms: tuple[float, ...] = field(repr=False)
+
+    @classmethod
+    def from_samples(cls, samples_ms: Iterable[float]) -> "LatencyStats":
+        """Summarize timings in milliseconds.
+
+        Args:
+            samples_ms: A nonempty iterable of finite, nonnegative durations.
+
+        Returns:
+            Latency statistics and the original sample sequence.
+
+        Raises:
+            ValueError: The sequence is empty or contains an invalid duration.
+        """
+        samples = tuple(float(sample) for sample in samples_ms)
+        if not samples or any(not math.isfinite(sample) or sample < 0 for sample in samples):
+            raise ValueError("Latency samples must be nonempty, finite and nonnegative")
+        ordered = sorted(samples)
+        count = len(ordered)
+        return cls(
+            median_ms=statistics.median(samples),
+            mean_ms=statistics.fmean(samples),
+            std_ms=statistics.pstdev(samples),
+            p90_ms=ordered[min(int(count * 0.90), count - 1)],
+            p99_ms=ordered[min(int(count * 0.99), count - 1)],
+            min_ms=ordered[0],
+            max_ms=ordered[-1],
+            samples_ms=samples,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return statistics and samples as JSON-ready data."""
+        result = asdict(self)
+        result["samples_ms"] = list(self.samples_ms)
+        return result
+
+
+def _timing_device(requested: torch.device | str) -> tuple[torch.device, ModuleType | None]:
+    """Resolve the device and runtime used to record timing events."""
+    if isinstance(requested, str) and requested.split(":", 1)[0] == "npu":
+        import_module("torch_npu")
+    device = torch.device(requested)
+    if device.type == "cpu":
+        return device, None
+    if device.type not in ("cuda", "npu"):
+        raise ValueError(f"Unsupported timing device {device}. Expected CPU, CUDA or Ascend NPU")
+    runtime = getattr(torch, device.type)
+    if device.index is None:
+        device = torch.device(device.type, runtime.current_device())
+    return device, runtime
+
+
+def _unlimited_recompile_patch() -> AbstractContextManager[None]:
+    """Raise Dynamo's recompile caps during the compilation forwards."""
+    config = torch._dynamo.config
+    # PyTorch 2.6 uses the cache_size_limit names.
+    overrides = {
+        name if hasattr(config, name) else legacy: _UNLIMITED_RECOMPILES
+        for name, legacy in (
+            ("recompile_limit", "cache_size_limit"),
+            ("accumulated_recompile_limit", "accumulated_cache_size_limit"),
+        )
+    }
+    return config.patch(**overrides)
+
+
+@contextmanager
+def _fail_on_late_compile() -> Iterator[None]:
+    """Reject compilation during warmup or timing."""
+    with torch.compiler.set_stance("fail_on_recompile"):
+        try:
+            yield
+        except RuntimeError as error:
+            if not str(error).startswith("Detected recompile when torch.compile stance"):
+                raise
+            raise RuntimeError(
+                f"PyTorch requested compilation after {_COMPILE_ITERS} preparation forwards. "
+                "The compiled latency measurement would be unreliable. "
+                "Set compile=False to measure eager latency."
+            ) from error
+
+
+@functools.cache
+def _default_npu_backend() -> Callable[..., Any]:
+    """Load and reuse the default TorchAir backend."""
+    try:
+        torchair = import_module("torchair")
+    except ModuleNotFoundError as error:
+        if error.name != "torchair":
+            raise
+        raise RuntimeError(
+            "Compiled Ascend latency requires a compatible TorchAir installation"
+        ) from error
+    return torchair.get_npu_backend()
+
+
+def _compile_options(device: torch.device, options: dict[str, Any] | None) -> dict[str, Any]:
+    """Select TorchAir for Ascend unless the caller supplied a backend."""
+    options = dict(options or {})
+    if device.type == "npu" and "backend" not in options:
+        options["backend"] = _default_npu_backend()
+    return options
+
+
+def measure_module_latency(
+    target: Callable[..., Any],
+    input_args: Any,
+    device: torch.device | str,
+    input_kwargs: dict[str, Any] | None = None,
+    *,
+    repetitions: int = 20,
+    warmup: int = 5,
+    compile: bool = True,
+    compile_kwargs: dict[str, Any] | None = None,
+) -> LatencyStats:
+    """Measure the latency of one forward per sample for the supplied batch.
+
+    Set the model mode and place the model and inputs before calling.
+    Forward calls run in inference mode, and their side effects persist.
+    Timing starts after compilation preparation and warmup.
+
+    Args:
+        target: A prepared module or callable.
+        input_args: One positional input or a tuple/list of positional inputs.
+        device: Device used for timing. CUDA and Ascend use its current stream.
+        input_kwargs: Optional keyword arguments.
+        repetitions: Positive number of timed forwards.
+        warmup: Number of untimed forwards after any compilation preparation.
+        compile: Enable `torch.compile` and three untimed preparation forwards.
+            Ascend defaults to TorchAir. Set `False` for eager timing if
+            TorchAir reports a compatibility error.
+        compile_kwargs: Keyword arguments forwarded to `torch.compile` when
+            compilation is enabled. On Ascend, pass `{"backend": "inductor"}`
+            to use Inductor with its default Triton compilation mode. If
+            `TORCHINDUCTOR_NPU_BACKEND` is set, it must be `default` for Triton.
+
+    Returns:
+        Forward durations in sampling order and their summary statistics in milliseconds.
+
+    Raises:
+        ValueError: The timing options or timing device are invalid.
+        RuntimeError: The requested accelerator or default NPU backend is
+            unavailable, or the forward requires compilation after preparation.
+    """
+    if type(repetitions) is not int or repetitions <= 0 or type(warmup) is not int or warmup < 0:
+        raise ValueError("Require positive repetitions and nonnegative warmup")
+    if type(compile) is not bool:
+        raise ValueError("compile must be a boolean")
+
+    args, kwargs = call_arguments(input_args, input_kwargs)
+    device, runtime = _timing_device(device)
+    forward = target
+    with torch.inference_mode(), ExitStack() as timing:
+        if compile:
+            with _unlimited_recompile_patch():
+                forward = torch.compile(target, **_compile_options(device, compile_kwargs))
+                # A forward can initialize state and invalidate the first graph.
+                for _ in range(_COMPILE_ITERS):
+                    forward(*args, **kwargs)
+                if runtime is not None:
+                    runtime.synchronize(device)
+            timing.enter_context(_fail_on_late_compile())
+        for _ in range(warmup):
+            forward(*args, **kwargs)
+
+        events = []
+        if runtime is not None:
+            stream = runtime.current_stream(device)
+            events = [
+                (runtime.Event(enable_timing=True), runtime.Event(enable_timing=True))
+                for _ in range(repetitions)
+            ]
+            # Event handles may be allocated lazily by the first record.
+            for start, end in events:
+                start.record(stream)
+                end.record(stream)
+            runtime.synchronize(device)
+
+        gc_enabled = gc.isenabled()
+        gc.disable()
+        try:
+            if runtime is None:
+                samples = []
+                for _ in range(repetitions):
+                    start_time = time.perf_counter()
+                    forward(*args, **kwargs)
+                    samples.append((time.perf_counter() - start_time) * 1000.0)
+            else:
+                for start, end in events:
+                    start.record(stream)
+                    forward(*args, **kwargs)
+                    end.record(stream)
+                runtime.synchronize(device)
+                samples = [start.elapsed_time(end) for start, end in events]
+        finally:
+            if gc_enabled:
+                gc.enable()
+            else:
+                gc.disable()
+
+    return LatencyStats.from_samples(samples)

torch_prisma-0.1.0.dist-info/METADATA

@@ -0,0 +1,72 @@
+Metadata-Version: 2.4
+Name: torch-prisma
+Version: 0.1.0
+Summary: Measure PyTorch inference latency, MACs, FLOPs, and parameters
+Keywords: pytorch,inference,benchmarking,latency,model-complexity,macs,flops,cuda,ascend
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 6 - Mature
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: System :: Benchmark
+Requires-Dist: torch>=2.6
+Requires-Python: >=3.10
+Project-URL: Documentation, https://github.com/ZaberKo/torch-prisma/blob/main/docs/index.md
+Project-URL: Issues, https://github.com/ZaberKo/torch-prisma/issues
+Project-URL: Repository, https://github.com/ZaberKo/torch-prisma
+Description-Content-Type: text/markdown
+
+# Torch Prisma
+
+Measure PyTorch inference latency, parameters, MACs, and FLOPs on CPU, CUDA, and Huawei Ascend.
+
+## Installation
+
+Requires Python 3.10+ and PyTorch 2.6+. Install in a virtual environment.
+
+```bash
+uv pip install torch-prisma
+```
+
+## Usage
+
+Latency uses `torch.compile` by default. Set `compile=False` for eager timing.
+
+```python
+import torch
+from torch import nn
+
+from torch_prisma import calculate_model_complexity, measure_module_latency
+
+# For inference, set the model to evaluation mode and place it and the inputs on the device.
+device = "cpu"
+model = nn.Linear(8, 4).eval().to(device)
+x = torch.randn(2, 8, device=device)
+
+complexity = calculate_model_complexity(model, x)
+latency = measure_module_latency(model, x, device=device)
+
+print(complexity.params, complexity.macs, complexity.flops)
+# Operations with incomplete counts, grouped by MACs and FLOPs.
+print(complexity.uncovered_ops)
+print(latency.median_ms)
+```
+
+[Measurement API](https://github.com/ZaberKo/torch-prisma/blob/main/docs/measurement.md) · [Latency benchmarks](https://github.com/ZaberKo/torch-prisma/blob/main/docs/benchmarks.md) · [Huawei Ascend](https://github.com/ZaberKo/torch-prisma/blob/main/docs/ascend.md)
+
+## Development
+
+```bash
+uv sync
+uv run pre-commit install
+uv run pytest
+uv build
+```
+
+Preview the documentation with `uv run --group docs zensical serve`.
+
+## License
+
+Torch Prisma is distributed under the [MIT License](https://github.com/ZaberKo/torch-prisma/blob/main/LICENSE).

torch_prisma-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+torch_prisma/__init__.py,sha256=41t9t0QybdB6hCurDsy2N1Dt0iyj2EW6SohQAqor6YY,468
+torch_prisma/complexity.py,sha256=t7R8lVFzm3ShKROcs4_2VFjDc6kIzbhsDL0iKgr0V8g,12100
+torch_prisma/execution.py,sha256=VDy9r15TZ2KPPUoYo3IuSuj_Dy8FxHKBQ7mM9v2pRzQ,680
+torch_prisma/formulas.py,sha256=URhkFz04AZ1BgRJcXYSRh8MmyFwgGCKbf_eopi9HEys,18731
+torch_prisma/latency.py,sha256=q1JAAfZf-SqfqKyTW7-uWwOBT6KKps-dcJMlbZgwP7s,9411
+torch_prisma-0.1.0.dist-info/licenses/LICENSE,sha256=X12NtdOZ6qnbw8nCphqDwNbZQE4A1ErhkOEqz9MaFaU,1064
+torch_prisma-0.1.0.dist-info/WHEEL,sha256=YUH1mBqsx8Dh2cQG2rlcuRYUhJddG9iClegy4IgnHik,79
+torch_prisma-0.1.0.dist-info/METADATA,sha256=vkK_ipBI5pQMkCgugpsPH1zqtXmqqfSepp-ViQKxag0,2376
+torch_prisma-0.1.0.dist-info/RECORD,,

torch_prisma-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.9.11
+Root-Is-Purelib: true
+Tag: py3-none-any

torch_prisma-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 ZaberKo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.