31 lines
870 B
Python
31 lines
870 B
Python
import copy
|
|
import dataclasses
|
|
from contextlib import contextmanager
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class CompilationCounter:
|
|
num_graphs_seen: int = 0
|
|
# including the splitting ops
|
|
num_piecewise_graphs_seen: int = 0
|
|
# not including the splitting ops
|
|
num_piecewise_capturable_graphs_seen: int = 0
|
|
num_inductor_compilations: int = 0
|
|
num_cudagraph_caputured: int = 0
|
|
|
|
def clone(self) -> "CompilationCounter":
|
|
return copy.deepcopy(self)
|
|
|
|
@contextmanager
|
|
def expect(self, **kwargs):
|
|
old = self.clone()
|
|
yield
|
|
for k, v in kwargs.items():
|
|
assert getattr(self, k) - getattr(old, k) == v, (
|
|
f"{k} not as expected, before it is {getattr(old, k)}"
|
|
f", after it is {getattr(self, k)}, "
|
|
f"expected diff is {v}")
|
|
|
|
|
|
compilation_counter = CompilationCounter()
|