62 lines
1.3 KiB
Python
62 lines
1.3 KiB
Python
from triton._C.libproton import proton as libproton
|
|
from .flags import get_profiling_on
|
|
from functools import wraps
|
|
|
|
|
|
class state:
|
|
"""
|
|
A context manager and decorator for entering and exiting a state.
|
|
|
|
Usage:
|
|
context manager:
|
|
```python
|
|
with proton.state("test0"):
|
|
foo[1,](x, y)
|
|
```
|
|
|
|
decorator:
|
|
```python
|
|
@proton.state("test0")
|
|
def foo(x, y):
|
|
...
|
|
```
|
|
|
|
Args:
|
|
name (str): The name of the state.
|
|
"""
|
|
|
|
def __init__(self, name: str) -> None:
|
|
self.name = name
|
|
|
|
def __enter__(self):
|
|
if not get_profiling_on():
|
|
return self
|
|
libproton.enter_state(self.name)
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
if not get_profiling_on():
|
|
return
|
|
libproton.exit_state()
|
|
|
|
def __call__(self, func):
|
|
|
|
@wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
if get_profiling_on():
|
|
libproton.enter_state(self.name)
|
|
ret = func(*args, **kwargs)
|
|
if get_profiling_on():
|
|
libproton.exit_state()
|
|
return ret
|
|
|
|
return wrapper
|
|
|
|
|
|
def enter_state(name: str) -> None:
|
|
libproton.enter_state(name)
|
|
|
|
|
|
def exit_state() -> None:
|
|
libproton.exit_state()
|