| """Opt-in per-phase latency profile for :func:`deg.infer.fill`. |
| |
| Off by default (zero cost: one ContextVar read per phase). Inside ``with Profile() as p`` |
| every :func:`phase` synchronizes CUDA at its edges and accumulates wall time, call count |
| and the counters passed to :func:`count` (forward calls, rows, tokens). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import time |
| from collections.abc import Iterator |
| from contextlib import contextmanager |
| from contextvars import ContextVar |
| from dataclasses import dataclass, field |
| from typing import Any |
|
|
| import torch |
|
|
| _ACTIVE: ContextVar[Profile | None] = ContextVar("deg_profile", default=None) |
|
|
|
|
| @dataclass |
| class Profile: |
| seconds: dict[str, float] = field(default_factory=dict) |
| calls: dict[str, int] = field(default_factory=dict) |
| counters: dict[str, int] = field(default_factory=dict) |
| _token: Any = None |
|
|
| def __enter__(self) -> Profile: |
| self._token = _ACTIVE.set(self) |
| return self |
|
|
| def __exit__(self, *exc: object) -> None: |
| _ACTIVE.reset(self._token) |
|
|
| def to_json(self) -> dict[str, Any]: |
| return {"seconds": dict(self.seconds), "calls": dict(self.calls), |
| "counters": dict(self.counters)} |
|
|
|
|
| def _sync() -> None: |
| if torch.cuda.is_available() and torch.cuda.is_initialized(): |
| torch.cuda.synchronize() |
|
|
|
|
| @contextmanager |
| def phase(name: str) -> Iterator[None]: |
| prof = _ACTIVE.get() |
| if prof is None: |
| yield |
| return |
| _sync() |
| t0 = time.perf_counter() |
| try: |
| yield |
| finally: |
| _sync() |
| prof.seconds[name] = prof.seconds.get(name, 0.0) + time.perf_counter() - t0 |
| prof.calls[name] = prof.calls.get(name, 0) + 1 |
|
|
|
|
| def count(**values: int) -> None: |
| prof = _ACTIVE.get() |
| if prof is not None: |
| for k, v in values.items(): |
| prof.counters[k] = prof.counters.get(k, 0) + int(v) |
|
|
