1# mypy: allow-untyped-defs 2r"""This package adds support for NVIDIA Tools Extension (NVTX) used in profiling.""" 3 4from contextlib import contextmanager 5 6 7try: 8 from torch._C import _nvtx 9except ImportError: 10 11 class _NVTXStub: 12 @staticmethod 13 def _fail(*args, **kwargs): 14 raise RuntimeError( 15 "NVTX functions not installed. Are you sure you have a CUDA build?" 16 ) 17 18 rangePushA = _fail 19 rangePop = _fail 20 markA = _fail 21 22 _nvtx = _NVTXStub() # type: ignore[assignment] 23 24__all__ = ["range_push", "range_pop", "range_start", "range_end", "mark", "range"] 25 26 27def range_push(msg): 28 """ 29 Push a range onto a stack of nested range span. Returns zero-based depth of the range that is started. 30 31 Args: 32 msg (str): ASCII message to associate with range 33 """ 34 return _nvtx.rangePushA(msg) 35 36 37def range_pop(): 38 """Pop a range off of a stack of nested range spans. Returns the zero-based depth of the range that is ended.""" 39 return _nvtx.rangePop() 40 41 42def range_start(msg) -> int: 43 """ 44 Mark the start of a range with string message. It returns an unique handle 45 for this range to pass to the corresponding call to rangeEnd(). 46 47 A key difference between this and range_push/range_pop is that the 48 range_start/range_end version supports range across threads (start on one 49 thread and end on another thread). 50 51 Returns: A range handle (uint64_t) that can be passed to range_end(). 52 53 Args: 54 msg (str): ASCII message to associate with the range. 55 """ 56 return _nvtx.rangeStartA(msg) 57 58 59def range_end(range_id) -> None: 60 """ 61 Mark the end of a range for a given range_id. 62 63 Args: 64 range_id (int): an unique handle for the start range. 65 """ 66 _nvtx.rangeEnd(range_id) 67 68 69def mark(msg): 70 """ 71 Describe an instantaneous event that occurred at some point. 72 73 Args: 74 msg (str): ASCII message to associate with the event. 75 """ 76 return _nvtx.markA(msg) 77 78 79@contextmanager 80def range(msg, *args, **kwargs): 81 """ 82 Context manager / decorator that pushes an NVTX range at the beginning 83 of its scope, and pops it at the end. If extra arguments are given, 84 they are passed as arguments to msg.format(). 85 86 Args: 87 msg (str): message to associate with the range 88 """ 89 range_push(msg.format(*args, **kwargs)) 90 try: 91 yield 92 finally: 93 range_pop() 94