counted-float 1.0.1__py3-none-any.whl

counted_float/_core/counting/_context_managers.py

@@ -0,0 +1,95 @@
+"""
+Context manager that conveniently allows collection of flop counts for the enclosed code block,
+as well as providing .pause() and .resume() methods to control flop counting.
+"""
+
+from counted_float._core.counting._global_counter import GLOBAL_COUNTER
+from counted_float._core.counting.models import FlopCounts
+
+
+# =================================================================================================
+#  FlopCountingContext
+# =================================================================================================
+class FlopCountingContext:
+    """
+    Context manager that can be used to count FLOP operations in a block of code.  Only floating-point
+    operations of CountedFloat objects are counted.  So make sure all math uses this type.
+
+    Flops need to be registered by either of the following:
+      - calls to register_flops(...)
+      - using CountedFloat() objects in the computations
+
+    LIMITATIONS:
+        - this context manager is not thread-safe
+        - not _all_ floating-point operations are counted, see the docs for more details.
+    """
+
+    # -------------------------------------------------------------------------
+    #  Constructor
+    # -------------------------------------------------------------------------
+    def __init__(self):
+        # Active/inactive flag  (toggled by __enter__ and __exit__ + by pause() and resume() methods)
+        # When inactive:
+        #   - current count == self.__cnt_subtotal
+        # When active:
+        #   - current count == GLOBAL_COUNTER - self.__cnt_start_snapshot
+        self.__active: bool = False
+
+        # flop count bookkeeping
+        self.__cnt_subtotal: FlopCounts = FlopCounts()
+        self.__cnt_start_snapshot: FlopCounts = FlopCounts()
+
+    # -------------------------------------------------------------------------
+    #  Properties
+    # -------------------------------------------------------------------------
+    def is_active(self) -> bool:
+        return self.__active
+
+    def flop_counts(self) -> FlopCounts:
+        """Returns current total flop count for this context manager.  See constructor comments for details."""
+        if self.__active:
+            return GLOBAL_COUNTER.flop_counts() - self.__cnt_start_snapshot
+        else:
+            return self.__cnt_subtotal.copy()
+
+    # -------------------------------------------------------------------------
+    #  Pause/Resume
+    # -------------------------------------------------------------------------
+    def pause(self):
+        if self.__active:
+            self.__cnt_subtotal = self.flop_counts()
+            self.__cnt_start_snapshot = FlopCounts()
+            self.__active = False
+
+    def resume(self):
+        if not self.__active:
+            self.__cnt_start_snapshot = GLOBAL_COUNTER.flop_counts() - self.__cnt_subtotal
+            self.__cnt_subtotal = FlopCounts()
+            self.__active = True
+
+    # -------------------------------------------------------------------------
+    #  Context manager interface
+    # -------------------------------------------------------------------------
+    def __enter__(self):
+        self.resume()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.pause()
+
+
+# =================================================================================================
+#  PauseFlopCounting
+# =================================================================================================
+class PauseFlopCounting:
+    """
+    Context manager that pauses flop counting for the enclosed code block.  This acts globally, across all
+    active FlopCountingContext instances.
+    """
+
+    def __enter__(self):
+        GLOBAL_COUNTER.pause()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        GLOBAL_COUNTER.resume()

