torch-prisma 0.1.0__tar.gz

torch_prisma-0.1.0/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.

torch_prisma-0.1.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,52 @@
+# 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/docs/ascend.md

@@ -0,0 +1,53 @@
+# Huawei Ascend
+
+## Runtime setup
+
+Install the Ascend driver, CANN, and compatible PyTorch and TorchNPU versions. Initialize CANN before starting Python. The default compiled timing backend, TorchAir, is bundled in supported TorchNPU distributions. [Ascend compatibility matrix](https://github.com/Ascend/pytorch/blob/master/COMPATIBILITY.en.md) · [TorchAir](https://github.com/Ascend/torchair)
+
+## Measuring latency
+
+```python
+import torch
+import torch_npu
+from torch import nn
+
+from torch_prisma import measure_module_latency
+
+device = torch.device("npu:0")
+model = nn.Linear(128, 64, device=device).eval()
+inputs = torch.randn(8, 128, device=device)
+
+print(measure_module_latency(model, inputs, device).median_ms)
+```
+
+TorchAir is the default compiler backend on Ascend. If it raises a compatibility error, use `compile=False` for eager timing.
+
+To configure TorchAir, create a `CompilerConfig` and pass the resulting backend to the measurement function. This example enables graph dumps.
+
+```python
+import torchair
+from torchair.configs.compiler_config import CompilerConfig
+
+config = CompilerConfig()
+config.debug.graph_dump.type = "pbtxt"
+backend = torchair.get_npu_backend(compiler_config=config)
+latency = measure_module_latency(model, inputs, device, compile_kwargs={"backend": backend})
+```
+
+See the [TorchAir documentation](https://github.com/Ascend/torchair) for configuration options.
+
+## Triton Ascend
+
+To use [Triton Ascend](https://github.com/triton-lang/triton-ascend) through PyTorch Inductor, install a compatible Triton Ascend, CANN, PyTorch, and TorchNPU combination. With `compile=True` (the default), select the Inductor backend explicitly:
+
+```python
+latency = measure_module_latency(model, inputs, device, compile_kwargs={"backend": "inductor"})
+```
+
+TorchNPU uses Triton as its default Inductor compilation mode. If `TORCHINDUCTOR_NPU_BACKEND` has been changed, set it to `default` before the first compilation to select Triton. See the [TorchNPU backend documentation](https://github.com/Ascend/pytorch/blob/master/docs/zh/api/environment_variable/inductor/TORCHINDUCTOR_NPU_BACKEND.md).
+
+For a callable that launches a Triton kernel directly, use eager timing:
+
+```python
+latency = measure_module_latency(triton_callable, (x, y, output), device, compile=False)
+```

torch_prisma-0.1.0/docs/benchmarks.md

@@ -0,0 +1,73 @@
+# Latency benchmarks
+
+These measurements examine variation across independent process starts. Each process calls `measure_module_latency()` once. Each supported combination has 16 results: eight on each of two physical cards. The two cards were tested concurrently, with case order shuffled. Results retain all timed samples, including outliers.
+
+Each test uses one accelerator, one PyTorch intra-op thread, randomly initialized TorchVision weights, and a fixed FP32 input of shape `(1, 3, 224, 224)`. The model is in evaluation mode, and weights and input use a fixed seed. The warmup and sample counts are shown below. Compiled measurements also use three preparation forwards. Inductor's disk cache was reused between processes separately for each card.
+
+The latency columns are the averages of the 16 returned `mean_ms` and `median_ms` values. CV means *coefficient of variation*: the sample standard deviation divided by the average, expressed as a percentage. Each CV pair reports the variation of the returned mean and median, respectively. Combined CV uses all 16 processes; card CV uses the eight processes on that card. The Ascend test used cards 6 and 7; the 3080 Ti test used cards 0 and 1.
+
+## Ascend 910B2
+
+The Ascend environment used PyTorch 2.10.0, TorchVision 0.25.0, TorchNPU 2.10.0.post7, Triton Ascend 3.6.0, CANN 9.2.0-beta.2, and driver 26.1.1. Inductor used Triton mode and four compiler threads. TorchAir is the default Ascend compiler backend.
+
+| Model | Execution | Warmup / samples | Mean (ms) | Median (ms) | Combined CV | Card 6 CV | Card 7 CV |
+| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
+| MobileNetV3-Small | Eager | 5 / 20 | 6.895 | 6.804 | 4.00% / 3.12% | 2.92% / 2.51% | 5.04% / 3.72% |
+| MobileNetV3-Small | Eager | 20 / 100 | 6.612 | 6.559 | 4.65% / 4.40% | 4.64% / 4.75% | 4.96% / 4.20% |
+| MobileNetV3-Small | Eager | 50 / 20 | 6.365 | 6.343 | 3.86% / 3.49% | 3.72% / 3.00% | 3.74% / 3.77% |
+| MobileNetV3-Small | TorchAir | 5 / 20 | 1.724 | 1.704 | 3.66% / 3.77% | 3.89% / 4.05% | 3.55% / 3.62% |
+| MobileNetV3-Small | TorchAir | 20 / 100 | 1.744 | 1.739 | 4.28% / 4.29% | 4.28% / 4.30% | 4.47% / 4.50% |
+| MobileNetV3-Small | TorchAir | 50 / 20 | 1.748 | 1.723 | 4.46% / 4.37% | 4.87% / 4.51% | 4.19% / 4.42% |
+| MobileNetV3-Small | Triton Inductor | 5 / 20 | 7.488 | 7.477 | 4.85% / 4.98% | 5.60% / 5.77% | 3.37% / 3.52% |
+| MobileNetV3-Small | Triton Inductor | 20 / 100 | 7.400 | 7.374 | 3.23% / 3.12% | 2.28% / 2.15% | 3.81% / 3.76% |
+| MobileNetV3-Small | Triton Inductor | 50 / 20 | 7.303 | 7.293 | 3.91% / 3.88% | 4.09% / 3.95% | 4.01% / 4.07% |
+| ResNet-50 | Eager | 5 / 20 | 7.481 | 7.451 | 6.15% / 6.15% | 7.04% / 7.14% | 5.60% / 5.47% |
+| ResNet-50 | Eager | 20 / 100 | 7.235 | 7.196 | 4.38% / 4.27% | 3.69% / 3.71% | 5.13% / 4.97% |
+| ResNet-50 | Eager | 50 / 20 | 7.181 | 7.162 | 3.21% / 3.25% | 3.07% / 3.28% | 3.38% / 3.31% |
+| ResNet-50 | TorchAir | 5 / 20 | 2.378 | 2.361 | 4.05% / 4.07% | 3.98% / 4.02% | 4.38% / 4.40% |
+| ResNet-50 | TorchAir | 20 / 100 | 2.375 | 2.371 | 3.77% / 3.76% | 4.26% / 4.25% | 0.37% / 0.36% |
+| ResNet-50 | TorchAir | 50 / 20 | 2.359 | 2.338 | 4.03% / 4.12% | 4.24% / 4.33% | 4.09% / 4.20% |
+| ResNet-50 | Triton Inductor | 5 / 20 | 12.365 | 12.324 | 0.29% / 0.28% | 0.20% / 0.22% | 0.34% / 0.29% |
+| ResNet-50 | Triton Inductor | 20 / 100 | 12.329 | 12.321 | 0.35% / 0.35% | 0.19% / 0.19% | 0.25% / 0.25% |
+| ResNet-50 | Triton Inductor | 50 / 20 | 12.376 | 12.321 | 0.32% / 0.35% | 0.13% / 0.24% | 0.18% / 0.21% |
+| ViT-B/16 | Eager | 5 / 20 | 8.810 | 8.799 | 4.79% / 4.49% | 3.05% / 2.64% | 6.19% / 5.90% |
+| ViT-B/16 | Eager | 20 / 100 | 8.344 | 8.311 | 6.71% / 6.14% | 9.26% / 8.45% | 3.28% / 3.08% |
+| ViT-B/16 | Eager | 50 / 20 | 8.278 | 8.253 | 4.42% / 4.39% | 5.82% / 5.72% | 2.72% / 2.69% |
+| ViT-B/16 | Triton Inductor | 5 / 20 | 10.945 | 10.899 | 2.51% / 2.43% | 2.02% / 1.94% | 2.98% / 2.88% |
+| ViT-B/16 | Triton Inductor | 20 / 100 | 10.578 | 10.553 | 3.05% / 3.15% | 2.44% / 2.44% | 3.48% / 3.66% |
+| ViT-B/16 | Triton Inductor | 50 / 20 | 10.672 | 10.643 | 5.36% / 5.21% | 5.91% / 5.74% | 5.01% / 4.89% |
+
+ViT-B/16 with TorchAir failed during compilation preparation on both cards with `only supported softmax_layout=TND`; no latency statistics are available for that combination.
+
+## GeForce RTX 3080 Ti
+
+The 3080 Ti environment used PyTorch 2.15.0, TorchVision 0.29.2, CUDA runtime 13.2, Triton 3.9.0, and NVIDIA driver 595.71.05. Inductor used four compiler threads. Software versions differ between devices, so these results are not a controlled cross-device performance comparison.
+
+| Model | Execution | Warmup / samples | Mean (ms) | Median (ms) | Combined CV | Card 0 CV | Card 1 CV |
+| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
+| MobileNetV3-Small | Eager | 5 / 20 | 4.302 | 4.267 | 11.07% / 10.99% | 13.12% / 12.99% | 2.31% / 1.34% |
+| MobileNetV3-Small | Eager | 20 / 100 | 4.399 | 4.383 | 12.16% / 12.36% | 11.37% / 11.69% | 13.54% / 13.61% |
+| MobileNetV3-Small | Eager | 50 / 20 | 4.149 | 4.127 | 6.94% / 7.10% | 2.00% / 1.86% | 9.47% / 9.66% |
+| MobileNetV3-Small | Inductor | 5 / 20 | 2.214 | 2.209 | 14.75% / 14.60% | 15.17% / 14.82% | 15.14% / 15.15% |
+| MobileNetV3-Small | Inductor | 20 / 100 | 2.115 | 2.114 | 11.96% / 13.02% | 14.83% / 15.83% | 0.67% / 0.81% |
+| MobileNetV3-Small | Inductor | 50 / 20 | 2.046 | 2.038 | 8.14% / 8.22% | 1.54% / 1.30% | 11.18% / 11.33% |
+| ResNet-50 | Eager | 5 / 20 | 5.197 | 5.106 | 10.02% / 10.16% | 8.87% / 9.16% | 11.21% / 11.41% |
+| ResNet-50 | Eager | 20 / 100 | 4.884 | 4.840 | 1.46% / 1.42% | 1.36% / 1.36% | 1.57% / 1.38% |
+| ResNet-50 | Eager | 50 / 20 | 5.006 | 4.981 | 8.10% / 8.31% | 8.33% / 8.61% | 8.44% / 8.59% |
+| ResNet-50 | Inductor | 5 / 20 | 2.656 | 2.612 | 15.85% / 13.01% | 10.20% / 10.24% | 18.49% / 14.66% |
+| ResNet-50 | Inductor | 20 / 100 | 2.674 | 2.658 | 13.58% / 13.51% | 10.56% / 10.48% | 14.70% / 14.58% |
+| ResNet-50 | Inductor | 50 / 20 | 2.613 | 2.590 | 15.67% / 14.44% | 10.95% / 11.09% | 19.15% / 17.22% |
+| ViT-B/16 | Eager | 5 / 20 | 4.522 | 4.505 | 7.30% / 7.23% | 6.22% / 5.73% | 8.07% / 8.25% |
+| ViT-B/16 | Eager | 20 / 100 | 4.224 | 4.151 | 7.89% / 8.41% | 2.35% / 1.45% | 10.07% / 10.50% |
+| ViT-B/16 | Eager | 50 / 20 | 4.284 | 4.249 | 8.91% / 9.10% | 10.22% / 10.27% | 7.66% / 8.15% |
+| ViT-B/16 | Inductor | 5 / 20 | 4.655 | 4.670 | 1.43% / 0.17% | 1.79% / 0.23% | 0.09% / 0.07% |
+| ViT-B/16 | Inductor | 20 / 100 | 4.315 | 4.240 | 1.62% / 0.98% | 1.60% / 0.18% | 1.04% / 0.34% |
+| ViT-B/16 | Inductor | 50 / 20 | 4.282 | 4.240 | 1.11% / 0.82% | 1.51% / 0.46% | 0.38% / 0.29% |
+
+## Interpreting the defaults
+
+The default is five warmup forwards and 20 timed forwards. On Ascend, the combined mean and median CVs at 5/20 ranged from 0.28% to 6.15% across the measured combinations. Increasing warmup or sample counts did not consistently reduce variation across all models and backends.
+
+On the 3080 Ti, the combined median CV at 5/20 ranged from 0.17% to 14.60%. For ViT-B/16 with Inductor, the returned median averaged 4.670 ms at 5/20 and 4.240 ms at 50/20, even though the 5/20 median CV was only 0.17%. A low CV therefore does not by itself establish that the measurement has reached a representative latency. Increasing warmup to 50 did not reduce CV for every measured combination.
+
+These results do not establish either 5/20 or 50/20 as a reliable setting for every workload or device. Compare settings for the workload and environment being measured, and inspect the samples when the mean and median differ substantially.

torch_prisma-0.1.0/docs/index.md

@@ -0,0 +1,41 @@
+# Torch Prisma
+
+Torch Prisma measures PyTorch inference latency, parameters, multiply-accumulates (MACs), and floating-point operations (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)
+inputs = torch.randn(2, 8, device=device)
+
+complexity = calculate_model_complexity(model, inputs)
+latency = measure_module_latency(model, inputs, 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)
+```
+
+## Guides
+
+- [Measurement API and conventions](measurement.md)
+- [Latency benchmarks](benchmarks.md)
+- [Huawei Ascend](ascend.md)

torch_prisma-0.1.0/docs/measurement.md

@@ -0,0 +1,87 @@
+# Measurement API
+
+Import the measurement functions from `torch_prisma`. Set the model's training or evaluation mode and place the model and inputs on their execution devices before measuring. Forward side effects persist.
+
+## `measure_module_latency`
+
+Measures one forward per sample for the entire batch in inference mode. CPU uses a wall clock. CUDA and Ascend use events on the selected device's current stream.
+
+**Parameters**
+
+- `target` — Prepared module or callable.
+- `input_args` — Positional input or tuple/list of positional inputs. Use `(items,)` to pass a list as one argument.
+- `device` — Timing device.
+- `input_kwargs` — Optional keyword arguments.
+- `repetitions` — Number of timed forwards. Defaults to 20.
+- `warmup` — Number of untimed forwards after compilation preparation. Defaults to 5.
+- `compile` — Enable `torch.compile`. Defaults to `True`.
+- `compile_kwargs` — Keyword arguments passed to `torch.compile` when enabled. On Ascend, pass `{"backend": "inductor"}` to use Triton through Inductor. See the [Ascend guide](ascend.md#triton-ascend).
+
+**Returns**
+
+A `LatencyStats` result in milliseconds. Its `samples_ms` preserve sampling order. The remaining fields are `mean_ms`, `median_ms`, `std_ms` (population standard deviation), `min_ms`, `max_ms`, `p90_ms`, and `p99_ms`. Percentiles select the sample at index `min(int(n * q), n - 1)` after sorting, with `q` equal to 0.90 or 0.99. Use `LatencyStats.from_samples()` to summarize external timings and `to_dict()` to serialize the result.
+
+**Raises**
+
+- `ValueError` for invalid timing options or an unsupported timing device.
+- `RuntimeError` when the accelerator or default Ascend backend is unavailable, or compilation is required after preparation.
+
+Compilation runs three untimed forwards before warmup. Set `compile=False` for eager timing. See [Ascend setup](ascend.md) and [latency benchmarks](benchmarks.md).
+
+## `calculate_model_complexity`
+
+Counts operations for one eager forward without gradients. Counting requests math scaled dot-product attention and disables the multihead attention fast path.
+
+**Parameters**
+
+- `target` — Prepared module or callable.
+- `input_args` — Positional input or tuple/list of positional inputs.
+- `input_kwargs` — Optional keyword arguments.
+
+**Returns**
+
+A `ModelComplexity` result for the supplied batch. `params` counts unique registered parameter elements. `macs` and `flops` are known operation subtotals. `uncovered_ops` lists visible operations with incomplete counts, grouped by MACs and FLOPs. A plain callable reports zero parameters. Use `to_dict()` to serialize the result.
+
+Each real floating-point dense MAC contributes two FLOPs. Supported arithmetic outside matrix products adds to FLOPs. Comparisons, ReLU, forward max pooling, indexing, and data movement contribute zero FLOPs. Integer matrix work contributes MACs only. Complex matrix work contributes MACs and a FLOP coverage gap.
+
+## `count_parameters`
+
+Counts the elements in unique registered parameter objects without executing a forward.
+
+**Parameters**
+
+- `model` — Module whose parameters are counted.
+
+**Returns**
+
+The parameter count as an integer. Distinct parameters sharing storage count separately.
+
+## `count_operations`
+
+Counts operations for a prepared zero-argument eager callable.
+
+**Parameters**
+
+- `forward` — Callable that executes the forward to count.
+
+**Returns**
+
+An `OperationCounts` result containing `macs`, `flops`, and `uncovered_ops`.
+
+```python
+from torch_prisma import count_operations
+
+operations = count_operations(lambda: model(inputs))
+```
+
+## Nominal FLOP formulas
+
+Each transcendental operation counts as one FLOP.
+
+- Addition, subtraction, multiplication, and division contribute one FLOP per output element. A nonunit add/subtract scale contributes an additional multiplication.
+- Sum over a nonempty group of `K` elements contributes `K - 1` additions. Mean adds one division.
+- Softmax over a nonempty group of `K` elements contributes `4K - 1` FLOPs.
+- LayerNorm and GroupNorm contribute `5K + 2` FLOPs per group of `K` elements, plus `K` for each affine weight or bias.
+- BatchNorm with stored statistics contributes two FLOPs per channel and two per input element, plus one per element for each affine weight or bias.
+- Matrix product bias and scaling contribute an addition per output element when the bias term is active, plus each applicable nonunit scalar multiplication.
+- Convolution bias contributes one additional FLOP per output element.

torch_prisma-0.1.0/pyproject.toml

@@ -0,0 +1,75 @@
+[build-system]
+requires = ["uv_build>=0.9.11,<0.10.0"]
+build-backend = "uv_build"
+
+[project]
+name = "torch-prisma"
+version = "0.1.0"
+description = "Measure PyTorch inference latency, MACs, FLOPs, and parameters"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+dependencies = ["torch>=2.6"]
+keywords = [
+    "pytorch",
+    "inference",
+    "benchmarking",
+    "latency",
+    "model-complexity",
+    "macs",
+    "flops",
+    "cuda",
+    "ascend",
+]
+classifiers = [
+    "Development Status :: 6 - Mature",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: System :: Benchmark",
+]
+
+[project.urls]
+Repository = "https://github.com/ZaberKo/torch-prisma"
+Documentation = "https://github.com/ZaberKo/torch-prisma/blob/main/docs/index.md"
+Issues = "https://github.com/ZaberKo/torch-prisma/issues"
+
+[dependency-groups]
+dev = [
+    "pre-commit; python_version >= '3.11'",
+    "pytest>=8",
+    "ruff",
+]
+docs = ["zensical"]
+
+[tool.uv.dependency-groups]
+docs = { requires-python = ">=3.11" }
+
+[tool.uv.build-backend]
+module-root = ""
+source-include = ["docs/**", "tests/**"]
+
+[tool.pytest.ini_options]
+minversion = "8.0"
+testpaths = ["tests"]
+markers = ["cuda: requires a CUDA runner", "ascend: requires an Ascend and TorchAir runner"]
+addopts = ["-ra", "--strict-config", "--strict-markers", "--import-mode=importlib"]
+
+[tool.ruff]
+target-version = "py310"
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B", "SIM", "C4", "RUF", "D"]
+ignore = ["E501"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**" = ["D"]
+
+[tool.ruff.format]
+docstring-code-format = true

torch_prisma-0.1.0/tests/test_accelerators.py

@@ -0,0 +1,145 @@
+import os
+import sys
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+
+import torch_prisma.latency as latency
+from torch_prisma import calculate_model_complexity, measure_module_latency
+
+
+def test_ascend_compiler_and_event_wiring(monkeypatch):
+    device = SimpleNamespace(type="npu")
+    stream = object()
+    operations = []
+    forward_limits = []
+    config = torch._dynamo.config
+    limit_name = "recompile_limit" if hasattr(config, "recompile_limit") else "cache_size_limit"
+    accumulated_limit_name = f"accumulated_{limit_name}"
+    original_limits = getattr(config, limit_name), getattr(config, accumulated_limit_name)
+
+    class Event:
+        def __init__(self, *, enable_timing):
+            assert enable_timing
+
+        def record(self, selected_stream):
+            assert selected_stream is stream
+            operations.append("record")
+
+        def elapsed_time(self, end):
+            assert isinstance(end, Event)
+            return 0.25
+
+    def current_stream(selected):
+        assert selected is device
+        return stream
+
+    def synchronize(selected):
+        assert selected is device
+        operations.append("sync")
+
+    runtime = SimpleNamespace(Event=Event, current_stream=current_stream, synchronize=synchronize)
+
+    def timing_device(requested):
+        assert requested == "npu:0"
+        return device, runtime
+
+    def compile_fn(fn, *, backend, fullgraph):
+        assert backend == "torchair-test" and fullgraph
+        return fn
+
+    def forward(x):
+        operations.append("forward")
+        forward_limits.append(
+            (getattr(config, limit_name), getattr(config, accumulated_limit_name))
+        )
+        return x
+
+    backend_factory = Mock(return_value="torchair-test")
+    monkeypatch.setitem(sys.modules, "torchair", SimpleNamespace(get_npu_backend=backend_factory))
+    monkeypatch.setattr(latency, "_timing_device", timing_device)
+    monkeypatch.setattr(torch, "compile", compile_fn)
+    latency._default_npu_backend.cache_clear()
+    try:
+        assert latency._compile_options(device, {"backend": "custom"}) == {"backend": "custom"}
+        backend_factory.assert_not_called()
+        result = measure_module_latency(
+            forward,
+            torch.ones(1),
+            "npu:0",
+            repetitions=2,
+            warmup=1,
+            compile_kwargs={"fullgraph": True},
+        )
+        assert latency._compile_options(device, None) == {"backend": "torchair-test"}
+        backend_factory.assert_called_once_with()
+    finally:
+        latency._default_npu_backend.cache_clear()
+    assert result.samples_ms == (0.25, 0.25)
+    assert operations == (
+        ["forward", "forward", "forward", "sync", "forward"]
+        + ["record"] * 4
+        + ["sync", "record", "forward", "record", "record", "forward", "record", "sync"]
+    )
+    raised_limits = (2**31 - 1, 2**31 - 1)
+    assert forward_limits == [raised_limits] * 3 + [original_limits] * 3
+    assert (getattr(config, limit_name), getattr(config, accumulated_limit_name)) == original_limits
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(os.environ.get("TORCH_PRISMA_TEST_DEVICE") != "cuda", reason="Requires CUDA")
+def test_cuda_measurements():
+    assert torch.cuda.is_available()
+    device = torch.device("cuda", torch.cuda.current_device())
+    model = nn.Linear(4, 3, device=device).eval()
+    x = torch.ones(2, 4, device=device)
+    stats = measure_module_latency(model, x, device, warmup=1, repetitions=3, compile=False)
+    assert len(stats.samples_ms) == 3
+    assert stats.median_ms > 0
+    complexity = calculate_model_complexity(model, x)
+    assert (complexity.params, complexity.macs, complexity.flops) == (15, 24, 54)
+
+
+@pytest.mark.ascend
+@pytest.mark.skipif(
+    os.environ.get("TORCH_PRISMA_TEST_DEVICE") != "ascend", reason="Requires Ascend"
+)
+def test_ascend_measurements():
+    __import__("torch_npu")
+    assert torch.npu.is_available()
+    device = torch.device("npu", torch.npu.current_device())
+    model = nn.Linear(4, 3, device=device).eval()
+    x = torch.ones(2, 4, device=device)
+    for compile in (False, True):
+        stats = measure_module_latency(model, x, device, warmup=2, repetitions=3, compile=compile)
+        assert len(stats.samples_ms) == 3
+        assert stats.median_ms > 0
+
+    complexity = calculate_model_complexity(model, x)
+    assert complexity.params == 15
+    assert complexity.macs == 24 or complexity.uncovered_ops["macs"]
+    assert complexity.flops == 54 or complexity.uncovered_ops["flops"]
+
+
+@pytest.mark.ascend
+@pytest.mark.skipif(
+    os.environ.get("TORCH_PRISMA_TEST_ASCEND_INDUCTOR") != "1",
+    reason="Requires Ascend with Triton Inductor",
+)
+def test_ascend_triton_inductor_measurement():
+    __import__("torch_npu")
+    device = torch.device("npu", torch.npu.current_device())
+    x = torch.randn(8192, device=device)
+
+    def forward(x):
+        return torch.sin(x) + x * x
+
+    stats = measure_module_latency(
+        forward, x, device, warmup=2, repetitions=3, compile_kwargs={"backend": "inductor"}
+    )
+    assert len(stats.samples_ms) == 3
+    assert stats.median_ms > 0
+    torch.testing.assert_close(torch.compile(forward, backend="inductor")(x), forward(x))