counted-float 1.0.1__py3-none-any.whl

counted_float/__init__.py

@@ -0,0 +1,39 @@
+import counted_float.config as config
+
+from ._core.counting import BuiltInData, CountedFloat, FlopCountingContext, PauseFlopCounting
+from ._core.counting.models import (
+    FlopCounts,
+    FlopsBenchmarkDurations,
+    FlopsBenchmarkResults,
+    FlopType,
+    FlopWeights,
+    FPUInstruction,
+    SystemInfo,
+)
+
+__all__ = [
+    "config",
+    "CountedFloat",
+    "FlopCountingContext",
+    "FlopCounts",
+    "FlopsBenchmarkDurations",
+    "FlopsBenchmarkResults",
+    "FlopType",
+    "FlopWeights",
+    "FPUInstruction",
+    "PauseFlopCounting",
+    "SystemInfo",
+]
+
+# -------------------------------------------------------------------------
+#  Optional [benchmarking] dependencies
+# -------------------------------------------------------------------------
+from ._core._optional_deps import FLAG_BENCHMARK_DEPS
+
+if FLAG_BENCHMARK_DEPS:
+    import counted_float.benchmarking as benchmarking
+
+    __all__.append("benchmarking")
+
+# delete variable again, we don't want to expose this
+del FLAG_BENCHMARK_DEPS

counted_float/_core/_backwards_compat.py

@@ -0,0 +1,6 @@
+import sys
+
+if sys.version_info >= (3, 11):
+    from enum import StrEnum
+else:
+    from backports.strenum import StrEnum

counted_float/_core/_optional_deps.py

@@ -0,0 +1,49 @@
+"""
+Module implementing boolean flags indicating whether optional dependencies are available & decorators to
+safeguard execution of functions that require these dependencies.
+"""
+
+import functools
+from typing import Callable
+
+# -------------------------------------------------------------------------
+#  Flags
+# -------------------------------------------------------------------------
+try:
+    import numba
+    import numpy
+    import psutil
+
+    FLAG_BENCHMARK_DEPS = True
+except ImportError:
+    FLAG_BENCHMARK_DEPS = False
+
+
+# -------------------------------------------------------------------------
+#  Decorators
+# -------------------------------------------------------------------------
+def requires_benchmark_deps(fun: Callable) -> Callable:
+    """
+    Decorator to mark a function that requires some deps to be installed.
+
+    USAGE:
+
+        @requires_benchmark_deps
+        def fun_that_needs_numba_or_numpy_or_psutil():
+            pass
+
+    This will ensure that the function - when it is called for the first time - will check if the required
+    dependencies are available.  If not, an ImportError will be raised with an informative message.
+    """
+
+    @functools.wraps(fun)
+    def wrapped_fun(*args, **kwargs):
+        if not FLAG_BENCHMARK_DEPS:
+            raise ImportError(
+                "This function requires counted_float to be installed with benchmark optional dependencies. "
+                + "Please install this package as counted_float[benchmark]."
+            )
+        else:
+            return fun(*args, **kwargs)
+
+    return wrapped_fun

counted_float/_core/benchmarking/__init__.py

@@ -0,0 +1,10 @@
+from counted_float._core._optional_deps import requires_benchmark_deps
+from counted_float._core.counting.models import FlopsBenchmarkResults
+
+from ._flops_benchmark_suite import FlopsBenchmarkSuite
+
+
+@requires_benchmark_deps
+def run_flops_benchmark() -> FlopsBenchmarkResults:
+    """Run the flops benchmark suite with default settings returns a FlopsBenchmarkResults object."""
+    return FlopsBenchmarkSuite().run()

counted_float/_core/benchmarking/_flops_benchmark_suite.py