counted_float/_core/counting/_counted_float.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import math
+
+from ._global_counter import GLOBAL_COUNTER
+from .models import FlopCounts
+
+
+class CountedFloat(float):
+    # -------------------------------------------------------------------------
+    #  FLOP COUNTING
+    # -------------------------------------------------------------------------
+    @classmethod
+    def get_global_flop_counts(cls) -> FlopCounts:
+        """
+        Returns the global FLOP counts for all CountedFloat instances.
+        """
+        return GLOBAL_COUNTER.flop_counts()
+
+    # -------------------------------------------------------------------------
+    #  CONSTRUCTOR
+    # -------------------------------------------------------------------------
+    def __new__(cls, value: float):
+        self = super().__new__(cls, value)
+        return self
+
+    def __str__(self):
+        return self.__repr__()
+
+    def __repr__(self):
+        return f"CountedFloat({super().__repr__()})"
+
+    def __hash__(self):
+        return super().__hash__()
+
+    # -------------------------------------------------------------------------
+    #  OVERLOADED MATH OPERATIONS
+    # -------------------------------------------------------------------------
+    def __abs__(self) -> CountedFloat:
+        """abs(x)"""
+        GLOBAL_COUNTER.incr_abs()
+        return CountedFloat(super().__abs__())
+
+    def __neg__(self) -> CountedFloat:
+        """-x"""
+        GLOBAL_COUNTER.incr_minus()
+        return CountedFloat(super().__neg__())
+
+    def __eq__(self, other) -> bool:
+        """x==other or other==x"""
+        if isinstance(other, int) and other == 0:
+            GLOBAL_COUNTER.incr_cmp_zero()
+        else:
+            GLOBAL_COUNTER.incr_equals()
+        return super().__eq__(other)
+
+    def __ne__(self, other) -> bool:
+        """x!=other or other!=x"""
+        if isinstance(other, int) and other == 0:
+            GLOBAL_COUNTER.incr_cmp_zero()
+        else:
+            GLOBAL_COUNTER.incr_equals()
+        return super().__ne__(other)
+
+    def __lt__(self, other):
+        """x<other"""
+        if isinstance(other, int) and other == 0:
+            GLOBAL_COUNTER.incr_cmp_zero()
+        else:
+            GLOBAL_COUNTER.incr_lte()
+        return super().__lt__(other)
+
+    def __le__(self, other):
+        """x<=other"""
+        if isinstance(other, int) and other == 0:
+            GLOBAL_COUNTER.incr_cmp_zero()
+        else:
+            GLOBAL_COUNTER.incr_lte()
+        return super().__le__(other)
+
+    def __gt__(self, other):
+        """x>other"""
+        if isinstance(other, int) and other == 0:
+            GLOBAL_COUNTER.incr_cmp_zero()
+        else:
+            GLOBAL_COUNTER.incr_gte()
+        return super().__gt__(other)
+
+    def __ge__(self, other):
+        """x>=other"""
+        if isinstance(other, int) and other == 0:
+            GLOBAL_COUNTER.incr_cmp_zero()
+        else:
+            GLOBAL_COUNTER.incr_gte()
+        return super().__ge__(other)
+
+    def __round__(self, n=None) -> int:
+        """round(x, n)"""
+        if n:
+            raise ValueError("only n==None or n==0 are supported in a CountedFloat")
+        GLOBAL_COUNTER.incr_rnd()  # assuming n=0, otherwise we can't reliably count the flops
+        return super().__round__()
+
+    def __floor__(self) -> int:
+        """math.floor(x)"""
+        GLOBAL_COUNTER.incr_rnd()
+        return super().__floor__()
+
+    def __ceil__(self) -> int:
+        """math.ceil(x)"""
+        GLOBAL_COUNTER.incr_rnd()
+        return super().__ceil__()
+
+    def __add__(self, other) -> CountedFloat:
+        """x+other"""
+        GLOBAL_COUNTER.incr_add()
+        return CountedFloat(super().__add__(other))
+
+    def __radd__(self, other) -> CountedFloat:
+        """other+x"""
+        GLOBAL_COUNTER.incr_add()
+        return CountedFloat(super().__radd__(other))
+
+    def __sub__(self, other) -> CountedFloat:
+        """x-other"""
+        GLOBAL_COUNTER.incr_sub()
+        return CountedFloat(super().__sub__(other))
+
+    def __rsub__(self, other) -> CountedFloat:
+        """other-x"""
+        GLOBAL_COUNTER.incr_sub()
+        return CountedFloat(super().__rsub__(other))
+
+    def __mul__(self, other) -> CountedFloat:
+        """x*other or other*x"""
+        GLOBAL_COUNTER.incr_mul()
+        return CountedFloat(super().__mul__(other))
+
+    def __rmul__(self, other) -> CountedFloat:
+        """other*x"""
+        GLOBAL_COUNTER.incr_mul()
+        return CountedFloat(super().__rmul__(other))
+
+    def __truediv__(self, other) -> CountedFloat:
+        """x/other"""
+        GLOBAL_COUNTER.incr_div()
+        return CountedFloat(super().__truediv__(other))
+
+    def __rtruediv__(self, other) -> CountedFloat:
+        """other/x"""
+        GLOBAL_COUNTER.incr_div()
+        return CountedFloat(super().__rtruediv__(other))
+
+    def __pow__(self, other) -> CountedFloat:
+        """x**other"""
+        if isinstance(other, int) and other == 2:
+            GLOBAL_COUNTER.incr_mul()  # x^2 = x*x
+        else:
+            GLOBAL_COUNTER.incr_pow()
+        return CountedFloat(super().__pow__(other))
+
+    def __rpow__(self, other) -> CountedFloat:
+        """other**x"""
+        if isinstance(other, int) and other == 2:
+            GLOBAL_COUNTER.incr_pow2()
+        else:
+            GLOBAL_COUNTER.incr_pow()
+        return CountedFloat(super().__rpow__(other))
+
+
+# -------------------------------------------------------------------------
+#  override some methods of math module
+# -------------------------------------------------------------------------
+original_math_sqrt = math.sqrt
+original_math_log2 = math.log2
+
+
+def math_sqrt(x: float) -> float | CountedFloat:
+    if isinstance(x, CountedFloat):
+        GLOBAL_COUNTER.incr_sqrt()
+        return CountedFloat(original_math_sqrt(x))
+    else:
+        return original_math_sqrt(x)
+
+
+def math_log2(x: float) -> float | CountedFloat:
+    if isinstance(x, CountedFloat):
+        GLOBAL_COUNTER.incr_log2()
+        return CountedFloat(original_math_log2(x))
+    else:
+        return original_math_log2(x)
+
+
+def math_pow(x: float, y: float) -> float | CountedFloat:
+    return x**y
+
+
+# override math module methods
+math.sqrt = math_sqrt
+math.log2 = math_log2
+math.pow = math_pow

