import contextlib
import functools
from cupy import cuda
@contextlib.contextmanager
[docs]def time_range(message, color_id=None, argb_color=None):
"""A context manager to describe the enclosed block as a nested range
>>> from cupy import prof
>>> with cupy.prof.time_range('some range in green', color_id=0):
... # do something you want to measure
... pass
Args:
message: Name of a range.
color_id: range color ID
argb_color: range color in ARGB (e.g. 0xFF00FF00 for green)
.. seealso:: :func:`cupy.cuda.nvtx.RangePush`
:func:`cupy.cuda.nvtx.RangePop`
"""
if not cuda.nvtx_enabled:
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')
if argb_color is not None:
cuda.nvtx.RangePushC(message, argb_color)
else:
if color_id is None:
color_id = -1
cuda.nvtx.RangePush(message, color_id)
try:
yield
finally:
cuda.nvtx.RangePop()
[docs]class TimeRangeDecorator(object):
"""Decorator to mark function calls with range in NVIDIA profiler
Decorated function calls are marked as ranges in NVIDIA profiler timeline.
>>> from cupy import prof
>>> @cupy.prof.TimeRangeDecorator()
... def function_to_profile():
... pass
Args:
message (str): Name of a range, default use ``func.__name__``.
color_id: range color ID
argb_color: range color in ARGB (e.g. 0xFF00FF00 for green)
.. seealso:: :func:`cupy.nvtx.range`
:func:`cupy.cuda.nvtx.RangePush`
:func:`cupy.cuda.nvtx.RangePop`
"""
def __init__(self, message=None, color_id=None, argb_color=None):
if not cuda.nvtx_enabled:
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
def __enter__(self):
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):
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