Blame for axis_numeric/gpu_exec.py

"""Optional CUDA critical section. The line process binds the shared gate."""

from __future__ import annotations

from collections.abc import Callable
from typing import TypeVar

T = TypeVar("T")
Runner = Callable[[Callable[[], T]], tuple[T, float]]

_runner: Runner | None = None


def bind_gpu_runner(runner: Runner) -> None:
    global _runner
    _runner = runner


def gpu_call(fn: Callable[[], T]) -> T:
    if _runner is None:
        return fn()
    value, _wait_ms = _runner(fn)
    return value