counted_float/_core/counting/_global_counter.py

@@ -0,0 +1,98 @@
+from counted_float._core.counting.models import FlopCounts
+
+
+class GlobalFlopCounter:
+    """
+    Global counter for FLOP operations.  Essentially this class wraps around a FlopCounts object,
+    limiting access to its fields (only allowing incrementing them) and providing a way to access copies of the counts.
+    On top of this, the class allows pausing and resuming counting globally.
+    """
+
+    # -------------------------------------------------------------------------
+    #  Constructor
+    # -------------------------------------------------------------------------
+    def __init__(self):
+        self.__counts = FlopCounts()
+        self.__incr = 1  # 1 if enabled, 0 if paused
+
+    # -------------------------------------------------------------------------
+    #  Pause / Resume / Status API
+    # -------------------------------------------------------------------------
+    def pause(self):
+        self.__incr = 0
+
+    def resume(self):
+        self.__incr = 1
+
+    def reset(self):
+        self.__counts.reset()
+        self.resume()
+
+    def is_active(self) -> bool:
+        return self.__incr > 0
+
+    def flop_counts(self) -> FlopCounts:
+        return self.__counts.copy()
+
+    def total_count(self) -> int:
+        """Shorthand for self.flop_counts().total_count()"""
+        return self.__counts.total_count()
+
+    def __getattr__(self, item):
+        # provide shorthand access to the counts
+        if item in FlopCounts.field_names():
+            return getattr(self.__counts, item)
+        else:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")
+
+    # -------------------------------------------------------------------------
+    #  Incrementing counts
+    # -------------------------------------------------------------------------
+    def incr_abs(self):
+        self.__counts.ABS += self.__incr
+
+    def incr_minus(self):
+        self.__counts.MINUS += self.__incr
+
+    def incr_equals(self):
+        self.__counts.EQUALS += self.__incr
+
+    def incr_gte(self):
+        self.__counts.GTE += self.__incr
+
+    def incr_lte(self):
+        self.__counts.LTE += self.__incr
+
+    def incr_cmp_zero(self):
+        self.__counts.CMP_ZERO += self.__incr
+
+    def incr_rnd(self):
+        self.__counts.RND += self.__incr
+
+    def incr_add(self):
+        self.__counts.ADD += self.__incr
+
+    def incr_sub(self):
+        self.__counts.SUB += self.__incr
+
+    def incr_mul(self):
+        self.__counts.MUL += self.__incr
+
+    def incr_div(self):
+        self.__counts.DIV += self.__incr
+
+    def incr_sqrt(self):
+        self.__counts.SQRT += self.__incr
+
+    def incr_pow2(self):
+        self.__counts.POW2 += self.__incr
+
+    def incr_log2(self):
+        self.__counts.LOG2 += self.__incr
+
+    def incr_pow(self):
+        self.__counts.POW += self.__incr
+
+
+# --- global variable through which we access the global counter ---
+GLOBAL_COUNTER = GlobalFlopCounter()

counted_float/_core/counting/config/__init__.py