@@ -0,0 +1,189 @@
+import platform
+
+import psutil
+
+from counted_float._core._optional_deps import requires_benchmark_deps
+from counted_float._core.counting.models import (
+    BenchmarkSettings,
+    FlopsBenchmarkDurations,
+    FlopsBenchmarkResults,
+    FlopType,
+    Quantiles,
+    SystemInfo,
+)
+
+from ._flops_micro_benchmark import FlopsMicroBenchmark
+
+
+class FlopsBenchmarkSuite:
+    # -------------------------------------------------------------------------
+    #  Main API
+    # -------------------------------------------------------------------------
+    def run(
+        self,
+        array_size: int = 1000,
+        n_runs_total: int = 30,
+        n_runs_warmup: int = 10,
+        n_seconds_per_run_target: float = 0.5,
+    ) -> FlopsBenchmarkResults:
+        """
+        Run entire flops benchmarking suite and return the results as a FlopsBenchmarkResults object.
+        """
+
+        # run actual benchmarks
+        benchmarks = self.get_flops_benchmarking_suite(size=array_size)
+        results_dict: dict[FlopType | None, Quantiles] = {
+            flop_type: benchmark.run_many(
+                n_runs_total=n_runs_total,
+                n_runs_warmup=n_runs_warmup,
+                n_seconds_per_run_target=n_seconds_per_run_target,
+            ).summary_stats()
+            for flop_type, benchmark in benchmarks.items()
+        }
+
+        # put results in appropriate format & return
+        return FlopsBenchmarkResults(
+            system_info=SystemInfo(
+                platform_processor=platform.processor(),
+                platform_machine=platform.machine(),
+                platform_system=platform.system(),
+                platform_release=platform.release(),
+                platform_python_version=platform.python_version(),
+                platform_python_implementation=platform.python_implementation(),
+                platform_python_compiler=platform.python_compiler(),
+                psutil_cpu_count_logical=psutil.cpu_count(logical=True),
+                psutil_cpu_count_physical=psutil.cpu_count(logical=False),
+            ),
+            benchmark_settings=BenchmarkSettings(
+                array_size=array_size,
+                n_runs_total=n_runs_total,
+                n_runs_warmup=n_runs_warmup,
+                n_seconds_per_run_target=n_seconds_per_run_target,
+            ),
+            results_ns=FlopsBenchmarkDurations(
+                baseline=results_dict[None],
+                flops={flop_type: results_dict[flop_type] for flop_type in FlopType},
+            ),
+        )
+
+    # -------------------------------------------------------------------------
+    #  Static methods
+    # -------------------------------------------------------------------------
+    @staticmethod
+    @requires_benchmark_deps
+    def get_flops_benchmarking_suite(size: int) -> dict[FlopType | None, FlopsMicroBenchmark]:
+        """
+        Returns a benchmark for each FlopType + None (=baseline test), of requested array size.
+        """
+
+        # --- late import of optional dependencies --------
+        import numba
+        import numpy as np
+
+        # --- define all test functions -------------------
+        @numba.njit(parallel=False)
+        def baseline(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            """baseline benchmark to measure the overhead of the benchmarking framework + iteration"""
+            for i in range(n):
+                pass
+
+        @numba.njit(parallel=False)
+        def flop_abs(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = abs(in_f1[i])
+
+        @numba.njit(parallel=False)
+        def flop_minus(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = -in_f1[i]
+
+        @numba.njit(parallel=False)
+        def flop_equals(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_i[i] = in_f1[i] == in_f2[i]  # assign to integer output, to avoid unnecessary conversion overhead
+
+        @numba.njit(parallel=False)
+        def flop_gte(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_i[i] = in_f1[i] >= in_f2[i]  # assign to integer output, to avoid unnecessary conversion overhead
+
+        @numba.njit(parallel=False)
+        def flop_lte(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_i[i] = in_f1[i] <= in_f2[i]  # assign to integer output, to avoid unnecessary conversion overhead
+
+        @numba.njit(parallel=False)
+        def flop_gte_zero(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_i[i] = in_f1[i] >= 0.0  # assign to integer output, to avoid unnecessary conversion overhead
+
+        @numba.njit(parallel=False)
+        def flop_rnd(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_i[i] = np.ceil(in_f1[i])
+
+        @numba.njit(parallel=False)
+        def flop_add(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = in_f1[i] + in_f2[i]
+
+        @numba.njit(parallel=False)
+        def flop_sub(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = in_f1[i] - in_f2[i]
+
+        @numba.njit(parallel=False)
+        def flop_mul(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = in_f1[i] * in_f2[i]
+
+        @numba.njit(parallel=False)
+        def flop_div(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = in_f1[i] / in_f2[i]
+
+        @numba.njit(parallel=False)
+        def flop_sqrt(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = np.sqrt(in_f1[i])
+
+        @numba.njit(parallel=False)
+        def flop_pow2(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = 2 ** in_f1[i]
+
+        @numba.njit(parallel=False)
+        def flop_log2(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = np.log2(in_f1[i])
+
+        @numba.njit(parallel=False)
+        def flop_pow(n: int, in_f1: np.ndarray, in_f2: np.ndarray, out_f: np.ndarray, out_i: np.ndarray):
+            for i in range(n):
+                out_f[i] = in_f1[i] ** in_f2[i]
+
+        # --- return in appropriate format ----------------
+        return {
+            key: FlopsMicroBenchmark(name=name, f=f, size=size)
+            for key, name, f in [
+                (key, key.long_name() if key else "baseline", f)
+                for key, f in [
+                    (None, baseline),
+                    (FlopType.ABS, flop_abs),
+                    (FlopType.CMP_ZERO, flop_gte_zero),
+                    (FlopType.RND, flop_rnd),
+                    (FlopType.MINUS, flop_minus),
+                    (FlopType.EQUALS, flop_equals),
+                    (FlopType.GTE, flop_gte),
+                    (FlopType.LTE, flop_lte),
+                    (FlopType.ADD, flop_add),
+                    (FlopType.SUB, flop_sub),
+                    (FlopType.MUL, flop_mul),
+                    (FlopType.SQRT, flop_sqrt),
+                    (FlopType.DIV, flop_div),
+                    (FlopType.POW2, flop_pow2),
+                    (FlopType.LOG2, flop_log2),
+                    (FlopType.POW, flop_pow),
+                ]
+            ]
+        }

counted_float/_core/benchmarking/_flops_micro_benchmark.py

@@ -0,0 +1,64 @@
+from typing import Callable
+
+from counted_float._core._optional_deps import requires_benchmark_deps
+
+from ._micro_benchmark import MicroBenchmark
+
+
+class FlopsMicroBenchmark(MicroBenchmark):
+    """
+    Base class for benchmark that checks speed of a certain type of floating point operation.
+
+    This is set up as follows:
+      - we configure the benchmark with a 'size' and a function 'f'
+      - we prepare the inputs: 2 1D numpy arrays of size 'size': in_f1, in_f2
+         - initialized as random floating point numbers in range [0, 10]
+      - we prepare the output arrays: of size 'size': out_f, out_i
+         - 1 output array per type of result: float, int
+         - initialized with zeros
+      - the function f will loop over a&b and write the result to c
+         - the function should be implemented such that it does not use vectorized operations, we want to avoid
+             using vectorized CPU instructions (AVX, etc...): we want to measure the speed of the regular, scalar
+             operations
+         - we numba.jit the function, to make sure it is compiled to machine code, to avoid Python overhead to dominate
+             the benchmark
+         - 'size' should be chosen large enough to avoid the overhead of the benchmarking framework to be noticeable
+         - 'size' should be chosen small enough to fit in the CPU cache, to avoid being RAM-bandwidth limited
+
+    Despite all these measures, it is not expected that the benchmark will be fully accurate in absolute terms.
+    In real-life the speed of execution of floating point operations on a CPU will also be influenced by branching,
+    memory access patterns, etc...
+    The main goal is to find a reasonably accurate estimate of the relative cost of the main types of floating point
+    operations, so we can make representative estimates of the number of FLOPS executed by instrumented algorithms.
+    """
+
+    @requires_benchmark_deps
+    def __init__(self, name: str, f: Callable, size: int):
+        import numpy as np  # late import of optional dependency
+
+        super().__init__(name=name)
+        self.size = size
+        self.f = f
+        self.n_operations = 0
+        # input arrays
+        self.in_f1: np.ndarray = np.zeros(size, dtype=float)
+        self.in_f2: np.ndarray = np.zeros(size, dtype=float)
+        # output arrays
+        self.out_f: np.ndarray = np.zeros(size, dtype=float)
+        self.out_i: np.ndarray = np.zeros(size, dtype=int)
+
+    def _prepare_benchmark(self, n_operations: int):
+        import numpy as np  # late import of optional dependency
+
+        self.n_operations = n_operations
+        # input arrays
+        self.in_f1 = 10 * np.random.rand(self.size)
+        self.in_f2 = 10 * np.random.rand(self.size)
+        # output arrays
+        self.out_f: np.ndarray = np.full(self.size, 0.0, dtype=float)
+        self.out_i: np.ndarray = np.full(self.size, 0, dtype=int)
+
+    def _run_benchmark(self):
+        # repeat 'f' n_operations times, each time on the same data
+        for _ in range(self.n_operations):
+            self.f(self.size, self.in_f1, self.in_f2, self.out_f, self.out_i)

counted_float/_core/benchmarking/_micro_benchmark.py

@@ -0,0 +1,113 @@
+from abc import ABC, abstractmethod
+
+from counted_float._core._optional_deps import requires_benchmark_deps
+
+from ._models import MicroBenchmarkResult, SingleRunResult
+from ._time_utils import Timer, format_time_durations
+
+
+# =================================================================================================
+#  Base class for micro-benchmarks
+# =================================================================================================
+class MicroBenchmark(ABC):
+    """
+    Base class for micro-benchmarks, where a child class needs to implement the following methods:
+      _prepare_benchmark --> prepares the benchmark_runs (e.g. sets up data); is called once before each _run_benchmark
+                                and time spent here is not counted.
+      _run_benchmark     --> runs the benchmark_runs; is called multiple times and time spent here is counted.
+
+    NOTE: this essentially offers similar functionality as the python-builtin timeit module, but with some benefits:
+            - automatic resizing of the benchmark_runs to achieve a target time per run
+            - automatic warmup_runs runs
+            - robust computation of median run time and ± range based on quantiles, rather than mean & std.
+                  (=more robust to outliers)
+    """
+
+    MAX_N_OPERATIONS_FACTOR = 10  # never adjust n_operations by more than this factor (up or down)
+
+    @requires_benchmark_deps
+    def __init__(self, name: str):
+        self.name = name
+
+    def run_many(
+        self, n_runs_total: int = 20, n_runs_warmup: int = 5, n_seconds_per_run_target: float = 0.5
+    ) -> MicroBenchmarkResult:
+        """
+        Runs the MicroBenchmark multiple times (warmup_runs & actual test runs), with provided parameters and returns
+        (q25, q50, q75) quantiles of run times in nanoseconds.
+        :param n_runs_total: (int, default=20) total number of benchmark_runs runs
+        :param n_runs_warmup: (int, default=5) number of warmup_runs runs
+                                             - the first n_runs_warmup of n_runs_total are not included in timing stats
+                                             - warmup_runs runs serve to initialize n_operations & get processor, cache, ...
+                                                 in a stable, representative state
+        :param n_seconds_per_run_target: (float, default=0.5) target time (sec) per benchmark_runs run (prepare + run).
+                                             n_operations will be iteratively adjusted to achieve this target time.
+        """
+
+        print(f"{self.name.ljust(35)}: ", end="")
+
+        # repeat benchmark_runs n_runs_total times
+        n_operations = 1  # start with a benchmark_runs of 1 operation and scale up as needed
+        warmup_runs: list[SingleRunResult] = []
+        benchmark_runs: list[SingleRunResult] = []
+        for i in range(n_runs_total):
+            # --- run benchmark_runs ---
+            with Timer() as t:
+                single_run_result = self.run_once(n_operations)
+            t_tot_seconds = t.t_elapsed_sec()  # total time (sec) of prepare + run
+
+            # --- capture result ---
+            if i < n_runs_warmup:
+                # warmup run that doesn't count
+                print("w", end="", flush=True)
+                warmup_runs.append(single_run_result)
+            else:
+                # benchmark run that does count
+                print(".", end="", flush=True)
+                benchmark_runs.append(single_run_result)
+
+            # --- adjust n_ops ---
+            n_ops_min = max(1, int(n_operations / self.MAX_N_OPERATIONS_FACTOR))
+            n_ops_max = int(n_operations * self.MAX_N_OPERATIONS_FACTOR)
+            n_operations = max(n_ops_min, min(n_ops_max, int(n_operations * n_seconds_per_run_target / t_tot_seconds)))
+
+        # final results
+        benchmark_result = MicroBenchmarkResult(
+            warmup_runs=warmup_runs,
+            benchmark_runs=benchmark_runs,
+        )
+        stats = benchmark_result.summary_stats()
+        print(f"   {format_time_durations(nsec_q25=stats.q25, nsec_q50=stats.q50, nsec_q75=stats.q75)}")
+
+        # return quantiles
+        return benchmark_result
+
+    def run_once(self, n_operations: int) -> SingleRunResult:
+        """Runs benchmark_runs once for a given # of operations and returns time in nanoseconds per operation"""
+
+        # prepare
+        self._prepare_benchmark(n_operations)
+
+        # run
+        with Timer() as t:
+            self._run_benchmark()
+
+        # report
+        return SingleRunResult(
+            n_operations=n_operations,
+            t_nsecs=t.t_elapsed_nsec(),
+        )
+
+    @abstractmethod
+    def _prepare_benchmark(self, n_operations: int):
+        """
+        Prepare benchmark_runs (e.g. set up data) based on requested number of operations.  This argument is adjusted each
+          run by the MicroBenchmarkRunner class to ensure that the benchmark_runs runs for a reasonable amount of time
+            (e.g. 1 second per run).
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def _run_benchmark(self):
+        """Run benchmark_runs.  This method is called multiple times and the time spent here is measured."""
+        raise NotImplementedError()

counted_float/_core/benchmarking/_models.py

@@ -0,0 +1,31 @@
+from counted_float._core.counting.models import MyBaseModel, Quantiles
+
+
+class SingleRunResult(MyBaseModel):
+    """Result of a single run of our micro-benchmark_runs for a give # of operations."""
+
+    n_operations: int
+    t_nsecs: float
+
+    def nsecs_per_op(self) -> float:
+        return self.t_nsecs / self.n_operations
+
+
+class MicroBenchmarkResult(MyBaseModel):
+    """Results of all runs in the micro-benchmark_runs (warmup_runs + actual benchmark_runs runs)."""
+
+    warmup_runs: list[SingleRunResult]
+    benchmark_runs: list[SingleRunResult]
+
+    def get_nsec_per_op_quantile(self, q: float) -> float:
+        """Returns a specific quantile of all results in the 'benchmark_runs' category expressed as nsec/op."""
+        import numpy as np  # late import of optional dependency
+
+        return float(np.quantile([el.nsecs_per_op() for el in self.benchmark_runs], q))
+
+    def summary_stats(self) -> Quantiles:
+        return Quantiles(
+            q25=self.get_nsec_per_op_quantile(q=0.25),
+            q50=self.get_nsec_per_op_quantile(q=0.50),
+            q75=self.get_nsec_per_op_quantile(q=0.75),
+        )

counted_float/_core/benchmarking/_time_utils.py

@@ -0,0 +1,70 @@
+import time
+
+
+# =================================================================================================
+#  Formatting
+# =================================================================================================
+def format_time_durations(nsec_q25: float, nsec_q50: float, nsec_q75: float) -> str:
+    if nsec_q50 < 1e3:
+        formatter = _format_nsec_as_ns
+    elif nsec_q50 < 1e6:
+        formatter = _format_nsec_as_us
+    elif nsec_q50 < 1e9:
+        formatter = _format_nsec_as_ms
+    else:
+        formatter = _format_nsec_as_s
+
+    return f"{formatter(nsec_q50)} ± {formatter((nsec_q75 - nsec_q25) / 2)} / operation"
+
+
+def _format_nsec_as_ns(nsec: float) -> str:
+    return f"{nsec:7.2f} ns"
+
+
+def _format_nsec_as_us(nsec: float) -> str:
+    return f"{nsec / 1e3:7.2f} µs"
+
+
+def _format_nsec_as_ms(nsec: float) -> str:
+    return f"{nsec / 1e6:7.2f} ms"
+
+
+def _format_nsec_as_s(nsec: float) -> str:
+    return f"{nsec / 1e9:7.2f} s"
+
+
+# =================================================================================================
+#  Timing
+# =================================================================================================
+class Timer:
+    """
+    Class acting as a context manager to measure time elapsed.  Elapsed time can be retrieved using t_elapsed().
+    """
+
+    # --- constructor -------------------------------------
+    def __init__(self):
+        self._start: float | None = None
+        self._end: float | None = None
+
+    # --- context manager ---------------------------------
+    def __enter__(self):
+        self._start = time.perf_counter_ns()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._end = time.perf_counter_ns()
+
+    # --- extract results ---------------------------------
+    def t_elapsed_nsec(self) -> float:
+        if self._start is None:
+            raise RuntimeError("Timer has not been started.")
+
+        if self._end is None:
+            # timer still running
+            return time.perf_counter_ns() - self._start
+        else:
+            # timer finished
+            return self._end - self._start
+
+    def t_elapsed_sec(self) -> float:
+        return self.t_elapsed_nsec() / 1e9

counted_float/_core/counting/__init__.py

@@ -0,0 +1,3 @@
+from ._builtin_data import BuiltInData
+from ._context_managers import FlopCountingContext, PauseFlopCounting
+from ._counted_float import CountedFloat

counted_float/_core/counting/_builtin_data.py

@@ -0,0 +1,25 @@
+from importlib.resources import files
+
+from counted_float._core.counting.models import FlopsBenchmarkResults, InstructionLatencies
+
+
+class BuiltInData:
+    """
+    A class that provides access to built-in data for the counted_float package.
+    """
+
+    @classmethod
+    def benchmarks(cls) -> dict[str, FlopsBenchmarkResults]:
+        return {
+            file.stem: FlopsBenchmarkResults.model_validate_json(file.read_text())
+            for file in files("counted_float._core.data.benchmarks").iterdir()
+            if file.is_file() and file.name.endswith(".json")
+        }
+
+    @classmethod
+    def specs(cls) -> dict[str, InstructionLatencies]:
+        return {
+            file.stem: InstructionLatencies.model_validate_json(file.read_text())
+            for file in files("counted_float._core.data.specs").iterdir()
+            if file.is_file() and file.name.endswith(".json")
+        }