1# mypy: allow-untyped-defs 2from contextlib import contextmanager 3 4 5try: 6 from torch._C import _itt 7except ImportError: 8 9 class _ITTStub: 10 @staticmethod 11 def _fail(*args, **kwargs): 12 raise RuntimeError( 13 "ITT functions not installed. Are you sure you have a ITT build?" 14 ) 15 16 @staticmethod 17 def is_available(): 18 return False 19 20 rangePush = _fail 21 rangePop = _fail 22 mark = _fail 23 24 _itt = _ITTStub() # type: ignore[assignment] 25 26 27__all__ = ["is_available", "range_push", "range_pop", "mark", "range"] 28 29 30def is_available(): 31 """ 32 Check if ITT feature is available or not 33 """ 34 return _itt.is_available() 35 36 37def range_push(msg): 38 """ 39 Pushes a range onto a stack of nested range span. Returns zero-based 40 depth of the range that is started. 41 42 Arguments: 43 msg (str): ASCII message to associate with range 44 """ 45 return _itt.rangePush(msg) 46 47 48def range_pop(): 49 """ 50 Pops a range off of a stack of nested range spans. Returns the 51 zero-based depth of the range that is ended. 52 """ 53 return _itt.rangePop() 54 55 56def mark(msg): 57 """ 58 Describe an instantaneous event that occurred at some point. 59 60 Arguments: 61 msg (str): ASCII message to associate with the event. 62 """ 63 return _itt.mark(msg) 64 65 66@contextmanager 67def range(msg, *args, **kwargs): 68 """ 69 Context manager / decorator that pushes an ITT range at the beginning 70 of its scope, and pops it at the end. If extra arguments are given, 71 they are passed as arguments to msg.format(). 72 73 Args: 74 msg (str): message to associate with the range 75 """ 76 range_push(msg.format(*args, **kwargs)) 77 try: 78 yield 79 finally: 80 range_pop() 81