flopscope 0.2.0__py3-none-any.whl

benchmarks/_contractions.py

@@ -0,0 +1,291 @@
+"""Benchmark BLAS contraction operations."""
+
+from __future__ import annotations
+
+import statistics
+
+from benchmarks._perf import measure_flops
+
+CONTRACTION_OPS: list[str] = [
+    "dot",
+    "matmul",
+    "inner",
+    "vdot",
+    "vecdot",
+    "outer",
+    "tensordot",
+    "kron",
+    "einsum",
+]
+
+_FORMULA_STRINGS: dict[str, str] = {
+    "dot": "MNK",
+    "matmul": "MNK",
+    "inner": "N (a.size)",
+    "vdot": "N (a.size)",
+    "vecdot": "batch * K (output_size * contracted_axis)",
+    "outer": "M*N",
+    "tensordot": "product of free * contracted dims",
+    "kron": "numel(output)",
+    "einsum": "α/M model (FMA=2 textbook)",
+}
+
+_BENCHMARK_SIZE_STRINGS: dict[str, str] = {
+    "dot": "A: (512,512), B: (512,512)",
+    "matmul": "A: (512,512), B: (512,512)",
+    "inner": "a: (1000000,), b: (1000000,)",
+    "vdot": "a: (1000000,), b: (1000000,)",
+    "vecdot": "A: (1000,512), B: (1000,512)",
+    "outer": "a: (5000,), b: (5000,)",
+    "tensordot": "A: (64,64,64), B: (64,64,64), axes=1",
+    "kron": "A: (64,64), B: (64,64)",
+    "einsum": "A: (512,512), B: (512,512), subscripts='ij,jk->ik'",
+}
+
+
+def _analytical_cost(op: str, **kwargs: int) -> int:
+    """Return analytical FLOP count for the benchmark configuration.
+
+    Parameters
+    ----------
+    op : str
+        Operation name (e.g. ``"dot"``).
+    **kwargs : int
+        Shape parameters used by each formula.
+
+    Returns
+    -------
+    int
+        Analytical FLOP count.
+    """
+    costs: dict[str, int] = {
+        # dot: 2D matrix multiply A(512,512) @ B(512,512), FMA=2 textbook
+        "dot": 512 * 512 * 512,
+        # matmul: identical to dot for 2D
+        "matmul": 512 * 512 * 512,
+        # inner: dot product of two 1M-element vectors.
+        # Runtime charges a.size — matches flopscope's convention (FMA=2, but
+        # a.size is pointwise-shaped so the FMA off-by-one doesn't apply here).
+        "inner": 1_000_000,
+        # vdot: same as inner for 1D real inputs.
+        # Runtime charges a.size (FMA=2, pointwise-shaped — no off-by-one).
+        "vdot": 1_000_000,
+        # vecdot: batched dot product A(1000,512) . B(1000,512)
+        # Output (1000,) with contracted axis 512.
+        # Runtime charges result.size * contracted = 1000 * 512 (FMA=2 textbook).
+        "vecdot": 1000 * 512,
+        # outer: outer product of two 5000-element vectors
+        "outer": 5000 * 5000,
+        # tensordot: A(64,64,64) . B(64,64,64) axes=1 -> contract last of A with first of B
+        "tensordot": 64**5,
+        # kron: Kronecker product A(64,64) x B(64,64)
+        "kron": 64**4,
+        # einsum: 'ij,jk->ik' is matrix multiply (512,512)x(512,512), FMA=2 textbook
+        "einsum": 512 * 512 * 512,
+    }
+    return costs[op]
+
+
+def benchmark_contractions(
+    dtype: str = "float64",
+    repeats: int = 10,
+) -> tuple[dict[str, float], dict[str, dict]]:
+    """Benchmark contraction ops, returning raw measurement per analytical FLOP.
+
+    In perf mode this is actual FP ops / analytical FLOPs (correction factor).
+    In timing mode this is nanoseconds / analytical FLOPs (same units as
+    pointwise -- the runner normalizes against baseline to get relative weights).
+
+    Parameters
+    ----------
+    dtype : str
+        NumPy dtype string.
+    repeats : int
+        Number of repetitions per measurement.
+
+    Returns
+    -------
+    tuple[dict[str, float], dict[str, dict]]
+        A pair of (alphas, details). ``alphas`` maps op name to median
+        raw measurement per analytical FLOP. ``details`` maps op name to
+        a dict of raw benchmark metadata.
+    """
+    results: dict[str, float] = {}
+    details: dict[str, dict] = {}
+
+    for op in CONTRACTION_OPS:
+        dist_values: list[float] = []
+        dist_raw_totals: list[int] = []
+
+        # --- Build setups and bench code per op ---
+
+        if op in ("dot", "matmul", "einsum"):
+            # Two 512x512 matrices
+            setups = [
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.standard_normal((512, 512)).astype(np.{dtype}); "
+                    f"B = rng.standard_normal((512, 512)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(0.01, 100, size=(512, 512)).astype(np.{dtype}); "
+                    f"B = rng.uniform(0.01, 100, size=(512, 512)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(-1000, 1000, size=(512, 512)).astype(np.{dtype}); "
+                    f"B = rng.uniform(-1000, 1000, size=(512, 512)).astype(np.{dtype})"
+                ),
+            ]
+            if op == "dot":
+                bench = "np.dot(A, B)"
+            elif op == "matmul":
+                bench = "np.matmul(A, B)"
+            else:  # einsum
+                bench = "np.einsum('ij,jk->ik', A, B)"
+
+        elif op in ("inner", "vdot"):
+            # Two 1M-element vectors — large enough for BLAS ddot FMA to dominate
+            # over per-call overhead (10K was too small, overhead inflated alpha)
+            vec_n = 1_000_000
+            setups = [
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"a = rng.standard_normal({vec_n}).astype(np.{dtype}); "
+                    f"b = rng.standard_normal({vec_n}).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"a = rng.uniform(0.01, 100, size={vec_n}).astype(np.{dtype}); "
+                    f"b = rng.uniform(0.01, 100, size={vec_n}).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"a = rng.uniform(-1000, 1000, size={vec_n}).astype(np.{dtype}); "
+                    f"b = rng.uniform(-1000, 1000, size={vec_n}).astype(np.{dtype})"
+                ),
+            ]
+            bench = f"np.{op}(a, b)"
+
+        elif op == "vecdot":
+            # Batched dot: A(1000,512), B(1000,512) -- NumPy 2.x only
+            setups = [
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.standard_normal((1000, 512)).astype(np.{dtype}); "
+                    f"B = rng.standard_normal((1000, 512)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(0.01, 100, size=(1000, 512)).astype(np.{dtype}); "
+                    f"B = rng.uniform(0.01, 100, size=(1000, 512)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(-1000, 1000, size=(1000, 512)).astype(np.{dtype}); "
+                    f"B = rng.uniform(-1000, 1000, size=(1000, 512)).astype(np.{dtype})"
+                ),
+            ]
+            bench = "np.vecdot(A, B)"
+
+        elif op == "outer":
+            # Two 5000-element vectors
+            setups = [
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"a = rng.standard_normal(5000).astype(np.{dtype}); "
+                    f"b = rng.standard_normal(5000).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"a = rng.uniform(0.01, 100, size=5000).astype(np.{dtype}); "
+                    f"b = rng.uniform(0.01, 100, size=5000).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"a = rng.uniform(-1000, 1000, size=5000).astype(np.{dtype}); "
+                    f"b = rng.uniform(-1000, 1000, size=5000).astype(np.{dtype})"
+                ),
+            ]
+            bench = "np.outer(a, b)"
+
+        elif op == "tensordot":
+            # Two (64,64,64) tensors, axes=1
+            setups = [
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.standard_normal((64, 64, 64)).astype(np.{dtype}); "
+                    f"B = rng.standard_normal((64, 64, 64)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(0.01, 100, size=(64, 64, 64)).astype(np.{dtype}); "
+                    f"B = rng.uniform(0.01, 100, size=(64, 64, 64)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(-1000, 1000, size=(64, 64, 64)).astype(np.{dtype}); "
+                    f"B = rng.uniform(-1000, 1000, size=(64, 64, 64)).astype(np.{dtype})"
+                ),
+            ]
+            bench = "np.tensordot(A, B, axes=1)"
+
+        elif op == "kron":
+            # Two (64,64) matrices
+            setups = [
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.standard_normal((64, 64)).astype(np.{dtype}); "
+                    f"B = rng.standard_normal((64, 64)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(0.01, 100, size=(64, 64)).astype(np.{dtype}); "
+                    f"B = rng.uniform(0.01, 100, size=(64, 64)).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; rng = np.random.default_rng(42); "
+                    f"A = rng.uniform(-1000, 1000, size=(64, 64)).astype(np.{dtype}); "
+                    f"B = rng.uniform(-1000, 1000, size=(64, 64)).astype(np.{dtype})"
+                ),
+            ]
+            bench = "np.kron(A, B)"
+
+        else:
+            continue  # pragma: no cover
+
+        analytical = _analytical_cost(op)
+
+        for setup in setups:
+            # For vecdot, wrap in try/except since it's NumPy 2.x only
+            if op == "vecdot":
+                try:
+                    result = measure_flops(setup, bench, repeats=repeats)
+                except (RuntimeError, AttributeError):
+                    continue
+            else:
+                try:
+                    result = measure_flops(setup, bench, repeats=repeats)
+                except RuntimeError:
+                    continue
+
+            measured = result.total_flops / repeats
+            dist_values.append(measured / analytical if analytical else 0.0)
+            dist_raw_totals.append(result.total_flops)
+
+        if dist_values:
+            results[op] = statistics.median(dist_values)
+            details[op] = {
+                "category": "counted_custom",
+                "measurement_mode": "blas",
+                "analytical_formula": _FORMULA_STRINGS.get(op, ""),
+                "analytical_flops": analytical,
+                "benchmark_size": _BENCHMARK_SIZE_STRINGS.get(op, ""),
+                "bench_code": bench,
+                "repeats": repeats,
+                "perf_instructions_total": dist_raw_totals,
+                "distribution_alphas": dist_values,
+            }
+
+    return results, details