@@ -0,0 +1,6 @@
+from ._config import get_flop_weights, set_flop_weights
+from ._defaults import (
+    get_default_consensus_flop_weights,
+    get_default_empirical_flop_weights,
+    get_default_theoretical_flop_weights,
+)

counted_float/_core/counting/config/_config.py

@@ -0,0 +1,54 @@
+from ..models import FlopWeights
+from ._defaults import get_default_consensus_flop_weights
+
+
+# =================================================================================================
+#  Central class for weight configuration
+# =================================================================================================
+class Config:
+    """Class to hold configuration settings for the counted_float package."""
+
+    # -------------------------------------------------------------------------
+    #  Internal State
+    # -------------------------------------------------------------------------
+
+    # these are the weights that are used to calculate weighted flop counts; update with set_flop_weights(...)
+    __weights: FlopWeights = get_default_consensus_flop_weights()
+
+    # -------------------------------------------------------------------------
+    #  Configuration Methods
+    # -------------------------------------------------------------------------
+    @classmethod
+    def set_flop_weights(cls, weights: FlopWeights):
+        """
+        Set the weights for the flops used in the package.  These weights will be used in any calculation of
+        weighted flops, going forward.
+        :param weights: FlopWeights instance containing the weights.
+        """
+        cls.__weights = weights
+
+    @classmethod
+    def get_flop_weights(cls) -> FlopWeights:
+        """
+        Get the currently configured flop weights.
+        """
+        return cls.__weights.model_copy()
+
+
+# =================================================================================================
+#  Functional accessors
+# =================================================================================================
+def set_flop_weights(weights: FlopWeights):
+    """
+    Set the weights for the flops used in the package.  These weights will be used in any calculation of
+    weighted flops, going forward.
+    :param weights: FlopWeights instance containing the weights.
+    """
+    Config.set_flop_weights(weights)
+
+
+def get_flop_weights() -> FlopWeights:
+    """
+    Get the currently configured flop weights.
+    """
+    return Config.get_flop_weights()

counted_float/_core/counting/config/_defaults.py

@@ -0,0 +1,49 @@
+from functools import cache
+
+from counted_float._core.counting._builtin_data import BuiltInData
+
+from ..models import FlopWeights
+
+
+@cache
+def get_default_consensus_flop_weights(rounded: bool = True) -> FlopWeights:
+    """
+    Get the default CONSENSUS flop weights.
+    Computed as the geo-mean of the unrounded empirical and theoretical weights, rounded to the nearest integer.
+    """
+    weights = FlopWeights.as_geo_mean(
+        [
+            get_default_empirical_flop_weights(rounded=False),
+            get_default_theoretical_flop_weights(rounded=False),
+        ]
+    )
+    if rounded:
+        return weights.round()
+    else:
+        return weights
+
+
+@cache
+def get_default_empirical_flop_weights(rounded: bool = True) -> FlopWeights:
+    """
+    Get the default EMPIRICAL flop weights.
+    Computed as the geo-mean of flop weights estimated from built-in benchmark results.
+    """
+    weights = FlopWeights.as_geo_mean([v.flop_weights for v in BuiltInData.benchmarks().values()])
+    if rounded:
+        return weights.round()
+    else:
+        return weights
+
+
+@cache
+def get_default_theoretical_flop_weights(rounded: bool = True) -> FlopWeights:
+    """
+    Get the default THEORETICAL flop weights.
+    Computed as the geo-mean of flop weights estimated from built-in instruction latency analyse.
+    """
+    weights = FlopWeights.as_geo_mean([v.flop_weights for v in BuiltInData.specs().values()])
+    if rounded:
+        return weights.round()
+    else:
+        return weights

counted_float/_core/counting/models/__init__.py

@@ -0,0 +1,13 @@
+from ._base import MyBaseModel
+from ._flop_counts import FlopCounts
+from ._flop_type import FlopType
+from ._flop_weights import FlopWeights
+from ._flops_benchmark_result import (
+    BenchmarkSettings,
+    FlopsBenchmarkDurations,
+    FlopsBenchmarkResults,
+    Quantiles,
+    SystemInfo,
+)
+from ._fpu_instruction import FPUInstruction
+from ._fpu_specs import InstructionLatencies

counted_float/_core/counting/models/_base.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class MyBaseModel(BaseModel):
+    def show(self):
+        print(self.model_dump_json(indent=4))

