Add optional GPU execution via CuPy

README.md

@@ -21,6 +21,31 @@
 
 **flopscope** is a drop-in replacement for a subset of NumPy that counts floating-point operations as you compute. Algorithms submitted to the ARC Whitebox Estimation Challenge are scored by their analytical FLOP cost, not wall-clock time, so researchers can focus on **algorithmic innovation** rather than hardware tuning. Every arithmetic call deducts from a fixed budget; exceed it and execution stops immediately.
 
+## Optional GPU execution
+
+flopscope can optionally execute selected backend calls with CuPy while
+preserving the public CPU `flopscope.numpy.ndarray` contract and analytical
+FLOP accounting. Install the optional extra and enable GPU execution locally:
+
+```sh
+pip install "flopscope[gpu]"
+FLOPSCOPE_GPU=1 python your_script.py
+```
+
+or:
+
+```python
+import flopscope as flops
+
+flops.configure_gpu(True)
+print(flops.gpu_status())
+```
+
+Results are copied back to CPU after each offloaded operation. The default gate
+keeps tiny or transfer-heavy calls on CPU with
+`FLOPSCOPE_GPU_MIN_FLOPS=5000000`,
+`FLOPSCOPE_GPU_MIN_FLOPS_PER_BYTE=0.05`, and no transfer-byte floor.
+
 ## Why flopscope?
 
 <table>

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
 
 [project.optional-dependencies]
 server = ["flopscope-server==0.8.0rc4"]
+gpu = ["cupy-cuda12x>=13.0.0"]
 dev = [
     "commitizen>=4.0",
     "gitlint>=0.19",

src/flopscope/__init__.py

@@ -32,7 +32,7 @@
 
 from flopscope._registry import REGISTRY_META as _REGISTRY_META
 
-__version__ = f"0.8.0rc4+np{_np.__version__}"
+__version__ = f"0.8.0rc4+gpu.np{_np.__version__}"
 __numpy_version__ = _np.__version__
 __numpy_pinned__ = _REGISTRY_META["numpy_version"]
 __numpy_supported__ = _REGISTRY_META.get("numpy_supported", ">=2.0.0,<2.5.0")
@@ -61,6 +61,12 @@
 )
 from flopscope._config import configure  # noqa: F401,E402
 from flopscope._display import budget_live, budget_summary  # noqa: F401,E402
+from flopscope._gpu import (  # noqa: F401,E402
+    configure_gpu,
+    gpu_available,
+    gpu_enabled,
+    gpu_status,
+)
 
 # --- Module base class ---
 from flopscope._module import Module  # noqa: F401,E402
@@ -266,9 +272,13 @@ def tier2_reduction_cost(a, axis=None, *, dense_per_output_cost=None):
     "budget_summary_dict",
     "clear_cache",
     "configure",
+    "configure_gpu",
     "einsum_accumulation_cost",
     "einsum_cache_info",
     "einsum_clear_caches",
+    "gpu_available",
+    "gpu_enabled",
+    "gpu_status",
     "is_symmetric",
     "namespace",
     "numpy",

src/flopscope/_budget.py

@@ -214,6 +214,21 @@ def _call_numpy(fn: Any, *args: Any, **kwargs: Any) -> Any:
     """
     t0 = time.perf_counter()
     try:
+        from flopscope._gpu import maybe_call_gpu
+
+        flop_cost = None
+        budget = get_active_budget()
+        timer = budget._current_op_timer if budget is not None else None
+        op_index = getattr(timer, "_op_index", None)
+        if op_index is not None:
+            try:
+                flop_cost = budget._op_log[op_index].flop_cost  # type: ignore[union-attr]
+            except (IndexError, TypeError):
+                flop_cost = None
+
+        used_gpu, result = maybe_call_gpu(fn, *args, flop_cost=flop_cost, **kwargs)
+        if used_gpu:
+            return result
         return fn(*args, **kwargs)
     finally:
         d = time.perf_counter() - t0