benchmarks/_fft.py

@@ -0,0 +1,198 @@
+"""Benchmark FFT operations."""
+
+from __future__ import annotations
+
+import math
+import statistics
+
+from benchmarks._perf import measure_flops
+
+FFT_OPS: list[str] = [
+    "fft.fft",
+    "fft.ifft",
+    "fft.rfft",
+    "fft.irfft",
+    "fft.fft2",
+    "fft.ifft2",
+    "fft.rfft2",
+    "fft.irfft2",
+    "fft.fftn",
+    "fft.ifftn",
+    "fft.rfftn",
+    "fft.irfftn",
+    "fft.hfft",
+    "fft.ihfft",
+]
+
+_RFFT_OPS = {
+    "fft.rfft",
+    "fft.irfft",
+    "fft.rfft2",
+    "fft.irfft2",
+    "fft.rfftn",
+    "fft.irfftn",
+}
+
+_FORMULA_STRINGS: dict[str, str] = {
+    "fft.fft": "5*n*ceil(log2(n))",
+    "fft.ifft": "5*n*ceil(log2(n))",
+    "fft.rfft": "5*(n/2)*ceil(log2(n))",
+    "fft.irfft": "5*(n/2)*ceil(log2(n))",
+    "fft.fft2": "5*n*ceil(log2(n))",
+    "fft.ifft2": "5*n*ceil(log2(n))",
+    "fft.rfft2": "5*(n/2)*ceil(log2(n))",
+    "fft.irfft2": "5*(n/2)*ceil(log2(n))",
+    "fft.fftn": "5*n*ceil(log2(n))",
+    "fft.ifftn": "5*n*ceil(log2(n))",
+    "fft.rfftn": "5*(n/2)*ceil(log2(n))",
+    "fft.irfftn": "5*(n/2)*ceil(log2(n))",
+    "fft.hfft": "5*n*ceil(log2(n))",
+    "fft.ihfft": "5*(n/2)*ceil(log2(n))",
+}
+
+
+def _ceil_log2(n: int) -> int:
+    """Return ceil(log2(n)), minimum 1."""
+    if n <= 1:
+        return 1
+    return math.ceil(math.log2(n))
+
+
+def _analytical_cost(op_name: str, n: int) -> int:
+    """Return the analytical FLOP count for an FFT operation.
+
+    Parameters
+    ----------
+    op_name : str
+        Operation name (e.g. ``"fft.fft"``).
+    n : int
+        Input size.
+
+    Returns
+    -------
+    int
+        Analytical FLOP count.
+    """
+    cl2 = _ceil_log2(n)
+    if op_name in _RFFT_OPS:
+        return 5 * (n // 2) * cl2
+    return 5 * n * cl2
+
+
+def benchmark_fft(
+    n: int = 2**20,
+    dtype: str = "float64",
+    repeats: int = 10,
+) -> tuple[dict[str, float], dict[str, dict]]:
+    """Benchmark FFT ops, returning correction factors (measured / analytical).
+
+    Parameters
+    ----------
+    n : int
+        Input size.
+    dtype : str
+        NumPy dtype string.
+    repeats : int
+        Number of repetitions per measurement.
+
+    Returns
+    -------
+    tuple[dict[str, float], dict[str, dict]]
+        A pair of (alphas, details). ``alphas`` maps op name to median
+        correction factor. ``details`` maps op name to a dict of raw
+        benchmark metadata.
+    """
+    results: dict[str, float] = {}
+    details: dict[str, dict] = {}
+
+    # 2D/nD ops use sqrt(n) x sqrt(n)
+    side = int(math.isqrt(n))
+
+    for op in FFT_OPS:
+        dist_values: list[float] = []
+        dist_raw_totals: list[int] = []
+
+        # Determine dimensionality
+        is_2d = "2" in op.split(".")[-1]
+        is_nd = op.endswith("n") or op.endswith("fftn")
+        is_multi = is_2d or is_nd
+
+        if is_multi:
+            shape_str = f"({side}, {side})"
+            effective_n = side * side
+            benchmark_size = f"x: ({side},{side})"
+        else:
+            shape_str = f"({n},)"
+            effective_n = n
+            benchmark_size = f"x: ({n},)"
+
+        # Determine input type needed
+        short = op.split(".")[-1]
+
+        # irfft variants need rfft output (complex, half-size)
+        needs_rfft_input = short in ("irfft", "irfft2", "irfftn")
+        # ifft variants and hfft need complex input
+        needs_complex = short in ("ifft", "ifft2", "ifftn", "hfft")
+
+        setups = []
+        if needs_rfft_input:
+            # Generate input by applying rfft to real data
+            rfft_func = short.replace("i", "", 1)  # irfft -> rfft
+            setups = [
+                f"import numpy as np; _r = np.random.default_rng(42).standard_normal({shape_str}).astype(np.{dtype}); x = np.fft.{rfft_func}(_r)",
+                f"import numpy as np; _r = np.random.default_rng(43).uniform(-1, 1, size={shape_str}).astype(np.{dtype}); x = np.fft.{rfft_func}(_r)",
+                f"import numpy as np; _r = np.random.default_rng(99).standard_normal({shape_str}).astype(np.{dtype}) * 100; x = np.fft.{rfft_func}(_r)",
+            ]
+        elif needs_complex:
+            setups = [
+                (
+                    f"import numpy as np; x = np.random.default_rng(42).standard_normal({shape_str}).astype(np.{dtype}) "
+                    f"+ 1j * np.random.default_rng(43).standard_normal({shape_str}).astype(np.{dtype})"
+                ),
+                (
+                    f"import numpy as np; t = np.linspace(0, 2*np.pi, {effective_n}).reshape({shape_str}).astype(np.{dtype}); "
+                    f"x = np.sin(t) + 1j * np.cos(t)"
+                ),
+                (
+                    f"import numpy as np; x = (np.random.default_rng(99).uniform(-1, 1, size={shape_str}).astype(np.{dtype}) "
+                    f"+ 1j * np.random.default_rng(100).uniform(-1, 1, size={shape_str}).astype(np.{dtype}))"
+                ),
+            ]
+        else:
+            # Real input: fft, rfft, fftn, rfftn, ihfft
+            setups = [
+                f"import numpy as np; x = np.random.default_rng(42).standard_normal({shape_str}).astype(np.{dtype})",
+                (
+                    f"import numpy as np; t = np.linspace(0, 2*np.pi, {effective_n}).reshape({shape_str}).astype(np.{dtype}); "
+                    f"x = np.sin(5*t) + 0.5*np.sin(13*t)"
+                ),
+                f"import numpy as np; x = np.random.default_rng(99).uniform(-1, 1, size={shape_str}).astype(np.{dtype})",
+            ]
+
+        bench = f"np.{op}(x)"
+        analytical = _analytical_cost(op, effective_n)
+
+        for setup in setups:
+            try:
+                result = measure_flops(setup, bench, repeats=repeats)
+            except RuntimeError:
+                continue
+            measured = result.total_flops / repeats
+            dist_values.append(measured / analytical if analytical else 0.0)
+            dist_raw_totals.append(result.total_flops)
+
+        if dist_values:
+            results[op] = statistics.median(dist_values)
+            details[op] = {
+                "category": "counted_custom",
+                "measurement_mode": "custom",
+                "analytical_formula": _FORMULA_STRINGS.get(op, ""),
+                "analytical_flops": analytical,
+                "benchmark_size": benchmark_size,
+                "bench_code": bench,
+                "repeats": repeats,
+                "perf_instructions_total": dist_raw_totals,
+                "distribution_alphas": dist_values,
+            }
+
+    return results, details

benchmarks/_impl_urls.py

@@ -0,0 +1,139 @@
+"""Map operation names to GitHub source code URLs.
+
+For each operation in flopscope, finds where the FLOP cost is charged in the
+source tree and constructs a GitHub permalink.
+
+Strategy (tried in order for each op):
+1. Literal ``budget.deduct("op_name", ...)`` call — works for hand-written ops.
+2. Factory registration like ``sin = _counted_unary(_np.sin, "sin")`` — works
+   for pointwise ops created via factory helpers.
+3. Falls back to file-level URL if only the file is identified.
+"""
+
+from __future__ import annotations
+
+import re
+import subprocess
+from pathlib import Path
+
+REPO_URL = "https://github.com/AIcrowd/flopscope/blob/main"
+SRC_ROOT = Path(__file__).resolve().parent.parent / "src" / "flopscope"
+
+
+def _grep(pattern: str, path: str | None = None) -> list[tuple[str, int, str]]:
+    """Run grep -rn and return list of (abs_path, line_no, line_text)."""
+    target = path or str(SRC_ROOT)
+    try:
+        result = subprocess.run(
+            ["grep", "-rn", "-E", pattern, target],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        return []
+
+    hits: list[tuple[str, int, str]] = []
+    if result.returncode == 0 and result.stdout.strip():
+        for line in result.stdout.strip().splitlines():
+            parts = line.split(":", 2)
+            if len(parts) >= 3:
+                try:
+                    hits.append((parts[0], int(parts[1]), parts[2]))
+                except ValueError:
+                    pass
+    return hits
+
+
+def _to_rel_path(abs_path: str) -> str:
+    """Convert absolute path to repo-relative path."""
+    repo_root = SRC_ROOT.parent.parent  # up from src/flopscope
+    try:
+        return str(Path(abs_path).relative_to(repo_root))
+    except ValueError:
+        return str(abs_path)
+
+
+def _find_deduct_line(op_name: str) -> tuple[str, int | None]:
+    """Find where the FLOP cost for *op_name* is charged in the source tree.
+
+    Returns ``(repo_relative_path, line_number)`` or ``("", None)``.
+    """
+    escaped = re.escape(op_name)
+
+    # --- Strategy 1: direct deduct("op_name" or deduct('op_name' --------
+    for quote in ('"', "'"):
+        hits = _grep(f"deduct\\({quote}{escaped}{quote}")
+        if hits:
+            abs_path, line_no, _ = hits[0]
+            return (_to_rel_path(abs_path), line_no)
+
+    # --- Strategy 2: factory registration like  sin = _counted_unary(..., "sin")
+    # Also handles _counted_binary, _counted_sampler, etc.
+    for quote in ('"', "'"):
+        # Match patterns like:
+        #   sin = _counted_unary(_np.sin, "sin")
+        #   random.rand = _counted_dims_sampler(_np.random.rand, "random.rand")
+        pattern = f"_counted_\\w+\\([^)]*,\\s*{quote}{escaped}{quote}"
+        hits = _grep(pattern)
+        if hits:
+            abs_path, line_no, _ = hits[0]
+            return (_to_rel_path(abs_path), line_no)
+
+    # --- Strategy 3: for reduction ops created via loops or dicts --------
+    # Search for the string literal "op_name" near a deduct or factory call
+    for quote in ('"', "'"):
+        pattern = f"{quote}{escaped}{quote}"
+        hits = _grep(pattern)
+        # Filter to source files (not test, not __pycache__, not data/)
+        src_hits = [
+            h
+            for h in hits
+            if "/src/flopscope/" in h[0]
+            and "__pycache__" not in h[0]
+            and "/data/" not in h[0]
+            and "_registry.py" not in h[0]
+            and "_docstrings.py" not in h[0]
+        ]
+        if src_hits:
+            abs_path, line_no, _ = src_hits[0]
+            return (_to_rel_path(abs_path), line_no)
+
+    return ("", None)
+
+
+def map_op_to_url(op_name: str) -> str:
+    """Return GitHub URL for the runtime cost implementation of an operation.
+
+    Returns an empty string if the source cannot be located.
+    """
+    rel_path, line_no = _find_deduct_line(op_name)
+    if not rel_path:
+        return ""
+    url = f"{REPO_URL}/{rel_path}"
+    if line_no is not None:
+        url += f"#L{line_no}"
+    return url
+
+
+def build_url_map(op_names: list[str] | None = None) -> dict[str, str]:
+    """Build URL map for all given operation names.
+
+    Parameters
+    ----------
+    op_names : list of str, optional
+        If *None*, reads all operation names from ``weights.json``.
+
+    Returns
+    -------
+    dict mapping operation name to GitHub URL (empty string if not found).
+    """
+    if op_names is None:
+        import json
+
+        weights_path = SRC_ROOT / "data" / "weights.json"
+        with open(weights_path) as f:
+            data = json.load(f)
+        op_names = list(data["weights"].keys())
+
+    return {op: map_op_to_url(op) for op in op_names}