/
_time_range.py
89 lines (73 loc) · 2.92 KB
/
_time_range.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import functools
from cupy import cuda
from cupy_backends.cuda.api import runtime
# Note: We use an (old-fashioned) custom object instead of
# @contextlib.contextmanager because for backward compatibility
# when used as a decorator this object needs to fetch the target
# function name.
class time_range:
"""Mark function calls with ranges using NVTX/rocTX. This object can be
used either as a decorator or a context manager.
When used as a decorator, the decorated function calls are marked as
ranges:
>>> from cupyx.profiler import time_range
>>> @time_range()
... def function_to_profile():
... pass
When used as a context manager, it describes the enclosed block as a nested
range:
>>> from cupyx.profiler import time_range
>>> with time_range('some range in green', color_id=0):
... # do something you want to measure
... pass
The marked ranges are visible in the profiler (such as nvvp, nsys-ui, etc)
timeline.
Args:
message (str): Name of a range. When used as a decorator, the default
is ``func.__name__``.
color_id: range color ID
argb_color: range color in ARGB (e.g. 0xFF00FF00 for green)
sync (bool): If ``True``, waits for completion of all outstanding
processing on GPU before calling :func:`cupy.cuda.nvtx.RangePush()`
or :func:`cupy.cuda.nvtx.RangePop()`
.. seealso:: :func:`cupy.cuda.nvtx.RangePush`,
:func:`cupy.cuda.nvtx.RangePop`
"""
def __init__(
self, message=None, color_id=None, argb_color=None, sync=False):
if not cuda.nvtx.available:
raise RuntimeError('nvtx is not installed')
if color_id is not None and argb_color is not None:
raise ValueError(
'Only either color_id or argb_color can be specified'
)
self.message = message
self.color_id = color_id if color_id is not None else -1
self.argb_color = argb_color
self.sync = sync
def __enter__(self):
if self.message is None:
raise ValueError(
'when used as a context manager, the message argument cannot '
'be None')
if self.sync:
runtime.deviceSynchronize()
if self.argb_color is not None:
cuda.nvtx.RangePushC(self.message, self.argb_color)
else:
cuda.nvtx.RangePush(self.message, self.color_id)
return self
def __exit__(self, exc_type, exc_value, traceback):
if self.sync:
runtime.deviceSynchronize()
cuda.nvtx.RangePop()
def _recreate_cm(self, message):
if self.message is None:
self.message = message
return self
def __call__(self, func):
@functools.wraps(func)
def inner(*args, **kwargs):
with self._recreate_cm(func.__name__):
return func(*args, **kwargs)
return inner