counted_float/_core/counting/models/_flop_counts.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import dataclasses
+
+from ._flop_type import FlopType
+from ._flop_weights import FlopWeights
+
+
+@dataclasses.dataclass(slots=True)
+class FlopCounts:
+    """
+    Class to keep track of flop counts per flop type.  The implementation is different from
+    the FlopWeights class, for two reasons:
+        - there's no need for (de)serialization, hence no usage of Pydantic
+        - we want to minimize overhead of flop counting, hence use no dict in favor of explicit fields per flop type
+    """
+
+    # --- Counting fields ---------------------------------
+    ABS: int = 0
+    MINUS: int = 0
+    EQUALS: int = 0
+    GTE: int = 0
+    LTE: int = 0
+    CMP_ZERO: int = 0
+    RND: int = 0
+    ADD: int = 0
+    SUB: int = 0
+    MUL: int = 0
+    DIV: int = 0
+    SQRT: int = 0
+    POW2: int = 0
+    LOG2: int = 0
+    POW: int = 0
+
+    # --- math --------------------------------------------
+    def __add__(self, other: FlopCounts) -> FlopCounts:
+        return FlopCounts(**{attr: getattr(self, attr) + getattr(other, attr) for attr in self.field_names()})
+
+    def __sub__(self, other: FlopCounts) -> FlopCounts:
+        return FlopCounts(**{attr: getattr(self, attr) - getattr(other, attr) for attr in self.field_names()})
+
+    # --- extract info ------------------------------------
+    def as_dict(self) -> dict[FlopType, int]:
+        """Return the flop counts as a dictionary with FlopType keys."""
+        return {flop_type: getattr(self, flop_type.name) for flop_type in FlopType}
+
+    def total_count(self) -> int:
+        """Sum of all flop counts."""
+        return sum(getattr(self, attr) for attr in self.field_names())
+
+    def total_weighted_cost(self, weights: FlopWeights | None = None) -> float:
+        """
+        Returns a weighted total count of all flops (counterpart of the unweighted total_count() method),
+        using the provided weights in the computations.
+        When omitted, the currently configured weights (see Config class) will be used.
+        """
+        if not weights:
+            from counted_float._core.counting.config import get_flop_weights
+
+            weights = get_flop_weights()
+
+        return sum([getattr(self, flop_type.name) * weights.weights[flop_type] for flop_type in FlopType])
+
+    # --- other -------------------------------------------
+    def reset(self):
+        """Reset all counts to 0"""
+        for attr in self.field_names():
+            setattr(self, attr, 0)
+
+    def copy(self) -> FlopCounts:
+        return FlopCounts(**dataclasses.asdict(self))
+
+    @classmethod
+    def field_names(cls) -> list[str]:
+        return [field.name for field in dataclasses.fields(cls)]

counted_float/_core/counting/models/_flop_type.py

@@ -0,0 +1,47 @@
+from counted_float._core._backwards_compat import StrEnum
+
+
+class FlopType(StrEnum):
+    """
+    Enum describing the different types of floating-point operations,
+    each of which are counted separately and can potentially have different weights.
+
+    Enum                Math                                    ~corresponding
+    Member              Operation                               x87 instruction(s)
+    -------             ----------                              -------------------
+
+    ABS                 abs(x)                                  FABS
+    MINUS               -x                                      FCHS
+    EQUALS              x == y                                  FCOM
+    GTE                 x >= y                                  FCOM
+    LTE                 x <= y                                  FCOM
+    CMP_ZERO            x >= 0                                  FTST
+    RND                 ceil(x) or floor(x)                     FRNDINT
+    ADD                 x + y                                   FADD
+    SUB                 x - y                                   FSUB
+    MUL                 x * y                                   FMUL
+    DIV                 x / y                                   FDIV
+    SQRT                sqrt(x)                                 FSQRT
+    POW2                2**x                                    > F2XM1
+    LOG2                log2(x)                                 FYLX2
+    POW                 x**y                                    > F2XM1 + FYLX2 + FMUL
+    """
+
+    ABS = "abs(x)"
+    MINUS = "-x"
+    EQUALS = "x==y"
+    GTE = "x>=y"
+    LTE = "x<=y"
+    CMP_ZERO = "x>=0"
+    RND = "round(x)"
+    ADD = "x+y"
+    SUB = "x-y"
+    MUL = "x*y"
+    DIV = "x/y"
+    SQRT = "sqrt(x)"
+    POW2 = "2^x"
+    LOG2 = "log2(x)"
+    POW = "x^y"
+
+    def long_name(self) -> str:
+        return f"FlopType.{self.name:<9}  [{self.value}]"